diff --git a/src/claude_agent_sdk/_internal/query.py b/src/claude_agent_sdk/_internal/query.py index e9869342a..3cdc9e2c0 100644 --- a/src/claude_agent_sdk/_internal/query.py +++ b/src/claude_agent_sdk/_internal/query.py @@ -156,6 +156,28 @@ def __init__( # a result that arrives while this set is non-empty must not close # stdin. self._inflight_tasks: set[str] = set() + # Subset of _inflight_tasks that has not yet been inflight during a + # processed result frame. A task leaves this set either when a + # result is processed while it is inflight (it has now had its + # "turn ended, I'm still running" result observed) or when it + # settles — in which case its exit was never covered by an observed + # result. See _task_settled_since_last_result below. + self._tasks_pending_first_result: set[str] = set() + # Set when a deferring task settles without ever having been seen + # inflight during a processed result frame; cleared the next time a + # result frame is evaluated. Covers the ordering gap where such a + # task settles and the turn's result frame both land with + # _inflight_tasks already empty at that result — closing stdin there + # would drop a continuation control request the settled task still + # owes (see #1190). A task that *was* seen inflight during an + # earlier result does not set this: that earlier result already + # covered the "turn ended, task still running" case, so the result + # that follows its settlement is the ordinary follow-up turn and + # should close stdin immediately, same as before this fix. + # The _read_messages finally block's unconditional + # _first_result_event.set() is the backstop if no further result + # frame ever arrives, so deferring here cannot hang the waiter. + self._task_settled_since_last_result = False # Set to the result's error text when the most recent message is a # result with is_error=True. Used to replace the generic "exit code 1" # ProcessError with the structured error the CLI already reported. @@ -332,6 +354,12 @@ async def _read_messages(self) -> None: if self._transcript_mirror_batcher is not None: await self._transcript_mirror_batcher.flush() if self._inflight_tasks: + # These tasks have now had a result observed while + # they were inflight; their eventual settlement is + # covered by *this* observation, so it must not defer + # a later result the way _task_settled_since_last_result + # does for a task that settles unobserved (#1190). + self._tasks_pending_first_result -= self._inflight_tasks # One turn ended, but background tasks are still # running and may need hook/SDK-MCP control responses # over stdin. Closing it now silently disables hooks @@ -344,6 +372,21 @@ async def _read_messages(self) -> None: "keeping stdin open", len(self._inflight_tasks), ) + elif self._task_settled_since_last_result: + # No task is in flight, but one reached a terminal + # status since the last result was evaluated — quite + # possibly in this very batch of frames, ahead of + # this result rather than because of it. That task's + # continuation turn may still owe a hook/SDK-MCP + # control round trip before the run truly ends + # (#1190). Keep stdin open through this result; the + # next one closes it, same as the plain-inflight case + # above. + logger.debug( + "Result received just after a task settled; " + "keeping stdin open for its continuation" + ) + self._task_settled_since_last_result = False else: self._first_result_event.set() if message.get("is_error"): @@ -869,14 +912,28 @@ def _track_task_lifecycle(self, message: dict[str, Any]) -> None: This is a mitigation, not a complete answer to #1088. An empty set means "nothing we know of is running", which is not the same as "the - run is over": a task that settles *before* the turn's result frame - leaves the set empty at that result, so stdin closes even though the - completion may still wake the parent for a continuation turn. No - ledger can close that gap, because the ledger cannot distinguish a - settled task whose continuation is pending from no work at all — that - needs a run-boundary signal from the CLI rather than an inference from - task bookkeeping. What this does fix is the common ordering, where the - task outlives the turn that spawned it. + run is over". The common ordering — the task outlives the turn that + spawned it — is fully handled: the result that ends that turn still + finds the task in ``_inflight_tasks`` and stdin stays open. + + A task that settles *before any result has ever been observed while + it was inflight* is handled too, but only for one extra result + rather than by distinguishing the run boundary outright (#1190). + ``_tasks_pending_first_result`` (see ``__init__``) tracks exactly + that condition; a settlement out of it sets + ``_task_settled_since_last_result``, which makes the very next + result frame — empty set or not — keep stdin open once more before + closing, on the chance that this is the pre-continuation result the + settled task is still owed a control round trip for. A task that + *has* already been seen inflight during a result does not set the + flag on settlement: that earlier result already covered "turn ended, + task still running", so the result that follows its settlement is + the ordinary next-turn result and closes stdin immediately, same as + before this mitigation. What no ledger can do — settled-flag or + inflight-count alike — is tell a run that is genuinely over from one + that will, after this one deferral, still need a further + continuation of its own; that needs a run-boundary signal from the + CLI rather than an inference from task bookkeeping. Only delegated agent work is tracked (``DEFERRING_TASK_TYPES``). A background *shell* — ``Bash(run_in_background=True)`` on a dev server or @@ -907,13 +964,29 @@ def _track_task_lifecycle(self, message: dict[str, Any]) -> None: if subtype == "task_started": if message.get("task_type") in DEFERRING_TASK_TYPES: self._inflight_tasks.add(task_id) + self._tasks_pending_first_result.add(task_id) elif subtype == "task_notification": - self._inflight_tasks.discard(task_id) + self._settle_task(task_id) elif subtype == "task_updated": patch = message.get("patch") status = patch.get("status") if isinstance(patch, dict) else None if status in TERMINAL_TASK_STATUSES: - self._inflight_tasks.discard(task_id) + self._settle_task(task_id) + + def _settle_task(self, task_id: str) -> None: + """Discards a task that reached a terminal status. + + If the task was never inflight during an already-processed result + frame, its exit is not yet covered by any result the caller has + seen — mark that the next result should be treated as possibly + pre-continuation rather than run-ending (see _track_task_lifecycle). + """ + if task_id not in self._inflight_tasks: + return + self._inflight_tasks.discard(task_id) + if task_id in self._tasks_pending_first_result: + self._tasks_pending_first_result.discard(task_id) + self._task_settled_since_last_result = True async def wait_for_result_and_end_input(self) -> None: """Wait for a run-ending result (if needed) then close stdin. diff --git a/tests/test_query.py b/tests/test_query.py index f6e98d194..32ce2b767 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -17,6 +17,7 @@ from claude_agent_sdk import ( AssistantMessage, ClaudeAgentOptions, + PermissionResultAllow, ResultMessage, create_sdk_mcp_server, query, @@ -494,6 +495,118 @@ async def mock_receive(): anyio.run(_test) + def test_task_settled_before_any_result_defers_close_for_late_control(self): + """A task that settles before its own result is ever observed keeps + stdin open through that result, so a late continuation control + request (e.g. a permission check) still gets a response (#1190). + + Unlike ``test_result_with_inflight_task_keeps_stdin_open`` above, + the task here never appears inflight during a processed result: the + very first result frame arrives *after* ``task_notification``. + Without the fix, that result would close stdin immediately and a + late control request would fail with "stdin already closed". + """ + + async def _test(): + server = _make_greet_server() + end_input_calls = [] + control_response_writes = [] + write_after_close = [] + + mock_transport = AsyncMock() + mock_transport.connect = AsyncMock() + mock_transport.close = AsyncMock() + mock_transport.is_ready = Mock(return_value=True) + + closed = False + + async def tracking_write(data): + if '"request_id":"late-control"' in data or ( + '"request_id": "late-control"' in data + ): + if closed: + write_after_close.append(True) + else: + control_response_writes.append(data) + + async def tracking_end_input(): + nonlocal closed + end_input_calls.append(True) + closed = True + + mock_transport.write = tracking_write + mock_transport.end_input = tracking_end_input + + closed_after_before_continuation_result = None + + async def mock_receive(): + nonlocal closed_after_before_continuation_result + yield dict(_TASK_STARTED) + yield dict(_TASK_NOTIFICATION) + yield _make_result("uuid-before-continuation") + for _ in range(20): + await anyio.sleep(0) + closed_after_before_continuation_result = bool(end_input_calls) + # The task's completion wakes a late continuation control + # request — e.g. a permission check for the tool call the + # subagent's result triggers next. + yield { + "type": "control_request", + "request_id": "late-control", + "request": { + "subtype": "can_use_tool", + "tool_name": "Read", + "input": {"file_path": "/tmp/late.txt"}, + "tool_use_id": "late-tool", + }, + } + for _ in range(20): + await anyio.sleep(0) + yield _make_result("uuid-final") + for _ in range(20): + await anyio.sleep(0) + + mock_transport.read_messages = mock_receive + + async def allow_tool(tool_name, tool_input, context): + return PermissionResultAllow(updated_input=tool_input) + + async def prompt_stream(): + yield { + "type": "user", + "message": {"role": "user", "content": "Hello"}, + } + + with ( + patch( + "claude_agent_sdk._internal.client.SubprocessCLITransport" + ) as mock_cls, + patch( + "claude_agent_sdk._internal.query.Query.initialize", + new_callable=AsyncMock, + ), + ): + mock_cls.return_value = mock_transport + + messages = [] + async for msg in query( + prompt=prompt_stream(), + options=ClaudeAgentOptions( + mcp_servers={"greeter": server}, + can_use_tool=allow_tool, + ), + ): + messages.append(msg) + + assert closed_after_before_continuation_result is False + assert len(control_response_writes) == 1 + assert write_after_close == [] + assert end_input_calls == [True] + results = [m for m in messages if isinstance(m, ResultMessage)] + assert len(results) == 2 + + anyio.run(_test) + def test_track_task_lifecycle_unit(self): """_track_task_lifecycle adds on start and clears only on terminal.""" transport = AsyncMock() @@ -531,6 +644,53 @@ def test_track_task_lifecycle_unit(self): q._track_task_lifecycle({"subtype": "task_started"}) assert q._inflight_tasks == set() + def test_settle_before_any_result_defers_the_next_result(self): + """A task that settles without ever being seen inflight during a + result must not let the very next result close stdin (#1190). + + ``task_started`` -> ``task_notification`` -> result is exactly the + ordering #1190 reports: the task's own turn result never arrived + while it was inflight, so its exit is not yet covered by any result + the caller has observed, and the CLI may still owe a continuation + control round trip after this result. + """ + transport = AsyncMock() + transport.is_ready = Mock(return_value=True) + q = Query(transport=transport, is_streaming_mode=False) + + q._track_task_lifecycle(dict(_TASK_STARTED)) + assert q._inflight_tasks == {"task-1"} + assert q._tasks_pending_first_result == {"task-1"} + assert q._task_settled_since_last_result is False + + q._track_task_lifecycle(dict(_TASK_NOTIFICATION)) + assert q._inflight_tasks == set() + assert q._tasks_pending_first_result == set() + assert q._task_settled_since_last_result is True + + def test_settle_after_a_seen_result_does_not_defer(self): + """A task settling after already being seen inflight during a result + does not set the deferral flag — that earlier result already + covered "turn ended, task still running", so the result that + follows settlement is the ordinary next-turn result (matches + ``test_result_with_inflight_task_keeps_stdin_open`` above, which + closes on that very result).""" + transport = AsyncMock() + transport.is_ready = Mock(return_value=True) + q = Query(transport=transport, is_streaming_mode=False) + + q._track_task_lifecycle(dict(_TASK_STARTED)) + assert q._tasks_pending_first_result == {"task-1"} + + # Simulates the "result" branch in _read_messages: a result was + # processed while the task was inflight. + q._tasks_pending_first_result -= q._inflight_tasks + assert q._tasks_pending_first_result == set() + + q._track_task_lifecycle(dict(_TASK_NOTIFICATION)) + assert q._inflight_tasks == set() + assert q._task_settled_since_last_result is False + def test_shell_and_monitor_tasks_never_defer_the_close(self): """Only delegated agent work defers the stdin close.