diff --git a/src/claude_agent_sdk/_internal/query.py b/src/claude_agent_sdk/_internal/query.py index 7aded9384..66f10f06e 100644 --- a/src/claude_agent_sdk/_internal/query.py +++ b/src/claude_agent_sdk/_internal/query.py @@ -17,6 +17,7 @@ from .._errors import ProcessError from ..types import ( + TERMINAL_TASK_STATUSES, PermissionMode, PermissionResultAllow, PermissionResultDeny, @@ -38,6 +39,22 @@ logger = logging.getLogger(__name__) +# Task types whose completion runs a follow-up turn, and which therefore may +# still need the control channel after the turn's result frame. +# +# This mirrors the set the CLI itself holds a result back for, which is +# narrower than its notion of "delegated agent work". The types left out are +# left out on purpose, and none of them is merely an oversight: +# - background shells and monitors run indefinitely by design, so deferring +# the close on one withholds it forever rather than briefly; +# - teammates are long-lived too — their status stays running for their whole +# lifetime, so they never settle the ledger; +# - remote agents can be long-running monitors the CLI likewise refuses to +# wait on. +# Anything added here must be a type that reliably reaches a terminal status, +# or it will hang the query (see Query._track_task_lifecycle). +DEFERRING_TASK_TYPES = frozenset({"local_agent", "local_workflow"}) + def _convert_hook_output_for_cli(hook_output: dict[str, Any]) -> dict[str, Any]: """Convert Python-safe field names to CLI-expected field names. @@ -128,8 +145,17 @@ def __init__( self._closed = False self._initialization_result: dict[str, Any] | None = None - # Track first result for proper stream closure with SDK MCP servers + # Set when a run-ending result arrives (a result frame with no tasks + # in flight) so the stdin-closing waiter can wake; see #1088 and + # _inflight_tasks below. Named for history — it once tracked the + # literal first result. self._first_result_event = anyio.Event() + # Task IDs of started-but-not-finished tasks. A result frame only ends + # one turn, not the run: background tasks keep running past it and + # still need stdin for hook/SDK-MCP control responses (see #1088), so + # a result that arrives while this set is non-empty must not close + # stdin. + self._inflight_tasks: set[str] = set() # 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. @@ -293,6 +319,11 @@ async def _read_messages(self) -> None: ) continue + # Track task lifecycle frames so results can tell "one turn + # ended" apart from "the run is done" (see #1088). + if msg_type == "system": + self._track_task_lifecycle(message) + # Track results for proper stream closure if msg_type == "result": # Flush pending transcript mirror entries before yielding @@ -300,7 +331,21 @@ async def _read_messages(self) -> None: # SessionStore being up to date for this turn. if self._transcript_mirror_batcher is not None: await self._transcript_mirror_batcher.flush() - self._first_result_event.set() + if 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 + # and fails SDK-MCP calls with "Stream closed" + # (#1088). Each task completion wakes the parent for + # a follow-up turn, so a later result frame arrives + # with no tasks in flight and closes stdin then. + logger.debug( + "Result received with %d task(s) in flight; " + "keeping stdin open", + len(self._inflight_tasks), + ) + else: + self._first_result_event.set() if message.get("is_error"): errors = message.get("errors") or [] self._last_error_result_text = "; ".join(errors) or str( @@ -806,19 +851,80 @@ async def stop_task(self, task_id: str) -> None: } ) + def _track_task_lifecycle(self, message: dict[str, Any]) -> None: + """Track in-flight tasks from ``system`` task lifecycle frames. + + ``task_started`` marks a task in flight; ``task_notification`` or a + ``task_updated`` patch with a terminal status clears it. Terminal + completion can arrive as either frame (not every terminal task emits + a notification), so both are handled; ``discard`` keeps the pair + idempotent. + + 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. + + Only delegated agent work is tracked (``DEFERRING_TASK_TYPES``). A + background *shell* — ``Bash(run_in_background=True)`` on a dev server or + ``tail -f`` — is also reported through these frames, but it may never + reach a terminal status, and the CLI in stream-json mode only exits on + stdin EOF. Tracking one would therefore withhold the close forever + rather than briefly: no terminal frame, no process exit, so not even the + reader's ``finally`` runs. Agent tasks are the ones whose completion + wakes the parent for the follow-up turn this relies on; shells and + monitors are bounded by the CLI's own post-close cleanup instead. + + ``background_tasks_changed`` is deliberately *not* consumed, in either + direction. Its payload is the live *background* set, while a subagent is + registered in the foreground and only flips to backgrounded later, + without a second ``task_started``. So the snapshot omits tracked work + that is still running: narrowing against it would drop an agent that + goes on to outlive its turn, which is the very close-too-early bug this + method exists to prevent. Widening from it is no better — the snapshot + spans every background task type and carries nothing marking an + observer agent, whose start and terminal frames are both suppressed, so + it could admit an id no later frame ever clears. The lifecycle frames + are the only self-consistent source here (see #1088). + """ + subtype = message.get("subtype") + task_id = message.get("task_id") + if not task_id: + return + if subtype == "task_started": + if message.get("task_type") in DEFERRING_TASK_TYPES: + self._inflight_tasks.add(task_id) + elif subtype == "task_notification": + self._inflight_tasks.discard(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) + async def wait_for_result_and_end_input(self) -> None: - """Wait for the first result (if needed) then close stdin. + """Wait for a run-ending result (if needed) then close stdin. If SDK MCP servers or hooks require bidirectional communication, - keeps stdin open until the first result arrives. The control protocol - requires stdin to remain open for the entire conversation, so no - timeout is applied. The event is guaranteed to fire: either when the - result message arrives, or in _read_messages' finally block if the - process exits early. + keeps stdin open until a result arrives with no tasks in flight. A + result frame ends one turn, not necessarily the run: background tasks + keep running past it and still need stdin for hook/SDK-MCP control + responses (see #1088). The control protocol requires stdin to remain + open for the entire conversation, so no timeout is applied. The event + is guaranteed to fire: either when a result message arrives with no + in-flight tasks (every task completion wakes the parent for a + follow-up turn, which ends in such a result), or in _read_messages' + finally block if the process exits early. """ if self.sdk_mcp_servers or self.hooks: logger.debug( - "Waiting for first result before closing stdin " + "Waiting for a run-ending result before closing stdin " f"(sdk_mcp_servers={len(self.sdk_mcp_servers)}, " f"has_hooks={bool(self.hooks)})" ) diff --git a/tests/test_query.py b/tests/test_query.py index 3ae65093d..a76201c7c 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -360,6 +360,317 @@ async def dummy_hook(input_data, tool_use_id, context): anyio.run(_test) +_TASK_STARTED = { + "type": "system", + "subtype": "task_started", + "task_id": "task-1", + "task_type": "local_agent", + "description": "background subagent", + "uuid": "uuid-ts1", + "session_id": "test", +} + +# A backgrounded shell (dev server, `tail -f`) is reported through the same +# frames but can run indefinitely, so it must never defer the stdin close. +_SHELL_TASK_STARTED = { + "type": "system", + "subtype": "task_started", + "task_id": "shell-1", + "task_type": "local_bash", + "description": "npm run dev", + "uuid": "uuid-ts2", + "session_id": "test", +} + +_TASK_NOTIFICATION = { + "type": "system", + "subtype": "task_notification", + "task_id": "task-1", + "status": "completed", + "output_file": "/tmp/task-1.output", + "summary": "done", + "uuid": "uuid-tn1", + "session_id": "test", +} + +_TASK_UPDATED_TERMINAL = { + "type": "system", + "subtype": "task_updated", + "task_id": "task-1", + "patch": {"status": "completed"}, +} + +_BACKGROUND_TASKS_CHANGED = { + "type": "system", + "subtype": "background_tasks_changed", + "tasks": [ + {"task_id": "task-1", "task_type": "local_agent"}, + {"task_id": "task-2", "task_type": "local_agent"}, + ], +} + + +def _make_result(uid): + return dict(_ASSISTANT_AND_RESULT[1], uuid=uid) + + +class TestStdinStaysOpenWithInflightTasks: + """A result frame with tasks in flight must not close stdin (#1088). + + Background tasks keep running past the turn's result frame and still + need stdin for hook/SDK-MCP control responses. Closing it there fails + those responses with "Stream closed" and silently disables hooks. + """ + + @pytest.mark.parametrize( + "drain_frame", + [_TASK_NOTIFICATION, _TASK_UPDATED_TERMINAL], + ids=["task_notification", "task_updated_terminal_patch"], + ) + def test_result_with_inflight_task_keeps_stdin_open(self, drain_frame): + """stdin stays open across an intermediate result, then closes on the + first result that arrives with no tasks in flight.""" + + async def _test(): + server = _make_greet_server() + end_input_calls = [] + + mock_transport = AsyncMock() + mock_transport.connect = AsyncMock() + mock_transport.close = AsyncMock() + mock_transport.write = AsyncMock() + mock_transport.is_ready = Mock(return_value=True) + + async def tracking_end_input(): + end_input_calls.append(True) + + mock_transport.end_input = tracking_end_input + + open_after_intermediate_result = None + closed_after_final_result = None + + async def mock_receive(): + nonlocal open_after_intermediate_result, closed_after_final_result + yield dict(_ASSISTANT_AND_RESULT[0]) + yield dict(_TASK_STARTED) + yield _make_result("uuid-r1") + # Let the stdin-closing waiter run if the first-result event + # (incorrectly) fired for the intermediate result. + for _ in range(20): + await anyio.sleep(0) + open_after_intermediate_result = not end_input_calls + yield dict(drain_frame) + yield _make_result("uuid-r2") + # Now the waiter should wake and close stdin before the + # stream ends (i.e. not merely via the reader's finally). + for _ in range(20): + await anyio.sleep(0) + closed_after_final_result = bool(end_input_calls) + + mock_transport.read_messages = mock_receive + + 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="Hello", + options=ClaudeAgentOptions(mcp_servers={"greeter": server}), + ): + messages.append(msg) + + assert open_after_intermediate_result is True + assert closed_after_final_result is 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() + 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"} + + # Non-terminal patch does not clear the task. + q._track_task_lifecycle( + { + "subtype": "task_updated", + "task_id": "task-1", + "patch": {"status": "running"}, + } + ) + assert q._inflight_tasks == {"task-1"} + + # Patch without a dict payload is ignored. + q._track_task_lifecycle( + {"subtype": "task_updated", "task_id": "task-1", "patch": None} + ) + assert q._inflight_tasks == {"task-1"} + + # Terminal patch clears it. + q._track_task_lifecycle(dict(_TASK_UPDATED_TERMINAL)) + assert q._inflight_tasks == set() + + # Draining an unknown/already-cleared task is a no-op, not an error. + q._track_task_lifecycle(dict(_TASK_NOTIFICATION)) + assert q._inflight_tasks == set() + + # Frames without a task_id are ignored. + q._track_task_lifecycle({"subtype": "task_started"}) + assert q._inflight_tasks == set() + + def test_shell_and_monitor_tasks_never_defer_the_close(self): + """Only delegated agent work defers the stdin close. + + A backgrounded shell can run forever, and the CLI only exits on stdin + EOF, so tracking one would withhold the close permanently rather than + briefly — the reader's ``finally`` never runs either (#1088). + """ + transport = AsyncMock() + transport.is_ready = Mock(return_value=True) + q = Query(transport=transport, is_streaming_mode=False) + + q._track_task_lifecycle(dict(_SHELL_TASK_STARTED)) + assert q._inflight_tasks == set() + + for task_type in ("monitor_mcp", "monitor_ws", "in_process_teammate"): + q._track_task_lifecycle( + { + "subtype": "task_started", + "task_id": f"{task_type}-1", + "task_type": task_type, + } + ) + assert q._inflight_tasks == set() + + # A start frame with no task_type at all is not assumed to be an agent. + q._track_task_lifecycle({"subtype": "task_started", "task_id": "unknown-1"}) + assert q._inflight_tasks == set() + + # Agent work still defers, alongside the ignored shell. + q._track_task_lifecycle(dict(_TASK_STARTED)) + assert q._inflight_tasks == {"task-1"} + + def test_background_snapshot_does_not_touch_the_ledger(self): + """background_tasks_changed is ignored in both directions (#1088). + + It reports the live *background* set, but a subagent is registered in + the foreground and only flips to backgrounded later without a second + ``task_started`` — so a tracked agent that is still running can be + absent from the snapshot entirely. + """ + transport = AsyncMock() + transport.is_ready = Mock(return_value=True) + q = Query(transport=transport, is_streaming_mode=False) + + # It cannot add: the snapshot spans every background task type and + # cannot distinguish an observer agent, whose bookends are suppressed. + q._track_task_lifecycle(dict(_BACKGROUND_TASKS_CHANGED)) + assert q._inflight_tasks == set() + + # Nor can it remove. A foreground subagent is tracked from its + # task_started but appears in no snapshot; narrowing against one would + # drop it, and when it is later auto-backgrounded no second + # task_started re-adds it — so the result frame would close stdin while + # it is still running, which is exactly the bug being fixed. + q._track_task_lifecycle(dict(_TASK_STARTED)) + q._track_task_lifecycle( + { + "type": "system", + "subtype": "background_tasks_changed", + "tasks": [{"task_id": "shell-1", "task_type": "local_bash"}], + } + ) + assert q._inflight_tasks == {"task-1"} + + # Not even an empty snapshot clears it; only a terminal frame does. + q._track_task_lifecycle( + {"type": "system", "subtype": "background_tasks_changed", "tasks": []} + ) + assert q._inflight_tasks == {"task-1"} + + q._track_task_lifecycle(dict(_TASK_NOTIFICATION)) + assert q._inflight_tasks == set() + + def test_never_ending_shell_does_not_wedge_stdin_open(self): + """A backgrounded shell that never finishes must not hang the query. + + Unlike the other streams in this file, this transport does not end on + its own: it models the real contract the stdin close depends on — the + CLI's stdout does not reach EOF until its stdin does. A tracked task + that never reaches a terminal status therefore cannot be rescued by + the reader's ``finally``, because that only runs once the process + exits (#1088). + """ + + async def _test(): + server = _make_greet_server() + stdin_closed = anyio.Event() + + mock_transport = AsyncMock() + mock_transport.connect = AsyncMock() + mock_transport.close = AsyncMock() + mock_transport.write = AsyncMock() + mock_transport.is_ready = Mock(return_value=True) + + async def end_input(): + stdin_closed.set() + + mock_transport.end_input = end_input + + async def mock_receive(): + yield dict(_ASSISTANT_AND_RESULT[0]) + yield dict(_SHELL_TASK_STARTED) + yield { + "type": "system", + "subtype": "background_tasks_changed", + "tasks": [{"task_id": "shell-1", "task_type": "local_bash"}], + } + yield _make_result("uuid-r1") + # The shell never exits, so no terminal frame ever arrives and + # the snapshot keeps listing it. Park until stdin closes. + await stdin_closed.wait() + + mock_transport.read_messages = mock_receive + + completed = False + 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 + + with anyio.move_on_after(10): + async for _ in query( + prompt="start the dev server", + options=ClaudeAgentOptions(mcp_servers={"greeter": server}), + ): + pass + completed = True + + assert stdin_closed.is_set(), "stdin was never closed" + assert completed, "query() did not terminate" + + anyio.run(_test) + + class TestAsyncIterablePromptWithSdkMcpServers: """Test that AsyncIterable prompts keep stdin open for SDK MCP servers."""