From b389b31eedf1541694f25c1c335933b7b3d34190 Mon Sep 17 00:00:00 2001 From: Steeve-Crypto Date: Thu, 24 Sep 2026 15:16:19 -0400 Subject: [PATCH] fix: surface ClientSession transport faults --- src/mcp/client/client.py | 2 ++ src/mcp/client/session.py | 2 ++ src/mcp/shared/jsonrpc_dispatcher.py | 17 ++++++++++ tests/client/test_session.py | 17 ++++++++++ tests/shared/test_jsonrpc_dispatcher.py | 42 +++++++++++++++++++++++++ 5 files changed, 80 insertions(+) diff --git a/src/mcp/client/client.py b/src/mcp/client/client.py index f921c7e30b..7f433684f5 100644 --- a/src/mcp/client/client.py +++ b/src/mcp/client/client.py @@ -166,6 +166,8 @@ async def handler(message: IncomingMessage) -> None: await user_handler(message) else: # Mirrors ClientSession's default handler (session._default_message_handler). + if isinstance(message, Exception): + raise message await anyio.lowlevel.checkpoint() return handler diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index a618112153..345d020f94 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -243,6 +243,8 @@ async def __call__(self, message: IncomingMessage) -> None: ... # pragma: no br async def _default_message_handler(message: IncomingMessage) -> None: + if isinstance(message, Exception): + raise message await anyio.lowlevel.checkpoint() diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index 87bdf31ceb..b42f875500 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/src/mcp/shared/jsonrpc_dispatcher.py @@ -536,6 +536,9 @@ async def _dispatch( are awaited; any other `await` would head-of-line block the read loop. """ if isinstance(item, Exception): + # Fail in-flight waiters with CONNECTION_CLOSED so they do not hang, + # but keep the receive loop open (unlike EOF / `_fan_out_closed`). + self._fan_out_transport_error(item) if self.on_stream_exception is None: logger.debug("transport yielded exception: %r", item) return @@ -698,6 +701,20 @@ def _fan_out_closed(self) -> None: pass self._pending.clear() + def _fan_out_transport_error(self, exc: Exception) -> None: + """Wake every pending `send_raw_request` waiter after a transport Exception item. + + Unlike `_fan_out_closed`, this does not mark the dispatcher closed: a single + Exception item is not EOF, and the receive loop must keep serving. + """ + error = ErrorData(code=CONNECTION_CLOSED, message=f"Transport error: {exc!r}") + for pending in self._pending.values(): + try: + pending.send.send_nowait(error) + except (anyio.WouldBlock, anyio.BrokenResourceError, anyio.ClosedResourceError): + pass + self._pending.clear() + async def _handle_request( self, req: JSONRPCRequest, diff --git a/tests/client/test_session.py b/tests/client/test_session.py index 6663fb47a2..f356da2279 100644 --- a/tests/client/test_session.py +++ b/tests/client/test_session.py @@ -1001,6 +1001,23 @@ async def handler(msg: object) -> None: assert "Failed to validate notification: notifications/progress" in caplog.text +@pytest.mark.anyio +async def test_default_message_handler_raises_on_transport_exception( + caplog: pytest.LogCaptureFixture, +): + """With no custom `message_handler`, a transport `Exception` is re-raised by the + default handler and logged via `_deliver_stream_exception` (SDK-defined). + Raw streams because only a transport can put an `Exception` item on the read stream.""" + async with raw_client_session() as (_session, to_client, from_client): + await to_client.send(ValueError("bad bytes")) + # Prove the receive loop kept serving after the default handler re-raised. + await to_client.send(SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=9, method="ping"))) + out = await from_client.receive() + assert isinstance(out.message, JSONRPCResponse) + assert out.message.id == 9 + assert "message_handler raised on transport exception" in caplog.text + + @pytest.mark.anyio async def test_raising_message_handler_on_transport_exception_costs_the_delivery_not_the_connection( caplog: pytest.LogCaptureFixture, diff --git a/tests/shared/test_jsonrpc_dispatcher.py b/tests/shared/test_jsonrpc_dispatcher.py index 9bee8b2c3b..d290f7ba8f 100644 --- a/tests/shared/test_jsonrpc_dispatcher.py +++ b/tests/shared/test_jsonrpc_dispatcher.py @@ -1594,6 +1594,48 @@ async def test_transport_exception_in_read_stream_is_logged_and_dropped(): s.close() +@pytest.mark.anyio +async def test_send_raw_request_raises_connection_closed_on_transport_exception(): + """A blocked send_raw_request is woken with CONNECTION_CLOSED when a transport + Exception item arrives; the receive loop stays healthy for a later request.""" + c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32) + s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32) + client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send) + on_request, on_notify = echo_handlers(Recorder()) + hiccup = ValueError("transport hiccup") + try: + async with anyio.create_task_group() as tg: + await tg.start(client.run, on_request, on_notify) + + async def caller() -> None: + with pytest.raises(MCPError) as exc: + await client.send_raw_request("ping", None) + assert exc.value.error.code == CONNECTION_CLOSED + assert "Transport error:" in exc.value.error.message + assert "transport hiccup" in exc.value.error.message + + tg.start_soon(caller) + # Let the outbound request register in `_pending` before the fault. + req = await c2s_recv.receive() + assert isinstance(req, SessionMessage) + assert isinstance(req.message, JSONRPCRequest) + await s2c_send.send(hiccup) + await anyio.sleep(0) + # Loop must still serve after the Exception item (not closed like EOF). + await s2c_send.send( + SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=99, method="t", params=None)) + ) + with anyio.fail_after(5): + resp = await c2s_recv.receive() + assert isinstance(resp, SessionMessage) + assert isinstance(resp.message, JSONRPCResponse) + assert resp.message.id == 99 + tg.cancel_scope.cancel() + finally: + for s in (c2s_send, c2s_recv, s2c_send, s2c_recv): + s.close() + + @pytest.mark.anyio async def test_on_stream_exception_observes_transport_exceptions(): """With an observer set, Exception items reach it instead of being dropped; the loop stays healthy."""