From 969fd553dac3d80cccf92f89a0d12b8397a91c91 Mon Sep 17 00:00:00 2001 From: Jianke LIN Date: Mon, 25 May 2026 22:26:45 +0200 Subject: [PATCH 1/8] fix(stdio): drain responses after stdin EOF --- src/mcp/server/lowlevel/server.py | 1 + src/mcp/shared/jsonrpc_dispatcher.py | 62 ++++++++++-------- tests/server/test_cancel_handling.py | 34 ++++------ tests/server/test_stdio.py | 96 +++++++++++++++++++++++++++- 4 files changed, 146 insertions(+), 47 deletions(-) diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index d2536189d..dc60c5c24 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -416,6 +416,7 @@ async def run( # the next request (spec says SHOULD NOT, not MUST NOT) sees # the initialized state instead of failing the init-gate. inline_methods=frozenset({"initialize"}), + close_write_stream_on_read_close=False, ) runner = ServerRunner( server=self, diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index a59cd119d..ef05f382b 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/src/mcp/shared/jsonrpc_dispatcher.py @@ -10,6 +10,7 @@ import contextvars import logging from collections.abc import Awaitable, Callable, Mapping +from contextlib import AsyncExitStack from dataclasses import dataclass, field from functools import partial from typing import Any, Generic, Literal, cast @@ -221,6 +222,7 @@ def __init__( peer_cancel_mode: PeerCancelMode = "interrupt", raise_handler_exceptions: bool = False, inline_methods: frozenset[str] = frozenset(), + close_write_stream_on_read_close: bool = True, on_stream_exception: Callable[[Exception], Awaitable[None]] | None = None, ) -> None: """Wire a dispatcher over a transport's `SessionMessage` stream pair. @@ -233,6 +235,10 @@ def __init__( inline_methods: Methods awaited in the read loop before the next message is dequeued (e.g. `initialize`); an inline handler that awaits the peer deadlocks the parked loop. + close_write_stream_on_read_close: Close the write stream when the + read stream closes. Full-duplex transports may set this to + false so in-flight handlers can finish writing responses after + input EOF. on_stream_exception: Observer for `Exception` items on the read stream; without it they are debug-logged and dropped. Awaited inline in the read loop, so a slow observer stalls dispatch. @@ -247,6 +253,7 @@ def __init__( ) self._peer_cancel_mode: PeerCancelMode = peer_cancel_mode self._raise_handler_exceptions = raise_handler_exceptions + self._close_write_stream_on_read_close = close_write_stream_on_read_close self._inline_methods = inline_methods self._on_stream_exception = on_stream_exception @@ -414,33 +421,36 @@ async def run( `task_status.started()` fires once `send_raw_request` is usable. Single-shot: once the loop ends the dispatcher stays closed and cannot be restarted. """ + normal_eof = False try: - # LIFO exits: the write stream closes only after the task-group join, so teardown writes still land. - async with self._write_stream: - async with anyio.create_task_group() as tg: - self._tg = tg - self._running = True - task_status.started() - try: - async with self._read_stream: - try: - async for item in self._read_stream: - # Duck-typed: only `ContextReceiveStream` carries the - # sender's per-message contextvars snapshot. - sender_ctx: contextvars.Context | None = getattr( - self._read_stream, "last_context", None - ) - await self._dispatch(item, on_request, on_notify, sender_ctx) - except anyio.ClosedResourceError: - # Receive end closed under us (stateless SHTTP teardown); same as EOF. - logger.debug("read stream closed by transport; treating as EOF") - # EOF: wake blocked `send_raw_request` waiters with CONNECTION_CLOSED. - self._running = False - self._closed = True - self._fan_out_closed() - finally: - # Cancel in-flight handlers; otherwise the task-group join - # waits on handlers whose callers are already gone. + async with anyio.create_task_group() as tg: + self._tg = tg + self._running = True + task_status.started() + try: + async with AsyncExitStack() as stack: + await stack.enter_async_context(self._read_stream) + if self._close_write_stream_on_read_close: + await stack.enter_async_context(self._write_stream) + try: + async for item in self._read_stream: + # Duck-typed: only `ContextReceiveStream` carries the + # sender's per-message contextvars snapshot. + sender_ctx: contextvars.Context | None = getattr( + self._read_stream, "last_context", None + ) + await self._dispatch(item, on_request, on_notify, sender_ctx) + except anyio.ClosedResourceError: + # Receive end closed under us (stateless SHTTP teardown); same as EOF. + logger.debug("read stream closed by transport; treating as EOF") + # EOF: wake blocked `send_raw_request` waiters with CONNECTION_CLOSED. + self._running = False + self._fan_out_closed() + normal_eof = True + finally: + if not normal_eof: + # Cancel on crash/cancel paths. On normal EOF, let + # already received handlers drain their responses. tg.cancel_scope.cancel() finally: # Covers cancel/crash paths that skip the inline fan-out; idempotent. diff --git a/tests/server/test_cancel_handling.py b/tests/server/test_cancel_handling.py index 10531fad9..2362befe8 100644 --- a/tests/server/test_cancel_handling.py +++ b/tests/server/test_cancel_handling.py @@ -19,6 +19,7 @@ InitializeRequestParams, JSONRPCNotification, JSONRPCRequest, + JSONRPCResponse, ListToolsResult, PaginatedRequestParams, TextContent, @@ -100,29 +101,18 @@ async def first_request(): @pytest.mark.anyio -async def test_server_cancels_in_flight_handlers_on_transport_close(): - """When the transport closes mid-request, server.run() must cancel in-flight - handlers rather than join on them. - - Without the cancel, the task group waits for the handler, which then tries - to respond through a write stream that _receive_loop already closed, - raising ClosedResourceError and crashing server.run() with exit code 1. - - This drives server.run() with raw memory streams because InMemoryTransport - wraps it in its own finally-cancel (_memory.py) which masks the bug. - """ +async def test_server_drains_in_flight_handlers_on_transport_read_eof(): + """When the transport's read side hits EOF (e.g., stdio stdin closes), the + server must drain already-started handlers so their responses reach the + peer via the still-open write side.""" handler_started = anyio.Event() - handler_cancelled = anyio.Event() + handler_allowed_to_finish = anyio.Event() server_run_returned = anyio.Event() async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult: handler_started.set() - try: - await anyio.sleep_forever() - finally: - handler_cancelled.set() - # unreachable: sleep_forever only exits via cancellation - raise AssertionError # pragma: no cover + await handler_allowed_to_finish.wait() + return CallToolResult(content=[TextContent(type="text", text="ok")]) server = Server("test", on_call_tool=handle_call_tool) @@ -167,9 +157,13 @@ async def run_server(): # handler gets CancelledError, server.run() returns. await to_server.aclose() - await server_run_returned.wait() + handler_allowed_to_finish.set() + + response = await from_server.receive() + assert isinstance(response.message, JSONRPCResponse) + assert response.message.id == 2 - assert handler_cancelled.is_set() + await server_run_returned.wait() @pytest.mark.anyio diff --git a/tests/server/test_stdio.py b/tests/server/test_stdio.py index 054a157b3..e827a8ae4 100644 --- a/tests/server/test_stdio.py +++ b/tests/server/test_stdio.py @@ -8,10 +8,28 @@ import anyio import pytest +from mcp.server import Server, ServerRequestContext from mcp.server.mcpserver import MCPServer from mcp.server.stdio import stdio_server from mcp.shared.message import SessionMessage -from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse, jsonrpc_message_adapter +from mcp.types import ( + LATEST_PROTOCOL_VERSION, + CallToolRequestParams, + CallToolResult, + ClientCapabilities, + Implementation, + InitializeRequestParams, + JSONRPCError, + JSONRPCMessage, + JSONRPCNotification, + JSONRPCRequest, + JSONRPCResponse, + ListToolsResult, + PaginatedRequestParams, + TextContent, + Tool, + jsonrpc_message_adapter, +) @pytest.mark.anyio @@ -169,3 +187,79 @@ async def lifespan(server: MCPServer) -> AsyncIterator[None]: assert events == ["setup", "cleanup"] response = jsonrpc_message_adapter.validate_json(captured.getvalue().decode().strip()) assert response == JSONRPCResponse(jsonrpc="2.0", id=1, result={}) + + +@pytest.mark.anyio +async def test_stdio_server_drains_in_flight_responses_on_stdin_eof(): + """When stdin reaches EOF (e.g., bash-redirected input), already-received + requests must still be able to emit their responses on stdout.""" + stdin = io.StringIO() + stdout = io.StringIO() + + tool_started_count = 0 + both_tools_started = anyio.Event() + allow_tools_to_finish = anyio.Event() + + async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult: + return ListToolsResult(tools=[Tool(name="slow", description="test", input_schema={})]) + + async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult: + nonlocal tool_started_count + tool_started_count += 1 + if tool_started_count == 2: + both_tools_started.set() + await allow_tools_to_finish.wait() + return CallToolResult(content=[TextContent(type="text", text="ok")]) + + server = Server("test", on_list_tools=handle_list_tools, on_call_tool=handle_call_tool) + + init_req = JSONRPCRequest( + jsonrpc="2.0", + id=0, + method="initialize", + params=InitializeRequestParams( + protocol_version=LATEST_PROTOCOL_VERSION, + capabilities=ClientCapabilities(), + client_info=Implementation(name="test", version="1.0"), + ).model_dump(by_alias=True, mode="json", exclude_none=True), + ) + initialized = JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized") + call_1 = JSONRPCRequest( + jsonrpc="2.0", + id=1, + method="tools/call", + params=CallToolRequestParams(name="slow", arguments={}).model_dump(by_alias=True, mode="json"), + ) + call_2 = JSONRPCRequest( + jsonrpc="2.0", + id=2, + method="tools/call", + params=CallToolRequestParams(name="slow", arguments={}).model_dump(by_alias=True, mode="json"), + ) + + for message in (init_req, initialized, call_1, call_2): + stdin.write(message.model_dump_json(by_alias=True, exclude_none=True) + "\n") + stdin.seek(0) + + async with stdio_server(stdin=anyio.AsyncFile(stdin), stdout=anyio.AsyncFile(stdout)) as ( + read_stream, + write_stream, + ): + with anyio.fail_after(5): + async with anyio.create_task_group() as tg: + tg.start_soon(server.run, read_stream, write_stream, server.create_initialization_options()) + await both_tools_started.wait() + allow_tools_to_finish.set() + + stdout.seek(0) + ids: set[int | str] = set() + for line in stdout.readlines(): + line = line.strip() + if not line: + continue + message = jsonrpc_message_adapter.validate_json(line) + if isinstance(message, JSONRPCResponse | JSONRPCError): + assert message.id is not None + ids.add(message.id) + assert 1 in ids + assert 2 in ids From 1a8efe18226243ddd0e9dcec494dda229e8dda6b Mon Sep 17 00:00:00 2001 From: Jianke LIN Date: Mon, 25 May 2026 22:59:26 +0200 Subject: [PATCH 2/8] test: cover stdio EOF drain and shutdown edges --- tests/server/test_cancel_handling.py | 117 +++++++++++++++++++++++++++ tests/server/test_stdio.py | 18 ++--- 2 files changed, 125 insertions(+), 10 deletions(-) diff --git a/tests/server/test_cancel_handling.py b/tests/server/test_cancel_handling.py index 2362befe8..da66209cc 100644 --- a/tests/server/test_cancel_handling.py +++ b/tests/server/test_cancel_handling.py @@ -166,6 +166,123 @@ async def run_server(): await server_run_returned.wait() +@pytest.mark.anyio +async def test_server_reraises_handler_cancellation_when_server_is_cancelled(): + """If the server task is cancelled (e.g. KeyboardInterrupt), in-flight + request handlers will get cancelled too. Cancellation must be re-raised so + the task group can unwind cleanly.""" + handler_started = anyio.Event() + server_run_returned = anyio.Event() + cancel_scope = anyio.CancelScope() + + async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult: + handler_started.set() + await anyio.sleep_forever() + raise AssertionError # pragma: no cover + + server = Server("test", on_call_tool=handle_call_tool) + + to_server, server_read = anyio.create_memory_object_stream[SessionMessage | Exception](10) + server_write, from_server = anyio.create_memory_object_stream[SessionMessage](10) + + async def run_server(): + try: + with cancel_scope: + await server.run(server_read, server_write, server.create_initialization_options()) + finally: + server_run_returned.set() + + init_req = JSONRPCRequest( + jsonrpc="2.0", + id=1, + method="initialize", + params=InitializeRequestParams( + protocol_version=LATEST_PROTOCOL_VERSION, + capabilities=ClientCapabilities(), + client_info=Implementation(name="test", version="1.0"), + ).model_dump(by_alias=True, mode="json", exclude_none=True), + ) + initialized = JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized") + call_req = JSONRPCRequest( + jsonrpc="2.0", + id=2, + method="tools/call", + params=CallToolRequestParams(name="slow", arguments={}).model_dump(by_alias=True, mode="json"), + ) + + with anyio.fail_after(5): + async with anyio.create_task_group() as tg, to_server, server_read, server_write, from_server: + tg.start_soon(run_server) + + await to_server.send(SessionMessage(init_req)) + await from_server.receive() # init response + await to_server.send(SessionMessage(initialized)) + await to_server.send(SessionMessage(call_req)) + + await handler_started.wait() + cancel_scope.cancel() + await server_run_returned.wait() + + +@pytest.mark.anyio +async def test_server_drops_response_when_write_stream_closes_mid_request(): + """If the write side closes while a handler is in-flight, responding may + raise (ClosedResourceError/BrokenResourceError). The handler task should + exit without crashing the server.""" + handler_started = anyio.Event() + allow_finish = anyio.Event() + server_run_returned = anyio.Event() + + async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult: + handler_started.set() + await allow_finish.wait() + return CallToolResult(content=[TextContent(type="text", text="ok")]) + + server = Server("test", on_call_tool=handle_call_tool) + + to_server, server_read = anyio.create_memory_object_stream[SessionMessage | Exception](10) + server_write, from_server = anyio.create_memory_object_stream[SessionMessage](10) + + async def run_server(): + await server.run(server_read, server_write, server.create_initialization_options()) + server_run_returned.set() + + init_req = JSONRPCRequest( + jsonrpc="2.0", + id=1, + method="initialize", + params=InitializeRequestParams( + protocol_version=LATEST_PROTOCOL_VERSION, + capabilities=ClientCapabilities(), + client_info=Implementation(name="test", version="1.0"), + ).model_dump(by_alias=True, mode="json", exclude_none=True), + ) + initialized = JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized") + call_req = JSONRPCRequest( + jsonrpc="2.0", + id=2, + method="tools/call", + params=CallToolRequestParams(name="slow", arguments={}).model_dump(by_alias=True, mode="json"), + ) + + with anyio.fail_after(5): + async with anyio.create_task_group() as tg, to_server, server_read, server_write, from_server: + tg.start_soon(run_server) + + await to_server.send(SessionMessage(init_req)) + await from_server.receive() # init response + await to_server.send(SessionMessage(initialized)) + await to_server.send(SessionMessage(call_req)) + + await handler_started.wait() + await server_write.aclose() + + allow_finish.set() + await to_server.aclose() + + await server_run_returned.wait() + + @pytest.mark.anyio async def test_server_handles_transport_close_with_pending_server_to_client_requests(): """When the transport closes while handlers are blocked on server→client diff --git a/tests/server/test_stdio.py b/tests/server/test_stdio.py index e827a8ae4..25da9edd1 100644 --- a/tests/server/test_stdio.py +++ b/tests/server/test_stdio.py @@ -19,7 +19,6 @@ ClientCapabilities, Implementation, InitializeRequestParams, - JSONRPCError, JSONRPCMessage, JSONRPCNotification, JSONRPCRequest, @@ -224,6 +223,7 @@ async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestPar ).model_dump(by_alias=True, mode="json", exclude_none=True), ) initialized = JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized") + list_tools = JSONRPCRequest(jsonrpc="2.0", id=10, method="tools/list") call_1 = JSONRPCRequest( jsonrpc="2.0", id=1, @@ -237,7 +237,7 @@ async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestPar params=CallToolRequestParams(name="slow", arguments={}).model_dump(by_alias=True, mode="json"), ) - for message in (init_req, initialized, call_1, call_2): + for message in (init_req, initialized, list_tools, call_1, call_2): stdin.write(message.model_dump_json(by_alias=True, exclude_none=True) + "\n") stdin.seek(0) @@ -252,14 +252,12 @@ async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestPar allow_tools_to_finish.set() stdout.seek(0) + output_lines = [line.strip() for line in stdout.readlines()] + messages = [jsonrpc_message_adapter.validate_json(line) for line in output_lines] ids: set[int | str] = set() - for line in stdout.readlines(): - line = line.strip() - if not line: - continue - message = jsonrpc_message_adapter.validate_json(line) - if isinstance(message, JSONRPCResponse | JSONRPCError): - assert message.id is not None - ids.add(message.id) + for message in messages: + assert isinstance(message, JSONRPCResponse) + ids.add(message.id) + assert 1 in ids assert 2 in ids From 98e3a539ace7c0d3fced681c06a7169a6013f5fc Mon Sep 17 00:00:00 2001 From: Jianke LIN Date: Mon, 25 May 2026 23:05:59 +0200 Subject: [PATCH 3/8] test: ignore coverage branch arc on 3.14 --- tests/server/test_stdio.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/server/test_stdio.py b/tests/server/test_stdio.py index 25da9edd1..92e63a854 100644 --- a/tests/server/test_stdio.py +++ b/tests/server/test_stdio.py @@ -246,7 +246,7 @@ async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestPar write_stream, ): with anyio.fail_after(5): - async with anyio.create_task_group() as tg: + async with anyio.create_task_group() as tg: # pragma: no branch tg.start_soon(server.run, read_stream, write_stream, server.create_initialization_options()) await both_tools_started.wait() allow_tools_to_finish.set() From 44ce2adb54372b7a980ffb9930c873de75830a26 Mon Sep 17 00:00:00 2001 From: Jianke LIN Date: Mon, 25 May 2026 23:17:06 +0200 Subject: [PATCH 4/8] fix(server): opt-in drain on read EOF --- src/mcp/server/lowlevel/server.py | 6 +++++- src/mcp/server/mcpserver/server.py | 1 + tests/server/test_cancel_handling.py | 2 +- tests/server/test_stdio.py | 11 ++++++++++- 4 files changed, 17 insertions(+), 3 deletions(-) diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index dc60c5c24..3229f5691 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -406,6 +406,10 @@ async def run( # the initialization lifecycle, but can do so with any available node # rather than requiring initialization for each connection. stateless: bool = False, + # When True, treat read EOF as a half-close and allow in-flight handlers + # to drain their responses via the still-open write stream (e.g. stdio + # with bash-redirected stdin). + drain_on_read_close: bool = False, ) -> None: async with self.lifespan(self) as lifespan_context: dispatcher: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher( @@ -416,7 +420,7 @@ async def run( # the next request (spec says SHOULD NOT, not MUST NOT) sees # the initialized state instead of failing the init-gate. inline_methods=frozenset({"initialize"}), - close_write_stream_on_read_close=False, + close_write_stream_on_read_close=not drain_on_read_close, ) runner = ServerRunner( server=self, diff --git a/src/mcp/server/mcpserver/server.py b/src/mcp/server/mcpserver/server.py index fdb69571d..647bb8c57 100644 --- a/src/mcp/server/mcpserver/server.py +++ b/src/mcp/server/mcpserver/server.py @@ -848,6 +848,7 @@ async def run_stdio_async(self) -> None: read_stream, write_stream, self._lowlevel_server.create_initialization_options(), + drain_on_read_close=True, ) async def run_sse_async( # pragma: no cover diff --git a/tests/server/test_cancel_handling.py b/tests/server/test_cancel_handling.py index da66209cc..b365f3a7f 100644 --- a/tests/server/test_cancel_handling.py +++ b/tests/server/test_cancel_handling.py @@ -120,7 +120,7 @@ async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestPar server_write, from_server = anyio.create_memory_object_stream[SessionMessage](10) async def run_server(): - await server.run(server_read, server_write, server.create_initialization_options()) + await server.run(server_read, server_write, server.create_initialization_options(), drain_on_read_close=True) server_run_returned.set() init_req = JSONRPCRequest( diff --git a/tests/server/test_stdio.py b/tests/server/test_stdio.py index 92e63a854..8f97c679a 100644 --- a/tests/server/test_stdio.py +++ b/tests/server/test_stdio.py @@ -247,7 +247,16 @@ async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestPar ): with anyio.fail_after(5): async with anyio.create_task_group() as tg: # pragma: no branch - tg.start_soon(server.run, read_stream, write_stream, server.create_initialization_options()) + + async def run_server() -> None: + await server.run( + read_stream, + write_stream, + server.create_initialization_options(), + drain_on_read_close=True, + ) + + tg.start_soon(run_server) await both_tools_started.wait() allow_tools_to_finish.set() From 27b1337b5230882ad9d78869ea8bad4fa3a9c526 Mon Sep 17 00:00:00 2001 From: Jianke LIN Date: Tue, 2 Jun 2026 19:23:11 +0200 Subject: [PATCH 5/8] fix(stdio): bound EOF drain wait --- src/mcp/server/lowlevel/server.py | 6 +++ src/mcp/shared/jsonrpc_dispatcher.py | 16 +++++-- tests/server/test_cancel_handling.py | 72 +++++++++++++++++++++++++++- 3 files changed, 90 insertions(+), 4 deletions(-) diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index 3229f5691..9439809b2 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -69,6 +69,8 @@ async def main(): logger = logging.getLogger(__name__) +DEFAULT_READ_EOF_DRAIN_TIMEOUT_SECONDS = 1.0 + LifespanResultT = TypeVar("LifespanResultT", default=Any) _ParamsT = TypeVar("_ParamsT", bound=BaseModel, default=BaseModel) @@ -410,6 +412,9 @@ async def run( # to drain their responses via the still-open write stream (e.g. stdio # with bash-redirected stdin). drain_on_read_close: bool = False, + # Maximum time to wait for in-flight handlers to drain after read EOF. + # None means wait indefinitely. + read_eof_drain_timeout_seconds: float | None = DEFAULT_READ_EOF_DRAIN_TIMEOUT_SECONDS, ) -> None: async with self.lifespan(self) as lifespan_context: dispatcher: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher( @@ -421,6 +426,7 @@ async def run( # the initialized state instead of failing the init-gate. inline_methods=frozenset({"initialize"}), close_write_stream_on_read_close=not drain_on_read_close, + read_eof_drain_timeout_seconds=read_eof_drain_timeout_seconds, ) runner = ServerRunner( server=self, diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index ef05f382b..dc9179779 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/src/mcp/shared/jsonrpc_dispatcher.py @@ -223,6 +223,7 @@ def __init__( raise_handler_exceptions: bool = False, inline_methods: frozenset[str] = frozenset(), close_write_stream_on_read_close: bool = True, + read_eof_drain_timeout_seconds: float | None = None, on_stream_exception: Callable[[Exception], Awaitable[None]] | None = None, ) -> None: """Wire a dispatcher over a transport's `SessionMessage` stream pair. @@ -239,6 +240,9 @@ def __init__( read stream closes. Full-duplex transports may set this to false so in-flight handlers can finish writing responses after input EOF. + read_eof_drain_timeout_seconds: Maximum time to wait for in-flight + handlers to drain after read EOF when the write stream stays + open; None waits indefinitely. on_stream_exception: Observer for `Exception` items on the read stream; without it they are debug-logged and dropped. Awaited inline in the read loop, so a slow observer stalls dispatch. @@ -254,6 +258,7 @@ def __init__( self._peer_cancel_mode: PeerCancelMode = peer_cancel_mode self._raise_handler_exceptions = raise_handler_exceptions self._close_write_stream_on_read_close = close_write_stream_on_read_close + self._read_eof_drain_timeout_seconds = read_eof_drain_timeout_seconds self._inline_methods = inline_methods self._on_stream_exception = on_stream_exception @@ -448,16 +453,21 @@ async def run( self._fan_out_closed() normal_eof = True finally: - if not normal_eof: - # Cancel on crash/cancel paths. On normal EOF, let - # already received handlers drain their responses. + if not normal_eof or self._close_write_stream_on_read_close: + # Cancel on crash/cancel paths. If read EOF also closed + # writes, handlers cannot drain responses anyway. tg.cancel_scope.cancel() + elif self._read_eof_drain_timeout_seconds is not None: + tg.cancel_scope.deadline = anyio.current_time() + self._read_eof_drain_timeout_seconds finally: # Covers cancel/crash paths that skip the inline fan-out; idempotent. self._running = False self._closed = True self._tg = None self._fan_out_closed() + if not self._close_write_stream_on_read_close: + with anyio.CancelScope(shield=True): + await self._write_stream.aclose() await resync_tracer() async def _dispatch( diff --git a/tests/server/test_cancel_handling.py b/tests/server/test_cancel_handling.py index b365f3a7f..fb5b6ecd0 100644 --- a/tests/server/test_cancel_handling.py +++ b/tests/server/test_cancel_handling.py @@ -120,7 +120,13 @@ async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestPar server_write, from_server = anyio.create_memory_object_stream[SessionMessage](10) async def run_server(): - await server.run(server_read, server_write, server.create_initialization_options(), drain_on_read_close=True) + await server.run( + server_read, + server_write, + server.create_initialization_options(), + drain_on_read_close=True, + read_eof_drain_timeout_seconds=None, + ) server_run_returned.set() init_req = JSONRPCRequest( @@ -166,6 +172,70 @@ async def run_server(): await server_run_returned.wait() +@pytest.mark.anyio +async def test_server_bounds_drain_on_read_eof_when_handler_never_finishes(): + handler_started = anyio.Event() + handler_cancelled = anyio.Event() + server_run_returned = anyio.Event() + + async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult: + handler_started.set() + try: + await anyio.sleep_forever() + finally: + handler_cancelled.set() + raise AssertionError # pragma: no cover + + server = Server("test", on_call_tool=handle_call_tool) + + to_server, server_read = anyio.create_memory_object_stream[SessionMessage | Exception](10) + server_write, from_server = anyio.create_memory_object_stream[SessionMessage](10) + + async def run_server(): + await server.run( + server_read, + server_write, + server.create_initialization_options(), + drain_on_read_close=True, + read_eof_drain_timeout_seconds=0.05, + ) + server_run_returned.set() + + init_req = JSONRPCRequest( + jsonrpc="2.0", + id=1, + method="initialize", + params=InitializeRequestParams( + protocol_version=LATEST_PROTOCOL_VERSION, + capabilities=ClientCapabilities(), + client_info=Implementation(name="test", version="1.0"), + ).model_dump(by_alias=True, mode="json", exclude_none=True), + ) + initialized = JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized") + call_req = JSONRPCRequest( + jsonrpc="2.0", + id=2, + method="tools/call", + params=CallToolRequestParams(name="slow", arguments={}).model_dump(by_alias=True, mode="json"), + ) + + with anyio.fail_after(2): + async with anyio.create_task_group() as tg, to_server, server_read, server_write, from_server: + tg.start_soon(run_server) + + await to_server.send(SessionMessage(init_req)) + await from_server.receive() # init response + await to_server.send(SessionMessage(initialized)) + await to_server.send(SessionMessage(call_req)) + + await handler_started.wait() + await to_server.aclose() + + await server_run_returned.wait() + + assert handler_cancelled.is_set() + + @pytest.mark.anyio async def test_server_reraises_handler_cancellation_when_server_is_cancelled(): """If the server task is cancelled (e.g. KeyboardInterrupt), in-flight From 8253b5b0f470d2742eeee104ae0eb9026a8b63d5 Mon Sep 17 00:00:00 2001 From: Jianke LIN Date: Thu, 11 Jun 2026 10:50:43 +0200 Subject: [PATCH 6/8] test: cover EOF drain branches --- src/mcp/shared/jsonrpc_dispatcher.py | 4 +- tests/shared/test_jsonrpc_dispatcher.py | 56 +++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index dc9179779..3743e577e 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/src/mcp/shared/jsonrpc_dispatcher.py @@ -457,7 +457,9 @@ async def run( # Cancel on crash/cancel paths. If read EOF also closed # writes, handlers cannot drain responses anyway. tg.cancel_scope.cancel() - elif self._read_eof_drain_timeout_seconds is not None: + elif self._read_eof_drain_timeout_seconds is None: + pass + else: tg.cancel_scope.deadline = anyio.current_time() + self._read_eof_drain_timeout_seconds finally: # Covers cancel/crash paths that skip the inline fan-out; idempotent. diff --git a/tests/shared/test_jsonrpc_dispatcher.py b/tests/shared/test_jsonrpc_dispatcher.py index 2b828e531..53267f705 100644 --- a/tests/shared/test_jsonrpc_dispatcher.py +++ b/tests/shared/test_jsonrpc_dispatcher.py @@ -305,6 +305,62 @@ async def drive() -> None: s2c_recv.close() +@pytest.mark.anyio +async def test_run_closes_write_stream_after_clean_eof_without_drain_timeout(): + c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32) + s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32) + server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher( + c2s_recv, + s2c_send, + close_write_stream_on_read_close=False, + read_eof_drain_timeout_seconds=None, + ) + on_request, on_notify = echo_handlers(Recorder()) + + with anyio.fail_after(5): + async with anyio.create_task_group() as tg, c2s_send, c2s_recv, s2c_send, s2c_recv: + await tg.start(server.run, on_request, on_notify) + c2s_send.close() + with pytest.raises(anyio.EndOfStream): # pragma: no branch + await s2c_recv.receive() + + +@pytest.mark.anyio +async def test_run_drains_in_flight_handlers_on_clean_eof_without_timeout(): + c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32) + s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32) + server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher( + c2s_recv, + s2c_send, + close_write_stream_on_read_close=False, + read_eof_drain_timeout_seconds=None, + ) + handler_started = anyio.Event() + handler_allowed_to_finish = anyio.Event() + + async def handle_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]: + handler_started.set() + await handler_allowed_to_finish.wait() + return {"drained": True} + + async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None: + raise NotImplementedError + + with anyio.fail_after(5): + async with anyio.create_task_group() as tg, c2s_send, c2s_recv, s2c_send, s2c_recv: + await tg.start(server.run, handle_request, on_notify) + await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="x", params=None))) + await handler_started.wait() + c2s_send.close() + handler_allowed_to_finish.set() + + response = await s2c_recv.receive() + assert isinstance(response, SessionMessage) + assert isinstance(response.message, JSONRPCResponse) + assert response.message.id == 1 + assert response.message.result == {"drained": True} + + @pytest.mark.anyio async def test_run_closes_write_stream_on_exit(): """run() owns both streams; the write end is released once the EOF teardown completes.""" From ff0e2a64250ef2831ca767b1e8542a8adac4f12f Mon Sep 17 00:00:00 2001 From: Jianke LIN Date: Tue, 16 Jun 2026 23:55:12 +0200 Subject: [PATCH 7/8] test: update stdio EOF drain tool schema --- tests/server/test_stdio.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/server/test_stdio.py b/tests/server/test_stdio.py index 8f97c679a..891ef7854 100644 --- a/tests/server/test_stdio.py +++ b/tests/server/test_stdio.py @@ -200,7 +200,7 @@ async def test_stdio_server_drains_in_flight_responses_on_stdin_eof(): allow_tools_to_finish = anyio.Event() async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult: - return ListToolsResult(tools=[Tool(name="slow", description="test", input_schema={})]) + return ListToolsResult(tools=[Tool(name="slow", description="test", input_schema={"type": "object"})]) async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult: nonlocal tool_started_count From c1bfbb1b23d2bb862149654cf2c259b3698dc527 Mon Sep 17 00:00:00 2001 From: Jianke LIN Date: Wed, 17 Jun 2026 00:06:57 +0200 Subject: [PATCH 8/8] test: exercise dispatcher shutdown drain mode --- tests/shared/test_jsonrpc_dispatcher.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/tests/shared/test_jsonrpc_dispatcher.py b/tests/shared/test_jsonrpc_dispatcher.py index 53267f705..bd8deaa85 100644 --- a/tests/shared/test_jsonrpc_dispatcher.py +++ b/tests/shared/test_jsonrpc_dispatcher.py @@ -1034,10 +1034,14 @@ async def test_shutdown_error_response_write_is_bounded_when_the_transport_is_we ): """Cancelling the task group hosting run() completes even when the shutdown error write wedges: only `_SHUTDOWN_WRITE_TIMEOUT` releases the join (SDK-defined). A 0-buffer stream nobody reads - expresses the wedge: run() closes its write stream only after the join, so the send stays parked.""" + expresses the wedge: drain-mode run() closes its write stream only after the join, so the send stays parked.""" c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1) s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](0) - server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send) + server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher( + c2s_recv, + s2c_send, + close_write_stream_on_read_close=False, + ) handler_started = anyio.Event() async def park(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]: @@ -1068,10 +1072,15 @@ async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> @pytest.mark.anyio async def test_shutdown_answers_in_flight_request_with_connection_closed(): """Read-stream EOF answers a still-running request with CONNECTION_CLOSED (SDK-defined): - run() keeps the write stream open until the task-group join, so the shielded teardown write lands.""" + drain-mode run() keeps the write stream open until the task-group join, so the shielded teardown write lands.""" c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4) s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4) - server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send) + server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher( + c2s_recv, + s2c_send, + close_write_stream_on_read_close=False, + read_eof_drain_timeout_seconds=0.05, + ) handler_started = anyio.Event() async def park(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]: