Serve the 2026-07-28 era over stdio and other stream-pair transports (#3038)

tests/server/test_stdio.py

@@ -7,7 +7,15 @@ from io import TextIOWrapper
 
 import anyio
 import pytest
-from mcp_types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse, jsonrpc_message_adapter
+from mcp_types import (
+    CLIENT_CAPABILITIES_META_KEY,
+    CLIENT_INFO_META_KEY,
+    PROTOCOL_VERSION_META_KEY,
+    JSONRPCMessage,
+    JSONRPCRequest,
+    JSONRPCResponse,
+    jsonrpc_message_adapter,
+)
 
 from mcp.server.mcpserver import MCPServer
 from mcp.server.stdio import stdio_server
@@ -96,32 +104,108 @@ async def test_stdio_server_invalid_utf8(monkeypatch: pytest.MonkeyPatch) -> Non
                 assert second.message == valid
 
 
-class _KeepOpenBytesIO(io.BytesIO):
-    """A BytesIO that survives its TextIOWrapper being closed.
+class _GatedStdin(io.RawIOBase):
+    """Raw stdin double: serves its frames, then blocks until released before EOF.
 
-    Lets the test read what was written after `run()` has torn the wrapper down.
+    A real stdio client keeps stdin open until it has read the responses it is
+    awaiting; an immediate EOF after the last frame races the dispatcher's
+    EOF-time cancellation of in-flight handlers (only inline-handled methods
+    would deterministically answer first). The blocked read sits in
+    `stdio_server`'s reader worker thread and unblocks on `release()`.
     """
 
+    name = "<gated-stdin>"
+
+    def __init__(self, payload: bytes) -> None:
+        self._pending = payload
+        self._released = threading.Event()
+
+    def readable(self) -> bool:
+        return True
+
+    def readinto(self, b: bytearray | memoryview) -> int:  # pyright: ignore[reportIncompatibleMethodOverride]
+        if self._pending:
+            n = min(len(b), len(self._pending))
+            b[:n] = self._pending[:n]
+            self._pending = self._pending[n:]
+            return n
+        # A missed release falls through to EOF after the bound; the caller's
+        # own response assertions then report what actually arrived.
+        self._released.wait(5)
+        return 0
+
+    def release(self) -> None:
+        self._released.set()
+
+
+class _NotifyingStdout(io.RawIOBase):
+    """Raw stdout double that counts newline-terminated lines and can be awaited on.
+
+    Survives wrapper close (`close()` is a no-op) so the test can read what was
+    written after `run()` has torn its TextIOWrapper down.
+    """
+
+    name = "<notifying-stdout>"
+
+    def __init__(self) -> None:
+        self._chunks: list[bytes] = []
+        self._lines = 0
+        self._cond = threading.Condition()
+
+    def writable(self) -> bool:
+        return True
+
+    def write(self, b: bytes | bytearray | memoryview) -> int:  # pyright: ignore[reportIncompatibleMethodOverride]
+        data = bytes(b)
+        with self._cond:
+            self._chunks.append(data)
+            self._lines += data.count(b"\n")
+            self._cond.notify_all()
+        return len(data)
+
+    def wait_for_lines(self, n: int, timeout: float = 5) -> bool:
+        with self._cond:
+            return self._cond.wait_for(lambda: self._lines >= n, timeout)
+
+    def getvalue(self) -> bytes:
+        with self._cond:
+            return b"".join(self._chunks)
+
     def close(self) -> None:
         pass
 
 
-def _run_stdio_bounded(server: MCPServer) -> None:
-    """Run the blocking `server.run("stdio")` in a daemon thread joined with a 5s bound.
+def _serve_stdio_and_collect(
+    monkeypatch: pytest.MonkeyPatch, server: MCPServer, frames: list[JSONRPCRequest], responses: int
+) -> list[JSONRPCMessage]:
+    """Serve `frames` over process stdio and return the parsed response lines.
 
-    `run()` creates its own event loop, so a sync test cannot arm `anyio.fail_after`;
-    the join timeout turns a run loop that never returns on stdin EOF into a red test
-    instead of a silent CI hang. An exception escaping `run()` still fails the test:
-    pytest's unhandled-thread warning is escalated by `filterwarnings = ["error"]`.
+    Runs the blocking `server.run("stdio")` in a daemon thread (it creates its
+    own event loop, so a sync test cannot arm `anyio.fail_after`) and signals
+    stdin EOF only after `responses` lines arrive on stdout - the way a real
+    client closes the pipe - so spawned in-flight handlers never race the
+    dispatcher's EOF cancellation. The join bound turns a run loop that never
+    returns on stdin EOF into a red test instead of a silent CI hang; an
+    exception escaping `run()` still fails the test via pytest's
+    unhandled-thread warning, escalated by `filterwarnings = ["error"]`.
     """
+    payload = "".join(f.model_dump_json(by_alias=True, exclude_none=True) + "\n" for f in frames).encode()
+    stdin = _GatedStdin(payload)
+    stdout = _NotifyingStdout()
+    monkeypatch.setattr(sys, "stdin", TextIOWrapper(stdin, encoding="utf-8"))
+    monkeypatch.setattr(sys, "stdout", TextIOWrapper(stdout, encoding="utf-8"))
 
     def target() -> None:
         server.run("stdio")
 
     thread = threading.Thread(target=target, daemon=True)
     thread.start()
+    arrived = stdout.wait_for_lines(responses)
+    stdin.release()
     thread.join(5)
     assert not thread.is_alive(), 'run("stdio") did not return after stdin EOF'
+    assert arrived, f"expected {responses} response line(s); stdout carried: {stdout.getvalue()!r}"
+    return [jsonrpc_message_adapter.validate_json(line) for line in stdout.getvalue().decode().splitlines()]
 
 
 def test_mcpserver_run_stdio_serves_until_stdin_closes(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -131,15 +215,10 @@ def test_mcpserver_run_stdio_serves_until_stdin_closes(monkeypatch: pytest.Monke
     rather than serving forever.
     """
     ping = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
-    stdin_bytes = io.BytesIO(ping.model_dump_json(by_alias=True, exclude_none=True).encode() + b"\n")
-    captured = _KeepOpenBytesIO()
-    monkeypatch.setattr(sys, "stdin", TextIOWrapper(stdin_bytes, encoding="utf-8"))
-    monkeypatch.setattr(sys, "stdout", TextIOWrapper(captured, encoding="utf-8"))
 
-    _run_stdio_bounded(MCPServer(name="RunStdioServer"))
+    responses = _serve_stdio_and_collect(monkeypatch, MCPServer(name="RunStdioServer"), [ping], 1)
 
-    response = jsonrpc_message_adapter.validate_json(captured.getvalue().decode().strip())
-    assert response == JSONRPCResponse(jsonrpc="2.0", id=1, result={})
+    assert responses == [JSONRPCResponse(jsonrpc="2.0", id=1, result={})]
 
 
 def test_mcpserver_run_stdio_runs_lifespan_cleanup_after_stdin_closes(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -159,13 +238,37 @@ def test_mcpserver_run_stdio_runs_lifespan_cleanup_after_stdin_closes(monkeypatc
             events.append("cleanup")
 
     ping = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
-    stdin_bytes = io.BytesIO(ping.model_dump_json(by_alias=True, exclude_none=True).encode() + b"\n")
-    captured = _KeepOpenBytesIO()
-    monkeypatch.setattr(sys, "stdin", TextIOWrapper(stdin_bytes, encoding="utf-8"))
-    monkeypatch.setattr(sys, "stdout", TextIOWrapper(captured, encoding="utf-8"))
 
-    _run_stdio_bounded(MCPServer(name="LifespanStdioServer", lifespan=lifespan))
+    server = MCPServer(name="LifespanStdioServer", lifespan=lifespan)
+    responses = _serve_stdio_and_collect(monkeypatch, server, [ping], 1)
 
     assert events == ["setup", "cleanup"]
-    response = jsonrpc_message_adapter.validate_json(captured.getvalue().decode().strip())
-    assert response == JSONRPCResponse(jsonrpc="2.0", id=1, result={})
+    assert responses == [JSONRPCResponse(jsonrpc="2.0", id=1, result={})]
+
+
+def test_mcpserver_run_stdio_serves_a_modern_connection(monkeypatch: pytest.MonkeyPatch) -> None:
+    """`MCPServer.run("stdio")` serves the modern era over process stdio.
+
+    A `server/discover` probe gets a DiscoverResult (no initialize handshake)
+    and a subsequent envelope-bearing request is served at the discovered
+    version - the wire exchange `Client(mode='auto')` drives against a stdio
+    server.
+    """
+    envelope = {
+        PROTOCOL_VERSION_META_KEY: "2026-07-28",
+        CLIENT_INFO_META_KEY: {"name": "probe", "version": "1.0"},
+        CLIENT_CAPABILITIES_META_KEY: {},
+    }
+    discover = JSONRPCRequest(jsonrpc="2.0", id=1, method="server/discover", params={"_meta": envelope})
+    tools = JSONRPCRequest(jsonrpc="2.0", id=2, method="tools/list", params={"_meta": envelope})
+
+    responses = _serve_stdio_and_collect(monkeypatch, MCPServer(name="ModernStdioServer"), [discover, tools], 2)
+
+    assert isinstance(responses[0], JSONRPCResponse) and responses[0].id == 1
+    assert "2026-07-28" in responses[0].result["supportedVersions"]
+    assert responses[0].result["serverInfo"]["name"] == "ModernStdioServer"
+    assert isinstance(responses[1], JSONRPCResponse) and responses[1].id == 2
+    # `resultType` is the modern-only wire field: its presence proves the
+    # request was served at the discovered version, not the handshake era.
+    assert responses[1].result["tools"] == []
+    assert responses[1].result["resultType"] == "complete"