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

Author: Jianke LIN

src/mcp/server/lowlevel/server.py

@@ -641,6 +641,15 @@ async def run(
         # but also make tracing exceptions much easier during testing and when using
         # in-process servers.
         raise_exceptions: bool = False,
+        # When True, the server is stateless and
+        # clients can perform initialization with any node. The client must still follow
+        # 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:
         """Serve a single connection over the given streams until the read side closes.
 
@@ -649,15 +658,20 @@ async def run(
         (the streamable-HTTP manager) call `serve_loop` directly instead.
         """
         async with self.lifespan(self) as lifespan_context:
-            await serve_loop(
-                self,
-                read_stream,
-                write_stream,
-                lifespan_state=lifespan_context,
-                init_options=initialization_options,
-                raise_exceptions=raise_exceptions,
-                close_write_stream_on_read_close=False,
-            )
+            try:
+                await serve_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()
 
     def streamable_http_app(
         self,

src/mcp/server/mcpserver/server.py

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

@@ -200,7 +200,7 @@ async def test_stdio_server_drains_in_flight_responses_on_stdin_eof():
     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={})])
+        return ListToolsResult(tools=[Tool(name="slow", description="test", input_schema={"type": "object"})])
 
     async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
         nonlocal tool_started_count
@@ -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()