fix(stdio): drain responses after stdin EOF

[Epochex]

Fixes #2678.

When stdin hits EOF (e.g. python server.py < payload.jsonl > response.jsonl), the client is done sending requests but is still reading stdout for results. On main, BaseSession._receive_loop() closes the write stream as soon as the read stream ends, and Server.run() cancels the handler task group on transport close. That combination can cancel in-flight tool handlers before they respond, dropping the last ools/call responses.

This change lets ServerSession keep the write stream open when the read side closes, so Server.run() can drain already-started handlers and flush their responses before shutdown.

Tests:

• uv run --frozen pytest tests/server/test_cancel_handling.py tests/server/test_stdio.py

[StantonMatt]

I ran this locally against the branch. The focused stdio/cancel tests pass:

TMPDIR=.codex-tmp/tmp UV_CACHE_DIR=.codex-tmp/uv-cache uv run --frozen pytest tests/server/test_cancel_handling.py tests/server/test_stdio.py -q

One shutdown edge looks worth tightening before merge. With drain_on_read_close=True, stdin EOF now waits for in-flight handlers to finish, but there is no bounded drain window. In a local memory-stream probe, I started a tools/call handler that does await anyio.sleep_forever(), sent initialize / notifications/initialized / tools/call, closed the input stream, and then waited for server.run(..., drain_on_read_close=True) to return. It did not return before anyio.fail_after(0.75).

That leaves redirected-stdin callers with the opposite failure mode: the previous branch can drop a response, while this one can keep the process alive forever if an accepted tool hangs after EOF. I think the drain path needs a small timeout or cancellation window so already-started handlers get a chance to flush, but EOF still remains a bounded shutdown signal.

[StantonMatt]

Rechecked the updated head 721dd43. The new read_eof_drain_timeout_seconds path covers the hanging-handler case I was worried about: in-flight handlers get a drain window, and the server cancels the task group if they never finish after read EOF.

Local validation:

• git diff --check origin/main...HEAD
• TMPDIR=/Users/mjstanton/Programming/open-source-maintainer4/.codex-tmp/python-sdk-tmp UV_CACHE_DIR=/Users/mjstanton/Programming/open-source-maintainer4/.codex-tmp/python-sdk-uv-cache uv run --frozen pytest tests/server/test_cancel_handling.py tests/server/test_stdio.py -q -> 9 passed

The visible CI matrix is also green on this head.

[StantonMatt]

The new read_eof_drain_timeout_seconds path addresses the bounded-shutdown concern I raised.

I re-ran the focused stdio/cancel suite on head 721dd43:

TMPDIR=.codex-tmp/tmp UV_CACHE_DIR=.codex-tmp/uv-cache uv run --frozen pytest tests/server/test_cancel_handling.py tests/server/test_stdio.py -q

That passed with 9 passed, including test_server_bounds_drain_on_read_eof_when_handler_never_finishes. git diff --check origin/main...HEAD is also clean, and the current GitHub Actions/Conformance/Zizmor checks are green.

[StantonMatt]

I retested the latest head (721dd43) against the shutdown case I called out earlier.

The new bounded drain path covers both sides of that tradeoff for me: test_server_drains_in_flight_handlers_on_transport_read_eof still lets already-started handlers flush responses after stdin EOF, and test_server_bounds_drain_on_read_eof_when_handler_never_finishes covers the hung-handler case so EOF stays bounded.

Local focused check passed:

TMPDIR=.codex-tmp/tmp UV_CACHE_DIR=.codex-tmp/uv-cache uv run --frozen pytest tests/server/test_cancel_handling.py tests/server/test_stdio.py -q

GitHub checks are green on the current head as well. No further concern from my earlier note.