Fix: keep stdin open through a result that lands right after a task settles (#1190)

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.

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.