fix(server): flush stdout and ensure line-buffering before restoring fd 1 (#3821)

Author: nikolas-sapa

src/mcp/server/stdio.py

@@ -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:
@@ -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]
@@ -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
@@ -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)

tests/conftest.py

@@ -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

tests/server/test_stdio.py

@@ -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,