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


Expand Down
32 changes: 27 additions & 5 deletions src/mcp/shared/jsonrpc_dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
87 changes: 87 additions & 0 deletions tests/shared/test_jsonrpc_dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Loading