Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/mcp/server/lowlevel/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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(
Expand Down
4 changes: 4 additions & 0 deletions src/mcp/server/mcpserver/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
19 changes: 16 additions & 3 deletions src/mcp/server/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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(
Expand All @@ -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()
Expand Down Expand Up @@ -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(
Expand All @@ -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)

Expand All @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
36 changes: 34 additions & 2 deletions src/mcp/shared/jsonrpc_dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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,
Expand Down Expand Up @@ -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.
Expand All @@ -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,
Expand Down
34 changes: 34 additions & 0 deletions tests/shared/test_jsonrpc_dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
Loading