diff --git a/src/mcp/shared/direct_dispatcher.py b/src/mcp/shared/direct_dispatcher.py index e17283afa2..c4b7b114df 100644 --- a/src/mcp/shared/direct_dispatcher.py +++ b/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: @@ -301,7 +304,10 @@ 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) + try: + await self._on_notify(dctx, method, params) + except Exception: + logger.exception("notification handler for %r raised", method) def create_direct_dispatcher_pair( diff --git a/tests/shared/test_dispatcher.py b/tests/shared/test_dispatcher.py index c6ebb401ff..afa3844670 100644 --- a/tests/shared/test_dispatcher.py +++ b/tests/shared/test_dispatcher.py @@ -216,6 +216,43 @@ 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, caplog: pytest.LogCaptureFixture +) -> None: + async def on_progress(progress: float, total: float | None, message: str | None) -> None: + raise RuntimeError("progress callback failed") + + 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 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_notification_handler_exception_does_not_reach_sender( + pair_factory: PairFactory, caplog: pytest.LogCaptureFixture +) -> None: + called = anyio.Event() + + async def on_notify(ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None) -> None: + called.set() + raise RuntimeError("notification handler failed") + + async with running_pair(pair_factory, server_on_notify=on_notify) as (client, *_): + with anyio.fail_after(5): + await client.notify("notifications/message", None) + await called.wait() + assert "notification handler for 'notifications/message' raised" in caplog.text + + @pytest.mark.anyio async def test_ctx_progress_is_noop_when_caller_supplied_no_callback(pair_factory: PairFactory): async def server_on_request(