fix(transport): keep reading stderr until the CLI exits on close

[alanhuangyoo]

With a stderr callback set, SubprocessCLITransport.close() stops reading stderr before it waits for the CLI to exit, so:

• Anything the CLI writes to stderr during the post-EOF grace period never reaches the callback.
• If it writes more than a pipe buffer's worth (16–64 KiB), the CLI blocks on the write and can't exit. It gets SIGTERM at 5s and SIGKILL at 10s, which interrupts the session-file flush the grace period was added for (Session file not flushed before subprocess termination in ClaudeSDKClient #625, fix: wait for graceful subprocess shutdown before SIGTERM #642).

close() cancelled the stderr task and closed the stream first, then closed stdin and waited on the process. This change keeps the reader running through that wait. Once the process has exited (or been terminated/killed), it gives the reader up to 1s to reach EOF, then cancels it as before, in case something else still holds the pipe open. Every await in the shielded scope is still bounded. Without a stderr callback, stderr isn't piped and nothing changes.

With a stand-in CLI that writes ~256 KiB to stderr after stdin EOF:

before: disconnect() took 10.01s; stderr lines delivered: 0   (child got SIGTERM, never flushed)
after:  all lines delivered, child exits 0 on its own

The new test test_close_reads_stderr_written_while_the_cli_shuts_down (asyncio and trio) spawns a real child that writes 4000 lines to stderr after stdin EOF, then asserts every line reached the callback, the child exited 0 rather than being terminated, and close() returned within 5s. On main it gets 0 lines and takes the terminate path. The full suite passes (1487 passed, 5 skipped), and ruff check, ruff format --check and mypy src/ are clean.

[tonydzi]

Hi — Mycroft here, TonyDzi's synthetic AI co-founder (I review shutdown paths for a living; the processes I watch die more gracefully than most of my hypotheses).

The reorder makes sense and the repro in the description is convincing: keeping the reader alive through the grace wait is the right fix for the full-pipe stall. One structural regression in the cancellation path, though, plus a small suggestion.

Moving the stderr cleanup after the wait() takes it off the path that the docstring's caveat describes.

The close() docstring already notes that the anyio shield doesn't defer a raw asyncio cancellation (asyncio.wait_for, Task.cancel() from outside). On main that was survivable for stderr, because the reader was cancelled before any await that could be interrupted. With this PR the stderr block runs only after the try/finally around self._process.wait() completes normally, and it is not itself in a finally. So a raw cancel that lands during the 5s grace wait propagates straight past it:

• _stderr_task is never cancelled and keeps running as a detached task,
• the user's stderr callback keeps firing after close() has already raised out to the caller,
• _stderr_stream is never aclose()d (fd held until the child dies).

Minimal probe (same harness as your new test: real child that ignores stdin EOF and keeps writing to stderr; close() wrapped in asyncio.ensure_future, cancelled after 0.5s, then 1s of observation):

PR 19ebce7:  close() raw-cancelled
             stderr_task done=False  lines_after_cancel=19  transport._stderr_task still set=True
main 37a52c9: close() raw-cancelled
             stderr_task done=True   lines_after_cancel=0   transport._stderr_task still set=False

Suggested shape: keep the 1s drain on the normal path, but move the teardown into a finally so an exception/cancellation unwinding through the wait still stops the reader — without awaiting during the unwind (a second raw cancel could interrupt that await anyway, and TaskHandle.cancel() is synchronous):

drained_normally = False
try:
    ...  # existing graceful wait / terminate / kill block
    drained_normally = True
finally:
    task = self._stderr_task
    if task is not None:
        if drained_normally and not task.done():
            with anyio.move_on_after(1):
                await task.wait()
        if not task.done():
            task.cancel()
        if drained_normally:
            with suppress(Exception):
                await task.wait()
    self._stderr_task = None
    if self._stderr_stream is not None and drained_normally:
        with suppress(Exception):
            await self._stderr_stream.aclose()
    self._stderr_stream = None

(On the exceptional path the reader's own finally still flushes the last partial line, so nothing the CLI already wrote is lost — it just stops delivering once close() has given up.)

Test: a sibling to test_close_reads_stderr_written_while_the_cli_shuts_down on the asyncio backend only — child that never exits on EOF and writes a line every 50ms, cancel close() via Task.cancel() mid-grace, assert task.done() and no new callback lines after the cancellation. The probe above has exactly that shape and flips between this branch and main, so such a test would pin the invariant the reorder moved. (On the exceptional path the stream object is dropped rather than aclose()d in my sketch; the fd goes when the transport is collected — if you'd rather close it eagerly there too, the reader's cancellation is the thing to wait on first.)

A side note that doesn't block anything: when the CLI's own children (stdio MCP servers, hooks) inherit its stderr, the pipe stays open after the CLI exits, so every close() in that setup will now spend the full extra 1s. That's a fine bound, just worth a sentence in the comment so nobody "optimizes" it away later.

— TonyDzi · this is one small piece of a multi-agent lab I run in public (agent consensus, persistent memory, fleet coordination): github.com/tonydzi

[alanhuangyoo]

@tonydzi Thanks, that's a real regression and I could reproduce it: with a raw Task.cancel() during the grace wait, the reader kept running on the previous head and stopped on main.

Fixed in 0daf8ae. close() now runs inside a small context manager that cancels the stderr reader, without awaiting, whenever the body unwinds with an exception. That covers a cancel in the grace wait, and also one during the stdin lock or the 1s drain. The stream is left for a later close() to aclose().

Added test_raw_cancel_during_close_grace_wait_stops_the_stderr_reader (asyncio only) along the lines you described: it fails on the previous head and passes on main and here. Also added a note about the extra 1s when a child process holds the stderr pipe open.

[tonydzi]

@alanhuangyoo I checked 0daf8ae locally. test_raw_cancel_during_close_grace_wait_stops_the_stderr_reader passes on the head. When I drop self._stop_stderr_reader_on_error() from the with line, it fails. So the test really guards the context manager. It also works on a second close(): _process isn't cleared on the raw-cancel path, so a later close() still reaches the _stderr_stream.aclose(). Looks good to me.

One nuance for the docstring, not a change request. Because the cancel isn't awaited, the reader's finally (emit(framer.flush())) runs on the next loop iteration, after close() has already re-raised. A partial last line can therefore reach the callback once after close gave up. I think that's the right trade (don't lose the tail), but "stop the reader" reads as "no more callbacks", so half a sentence would save someone a confusing bug report later.

— TonyDzi · github.com/tonydzi