fix(direct-dispatcher): contain callback failures

Author: luojiyin1987

src/mcp/shared/direct_dispatcher.py

@@ -30,6 +30,7 @@
 from mcp.shared._compat import resync_tracer
 from mcp.shared.dispatcher import (
     CallOptions,
+    DispatchContext,
     OnNotify,
     OnNotifyIntercept,
     OnRequest,
@@ -43,6 +44,30 @@
 
 logger = logging.getLogger(__name__)
 
+
+def _shielded_progress(fn: ProgressFnT) -> ProgressFnT:
+    """Wrap a progress callback so its failure does not fail the request."""
+
+    async def _wrapped(progress: float, total: float | None, message: str | None) -> None:
+        try:
+            await fn(progress, total, message)
+        except Exception:
+            logger.exception("progress callback raised")
+
+    return _wrapped
+
+
+def _contained_notify(fn: OnNotify) -> OnNotify:
+    """Wrap a notification handler so its failure does not reach the sender."""
+
+    async def _wrapped(dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None) -> None:
+        try:
+            await fn(dctx, method, params)
+        except Exception:
+            logger.exception("notification handler for %r raised", method)
+
+    return _wrapped
+
 __all__ = ["DirectDispatcher", "create_direct_dispatcher_pair"]
 
 DIRECT_TRANSPORT_KIND = "direct"
@@ -206,7 +231,7 @@ def _make_context(
             _back_request=lambda m, p, o: peer._dispatch_request(m, p, o),
             _back_notify=lambda m, p: peer._dispatch_notify(m, p),
             request_id=request_id,
-            _on_progress=on_progress,
+            _on_progress=_shielded_progress(on_progress) if on_progress is not None else None,
         )
 
     async def _wait_ready(self) -> None:
@@ -301,7 +326,7 @@ async def _dispatch_notify(self, method: str, params: Mapping[str, Any] | None)
             return
         assert self._on_notify is not None
         dctx = self._make_context()
-        await self._on_notify(dctx, method, params)
+        await _contained_notify(self._on_notify)(dctx, method, params)
 
 
 def create_direct_dispatcher_pair(

tests/shared/test_dispatcher.py

@@ -216,6 +216,35 @@ async def on_progress(progress: float, total: float | None, message: str | None)
     assert received == [(0.5, 1.0, "halfway")]
 
 
+@pytest.mark.anyio
+async def test_progress_callback_exception_does_not_fail_request(pair_factory: PairFactory):
+    async def server_on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        await ctx.progress(0.5)
+        return {"ok": True}
+
+    async def on_progress(progress: float, total: float | None, message: str | None) -> None:
+        raise RuntimeError("consumer failed")
+
+    async with running_pair(pair_factory, server_on_request=server_on_request) as (client, *_):
+        with anyio.fail_after(5):
+            result = await client.send_raw_request("tools/call", None, {"on_progress": on_progress})
+    assert result == {"ok": True}
+
+
+@pytest.mark.anyio
+async def test_notification_handler_exception_does_not_reach_sender(pair_factory: PairFactory):
+    async def on_notify(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> None:
+        raise RuntimeError("handler failed")
+
+    async with running_pair(pair_factory, client_on_notify=on_notify) as (client, *_):
+        with anyio.fail_after(5):
+            await client.notify("notifications/message", None)
+
+
 @pytest.mark.anyio
 async def test_ctx_progress_is_noop_when_caller_supplied_no_callback(pair_factory: PairFactory):
     async def server_on_request(