From dcbbf6dd4ae5ba142db24b837435f16060e92ac1 Mon Sep 17 00:00:00 2001 From: Tasneem Date: Fri, 18 Sep 2026 10:43:46 +0300 Subject: [PATCH] fix(client): prevent tool hang by failing pending requests on transport exception (Closes #1401) --- src/mcp/client/session.py | 9 +++ src/mcp/shared/jsonrpc_dispatcher.py | 32 +++++++-- tests/shared/test_jsonrpc_dispatcher.py | 87 +++++++++++++++++++++++++ 3 files changed, 123 insertions(+), 5 deletions(-) diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index a618112153..d5aadf7bf9 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -243,6 +243,15 @@ async def __call__(self, message: IncomingMessage) -> None: ... # pragma: no br async def _default_message_handler(message: IncomingMessage) -> None: + # Transport-level faults (e.g. an SSE read timeout) reach the handler as the + # `Exception` arm of `IncomingMessage`. The default used to only checkpoint, + # so these vanished with no trace. Log them at ERROR: in-flight requests are + # already failed by the dispatcher, but a fault with no request in flight + # (an idle SSE stream dropping) would otherwise leave no signal at all. + # Server notifications remain a no-op - a handler is optional for those. + if isinstance(message, Exception): + logger.error("transport error surfaced to message handler: %r", message) + return await anyio.lowlevel.checkpoint() diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index 87bdf31ceb..200f7e1e11 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/src/mcp/shared/jsonrpc_dispatcher.py @@ -536,6 +536,11 @@ async def _dispatch( are awaited; any other `await` would head-of-line block the read loop. """ if isinstance(item, Exception): + # A transport fault fails every in-flight request now, rather than + # leaving each waiter to block until its own `opts["timeout"]` + # elapses (or forever, if it has none). Done before the observer + # runs so a slow observer can't delay freeing the waiters. + self._fail_pending(item) if self.on_stream_exception is None: logger.debug("transport yielded exception: %r", item) return @@ -685,19 +690,36 @@ def _spawn( else: self._tg.start_soon(fn, *args) - def _fan_out_closed(self) -> None: - """Wake every pending `send_raw_request` waiter with `CONNECTION_CLOSED`. + def _fan_out_error(self, error: ErrorData) -> None: + """Wake every pending `send_raw_request` waiter with `error`. - Synchronous: callers may be inside a cancelled scope. Idempotent. + Synchronous: callers may be inside a cancelled scope. Idempotent - the + `_pending` table is cleared, so a second call is a no-op. A waiter that + already holds an outcome (buffer-of-1 full) keeps it: the `WouldBlock` + is swallowed rather than clobbering the real response with the signal. """ - closed = ErrorData(code=CONNECTION_CLOSED, message="Connection closed") for pending in self._pending.values(): try: - pending.send.send_nowait(closed) + pending.send.send_nowait(error) except (anyio.WouldBlock, anyio.BrokenResourceError, anyio.ClosedResourceError): pass self._pending.clear() + def _fan_out_closed(self) -> None: + """Wake every pending `send_raw_request` waiter with `CONNECTION_CLOSED` (EOF/shutdown).""" + self._fan_out_error(ErrorData(code=CONNECTION_CLOSED, message="Connection closed")) + + def _fail_pending(self, exc: Exception) -> None: + """Wake every pending waiter because the read stream yielded a transport fault. + + Reported as `CONNECTION_CLOSED` for parity with the EOF path (`send_raw_request` + raises `MCPError` either way), but the message carries `exc` so callers can see + what actually broke - e.g. an `httpx.ReadTimeout` from a low `sse_read_timeout` - + instead of the request hanging until its own timeout. The raw `exc` still reaches + any `on_stream_exception` observer untouched. + """ + self._fan_out_error(ErrorData(code=CONNECTION_CLOSED, message=f"Transport error: {exc!r}")) + async def _handle_request( self, req: JSONRPCRequest, diff --git a/tests/shared/test_jsonrpc_dispatcher.py b/tests/shared/test_jsonrpc_dispatcher.py index 9bee8b2c3b..4772da609b 100644 --- a/tests/shared/test_jsonrpc_dispatcher.py +++ b/tests/shared/test_jsonrpc_dispatcher.py @@ -2514,3 +2514,90 @@ async def call() -> None: for stream in (c2s_send, c2s_recv, s2c_send, s2c_recv): stream.close() assert result_box == [{"ok": True}] + + +@pytest.mark.anyio +async def test_transport_exception_fails_pending_request_without_hanging(): + """A read-stream fault wakes an in-flight `send_raw_request` with `CONNECTION_CLOSED`. + + Regression for the streamable-http hang: before this, a transport error (e.g. an SSE + read timeout) reached only the observer, so a request already waiting on its response + sat until its own `opts["timeout"]` elapsed. Now the dispatcher fails the waiter at once, + and the error message carries the transport exception so the caller can see the cause. + """ + c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4) + s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4) + + client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send) + on_request, on_notify = echo_handlers(Recorder()) + outcome: dict[str, BaseException] = {} + boom = TimeoutError("sse read timed out") + try: + async with anyio.create_task_group() as tg: + await tg.start(client.run, on_request, on_notify) + + async def call() -> None: + # No `timeout` opt: without the fix this would block forever. + try: + await client.send_raw_request("tools/call", {"name": "slow"}) + except BaseException as exc: # noqa: BLE001 - capture whatever the waiter raises + outcome["exc"] = exc + + tg.start_soon(call) + # Let the request register in `_pending` and park on its response. + with anyio.fail_after(5): + sent = await s2c_recv.receive() + assert isinstance(sent, SessionMessage) + assert isinstance(sent.message, JSONRPCRequest) + + # The transport now yields an exception instead of a response. + await c2s_send.send(boom) + with anyio.fail_after(5): + while "exc" not in outcome: + await anyio.sleep(0) + tg.cancel_scope.cancel() + finally: + for s in (c2s_send, c2s_recv, s2c_send, s2c_recv): + s.close() + + raised = outcome["exc"] + assert isinstance(raised, MCPError) + assert raised.error.code == CONNECTION_CLOSED + # The original transport exception is preserved in the message for debugging. + assert "sse read timed out" in raised.error.message + + +def test_fail_pending_reports_transport_exception_and_clears_pending(): + """White-box: `_fail_pending` wakes waiters with the exception detail, then empties `_pending`.""" + c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1) + s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1) + d: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send) + send, recv = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1) + d._pending[1] = _Pending(send=send, receive=recv) # pyright: ignore[reportPrivateUsage] + + d._fail_pending(TimeoutError("sse read timed out")) # pyright: ignore[reportPrivateUsage] + + signalled = recv.receive_nowait() + assert isinstance(signalled, ErrorData) + assert signalled.code == CONNECTION_CLOSED + assert "sse read timed out" in signalled.message + assert d._pending == {} # pyright: ignore[reportPrivateUsage] + for s in (c2s_send, c2s_recv, s2c_send, s2c_recv, send, recv): + s.close() + + +def test_fail_pending_keeps_existing_outcome_when_waiter_already_resolved(): + """White-box: a waiter that already holds a real result is not clobbered by the fault signal.""" + c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1) + s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1) + d: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send) + send, recv = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1) + d._pending[1] = _Pending(send=send, receive=recv) # pyright: ignore[reportPrivateUsage] + send.send_nowait({"real": "result"}) + + d._fail_pending(TimeoutError("sse read timed out")) # pyright: ignore[reportPrivateUsage] + + assert recv.receive_nowait() == {"real": "result"} + assert d._pending == {} # pyright: ignore[reportPrivateUsage] + for s in (c2s_send, c2s_recv, s2c_send, s2c_recv, send, recv): + s.close()