From 48e5f12af33fa44aaf99f76d4445d2f4b3171470 Mon Sep 17 00:00:00 2001 From: "767092677@qq.com" <767092677@qq.com> Date: Thu, 27 Aug 2026 16:25:21 +0800 Subject: [PATCH] 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 --- src/mcp/shared/jsonrpc_dispatcher.py | 28 ++++++++++---- tests/shared/test_jsonrpc_dispatcher.py | 51 ++++++++++++++++++++++++- 2 files changed, 70 insertions(+), 9 deletions(-) diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index 87bdf31ceb..5929254d88 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/src/mcp/shared/jsonrpc_dispatcher.py @@ -120,8 +120,8 @@ def cancelled_request_id_from_params(params: Mapping[str, Any] | None) -> Reques class _Pending: """An outbound request awaiting its response.""" - send: MemoryObjectSendStream[dict[str, Any] | ErrorData] - receive: MemoryObjectReceiveStream[dict[str, Any] | ErrorData] + send: MemoryObjectSendStream[dict[str, Any] | ErrorData | Exception] + receive: MemoryObjectReceiveStream[dict[str, Any] | ErrorData | Exception] on_progress: ProgressFnT | None = None @@ -329,6 +329,8 @@ async def send_raw_request( MCPError: Peer error response; `REQUEST_TIMEOUT` if `opts["timeout"]` elapsed; `CONNECTION_CLOSED` if the transport closed or the dispatcher shut down. + Exception: The read stream yielded an exception (transport + fault) while awaiting; re-raised as-is. RuntimeError: Called before `run()`. """ # Post-close sends get the same CONNECTION_CLOSED contract as in-flight waiters. @@ -363,7 +365,7 @@ async def send_raw_request( # buffer=1: a close signal can arrive before the waiter parks in receive(); # a WouldBlock later just means the waiter already has its one outcome. - send, receive = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1) + send, receive = anyio.create_memory_object_stream[dict[str, Any] | ErrorData | Exception](1) pending = _Pending(send=send, receive=receive, on_progress=on_progress) self._pending[pending_key] = pending @@ -442,6 +444,10 @@ async def send_raw_request( if isinstance(outcome, ErrorData): raise MCPError(code=outcome.code, message=outcome.message, data=outcome.data) + if isinstance(outcome, Exception): + # Read stream faulted mid-await: re-raise the transport's exception + # as-is so callers see the original type (e.g. httpx.ReadTimeout). + raise outcome return outcome async def notify( @@ -536,6 +542,9 @@ async def _dispatch( are awaited; any other `await` would head-of-line block the read loop. """ if isinstance(item, Exception): + # No response can arrive over a faulted transport: fail the pending + # waiters now instead of parking them until their timeout elapses. + self._fail_pending(item) if self.on_stream_exception is None: logger.debug("transport yielded exception: %r", item) return @@ -686,14 +695,19 @@ def _spawn( self._tg.start_soon(fn, *args) def _fan_out_closed(self) -> None: - """Wake every pending `send_raw_request` waiter with `CONNECTION_CLOSED`. + """Wake every pending `send_raw_request` waiter with `CONNECTION_CLOSED`.""" + self._fail_pending(ErrorData(code=CONNECTION_CLOSED, message="Connection closed")) - Synchronous: callers may be inside a cancelled scope. Idempotent. + def _fail_pending(self, outcome: ErrorData | Exception) -> None: + """Wake every pending `send_raw_request` waiter with `outcome`. + + `CONNECTION_CLOSED` on EOF, the transport's exception on a faulted + read stream. Synchronous: callers may be inside a cancelled scope. + Idempotent. """ - closed = ErrorData(code=CONNECTION_CLOSED, message="Connection closed") for pending in self._pending.values(): try: - pending.send.send_nowait(closed) + pending.send.send_nowait(outcome) except (anyio.WouldBlock, anyio.BrokenResourceError, anyio.ClosedResourceError): pass self._pending.clear() diff --git a/tests/shared/test_jsonrpc_dispatcher.py b/tests/shared/test_jsonrpc_dispatcher.py index 9bee8b2c3b..2a35ffad78 100644 --- a/tests/shared/test_jsonrpc_dispatcher.py +++ b/tests/shared/test_jsonrpc_dispatcher.py @@ -316,6 +316,53 @@ async def caller() -> None: s.close() +@pytest.mark.anyio +@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"], indirect=True) +async def test_send_raw_request_raises_transport_exception_yielded_mid_await(): + """A blocked send_raw_request is woken with the transport's own exception, not parked + until its timeout elapses; the dispatcher keeps serving once the stream recovers (#1401).""" + 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) + server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send) + release_first = anyio.Event() + + async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]: + # Park the first request so the caller is mid-await when the fault lands. + await release_first.wait() + return {"echoed": method, "params": {}} + + async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None: + raise NotImplementedError + + fault_consumed = anyio.Event() + + async def caller() -> None: + with pytest.raises(RuntimeError, match="transport fault"): + await client.send_raw_request("ping", None) + fault_consumed.set() + + try: + async with anyio.create_task_group() as tg: + await tg.start(client.run, *echo_handlers(Recorder())) + await tg.start(server.run, server_on_request, on_notify) + + tg.start_soon(caller) + await anyio.sleep(0) + # Fault the client's read side mid-await. The buffered send yields no + # checkpoint, so wait for the waiter to consume the fault first. + await s2c_send.send(RuntimeError("transport fault")) + await fault_consumed.wait() + release_first.set() # the parked first response arrives late and is dropped + # The stream stays open, so a later round-trip must still work. + assert await client.send_raw_request("ping", None) == {"echoed": "ping", "params": {}} + s2c_send.close() # EOF both read streams so run() loops exit and the tg joins + c2s_send.close() + finally: + for s in (c2s_send, c2s_recv, s2c_send, s2c_recv): + s.close() + + @pytest.mark.anyio async def test_run_returns_cleanly_when_read_stream_receive_end_is_closed(): """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(): 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) + send, recv = anyio.create_memory_object_stream[dict[str, Any] | ErrorData | Exception](1) d._pending[1] = _Pending(send=send, receive=recv) # pyright: ignore[reportPrivateUsage] recv.close() # waiter gone - send_nowait will raise BrokenResourceError 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(): 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) + send, recv = anyio.create_memory_object_stream[dict[str, Any] | ErrorData | Exception](1) d._pending[1] = _Pending(send=send, receive=recv) # pyright: ignore[reportPrivateUsage] send.send_nowait({"real": "result"}) d._fan_out_closed() # pyright: ignore[reportPrivateUsage]