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/client/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions src/mcp/client/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()


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