Add an on_stream_exception observer to the dispatcher

Transports yield Exception items on the read stream for connection faults
and parse errors; the dispatcher debug-logged and dropped them. An optional
observer now receives them (awaited in the read loop, contained so a raising
observer costs the item, not the connection). Unset keeps the old behavior.

Author: maxisbey

src/mcp/shared/jsonrpc_dispatcher.py

@@ -245,6 +245,7 @@ def __init__(
         peer_cancel_mode: PeerCancelMode = "interrupt",
         raise_handler_exceptions: bool = False,
         inline_methods: frozenset[str] = frozenset(),
+        on_stream_exception: Callable[[Exception], Awaitable[None]] | None = None,
     ) -> None: ...
     @overload
     def __init__(
@@ -256,6 +257,7 @@ def __init__(
         peer_cancel_mode: PeerCancelMode = "interrupt",
         raise_handler_exceptions: bool = False,
         inline_methods: frozenset[str] = frozenset(),
+        on_stream_exception: Callable[[Exception], Awaitable[None]] | None = None,
     ) -> None: ...
     def __init__(
         self,
@@ -266,6 +268,7 @@ def __init__(
         peer_cancel_mode: PeerCancelMode = "interrupt",
         raise_handler_exceptions: bool = False,
         inline_methods: frozenset[str] = frozenset(),
+        on_stream_exception: Callable[[Exception], Awaitable[None]] | None = None,
     ) -> None:
         self._read_stream = read_stream
         self._write_stream = write_stream
@@ -287,6 +290,11 @@ def __init__(
         # while inline will deadlock because the parked read loop cannot dequeue
         # the response.
         self._inline_methods = inline_methods
+        # Observer for Exception items the transport yields on the read stream
+        # (SSE/streamable-HTTP connection faults, stdio parse errors). Without
+        # it they are debug-logged and dropped. Awaited in the read loop and
+        # contained: a raising observer costs the item, not the connection.
+        self._on_stream_exception = on_stream_exception
 
         self._next_id = 0
         self._pending: dict[RequestId, _Pending] = {}
@@ -482,13 +490,19 @@ async def _dispatch(
     ) -> None:
         """Route one inbound item.
 
-        Everything here is `send_nowait` or `_spawn`; the only `await` is for
-        `inline_methods` requests, which deliberately block dequeuing until
-        handled. Any other `await` would let one slow message head-of-line
-        block the entire read loop.
+        Everything here is `send_nowait` or `_spawn`; the only `await`s are
+        `inline_methods` requests and the `on_stream_exception` observer,
+        which deliberately block dequeuing until handled. Any other `await`
+        would let one slow message head-of-line block the entire read loop.
         """
         if isinstance(item, Exception):
-            logger.debug("transport yielded exception: %r", item)
+            if self._on_stream_exception is None:
+                logger.debug("transport yielded exception: %r", item)
+                return
+            try:
+                await self._on_stream_exception(item)
+            except Exception:
+                logger.exception("on_stream_exception observer raised")
             return
         metadata = item.metadata
         msg = item.message

tests/shared/test_jsonrpc_dispatcher.py

@@ -912,6 +912,63 @@ async def test_transport_exception_in_read_stream_is_logged_and_dropped():
             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."""
+    c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
+    s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
+
+    seen: list[Exception] = []
+
+    async def observe(exc: Exception) -> None:
+        seen.append(exc)
+
+    server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send, on_stream_exception=observe)
+    on_request, on_notify = echo_handlers(Recorder())
+    hiccup = ValueError("transport hiccup")
+    try:
+        async with anyio.create_task_group() as tg:
+            await tg.start(server.run, on_request, on_notify)
+            await c2s_send.send(hiccup)
+            await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="t", params=None)))
+            with anyio.fail_after(5):
+                resp = await s2c_recv.receive()
+            assert isinstance(resp, SessionMessage)
+            assert isinstance(resp.message, JSONRPCResponse)
+            tg.cancel_scope.cancel()
+    finally:
+        for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
+            s.close()
+    assert seen == [hiccup]
+
+
+@pytest.mark.anyio
+async def test_on_stream_exception_observer_raising_is_contained(caplog: pytest.LogCaptureFixture):
+    """A raising observer costs the item, not the connection: it runs in the read loop itself."""
+    c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
+    s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
+
+    async def observe(exc: Exception) -> None:
+        raise RuntimeError("observer boom")
+
+    server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send, on_stream_exception=observe)
+    on_request, on_notify = echo_handlers(Recorder())
+    try:
+        async with anyio.create_task_group() as tg:
+            await tg.start(server.run, on_request, on_notify)
+            await c2s_send.send(ValueError("transport hiccup"))
+            await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="t", params=None)))
+            with anyio.fail_after(5):
+                resp = await s2c_recv.receive()
+            assert isinstance(resp, SessionMessage)
+            assert isinstance(resp.message, JSONRPCResponse)
+            tg.cancel_scope.cancel()
+    finally:
+        for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
+            s.close()
+    assert "on_stream_exception observer raised" in caplog.text
+
+
 @pytest.mark.anyio
 async def test_progress_notification_for_unknown_token_falls_through_to_on_notify():
     async with running_pair(jsonrpc_pair) as (client, _server, _crec, srec):