fix(query): don't close stdin on a result frame while tasks are in flight

[yayayouyou]

Summary

Fixes #1088.

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. Query.wait_for_result_and_end_input() closed stdin on the first result frame, which broke everything a still-running subagent needed afterwards:

• its SDK-MCP tool calls failed with exactly "Stream closed" (the error reported in SDK closes stdin on first result frame, breaking bidirectional control from nested Task/Agent subagents #1088), and
• its PreToolUse hooks were silently bypassed — built-in tools kept executing with no hook callback delivered, so deny-gate hooks stopped gating.

Reproduction

Reproduced end-to-end against the live CLI (2.1.206) with query() + a PreToolUse hook + an in-process SDK-MCP server whose tool sleeps 12s, called from a run_in_background: true Task. Observed timeline on main:

[19.3s] parent turn ends → ResultMessage #1 → SDK closes stdin   (task still running)
[25s+]  subagent's ToolSearch calls execute with NO PreToolUse callback (hooks silently bypassed)
[~25s]  subagent calls the SDK-MCP tool → tool_result is_error=true: "Stream closed" (×2 retries)
[47.9s] ResultMessage #2 arrives (task completion wakes the parent for a follow-up turn)

The second result frame is the key observation: result frames are per-turn, so closing stdin on the first one is the bug. (On current CLI, two sequential foreground subagents — the issue's literal scenario — no longer emit intermediate result frames; background tasks are the deterministic trigger.)

Fix

Track in-flight tasks from the task_started / task_notification / terminal task_updated lifecycle frames (using the existing TERMINAL_TASK_STATUSES, whose docstring already prescribes exactly this clearing behavior), and only treat a result frame as run-ending when no tasks are in flight.

Why not the issue's first suggestion (keep stdin open until close() whenever hooks/SDK-MCP are registered): the CLI in stream-json mode only exits on stdin EOF, so that would hang one-shot query() forever. The narrower rule preserves prompt closure:

• no background tasks → behavior unchanged (first result still closes stdin; existing tests cover this),
• background task in flight → stdin stays open; its completion wakes the parent for a follow-up turn that ends in another result frame, which closes stdin (this also makes chained background tasks work),
• early process exit → the reader's finally still unblocks the waiter, so no new hang mode is introduced.

Tests

• test_result_with_inflight_task_keeps_stdin_open (parametrized over both drain frames: task_notification and terminal task_updated patch) — asserts stdin stays open across an intermediate result and closes on the first result with no tasks in flight. Fails on current main, passes with the fix.
• test_track_task_lifecycle_unit — add/non-terminal/terminal/unknown-id/missing-id transitions.
• Full suite: 1059 passed, 5 skipped (asyncio + trio); ruff and mypy clean.
• E2E re-run of the reproduction against the live CLI with the fix: hook fires for the background subagent's calls, the SDK-MCP tool executes and returns its result, stdin closes exactly at the final result frame, run exits normally.

Used AI assistance; reviewed and tested by me.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

Fixes a streaming-mode stdin teardown bug where Query.wait_for_result_and_end_input() closed stdin on the first result frame even when background Tasks were still running, breaking subsequent hook / SDK-MCP control responses (issue #1088).

Changes:

• Track in-flight background task IDs from task_started / terminal task_notification / terminal task_updated frames.
• Treat a result frame as run-ending (and thus eligible to trigger stdin closure) only when no tasks are in flight.
• Add regression + unit tests covering the intermediate-result + background-task lifecycle.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
src/claude_agent_sdk/_internal/query.py	Track task lifecycle and defer stdin closure until a result arrives with no in-flight tasks.
tests/test_query.py	Add tests asserting stdin remains open across intermediate results while tasks are in flight, plus lifecycle unit coverage.

Comments suppressed due to low confidence (1)

src/claude_agent_sdk/_internal/query.py:877

• This debug log still says "Waiting for first result", but the code now waits for a run-ending result (a result with no tasks in flight). Adjusting the message will make logs match actual behavior during debugging.

            logger.debug(
                "Waiting for first result before closing stdin "
                f"(sdk_mcp_servers={len(self.sdk_mcp_servers)}, "
                f"has_hooks={bool(self.hooks)})"
            )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.