Max Isbey
98bd125ddd
The 2026-07-28 protocol did not work over stdio: subscriptions/listen was hard-refused, a legacy initialize arriving during an in-flight modern request was accepted and re-locked the connection, and a peer cancel produced a trailing "Request cancelled" frame. All three share one root cause: the connection's era was derived from which requests had completed instead of being decided once, in wire order, from how the client opened the connection. Replace serve_dual_era_loop and serve_loop with a single serve_stream driver that decides the era synchronously in the dispatcher's read loop, before the request body is spawned. initialize (or any envelope-less request) opens the legacy era, an enveloped request opens the modern era, server/discover is answered without pinning, and a stray leading notification opens nothing. A conflicting era claim on a committed connection is refused (-32022 or -32600) rather than silently switching. Cancel silence is structural: each request answers through a one-shot channel whose write target becomes a powerless void on a peer cancel, so there is no cancelled-check at any write site. The generic JSON-RPC dispatcher loses the code-0 cancel frame, the code-0 str(exc) catch-all, and the inline_methods knob, and documents that handlers are invoked synchronously in receive order with the returned awaitable as the body. Add a Posture enum (DUAL default, LEGACY_ONLY, MODERN_ONLY) on the Server and MCPServer constructors, honoured by the stream driver and the streamable-HTTP manager alike. Server.run(read, write) now stands alone, Server.lifespan() is a bound context manager, and serve_listener / newline_json_transport / close_subscriptions() give a straightforward path for custom transports. See docs/migration.md for the full list of observable changes.
459 lines
18 KiB
Python
459 lines
18 KiB
Python
import io
|
|
import sys
|
|
import threading
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
from io import TextIOWrapper
|
|
from pathlib import Path
|
|
|
|
import anyio
|
|
import anyio.abc
|
|
import pytest
|
|
from anyio.streams.buffered import BufferedByteReceiveStream
|
|
from mcp_types import (
|
|
CLIENT_CAPABILITIES_META_KEY,
|
|
CLIENT_INFO_META_KEY,
|
|
PROTOCOL_VERSION_META_KEY,
|
|
SERVER_INFO_META_KEY,
|
|
JSONRPCMessage,
|
|
JSONRPCRequest,
|
|
JSONRPCResponse,
|
|
jsonrpc_message_adapter,
|
|
)
|
|
from typing_extensions import Buffer
|
|
|
|
from mcp.server import serve_stream
|
|
from mcp.server.mcpserver import MCPServer
|
|
from mcp.server.stdio import newline_json_transport, stdio_server
|
|
from mcp.shared.message import SessionMessage
|
|
|
|
|
|
@pytest.fixture(params=["asyncio", "trio"])
|
|
def anyio_backend(request: pytest.FixtureRequest) -> str:
|
|
"""Run every async test in this module on both anyio backends; the sync
|
|
`run("stdio")` tests drive their own loop in a worker thread and take neither."""
|
|
return request.param
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _module_runner_lease() -> None:
|
|
"""Opt out of the shared per-module event loop: this module parametrizes `anyio_backend`."""
|
|
|
|
|
|
class _ByteStreamDouble(anyio.abc.ByteStream):
|
|
"""A minimal `ByteStream` over an inbound memory stream, recording what the framer writes.
|
|
|
|
`broken=True` makes every send fail the way a peer that vanished does, so the
|
|
framer's writer can be shown to end quietly instead of crashing the connection.
|
|
"""
|
|
|
|
def __init__(self, inbound: anyio.abc.ObjectReceiveStream[bytes], *, broken: bool = False) -> None:
|
|
self._inbound = inbound
|
|
self._broken = broken
|
|
self.sent: list[bytes] = []
|
|
self.closed = False
|
|
|
|
async def receive(self, max_bytes: int = 65536) -> bytes:
|
|
try:
|
|
return await self._inbound.receive()
|
|
except (anyio.EndOfStream, anyio.ClosedResourceError):
|
|
raise anyio.EndOfStream from None
|
|
|
|
async def send(self, item: bytes) -> None:
|
|
if self._broken:
|
|
raise anyio.BrokenResourceError
|
|
self.sent.append(item)
|
|
|
|
async def send_eof(self) -> None:
|
|
await self.aclose()
|
|
|
|
async def aclose(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_stdio_server_round_trips_messages_over_injected_streams() -> None:
|
|
"""stdio_server frames JSON-RPC messages as one line each in both directions.
|
|
|
|
Parses one message per stdin line and writes each outgoing message as exactly one
|
|
line, driven over injected in-process streams.
|
|
"""
|
|
stdin = io.StringIO()
|
|
stdout = io.StringIO()
|
|
|
|
messages = [
|
|
JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"),
|
|
JSONRPCResponse(jsonrpc="2.0", id=2, result={}),
|
|
]
|
|
|
|
for message in messages:
|
|
stdin.write(message.model_dump_json(by_alias=True, exclude_none=True) + "\n")
|
|
stdin.seek(0)
|
|
|
|
with anyio.fail_after(5):
|
|
async with stdio_server(stdin=anyio.AsyncFile(stdin), stdout=anyio.AsyncFile(stdout)) as (
|
|
read_stream,
|
|
write_stream,
|
|
):
|
|
async with read_stream:
|
|
received_messages: list[JSONRPCMessage] = []
|
|
for _ in range(2):
|
|
received = await read_stream.receive()
|
|
assert not isinstance(received, Exception)
|
|
received_messages.append(received.message)
|
|
|
|
assert received_messages[0] == JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
|
|
assert received_messages[1] == JSONRPCResponse(jsonrpc="2.0", id=2, result={})
|
|
|
|
responses = [
|
|
JSONRPCRequest(jsonrpc="2.0", id=3, method="ping"),
|
|
JSONRPCResponse(jsonrpc="2.0", id=4, result={}),
|
|
]
|
|
|
|
for response in responses:
|
|
await write_stream.send(SessionMessage(response))
|
|
await write_stream.aclose()
|
|
|
|
stdout.seek(0)
|
|
output_lines = stdout.readlines()
|
|
assert len(output_lines) == 2
|
|
|
|
received_responses = [jsonrpc_message_adapter.validate_json(line.strip()) for line in output_lines]
|
|
assert received_responses[0] == JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")
|
|
assert received_responses[1] == JSONRPCResponse(jsonrpc="2.0", id=4, result={})
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_stdio_server_invalid_utf8(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Non-UTF-8 stdin bytes surface as an in-stream exception without killing the stream.
|
|
|
|
Invalid bytes are replaced with U+FFFD, fail JSON parsing, and arrive as an in-stream
|
|
exception; subsequent valid messages are still processed.
|
|
"""
|
|
# \xff\xfe are invalid UTF-8 start bytes.
|
|
valid = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
|
|
raw_stdin = io.BytesIO(b"\xff\xfe\n" + valid.model_dump_json(by_alias=True, exclude_none=True).encode() + b"\n")
|
|
|
|
# Replace sys.stdin with a wrapper whose .buffer is our raw bytes, so that
|
|
# stdio_server()'s default path wraps it with errors='replace'.
|
|
monkeypatch.setattr(sys, "stdin", TextIOWrapper(raw_stdin, encoding="utf-8"))
|
|
monkeypatch.setattr(sys, "stdout", TextIOWrapper(io.BytesIO(), encoding="utf-8"))
|
|
|
|
with anyio.fail_after(5):
|
|
async with stdio_server() as (read_stream, write_stream):
|
|
await write_stream.aclose()
|
|
async with read_stream: # pragma: no branch
|
|
# First line: \xff\xfe -> U+FFFD U+FFFD -> JSON parse fails -> exception in stream
|
|
first = await read_stream.receive()
|
|
assert isinstance(first, Exception)
|
|
|
|
# Second line: valid message still comes through
|
|
second = await read_stream.receive()
|
|
assert isinstance(second, SessionMessage)
|
|
assert second.message == valid
|
|
|
|
|
|
class _GatedStdin(io.RawIOBase):
|
|
"""Raw stdin double: serves its frames, then blocks until released before EOF.
|
|
|
|
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: Buffer) -> int:
|
|
view = memoryview(b)
|
|
if self._pending:
|
|
n = min(len(view), len(self._pending))
|
|
view[: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: Buffer) -> int:
|
|
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 _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.
|
|
|
|
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:
|
|
"""`MCPServer.run("stdio")` serves over process stdio and returns at stdin EOF.
|
|
|
|
Answers a request over the process's stdio and returns when stdin reaches EOF,
|
|
rather than serving forever.
|
|
"""
|
|
ping = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
|
|
|
|
responses = _serve_stdio_and_collect(monkeypatch, MCPServer(name="RunStdioServer"), [ping], 1)
|
|
|
|
assert responses == [JSONRPCResponse(jsonrpc="2.0", id=1, result={})]
|
|
|
|
|
|
def test_mcpserver_run_stdio_runs_lifespan_cleanup_after_stdin_closes(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Code after `yield` in a lifespan runs when stdin EOF ends `run("stdio")`.
|
|
|
|
Regression lock for the issue #1027 shutdown chain: the run loop must end on
|
|
stdin EOF and unwind the lifespan rather than be killed before returning.
|
|
"""
|
|
events: list[str] = []
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(server: MCPServer) -> AsyncIterator[None]:
|
|
events.append("setup")
|
|
try:
|
|
yield
|
|
finally:
|
|
events.append("cleanup")
|
|
|
|
ping = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
|
|
|
|
server = MCPServer(name="LifespanStdioServer", lifespan=lifespan)
|
|
responses = _serve_stdio_and_collect(monkeypatch, server, [ping], 1)
|
|
|
|
assert events == ["setup", "cleanup"]
|
|
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})
|
|
|
|
server = MCPServer(name="ModernStdioServer", version="1.2.3")
|
|
responses = _serve_stdio_and_collect(monkeypatch, server, [discover, tools], 2)
|
|
|
|
assert isinstance(responses[0], JSONRPCResponse) and responses[0].id == 1
|
|
assert "2026-07-28" in responses[0].result["supportedVersions"]
|
|
# Server identity travels as the result `_meta` stamp, not a DiscoverResult
|
|
# body field (spec 2026-07-28, #3002).
|
|
assert responses[0].result["_meta"][SERVER_INFO_META_KEY] == {"name": "ModernStdioServer", "version": "1.2.3"}
|
|
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"
|
|
|
|
|
|
# --- newline_json_transport: the stdio framing over any byte stream --------------
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_custom_transport_over_a_unix_socket_needs_only_framing_plus_serve_stream(tmp_path: Path) -> None:
|
|
"""The whole custom-transport story: frame a real Unix socket with
|
|
`newline_json_transport` and hand the streams to `serve_stream` - a raw
|
|
JSON-RPC client on the other end gets a served answer, one frame per line."""
|
|
server = MCPServer(name="socket-server")
|
|
|
|
@server.tool()
|
|
def add(a: int, b: int) -> str:
|
|
"""Add two numbers."""
|
|
return str(a + b)
|
|
|
|
socket_path = tmp_path / "mcp.sock"
|
|
listener = await anyio.create_unix_listener(socket_path)
|
|
connection_served = anyio.Event()
|
|
|
|
async def handle(stream: anyio.abc.SocketStream) -> None:
|
|
async with stream, newline_json_transport(stream) as (read_stream, write_stream):
|
|
await serve_stream(server._lowlevel_server, read_stream, write_stream) # pyright: ignore[reportPrivateUsage]
|
|
connection_served.set() # the driver returned once the client closed its end
|
|
|
|
envelope = {
|
|
PROTOCOL_VERSION_META_KEY: "2026-07-28",
|
|
CLIENT_INFO_META_KEY: {"name": "socket-client", "version": "1.0"},
|
|
CLIENT_CAPABILITIES_META_KEY: {},
|
|
}
|
|
call = JSONRPCRequest(
|
|
jsonrpc="2.0",
|
|
id=1,
|
|
method="tools/call",
|
|
params={"name": "add", "arguments": {"a": 3, "b": 4}, "_meta": envelope},
|
|
)
|
|
line = b""
|
|
async with listener, anyio.create_task_group() as tg:
|
|
tg.start_soon(listener.serve, handle)
|
|
with anyio.fail_after(5):
|
|
async with await anyio.connect_unix(socket_path) as client:
|
|
await client.send(call.model_dump_json(by_alias=True, exclude_none=True).encode("utf-8") + b"\n")
|
|
line = await BufferedByteReceiveStream(client).receive_until(b"\n", 1_000_000)
|
|
# The client closed its socket: the driver sees EOF and ends the connection.
|
|
await connection_served.wait()
|
|
tg.cancel_scope.cancel()
|
|
|
|
answer = jsonrpc_message_adapter.validate_json(line, by_name=False)
|
|
assert isinstance(answer, JSONRPCResponse) and answer.id == 1
|
|
assert answer.result["content"][0]["text"] == "7"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_newline_json_transport_frames_lines_both_ways_and_survives_a_malformed_line() -> None:
|
|
"""One JSON-RPC message per line in both directions: a line that is not JSON-RPC
|
|
reaches the read stream as an exception item and the frame after it is still delivered
|
|
(one bad line costs one line, not the connection), outbound messages are written as
|
|
exactly one line each, and closing the byte stream stays the caller's job."""
|
|
peer_send, transport_side = anyio.create_memory_object_stream[bytes](8)
|
|
good = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
|
|
reply = JSONRPCResponse(jsonrpc="2.0", id=1, result={})
|
|
stream = _ByteStreamDouble(transport_side)
|
|
|
|
async with (
|
|
peer_send,
|
|
transport_side,
|
|
newline_json_transport(stream) as (read_stream, write_stream),
|
|
read_stream, # the driver's ends: a driver would own and close these
|
|
write_stream,
|
|
):
|
|
await peer_send.send(b"this is not json\n")
|
|
await peer_send.send(good.model_dump_json(by_alias=True, exclude_none=True).encode("utf-8") + b"\n")
|
|
with anyio.fail_after(5):
|
|
first = await read_stream.receive()
|
|
second = await read_stream.receive()
|
|
await write_stream.send(SessionMessage(reply))
|
|
await anyio.wait_all_tasks_blocked() # let the writer flush onto the byte stream
|
|
peer_send.close() # peer EOF: the framer ends its inbound stream
|
|
with anyio.fail_after(5):
|
|
assert [item async for item in read_stream] == []
|
|
|
|
assert isinstance(first, Exception)
|
|
assert isinstance(second, SessionMessage) and second.message == good
|
|
assert stream.sent == [reply.model_dump_json(by_alias=True, exclude_unset=True).encode("utf-8") + b"\n"]
|
|
assert not stream.closed # the framer never closes the byte stream itself
|
|
await stream.send_eof()
|
|
assert stream.closed
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_newline_json_transport_writer_ends_quietly_when_the_peer_is_gone() -> None:
|
|
"""A write onto a byte stream whose peer has vanished ends the writer without an
|
|
error escaping: outbound frames after that are undeliverable, not a crash."""
|
|
peer_send, transport_side = anyio.create_memory_object_stream[bytes](8)
|
|
stream = _ByteStreamDouble(transport_side, broken=True)
|
|
|
|
async with (
|
|
peer_send,
|
|
transport_side,
|
|
newline_json_transport(stream) as (read_stream, write_stream),
|
|
read_stream,
|
|
write_stream,
|
|
):
|
|
with anyio.fail_after(5):
|
|
await write_stream.send(SessionMessage(JSONRPCResponse(jsonrpc="2.0", id=1, result={})))
|
|
await anyio.wait_all_tasks_blocked() # the writer hits the broken send and returns
|
|
assert stream.sent == []
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_newline_json_transport_ends_the_read_side_on_an_oversized_frame() -> None:
|
|
"""A frame that overruns the bound cannot be resynchronised: the framer surfaces
|
|
the overrun as an exception item and then ends the inbound stream (EOF)."""
|
|
peer_send, transport_side = anyio.create_memory_object_stream[bytes](8)
|
|
|
|
async with (
|
|
peer_send,
|
|
transport_side,
|
|
newline_json_transport(_ByteStreamDouble(transport_side), max_frame_bytes=8) as (
|
|
read_stream,
|
|
write_stream,
|
|
),
|
|
read_stream, # the driver's ends: a driver would own and close these
|
|
write_stream,
|
|
):
|
|
await peer_send.send(b"x" * 32) # no newline within the 8-byte bound
|
|
with anyio.fail_after(5):
|
|
items = [item async for item in read_stream]
|
|
|
|
assert len(items) == 1 and isinstance(items[0], anyio.DelimiterNotFound)
|