From 51324a57fdc23b7aa17c142e733823440dc7d193 Mon Sep 17 00:00:00 2001 From: shihyayou Date: Fri, 10 Jul 2026 16:31:05 +0800 Subject: [PATCH 1/4] fix(query): don't close stdin on a result frame while tasks are in flight MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A result frame marks the end of one turn, not the end of the run: a background Task keeps running past it and still needs stdin for hook and SDK-MCP control responses. wait_for_result_and_end_input() closed stdin on the first result frame, so a still-running subagent's SDK-MCP tool calls failed with "Stream closed" and its PreToolUse hooks were silently bypassed (built-in tools kept executing without the hook callback). Track in-flight tasks from the task_started / task_notification / terminal task_updated lifecycle frames (via the existing TERMINAL_TASK_STATUSES) and only treat a result frame as run-ending when no tasks are in flight. Each task completion wakes the parent for a follow-up turn that ends in another result frame, so stdin still closes promptly — including for chained background tasks — and the reader's finally block still guarantees the waiter is unblocked if the process exits early. Fixes #1088. Used AI assistance; reviewed and tested by me (repro'd the "Stream closed" failure end-to-end against the live CLI before the fix and verified the same scenario passes after; regression tests fail on current main; ruff, mypy, and the full unit suite green). Co-Authored-By: Claude Opus 4.8 (1M context) --- src/claude_agent_sdk/_internal/query.py | 67 +++++++++-- tests/test_query.py | 150 ++++++++++++++++++++++++ 2 files changed, 210 insertions(+), 7 deletions(-) diff --git a/src/claude_agent_sdk/_internal/query.py b/src/claude_agent_sdk/_internal/query.py index 7aded9384..7f216c765 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, @@ -130,6 +131,12 @@ def __init__( # Track first result for proper stream closure with SDK MCP servers 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 +300,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 +312,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,15 +832,42 @@ 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. + """ + task_id = message.get("task_id") + if not task_id: + return + subtype = message.get("subtype") + if subtype == "task_started": + 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( diff --git a/tests/test_query.py b/tests/test_query.py index 3ae65093d..e9a7508ed 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -360,6 +360,156 @@ async def dummy_hook(input_data, tool_use_id, context): anyio.run(_test) +_TASK_STARTED = { + "type": "system", + "subtype": "task_started", + "task_id": "task-1", + "description": "background subagent", + "uuid": "uuid-ts1", + "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"}, +} + + +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() + + class TestAsyncIterablePromptWithSdkMcpServers: """Test that AsyncIterable prompts keep stdin open for SDK MCP servers.""" From 1e9b40438653ac60efa7c4a3f7c67be45f2feee0 Mon Sep 17 00:00:00 2001 From: shihyayou Date: Fri, 10 Jul 2026 16:53:43 +0800 Subject: [PATCH 2/4] docs(query): align stdin-closure comments with run-ending-result logic Address Copilot review on #1103: the _first_result_event comment and the debug log still said "first result", but the waiter now wakes on a run-ending result (a result frame with no tasks in flight). Wording only; no behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/claude_agent_sdk/_internal/query.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/claude_agent_sdk/_internal/query.py b/src/claude_agent_sdk/_internal/query.py index 7f216c765..7013e68eb 100644 --- a/src/claude_agent_sdk/_internal/query.py +++ b/src/claude_agent_sdk/_internal/query.py @@ -129,7 +129,10 @@ 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 @@ -871,7 +874,7 @@ async def wait_for_result_and_end_input(self) -> None: """ 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)})" ) From 7a7f069f4cc841d45168990bec69f3293b899e7d Mon Sep 17 00:00:00 2001 From: shihyayou Date: Wed, 22 Jul 2026 09:39:00 +0800 Subject: [PATCH 3/4] feat(query): self-heal in-flight task tracking via background_tasks_changed Consume the CLI's authoritative `background_tasks_changed` snapshot frame ({"type":"system","subtype":"background_tasks_changed","tasks":[{"task_id":...}]}) as a set-replacement of `_inflight_tasks`, in addition to the incremental task_started / task_notification / terminal task_updated tracking. This self-heals a missed start or terminal frame (dropped lines, older CLIs with different terminal statuses): tracking drift can then only *delay* the stdin close (bounded by the reader's finally), never wedge it open or close it early. Handled before the task_id guard since the snapshot has no top-level task_id; a malformed frame leaves live tracking untouched. Suggested by @b-randongee, who observed the frame emitted reliably on CLI 2.1.215 in a production incident (#1088). Used AI assistance; reviewed and tested by me (ruff/mypy clean, full unit suite green on asyncio and trio). Co-Authored-By: Claude Opus 4.8 --- src/claude_agent_sdk/_internal/query.py | 20 +++++++- tests/test_query.py | 67 +++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/src/claude_agent_sdk/_internal/query.py b/src/claude_agent_sdk/_internal/query.py index 7013e68eb..c171b2d23 100644 --- a/src/claude_agent_sdk/_internal/query.py +++ b/src/claude_agent_sdk/_internal/query.py @@ -843,11 +843,29 @@ def _track_task_lifecycle(self, message: dict[str, Any]) -> None: completion can arrive as either frame (not every terminal task emits a notification), so both are handled; ``discard`` keeps the pair idempotent. + + ``background_tasks_changed`` carries an authoritative snapshot of the + running background tasks (``{"tasks": [{"task_id": ...}, ...]}``) and is + applied as a set-replacement. This self-heals a missed ``task_started`` + or terminal frame (dropped lines, older CLIs with different terminal + statuses): tracking drift can then only *delay* the stdin close (bounded + by the reader's ``finally``), never wedge it open or close it early + (see #1088). """ + subtype = message.get("subtype") + if subtype == "background_tasks_changed": + tasks = message.get("tasks") + if isinstance(tasks, list): + self._inflight_tasks = { + task["task_id"] + for task in tasks + if isinstance(task, dict) and task.get("task_id") + } + return + task_id = message.get("task_id") if not task_id: return - subtype = message.get("subtype") if subtype == "task_started": self._inflight_tasks.add(task_id) elif subtype == "task_notification": diff --git a/tests/test_query.py b/tests/test_query.py index e9a7508ed..2cdbdf5e1 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -387,6 +387,15 @@ async def dummy_hook(input_data, tool_use_id, context): "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) @@ -509,6 +518,64 @@ def test_track_task_lifecycle_unit(self): q._track_task_lifecycle({"subtype": "task_started"}) assert q._inflight_tasks == set() + def test_track_task_lifecycle_background_snapshot(self): + """background_tasks_changed replaces the in-flight set, self-healing + missed start/terminal frames (#1088).""" + transport = AsyncMock() + transport.is_ready = Mock(return_value=True) + q = Query(transport=transport, is_streaming_mode=False) + + # The snapshot is authoritative: adopt exactly the listed tasks, even + # with no prior task_started frames for them. + q._track_task_lifecycle(dict(_BACKGROUND_TASKS_CHANGED)) + assert q._inflight_tasks == {"task-1", "task-2"} + + # A later snapshot with a task gone clears it even if its terminal + # frame was missed — the drift-healing this frame exists for. + q._track_task_lifecycle( + { + "type": "system", + "subtype": "background_tasks_changed", + "tasks": [{"task_id": "task-2"}], + } + ) + assert q._inflight_tasks == {"task-2"} + + # An empty snapshot means every background task is done. + q._track_task_lifecycle( + {"type": "system", "subtype": "background_tasks_changed", "tasks": []} + ) + assert q._inflight_tasks == set() + + # A malformed frame (missing / non-list tasks) must not wipe live + # tracking — better to delay the close than to close early. + q._inflight_tasks = {"task-9"} + q._track_task_lifecycle( + {"type": "system", "subtype": "background_tasks_changed"} + ) + assert q._inflight_tasks == {"task-9"} + q._track_task_lifecycle( + {"type": "system", "subtype": "background_tasks_changed", "tasks": None} + ) + assert q._inflight_tasks == {"task-9"} + # A truthy-but-non-list tasks value is ignored too (guards the + # isinstance(tasks, list) check against a regression to ``if tasks:``). + q._track_task_lifecycle( + {"type": "system", "subtype": "background_tasks_changed", "tasks": {}} + ) + assert q._inflight_tasks == {"task-9"} + + # Non-dict entries and entries without a task_id are skipped, not added. + q._inflight_tasks = set() + q._track_task_lifecycle( + { + "type": "system", + "subtype": "background_tasks_changed", + "tasks": [123, {"task_type": "local_agent"}, {"task_id": "task-3"}], + } + ) + assert q._inflight_tasks == {"task-3"} + class TestAsyncIterablePromptWithSdkMcpServers: """Test that AsyncIterable prompts keep stdin open for SDK MCP servers.""" From 3e74937b9738687f66da2dd99c1d2a5c17025e7a Mon Sep 17 00:00:00 2001 From: Ashwin Bhat Date: Wed, 22 Jul 2026 13:46:14 -0700 Subject: [PATCH 4/4] fix(query): only defer the stdin close for delegated agent tasks The in-flight ledger tracked every task type reported by task_started, including background shells. A shell such as a dev server or `tail -f` never reaches a terminal status, so the set never emptied and stdin was never closed. Because the CLI only exits on stdin EOF, the process then never exited either, so the reader's finally block never ran and query() hung instead of returning. Track only local_agent and local_workflow: the types whose completion wakes the parent for the follow-up turn this relies on. Shells, monitors, teammates and remote agents can all run indefinitely by design, and are bounded by the CLI's own cleanup once stdin closes. Also stop consuming background_tasks_changed. It reports the live background set, but a subagent is registered in the foreground and only flips to backgrounded later without emitting a second task_started, so a tracked agent that is still running can be absent from it. Narrowing against the snapshot therefore drops work that goes on to outlive its turn, and widening from it can admit an observer agent whose start and terminal frames are both suppressed and which no later frame clears. The lifecycle frames are the only self-consistent source. Adds a regression test whose transport stays open until end_input() is called; every other stream in the suite self-terminates, so none of them could detect a missing close. --- src/claude_agent_sdk/_internal/query.py | 68 +++++++--- tests/test_query.py | 170 ++++++++++++++++++------ 2 files changed, 182 insertions(+), 56 deletions(-) diff --git a/src/claude_agent_sdk/_internal/query.py b/src/claude_agent_sdk/_internal/query.py index c171b2d23..66f10f06e 100644 --- a/src/claude_agent_sdk/_internal/query.py +++ b/src/claude_agent_sdk/_internal/query.py @@ -39,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. @@ -844,30 +860,46 @@ def _track_task_lifecycle(self, message: dict[str, Any]) -> None: a notification), so both are handled; ``discard`` keeps the pair idempotent. - ``background_tasks_changed`` carries an authoritative snapshot of the - running background tasks (``{"tasks": [{"task_id": ...}, ...]}``) and is - applied as a set-replacement. This self-heals a missed ``task_started`` - or terminal frame (dropped lines, older CLIs with different terminal - statuses): tracking drift can then only *delay* the stdin close (bounded - by the reader's ``finally``), never wedge it open or close it early - (see #1088). + 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") - if subtype == "background_tasks_changed": - tasks = message.get("tasks") - if isinstance(tasks, list): - self._inflight_tasks = { - task["task_id"] - for task in tasks - if isinstance(task, dict) and task.get("task_id") - } - return - task_id = message.get("task_id") if not task_id: return if subtype == "task_started": - self._inflight_tasks.add(task_id) + 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": diff --git a/tests/test_query.py b/tests/test_query.py index 2cdbdf5e1..a76201c7c 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -364,11 +364,24 @@ async def dummy_hook(input_data, tool_use_id, context): "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", @@ -518,63 +531,144 @@ def test_track_task_lifecycle_unit(self): q._track_task_lifecycle({"subtype": "task_started"}) assert q._inflight_tasks == set() - def test_track_task_lifecycle_background_snapshot(self): - """background_tasks_changed replaces the in-flight set, self-healing - missed start/terminal frames (#1088).""" + 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) - # The snapshot is authoritative: adopt exactly the listed tasks, even - # with no prior task_started frames for them. + # 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 == {"task-1", "task-2"} + assert q._inflight_tasks == set() - # A later snapshot with a task gone clears it even if its terminal - # frame was missed — the drift-healing this frame exists for. + # 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": "task-2"}], + "tasks": [{"task_id": "shell-1", "task_type": "local_bash"}], } ) - assert q._inflight_tasks == {"task-2"} + assert q._inflight_tasks == {"task-1"} - # An empty snapshot means every background task is done. + # 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() - # A malformed frame (missing / non-list tasks) must not wipe live - # tracking — better to delay the close than to close early. - q._inflight_tasks = {"task-9"} - q._track_task_lifecycle( - {"type": "system", "subtype": "background_tasks_changed"} - ) - assert q._inflight_tasks == {"task-9"} - q._track_task_lifecycle( - {"type": "system", "subtype": "background_tasks_changed", "tasks": None} - ) - assert q._inflight_tasks == {"task-9"} - # A truthy-but-non-list tasks value is ignored too (guards the - # isinstance(tasks, list) check against a regression to ``if tasks:``). - q._track_task_lifecycle( - {"type": "system", "subtype": "background_tasks_changed", "tasks": {}} - ) - assert q._inflight_tasks == {"task-9"} + def test_never_ending_shell_does_not_wedge_stdin_open(self): + """A backgrounded shell that never finishes must not hang the query. - # Non-dict entries and entries without a task_id are skipped, not added. - q._inflight_tasks = set() - q._track_task_lifecycle( - { - "type": "system", - "subtype": "background_tasks_changed", - "tasks": [123, {"task_type": "local_agent"}, {"task_id": "task-3"}], - } - ) - assert q._inflight_tasks == {"task-3"} + 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: