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

Author: Epochex

src/mcp/server/lowlevel/server.py

@@ -406,6 +406,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,
     ) -> None:
         async with self.lifespan(self) as lifespan_context:
             dispatcher: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
@@ -416,7 +420,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,
+                close_write_stream_on_read_close=not drain_on_read_close,
             )
             runner = ServerRunner(
                 server=self,

src/mcp/server/mcpserver/server.py

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

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

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