fix(stdio): drain responses after stdin EOF

Author: Jianke LIN

src/mcp/server/lowlevel/server.py

@@ -656,6 +656,7 @@ async def run(
                 lifespan_state=lifespan_context,
                 init_options=initialization_options,
                 raise_exceptions=raise_exceptions,
+                close_write_stream_on_read_close=False,
             )
 
     def streamable_http_app(

src/mcp/server/runner.py

@@ -392,6 +392,7 @@ async def serve_loop(
     session_id: str | None = None,
     init_options: InitializationOptions | None = None,
     raise_exceptions: bool = False,
+    close_write_stream_on_read_close: bool = True,
 ) -> None:
     """Drive ``server`` in loop mode over a stream pair until the channel closes.
 
@@ -409,6 +410,7 @@ async def serve_loop(
         # next request (spec: 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=close_write_stream_on_read_close,
     )
     connection = Connection.for_loop(dispatcher, session_id=session_id)
     await serve_connection(

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
@@ -250,6 +251,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.
@@ -262,6 +264,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.
@@ -276,6 +282,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
         """Observer for ``Exception`` items on the read stream. Mutable so a session can
@@ -447,33 +454,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.

tests/server/test_cancel_handling.py

@@ -13,6 +13,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

@@ -7,8 +7,26 @@
 
 import anyio
 import pytest
-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,
+)
+
+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
@@ -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