Skip to content

Commit 48e5f12

Browse files
chaucerjclaude
andcommitted
Fail pending send_raw_request waiters when the read stream yields an exception
When a transport's read stream yields an exception item, pending send_raw_request waiters were only woken by the on_stream_exception observer (a no-op by default) -- they parked until their own timeout elapsed. Fan the raw exception out to every pending waiter and re-raise it as-is so callers see the transport's original exception type (e.g. httpx.ReadTimeout). _fan_out_closed semantics are unchanged (it now delegates to the generalized _fail_pending helper). Fixes #1401 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 6705402 commit 48e5f12

2 files changed

Lines changed: 70 additions & 9 deletions

File tree

src/mcp/shared/jsonrpc_dispatcher.py

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -120,8 +120,8 @@ def cancelled_request_id_from_params(params: Mapping[str, Any] | None) -> Reques
120120
class _Pending:
121121
"""An outbound request awaiting its response."""
122122

123-
send: MemoryObjectSendStream[dict[str, Any] | ErrorData]
124-
receive: MemoryObjectReceiveStream[dict[str, Any] | ErrorData]
123+
send: MemoryObjectSendStream[dict[str, Any] | ErrorData | Exception]
124+
receive: MemoryObjectReceiveStream[dict[str, Any] | ErrorData | Exception]
125125
on_progress: ProgressFnT | None = None
126126

127127

@@ -329,6 +329,8 @@ async def send_raw_request(
329329
MCPError: Peer error response; `REQUEST_TIMEOUT` if
330330
`opts["timeout"]` elapsed; `CONNECTION_CLOSED` if the
331331
transport closed or the dispatcher shut down.
332+
Exception: The read stream yielded an exception (transport
333+
fault) while awaiting; re-raised as-is.
332334
RuntimeError: Called before `run()`.
333335
"""
334336
# Post-close sends get the same CONNECTION_CLOSED contract as in-flight waiters.
@@ -363,7 +365,7 @@ async def send_raw_request(
363365

364366
# buffer=1: a close signal can arrive before the waiter parks in receive();
365367
# a WouldBlock later just means the waiter already has its one outcome.
366-
send, receive = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1)
368+
send, receive = anyio.create_memory_object_stream[dict[str, Any] | ErrorData | Exception](1)
367369
pending = _Pending(send=send, receive=receive, on_progress=on_progress)
368370
self._pending[pending_key] = pending
369371

@@ -442,6 +444,10 @@ async def send_raw_request(
442444

443445
if isinstance(outcome, ErrorData):
444446
raise MCPError(code=outcome.code, message=outcome.message, data=outcome.data)
447+
if isinstance(outcome, Exception):
448+
# Read stream faulted mid-await: re-raise the transport's exception
449+
# as-is so callers see the original type (e.g. httpx.ReadTimeout).
450+
raise outcome
445451
return outcome
446452

447453
async def notify(
@@ -536,6 +542,9 @@ async def _dispatch(
536542
are awaited; any other `await` would head-of-line block the read loop.
537543
"""
538544
if isinstance(item, Exception):
545+
# No response can arrive over a faulted transport: fail the pending
546+
# waiters now instead of parking them until their timeout elapses.
547+
self._fail_pending(item)
539548
if self.on_stream_exception is None:
540549
logger.debug("transport yielded exception: %r", item)
541550
return
@@ -686,14 +695,19 @@ def _spawn(
686695
self._tg.start_soon(fn, *args)
687696

688697
def _fan_out_closed(self) -> None:
689-
"""Wake every pending `send_raw_request` waiter with `CONNECTION_CLOSED`.
698+
"""Wake every pending `send_raw_request` waiter with `CONNECTION_CLOSED`."""
699+
self._fail_pending(ErrorData(code=CONNECTION_CLOSED, message="Connection closed"))
690700

691-
Synchronous: callers may be inside a cancelled scope. Idempotent.
701+
def _fail_pending(self, outcome: ErrorData | Exception) -> None:
702+
"""Wake every pending `send_raw_request` waiter with `outcome`.
703+
704+
`CONNECTION_CLOSED` on EOF, the transport's exception on a faulted
705+
read stream. Synchronous: callers may be inside a cancelled scope.
706+
Idempotent.
692707
"""
693-
closed = ErrorData(code=CONNECTION_CLOSED, message="Connection closed")
694708
for pending in self._pending.values():
695709
try:
696-
pending.send.send_nowait(closed)
710+
pending.send.send_nowait(outcome)
697711
except (anyio.WouldBlock, anyio.BrokenResourceError, anyio.ClosedResourceError):
698712
pass
699713
self._pending.clear()

tests/shared/test_jsonrpc_dispatcher.py

Lines changed: 49 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,53 @@ async def caller() -> None:
316316
s.close()
317317

318318

319+
@pytest.mark.anyio
320+
@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"], indirect=True)
321+
async def test_send_raw_request_raises_transport_exception_yielded_mid_await():
322+
"""A blocked send_raw_request is woken with the transport's own exception, not parked
323+
until its timeout elapses; the dispatcher keeps serving once the stream recovers (#1401)."""
324+
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
325+
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
326+
client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
327+
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
328+
release_first = anyio.Event()
329+
330+
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
331+
# Park the first request so the caller is mid-await when the fault lands.
332+
await release_first.wait()
333+
return {"echoed": method, "params": {}}
334+
335+
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
336+
raise NotImplementedError
337+
338+
fault_consumed = anyio.Event()
339+
340+
async def caller() -> None:
341+
with pytest.raises(RuntimeError, match="transport fault"):
342+
await client.send_raw_request("ping", None)
343+
fault_consumed.set()
344+
345+
try:
346+
async with anyio.create_task_group() as tg:
347+
await tg.start(client.run, *echo_handlers(Recorder()))
348+
await tg.start(server.run, server_on_request, on_notify)
349+
350+
tg.start_soon(caller)
351+
await anyio.sleep(0)
352+
# Fault the client's read side mid-await. The buffered send yields no
353+
# checkpoint, so wait for the waiter to consume the fault first.
354+
await s2c_send.send(RuntimeError("transport fault"))
355+
await fault_consumed.wait()
356+
release_first.set() # the parked first response arrives late and is dropped
357+
# The stream stays open, so a later round-trip must still work.
358+
assert await client.send_raw_request("ping", None) == {"echoed": "ping", "params": {}}
359+
s2c_send.close() # EOF both read streams so run() loops exit and the tg joins
360+
c2s_send.close()
361+
finally:
362+
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
363+
s.close()
364+
365+
319366
@pytest.mark.anyio
320367
async def test_run_returns_cleanly_when_read_stream_receive_end_is_closed():
321368
"""Iterating a closed receive end is EOF, not a crash (stateless SHTTP closes it during teardown)."""
@@ -1826,7 +1873,7 @@ def test_resolve_pending_drops_outcome_when_waiter_stream_already_closed():
18261873
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
18271874
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
18281875
d: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
1829-
send, recv = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1)
1876+
send, recv = anyio.create_memory_object_stream[dict[str, Any] | ErrorData | Exception](1)
18301877
d._pending[1] = _Pending(send=send, receive=recv) # pyright: ignore[reportPrivateUsage]
18311878
recv.close() # waiter gone - send_nowait will raise BrokenResourceError
18321879
d._resolve_pending(1, {"late": True}) # pyright: ignore[reportPrivateUsage]
@@ -1839,7 +1886,7 @@ def test_fan_out_closed_drops_signal_when_waiter_already_has_outcome():
18391886
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
18401887
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
18411888
d: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
1842-
send, recv = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1)
1889+
send, recv = anyio.create_memory_object_stream[dict[str, Any] | ErrorData | Exception](1)
18431890
d._pending[1] = _Pending(send=send, receive=recv) # pyright: ignore[reportPrivateUsage]
18441891
send.send_nowait({"real": "result"})
18451892
d._fan_out_closed() # pyright: ignore[reportPrivateUsage]

0 commit comments

Comments
 (0)