fix(stdio): bound EOF drain wait

Author: Epochex

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)
 
 
@@ -351,6 +353,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,
     ):
         async with AsyncExitStack() as stack:
             lifespan_context = await stack.enter_async_context(self.lifespan(self))
@@ -383,7 +388,14 @@ async def run(
                             raise_exceptions,
                         )
                 finally:
-                    if not drain_on_read_close:
+                    if drain_on_read_close:
+                        if read_eof_drain_timeout_seconds is not None:
+                            with anyio.move_on_after(read_eof_drain_timeout_seconds) as drain_scope:
+                                while session.has_in_flight_requests:
+                                    await anyio.sleep(0.01)
+                            if drain_scope.cancelled_caught:
+                                tg.cancel_scope.cancel()
+                    else:
                         # Transport closed: cancel in-flight handlers. Without this the
                         # TG join waits for them, and when they eventually try to
                         # respond they hit a closed write stream (the session's

src/mcp/shared/session.py

@@ -209,6 +209,11 @@ def __init__(
         self._exit_stack.push_async_callback(self._read_stream.aclose)
         self._exit_stack.push_async_callback(self._write_stream.aclose)
 
+    @property
+    def has_in_flight_requests(self) -> bool:
+        """Whether any received requests are still awaiting a response."""
+        return bool(self._in_flight)
+
     async def __aenter__(self) -> Self:
         self._task_group = anyio.create_task_group()
         await self._task_group.__aenter__()

tests/server/test_cancel_handling.py

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