Skip to content

Commit 88f9276

Browse files
author
Nova
committed
fix(stdio): drain pending responses before closing read stream on EOF
Closes #2678. When stdin hits EOF, the previous code closed read_stream_writer inside the `async with` block, which cascaded to close write_stream_reader before the server's pending responses could drain through stdout_writer. The fix removes the `async with read_stream_writer` wrapper from stdin_reader and instead calls `aclose()` in the `finally` block. This ensures: 1. All stdin lines are read and forwarded to the server 2. The read stream is closed promptly on EOF (signaling the server) 3. Buffered responses in write_stream_reader drain through stdout_writer before the task group exits All existing tests pass. Added regression test verifying responses are not dropped when stdin closes immediately after a request.
1 parent b478bff commit 88f9276

2 files changed

Lines changed: 72 additions & 9 deletions

File tree

src/mcp/server/stdio.py

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -48,18 +48,22 @@ async def stdio_server(stdin: anyio.AsyncFile[str] | None = None, stdout: anyio.
4848

4949
async def stdin_reader():
5050
try:
51-
async with read_stream_writer:
52-
async for line in stdin:
53-
try:
54-
message = types.jsonrpc_message_adapter.validate_json(line, by_name=False)
55-
except Exception as exc:
56-
await read_stream_writer.send(exc)
57-
continue
51+
async for line in stdin:
52+
try:
53+
message = types.jsonrpc_message_adapter.validate_json(line, by_name=False)
54+
except Exception as exc:
55+
await read_stream_writer.send(exc)
56+
continue
5857

59-
session_message = SessionMessage(message)
60-
await read_stream_writer.send(session_message)
58+
session_message = SessionMessage(message)
59+
await read_stream_writer.send(session_message)
6160
except anyio.ClosedResourceError: # pragma: no cover
6261
await anyio.lowlevel.checkpoint()
62+
finally:
63+
# Close the read stream to signal EOF to the server. Any pending
64+
# server responses are already buffered in write_stream_reader and
65+
# will drain through stdout_writer before the task group exits.
66+
await read_stream_writer.aclose()
6367

6468
async def stdout_writer():
6569
try:

tests/server/test_stdio_2678.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
"""Regression test for #2678: in-flight responses should not be dropped on stdin EOF.
2+
3+
When a server receives a request and stdin hits EOF while the server is still
4+
processing, the response must still be written to stdout. The fix closes
5+
read_stream_writer in stdin_reader's finally block so the server sees EOF and
6+
can flush pending writes before the task group exits.
7+
"""
8+
import io
9+
import sys
10+
import threading
11+
import time
12+
from io import TextIOWrapper
13+
14+
import anyio
15+
import pytest
16+
17+
from mcp.server.mcpserver import MCPServer
18+
from mcp.types import (
19+
JSONRPCRequest,
20+
JSONRPCResponse,
21+
jsonrpc_message_adapter,
22+
)
23+
24+
25+
class _KeepOpenBytesIO(io.BytesIO):
26+
"""A BytesIO that survives its TextIOWrapper being closed."""
27+
28+
def close(self) -> None:
29+
pass
30+
31+
32+
def _run_stdio_bounded(server: MCPServer, timeout: float = 5) -> None:
33+
def target() -> None:
34+
server.run("stdio")
35+
36+
thread = threading.Thread(target=target, daemon=True)
37+
thread.start()
38+
thread.join(timeout)
39+
assert not thread.is_alive(), "run('stdio') did not return after stdin EOF"
40+
41+
42+
def test_stdio_response_not_dropped_on_eof(monkeypatch: pytest.MonkeyPatch) -> None:
43+
"""Server response is written to stdout even when stdin closes right after the request.
44+
45+
Regression test for #2678: stdin EOF used to close read_stream_writer before
46+
the server could flush its response through stdout_writer.
47+
"""
48+
ping = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
49+
stdin_bytes = io.BytesIO(
50+
ping.model_dump_json(by_alias=True, exclude_none=True).encode() + b"\n"
51+
)
52+
captured = _KeepOpenBytesIO()
53+
monkeypatch.setattr(sys, "stdin", TextIOWrapper(stdin_bytes, encoding="utf-8"))
54+
monkeypatch.setattr(sys, "stdout", TextIOWrapper(captured, encoding="utf-8"))
55+
56+
_run_stdio_bounded(MCPServer(name="TestEOF"))
57+
58+
response = jsonrpc_message_adapter.validate_json(captured.getvalue().decode().strip())
59+
assert response == JSONRPCResponse(jsonrpc="2.0", id=1, result={})

0 commit comments

Comments
 (0)