fix(stdio): drain responses after stdin EOF

Author: Epochex

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,

src/mcp/shared/jsonrpc_dispatcher.py

@@ -23,6 +23,7 @@
 import contextvars
 import logging
 from collections.abc import Awaitable, Callable, Mapping
+from contextlib import AsyncExitStack
 from dataclasses import dataclass, field
 from typing import Any, Generic, Literal, TypeVar, cast, overload
 
@@ -226,6 +227,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,
     ) -> None: ...
     @overload
     def __init__(
@@ -237,6 +239,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,
     ) -> None: ...
     def __init__(
         self,
@@ -247,6 +250,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,
     ) -> None:
         self._read_stream = read_stream
         self._write_stream = write_stream
@@ -259,6 +263,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
         # Request methods handled inline in the read loop (awaited before the
         # next message is dequeued) instead of spawned concurrently. Use for
         # methods whose side effects must be observable to the next message,
@@ -400,13 +405,17 @@ async def run(
         `await tg.start(dispatcher.run, ...)` resumes when `send_raw_request`
         is usable.
         """
+        normal_eof = False
         try:
             async with anyio.create_task_group() as tg:
                 self._tg = tg
                 self._running = True
                 task_status.started()
                 try:
-                    async with self._read_stream, self._write_stream:
+                    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: `_context_streams.ContextReceiveStream`
@@ -425,14 +434,13 @@ async def run(
                     # (callers outside this task group) with CONNECTION_CLOSED.
                     self._running = False
                     self._fan_out_closed()
+                    normal_eof = True
                 finally:
-                    # Transport closed: cancel in-flight handlers. Without this
-                    # the task-group join waits for them, and a handler that
-                    # outlives its caller (its request timed out client-side, or
-                    # the client disconnected mid-call) would keep `run()` from
-                    # returning forever. Same behaviour as `Server.run()` before
-                    # the dispatcher rework.
-                    tg.cancel_scope.cancel()
+                    if not normal_eof:
+                        # Transport closed abnormally: cancel in-flight handlers.
+                        # On normal EOF, let already-received handlers drain
+                        # their responses before the task group exits.
+                        tg.cancel_scope.cancel()
         finally:
             # Covers the cancel/crash paths where the inline fan-out above is
             # never reached. Idempotent.

src/mcp/shared/session.py

@@ -148,14 +148,24 @@ def __init__(
         write_stream: WriteStream[SessionMessage],
         # If none, reading will never time out
         read_timeout_seconds: float | None = None,
+        # When True, closing/EOF on the read stream closes the write stream too.
+        #
+        # For full-duplex transports (e.g., stdio), an input EOF can be a
+        # half-close: the peer is done sending requests but still expects
+        # responses on the output stream. In that case, callers may opt out so
+        # in-flight handlers can drain their responses before shutdown.
+        close_write_stream_on_read_close: bool = True,
     ) -> None:
         self._read_stream = read_stream
         self._write_stream = write_stream
         self._response_streams = {}
         self._request_id = 0
         self._session_read_timeout_seconds = read_timeout_seconds
+        self._close_write_stream_on_read_close = close_write_stream_on_read_close
         self._progress_callbacks = {}
         self._exit_stack = AsyncExitStack()
+        self._exit_stack.push_async_callback(self._read_stream.aclose)
+        self._exit_stack.push_async_callback(self._write_stream.aclose)
 
     async def __aenter__(self) -> Self:
         self._task_group = anyio.create_task_group()
@@ -291,7 +301,10 @@ def _receive_notification_adapter(self) -> TypeAdapter[ReceiveNotificationT]:
         raise NotImplementedError
 
     async def _receive_loop(self) -> None:
-        async with self._read_stream, self._write_stream:
+        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 def _handle_session_message(message: SessionMessage) -> None:

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

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