Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 23 additions & 1 deletion src/mcp/server/stdio.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,12 @@ def _open_stdout_diversion() -> int:

def _restore_fd(fd: int, private_fd: int) -> bool:
"""Point fd back at the wire; the Windows handle rebind never affects the outcome."""
if fd == 1:
with suppress(Exception):
sys.stdout.flush()
if hasattr(sys.stdout, "reconfigure"):
with suppress(Exception):
sys.stdout.reconfigure(line_buffering=True)
try:
os.dup2(private_fd, fd)
except OSError:
Expand Down Expand Up @@ -123,6 +129,17 @@ def _claim_fd(
_claims[fd] = claim

def release() -> None:
if fd == 1:
with suppress(Exception):
sys.stdout.flush()
with suppress(Exception):
stream.flush()
if hasattr(sys.stdout, "reconfigure"):
with suppress(Exception):
sys.stdout.reconfigure(line_buffering=True)
if hasattr(stream, "reconfigure"):
with suppress(Exception):
stream.reconfigure(line_buffering=True)
if claim.private_fd is None or _restore_fd(fd, claim.private_fd):
with _claims_lock:
del _claims[fd]
Expand All @@ -138,6 +155,11 @@ def release() -> None:
except OSError:
return stream.buffer, release
try:
if fd == 1:
with suppress(Exception):
sys.stdout.flush()
with suppress(Exception):
stream.flush()
os.dup2(diversion_fd, fd)
except OSError:
# The divert did not land; ensure fd carries the wire (a Windows dup2 can
Expand Down Expand Up @@ -176,7 +198,7 @@ async def stdio_server(stdin: anyio.AsyncFile[str] | None = None, stdout: anyio.
stdin = anyio.wrap_file(_UnownedTextWrapper(stdin_buffer, encoding="utf-8", errors="replace"))
if not stdout:
stdout_buffer, restore_stdout = _claim_fd(1, sys.stdout, "wb", _open_stdout_diversion)
stdout = anyio.wrap_file(_UnownedTextWrapper(stdout_buffer, encoding="utf-8"))
stdout = anyio.wrap_file(_UnownedTextWrapper(stdout_buffer, encoding="utf-8", line_buffering=True))

read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
Expand Down
4 changes: 4 additions & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,10 @@ def blockbuster() -> Iterator[None]:
bb.functions["io.BufferedReader.read"].can_block_in(
"mcp/server/mcpserver/utilities/types.py", ("to_image_content", "to_audio_content")
)
# Stdio transport flushes stdout before restoring fd 1.
for name in ("os.write", "io.BufferedWriter.write", "io.TextIOWrapper.write"):
if name in bb.functions:
bb.functions[name].can_block_in("mcp/server/stdio.py", ("release", "_restore_fd"))
bb.activate()
try:
yield
Expand Down
40 changes: 40 additions & 0 deletions tests/server/test_stdio.py
Original file line number Diff line number Diff line change
Expand Up @@ -344,6 +344,46 @@ async def test_stdio_server_takes_stdout_off_the_descriptor_table_while_serving(
assert os.path.sameopenfile(1, out_w)


@pytest.mark.anyio
async def test_stdio_server_flushes_stdout_and_ensures_line_buffering_before_restore(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Unflushed writes to sys.stdout during stdio_server flush to diversion before fd 1 restore.

Regression for issue #3821: without a flush before restoring fd 1, buffered writes to
sys.stdout made while fd 1 was diverted leak onto the wire when flushed after restore,
corrupting client JSON-RPC streams. Also verifies line-buffering is ensured on restore.
"""
response = JSONRPCResponse(jsonrpc="2.0", id=1, result={})
with _pipe_planted_on_fd1(monkeypatch) as (out_r, out_w), _pipe_planted_on_fd2() as err_r:
# sys.stdout begins without line-buffering (pipe default)
assert getattr(sys.stdout, "line_buffering", None) is False

with anyio.fail_after(5):
async with stdio_server(stdin=anyio.AsyncFile(io.StringIO())) as (read_stream, write_stream):
read_stream.close()

# Buffered write to sys.stdout without explicit flush
sys.stdout.write("stray buffered log\n")

await write_stream.send(SessionMessage(response))
line = await _read_from(out_r)
assert jsonrpc_message_adapter.validate_json(line.decode().strip()) == response
await write_stream.aclose()

# The unflushed write must have been flushed to stderr before fd 1 was restored
assert await _read_from(err_r) == b"stray buffered log\n"

# After restore, line buffering is ensured on sys.stdout
assert getattr(sys.stdout, "line_buffering", None) is True

# And subsequent flush on sys.stdout does not write the stray log to the wire
sys.stdout.flush()
os.set_blocking(out_r, False)
with pytest.raises(BlockingIOError):
os.read(out_r, 4096)


@pytest.mark.anyio
async def test_stdio_server_diverts_stdout_to_the_null_device_when_stderr_is_unusable(
monkeypatch: pytest.MonkeyPatch,
Expand Down
Loading