fix(stdio): drain accepted responses before EOF shutdown

Signed-off-by: ace-trump-tech <ace-trump-tech@users.noreply.github.com>

Author: ace-trump-tech

src/mcp/server/lowlevel/server.py

@@ -699,6 +699,7 @@ async def run(
         # but also make tracing exceptions much easier during testing and when using
         # in-process servers.
         raise_exceptions: bool = False,
+        graceful_shutdown_timeout: float = 0,
     ) -> None:
         """Serve a single connection over the given streams until the read side closes.
 
@@ -716,6 +717,7 @@ async def run(
                 lifespan_state=lifespan_context,
                 init_options=initialization_options,
                 raise_exceptions=raise_exceptions,
+                graceful_shutdown_timeout=graceful_shutdown_timeout,
             )
 
     def streamable_http_app(

src/mcp/server/mcpserver/server.py

@@ -1069,6 +1069,10 @@ async def run_stdio_async(self) -> None:
                 read_stream,
                 write_stream,
                 self._lowlevel_server.create_initialization_options(),
+                # File-redirected stdin can reach EOF while accepted tool
+                # handlers are still producing responses. Give those writes
+                # a bounded drain window before cancelling the connection.
+                graceful_shutdown_timeout=0.5,
             )
 
     async def run_sse_async(  # pragma: no cover

src/mcp/server/runner.py

@@ -608,6 +608,7 @@ async def serve_dual_era_loop(
     session_id: str | None = None,
     init_options: InitializationOptions | None = None,
     raise_exceptions: bool = False,
+    graceful_shutdown_timeout: float = 0,
 ) -> None:
     """Drive `server` over a duplex stream pair, in the era the client opens with.
 
@@ -630,7 +631,8 @@ async def serve_dual_era_loop(
             )
             if opens_modern:
                 await _serve_modern_stream(
-                    server, replayed, write_stream, lifespan_state=lifespan_state, raise_exceptions=raise_exceptions
+                    server, replayed, write_stream, lifespan_state=lifespan_state,
+                    raise_exceptions=raise_exceptions, graceful_shutdown_timeout=graceful_shutdown_timeout
                 )
             else:
                 await _serve_legacy_stream(
@@ -641,6 +643,7 @@ async def serve_dual_era_loop(
                     session_id=session_id,
                     init_options=init_options,
                     raise_exceptions=raise_exceptions,
+                    graceful_shutdown_timeout=graceful_shutdown_timeout,
                 )
     finally:
         await write_stream.aclose()
@@ -723,6 +726,7 @@ async def _serve_legacy_stream(
     session_id: str | None,
     init_options: InitializationOptions | None,
     raise_exceptions: bool,
+    graceful_shutdown_timeout: float = 0,
 ) -> None:
     """Serve a 2025 handshake connection; enveloped requests are refused."""
     dispatcher: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
@@ -747,7 +751,11 @@ async def on_request(
         return await runner.on_request(dctx, method, params)
 
     try:
-        await dispatcher.run(on_request, runner.on_notify)
+        await dispatcher.run(
+            on_request,
+            runner.on_notify,
+            graceful_shutdown_timeout=graceful_shutdown_timeout,
+        )
     finally:
         await aclose_shielded(connection)
 
@@ -759,6 +767,7 @@ async def _serve_modern_stream(
     *,
     lifespan_state: LifespanT,
     raise_exceptions: bool,
+    graceful_shutdown_timeout: float = 0,
 ) -> None:
     """Serve a 2026-07-28 connection: every request carries its own envelope."""
     dispatcher: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
@@ -809,7 +818,11 @@ async def on_notify(dctx: DispatchContext[TransportContext], method: str, params
         finally:
             await aclose_shielded(connection)
 
-    await dispatcher.run(on_request, on_notify)
+    await dispatcher.run(
+        on_request,
+        on_notify,
+        graceful_shutdown_timeout=graceful_shutdown_timeout,
+    )
 
 
 async def serve_one(

src/mcp/shared/jsonrpc_dispatcher.py

@@ -307,6 +307,7 @@ def __init__(
         self._next_id = 0
         self._pending: dict[RequestId, _Pending] = {}
         self._in_flight: dict[RequestId, _InFlight[TransportT]] = {}
+        self._active_requests: set[anyio.Event] = set()
         self._on_notify_intercept: OnNotifyIntercept | None = None
         self._tg: anyio.abc.TaskGroup | None = None
         self._running = False
@@ -480,12 +481,15 @@ async def run(
         on_notify_intercept: OnNotifyIntercept | None = None,
         *,
         task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
+        graceful_shutdown_timeout: float = 0,
     ) -> None:
         """Drive the receive loop until the read stream closes.
 
         `task_status.started()` fires once `send_raw_request` is usable.
         Single-shot: once the loop ends the dispatcher stays closed and cannot be restarted.
         """
+        if graceful_shutdown_timeout < 0:
+            raise ValueError("graceful_shutdown_timeout must be non-negative")
         self._on_notify_intercept = on_notify_intercept
         try:
             # LIFO exits: the write stream closes only after the task-group join, so teardown writes still land.
@@ -511,6 +515,9 @@ async def run(
                         self._running = False
                         self._closed = True
                         self._fan_out_closed()
+                        if graceful_shutdown_timeout and self._active_requests:
+                            with anyio.move_on_after(graceful_shutdown_timeout):
+                                await self._wait_for_active_requests()
                     finally:
                         # Cancel in-flight handlers; otherwise the task-group join
                         # waits on handlers whose callers are already gone.
@@ -523,6 +530,15 @@ async def run(
             self._fan_out_closed()
             await resync_tracer()
 
+    async def _wait_for_active_requests(self) -> None:
+        """Wait for requests already accepted from a transport before cancellation."""
+        events = tuple(self._active_requests)
+        if not events:
+            return
+        async with anyio.create_task_group() as tg:
+            for event in events:
+                tg.start_soon(event.wait)
+
     async def _dispatch(
         self,
         item: SessionMessage | Exception,
@@ -585,6 +601,8 @@ async def _dispatch_request(
             _progress_token=progress_token,
         )
         scope = anyio.CancelScope()
+        completion = anyio.Event()
+        self._active_requests.add(completion)
         # TODO(maxisbey): duplicate ids blind-overwrite (v1/TS parity); revisit
         # rejecting with INVALID_REQUEST. Key coerced so a stringified
         # `notifications/cancelled` id still correlates.
@@ -596,14 +614,28 @@ async def _dispatch_request(
 
             async def _run_inline() -> None:
                 try:
-                    await self._handle_request(req, dctx, scope, on_request)
+                    await self._run_request(req, dctx, scope, on_request, completion)
                 finally:
                     done.set()
 
             self._spawn(_run_inline, sender_ctx=sender_ctx)
             await done.wait()
         else:
-            self._spawn(self._handle_request, req, dctx, scope, on_request, sender_ctx=sender_ctx)
+            self._spawn(self._run_request, req, dctx, scope, on_request, completion, sender_ctx=sender_ctx)
+
+    async def _run_request(
+        self,
+        req: JSONRPCRequest,
+        dctx: _JSONRPCDispatchContext[TransportT],
+        scope: anyio.CancelScope,
+        on_request: OnRequest,
+        completion: anyio.Event,
+    ) -> None:
+        try:
+            await self._handle_request(req, dctx, scope, on_request)
+        finally:
+            self._active_requests.discard(completion)
+            completion.set()
 
     def _dispatch_notification(
         self,

tests/shared/test_jsonrpc_dispatcher.py

@@ -369,6 +369,40 @@ async def drive() -> None:
     s2c_recv.close()
 
 
+@pytest.mark.anyio
+async def test_run_can_drain_in_flight_handlers_before_eof_shutdown():
+    """A stdio-style EOF may follow a piped request; accepted responses get a bounded drain window."""
+    c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
+    handler_started = anyio.Event()
+
+    async def slow(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
+        handler_started.set()
+        await anyio.sleep(0.01)
+        return {"ok": True}
+
+    async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
+        raise NotImplementedError
+
+    async def drive() -> None:
+        await server.run(slow, on_notify, graceful_shutdown_timeout=0.5)
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(drive)
+        await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="slow")))
+        await handler_started.wait()
+        c2s_send.close()
+        with anyio.fail_after(5):
+            response = await s2c_recv.receive()
+
+    assert isinstance(response, SessionMessage)
+    assert isinstance(response.message, JSONRPCResponse)
+    assert response.message.id == 1
+    assert response.message.result == {"ok": True}
+    s2c_recv.close()
+
+
 @pytest.mark.anyio
 async def test_run_closes_write_stream_on_exit():
     """run() owns both streams; the write end is released once the EOF teardown completes."""