diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index 6df5341e41..cad7dc70f0 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/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( diff --git a/src/mcp/server/mcpserver/server.py b/src/mcp/server/mcpserver/server.py index 1f3e863dbd..0cb433a744 100644 --- a/src/mcp/server/mcpserver/server.py +++ b/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 diff --git a/src/mcp/server/runner.py b/src/mcp/server/runner.py index 26e8efbe57..496e5bed3e 100644 --- a/src/mcp/server/runner.py +++ b/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( diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index 87bdf31ceb..4b13d33498 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/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, diff --git a/tests/shared/test_jsonrpc_dispatcher.py b/tests/shared/test_jsonrpc_dispatcher.py index 9bee8b2c3b..6659bb2cb9 100644 --- a/tests/shared/test_jsonrpc_dispatcher.py +++ b/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."""