fix(server): opt-in drain on read EOF

Author: Epochex

src/mcp/server/lowlevel/server.py

@@ -347,6 +347,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,
     ):
         async with AsyncExitStack() as stack:
             lifespan_context = await stack.enter_async_context(self.lifespan(self))
@@ -356,26 +360,35 @@ async def run(
                     write_stream,
                     initialization_options,
                     stateless=stateless,
+                    close_write_stream_on_read_close=not drain_on_read_close,
                 )
             )
 
             async with anyio.create_task_group() as tg:
-                async for message in session.incoming_messages:
-                    logger.debug("Received message: %s", message)
-
-                    if isinstance(message, RequestResponder) and message.context is not None:
-                        context = message.context
-                    else:
-                        context = contextvars.copy_context()
-
-                    context.run(
-                        tg.start_soon,
-                        self._handle_message,
-                        message,
-                        session,
-                        lifespan_context,
-                        raise_exceptions,
-                    )
+                try:
+                    async for message in session.incoming_messages:
+                        logger.debug("Received message: %s", message)
+
+                        if isinstance(message, RequestResponder) and message.context is not None:
+                            context = message.context
+                        else:
+                            context = contextvars.copy_context()
+
+                        context.run(
+                            tg.start_soon,
+                            self._handle_message,
+                            message,
+                            session,
+                            lifespan_context,
+                            raise_exceptions,
+                        )
+                finally:
+                    if not drain_on_read_close:
+                        # 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
+                        # _receive_loop closed it when the read stream ended).
+                        tg.cancel_scope.cancel()
 
     async def _handle_message(
         self,

src/mcp/server/mcpserver/server.py

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

src/mcp/server/session.py

@@ -80,8 +80,9 @@ def __init__(
         write_stream: WriteStream[SessionMessage],
         init_options: InitializationOptions,
         stateless: bool = False,
+        close_write_stream_on_read_close: bool = True,
     ) -> None:
-        super().__init__(read_stream, write_stream, close_write_stream_on_read_close=False)
+        super().__init__(read_stream, write_stream, close_write_stream_on_read_close=close_write_stream_on_read_close)
         self._stateless = stateless
         self._initialization_state = (
             InitializationState.Initialized if stateless else InitializationState.NotInitialized

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(

tests/server/test_stdio.py

@@ -170,7 +170,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()