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

Author: Jianke LIN

src/mcp/server/lowlevel/server.py

@@ -686,6 +686,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.
 
@@ -696,15 +705,20 @@ async def run(
         streamable-HTTP manager) call `serve_loop` directly instead.
         """
         async with self.lifespan(self) as lifespan_context:
-            await serve_dual_era_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_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()
 
     def streamable_http_app(
         self,

src/mcp/server/mcpserver/server.py

@@ -1010,6 +1010,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/runner.py

@@ -547,6 +547,7 @@ async def serve_dual_era_loop(
     session_id: str | None = None,
     init_options: InitializationOptions | None = None,
     raise_exceptions: bool = False,
+    close_write_stream_on_read_close: bool = True,
 ) -> None:
     """Drive `server` over a duplex stream pair, serving both protocol eras.
 
@@ -595,6 +596,7 @@ async def serve_dual_era_loop(
         # `server/discover` inline so the modern era lock commits before the
         # next pipelined message is read.
         inline_methods=frozenset({"initialize", "server/discover"}),
+        close_write_stream_on_read_close=close_write_stream_on_read_close,
     )
     loop_connection = Connection.for_loop(dispatcher, session_id=session_id)
     loop_runner = ServerRunner(server, loop_connection, lifespan_state, init_options=init_options)

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

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