fix(stdio): bound EOF drain wait

Author: Jianke LIN

src/mcp/server/lowlevel/server.py

@@ -72,6 +72,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)
@@ -695,6 +697,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:
         """Serve a single connection over the given streams until the read side closes.
 
@@ -705,20 +710,17 @@ async def run(
         streamable-HTTP manager) call `serve_loop` directly instead.
         """
         async with self.lifespan(self) as lifespan_context:
-            try:
-                await serve_dual_era_loop(
-                    self,
-                    read_stream,
-                    write_stream,
-                    lifespan_state=lifespan_context,
-                    init_options=initialization_options,
-                    raise_exceptions=raise_exceptions,
-                    session_id=None,
-                    close_write_stream_on_read_close=not drain_on_read_close,
-                )
-            finally:
-                if drain_on_read_close:
-                    await write_stream.aclose()
+            await serve_dual_era_loop(
+                self,
+                read_stream,
+                write_stream,
+                lifespan_state=lifespan_context,
+                init_options=initialization_options,
+                raise_exceptions=raise_exceptions,
+                session_id=None,
+                close_write_stream_on_read_close=not drain_on_read_close,
+                read_eof_drain_timeout_seconds=read_eof_drain_timeout_seconds,
+            )
 
     def streamable_http_app(
         self,

src/mcp/server/runner.py

@@ -415,6 +415,7 @@ async def serve_loop(
     init_options: InitializationOptions | None = None,
     raise_exceptions: bool = False,
     close_write_stream_on_read_close: bool = True,
+    read_eof_drain_timeout_seconds: float | None = None,
 ) -> None:
     """Drive ``server`` in handshake-only loop mode over a stream pair until the channel closes.
 
@@ -434,6 +435,7 @@ async def serve_loop(
         # state instead of failing the init-gate.
         inline_methods=frozenset({"initialize"}),
         close_write_stream_on_read_close=close_write_stream_on_read_close,
+        read_eof_drain_timeout_seconds=read_eof_drain_timeout_seconds,
     )
     connection = Connection.for_loop(dispatcher, session_id=session_id)
     await serve_connection(
@@ -548,6 +550,7 @@ async def serve_dual_era_loop(
     init_options: InitializationOptions | None = None,
     raise_exceptions: bool = False,
     close_write_stream_on_read_close: bool = True,
+    read_eof_drain_timeout_seconds: float | None = None,
 ) -> None:
     """Drive `server` over a duplex stream pair, serving both protocol eras.
 
@@ -597,6 +600,7 @@ async def serve_dual_era_loop(
         # next pipelined message is read.
         inline_methods=frozenset({"initialize", "server/discover"}),
         close_write_stream_on_read_close=close_write_stream_on_read_close,
+        read_eof_drain_timeout_seconds=read_eof_drain_timeout_seconds,
     )
     loop_connection = Connection.for_loop(dispatcher, session_id=session_id)
     loop_runner = ServerRunner(server, loop_connection, lifespan_state, init_options=init_options)

src/mcp/shared/jsonrpc_dispatcher.py

@@ -252,6 +252,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.
@@ -268,6 +269,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.
@@ -283,6 +287,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
         """Observer for ``Exception`` items on the read stream. Mutable so a session can
@@ -481,16 +486,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(

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_HANDSHAKE_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