Contain a raising progress callback on the in-process dispatcher (#3623)

Author: maxisbey

src/mcp/shared/direct_dispatcher.py

@@ -90,7 +90,10 @@ async def send_raw_request(
 
     async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None:
         if self._on_progress is not None:
-            await self._on_progress(progress, total, message)
+            try:
+                await self._on_progress(progress, total, message)
+            except Exception:
+                logger.exception("progress callback raised")
 
 
 class DirectDispatcher:

tests/shared/test_dispatcher.py

@@ -230,6 +230,28 @@ async def server_on_request(
     assert result == {"ok": True}
 
 
+@pytest.mark.anyio
+async def test_raising_on_progress_callback_is_logged_and_request_still_succeeds(
+    pair_factory: PairFactory, caplog: pytest.LogCaptureFixture
+):
+    """A caller's `on_progress` callback that raises is logged and does not fail the request (SDK-defined)."""
+
+    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("progress callback boom")
+
+    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}
+    assert "progress callback raised" in caplog.text
+
+
 @pytest.mark.anyio
 async def test_ctx_message_metadata_is_none_when_transport_attaches_nothing(pair_factory: PairFactory):
     """Plain requests carry no transport metadata, so handlers see `None`."""

tests/shared/test_jsonrpc_dispatcher.py

@@ -1449,25 +1449,6 @@ async def on_progress(progress: float, total: float | None, message: str | None)
         assert crec.notifications == []
 
 
-@pytest.mark.anyio
-async def test_progress_callback_exception_is_swallowed_and_logged(caplog: pytest.LogCaptureFixture):
-    """A user progress callback raising must not crash the dispatcher."""
-
-    async def boom(progress: float, total: float | None, message: str | None) -> None:
-        raise RuntimeError("progress callback boom")
-
-    async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
-        await ctx.progress(0.5)
-        return {"ok": True}
-
-    opts: CallOptions = {"on_progress": boom}
-    async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, *_):
-        with anyio.fail_after(5):
-            result = await client.send_raw_request("t", None, opts)
-    assert result == {"ok": True}
-    assert "progress callback raised" in caplog.text
-
-
 @pytest.mark.anyio
 async def test_inline_methods_are_handled_before_next_message_is_dequeued():
     """An `inline_methods` method runs to completion before the next message is dispatched."""