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.

[b-randongee]

Confirmed this fix's semantics end-to-end against CLI 2.1.215 from a production incident (details in #1088 (comment)): with in-flight-task tracking and close-only-on-quiescent-result, the previously-denied post-first-result SDK-MCP calls and hook callbacks succeed and the run closes cleanly.

One robustness suggestion: also consume the CLI's background_tasks_changed snapshot frame ({"type":"system","subtype":"background_tasks_changed","tasks":[...]}) as a set-replacement for _inflight_tasks. It's authoritative on 2.1.215 and self-heals a missed task_started/terminal frame, so tracking drift can only delay the close (still bounded by the reader's finally), never wedge stdin open or close it early. Worth it given the failure mode when tracking undercounts is exactly the silent hook-bypass this PR fixes.

[yayayouyou]

Thanks for taking the time to write this up — the production timeline (and especially the "presents as a phantom permissions problem" bit) is really valuable context. Glad the fix held up end-to-end on 2.1.215 🙏

The background_tasks_changed snapshot is a great call, so I went ahead and added it in 7a7f069: it's applied as a set-replacement of _inflight_tasks alongside the incremental start/terminal tracking, so a missed task_started/terminal frame now self-heals — drift can only delay the close (still bounded by the reader's finally), never wedge stdin open or close it early. It's handled before the task_id guard since the snapshot has no top-level task_id, and a malformed/absent tasks list leaves live tracking untouched rather than wiping it. Added unit coverage for the snapshot adoption, the self-heal, the empty-snapshot case, and the malformed shapes.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (main@b817bf5). Learn more about missing BASE report.
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1103   +/-   ##
=======================================
  Coverage        ?   90.03%           
=======================================
  Files           ?       23           
  Lines           ?     4154           
  Branches        ?        0           
=======================================
  Hits            ?     3740           
  Misses          ?      414           
  Partials        ?        0           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ashwin-ant]

Pushed 3e74937, quick summary of what changed.

The ledger was tracking every background task, including shells. A dev server or tail -f never finishes, so the set never emptied, stdin never closed, and query() hung. I reproduced that. Now it only tracks agent subagents and workflows, the ones that actually wake the parent when they finish.

I also removed the background_tasks_changed handling. Turned out it can miss a subagent that's still running, so using it to trim the set can drop live work and close stdin early. The task_started and terminal frames are more reliable on their own.

Added a test where the stream stays open until stdin actually closes, since none of the existing tests could catch a missing close.

Heads up this is a mitigation for #1088 rather than a complete fix. If a subagent happens to finish before the turn's result arrives, stdin still closes at that result. Fixing that fully needs a separate change. But this covers the common case.