diff --git a/src/mcp/client/client.py b/src/mcp/client/client.py index 2f467c2614..dad660790d 100644 --- a/src/mcp/client/client.py +++ b/src/mcp/client/client.py @@ -57,6 +57,7 @@ LoggingFnT, MessageHandlerFnT, SamplingFnT, + TransportExceptionHandlerFnT, ) from mcp.client.stdio import StdioServerParameters, stdio_client from mcp.client.streamable_http import streamable_http_client @@ -329,6 +330,14 @@ async def main(): message_handler: MessageHandlerFnT | None = None """Callback for handling raw messages.""" + transport_exception_handler: TransportExceptionHandlerFnT | None = None + """Callback for handling transport-level exceptions (timeouts, connection errors, etc.). + + When provided, this handler receives transport exceptions directly, allowing the caller + to propagate, log, or handle them as needed. If not provided, exceptions are delivered + to `message_handler` for backwards compatibility. + """ + client_info: Implementation | None = None """Client implementation info to send to server.""" @@ -442,6 +451,7 @@ async def _build_session(self, exit_stack: AsyncExitStack) -> ClientSession: logging_callback=self.logging_callback, log_level=self.log_level, message_handler=message_handler, + transport_exception_handler=self.transport_exception_handler, client_info=self.client_info, elicitation_callback=self.elicitation_callback, extensions=self._folded_extensions.ad, diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index f18cc0ef10..1cb0df4e6a 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -249,7 +249,13 @@ class MessageHandlerFnT(Protocol): async def __call__(self, message: IncomingMessage) -> None: ... # pragma: no branch +class TransportExceptionHandlerFnT(Protocol): + async def __call__(self, exc: Exception) -> None: ... # pragma: no branch + + async def _default_message_handler(message: IncomingMessage) -> None: + if isinstance(message, Exception): + logger.exception("Transport exception received: %s", message) await anyio.lowlevel.checkpoint() @@ -418,6 +424,7 @@ def __init__( list_roots_callback: ListRootsFnT | None = None, logging_callback: LoggingFnT | None = None, message_handler: MessageHandlerFnT | None = None, + transport_exception_handler: TransportExceptionHandlerFnT | None = None, client_info: types.Implementation | None = None, *, log_level: types.LoggingLevel | None = None, @@ -444,6 +451,7 @@ def __init__( self._logging_callback = logging_callback or _default_logging_callback self._log_level: types.LoggingLevel | None = log_level self._message_handler = message_handler or _default_message_handler + self._transport_exception_handler = transport_exception_handler self._tool_output_schemas: dict[str, dict[str, Any] | None] = {} # Compiled output-schema validators, derived from `_tool_output_schemas` and owned by # `_absorb_tool_listing`, which evicts a tool's entry whenever its schema changes. @@ -1501,7 +1509,10 @@ async def _on_stream_exception(self, exc: Exception) -> None: self._task_group.start_soon(self._deliver_stream_exception, exc) async def _deliver_stream_exception(self, exc: Exception) -> None: + # If a dedicated transport exception handler is provided, use it. + # Otherwise fall back to message_handler for backwards compatibility. + handler = self._transport_exception_handler or self._message_handler try: - await self._message_handler(exc) + await handler(exc) except Exception: - logger.exception("message_handler raised on transport exception") + logger.exception("transport exception handler raised") diff --git a/tests/client/test_session.py b/tests/client/test_session.py index 6663fb47a2..d01feb0eef 100644 --- a/tests/client/test_session.py +++ b/tests/client/test_session.py @@ -1026,7 +1026,7 @@ async def handler(msg: object) -> None: assert seen == [exc] assert isinstance(out.message, JSONRPCResponse) assert out.message.id == 9 - assert "message_handler raised on transport exception" in caplog.text + assert "transport exception handler raised" in caplog.text @pytest.mark.anyio @@ -1052,6 +1052,66 @@ async def handler(msg: object) -> None: await ponged.wait() +@pytest.mark.anyio +async def test_transport_exception_handler_receives_exceptions_separately_from_message_handler( + caplog: pytest.LogCaptureFixture, +): + """A dedicated `transport_exception_handler` receives transport exceptions, while + `message_handler` only receives server notifications (SDK-defined).""" + seen_messages: list[object] = [] + seen_exceptions: list[Exception] = [] + msg_delivered = anyio.Event() + exc_delivered = anyio.Event() + + async def message_handler(msg: object) -> None: + seen_messages.append(msg) + msg_delivered.set() + + async def transport_exception_handler(exc: Exception) -> None: + seen_exceptions.append(exc) + exc_delivered.set() + + async with raw_client_session( + message_handler=message_handler, + transport_exception_handler=transport_exception_handler, + ) as (_session, to_client, _from_client): + # Send a transport exception + exc = ValueError("transport timeout") + await to_client.send(exc) + await exc_delivered.wait() + + # Transport exception should only go to transport_exception_handler + assert seen_exceptions == [exc] + assert seen_messages == [] + + # Send a server notification + await to_client.send( + SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/tools/list_changed")) + ) + await msg_delivered.wait() + + # Server notification should only go to message_handler + assert len(seen_messages) == 1 + assert isinstance(seen_messages[0], type(types.ToolListChangedNotification())) + + +@pytest.mark.anyio +async def test_transport_exception_handler_fallback_to_message_handler(caplog: pytest.LogCaptureFixture): + """When no `transport_exception_handler` is provided, transport exceptions fall back + to `message_handler` for backwards compatibility (SDK-defined). The default + `message_handler` logs the exception.""" + async with raw_client_session() as (_session, to_client, _from_client): + exc = ValueError("bad bytes") + await to_client.send(exc) + # Give the handler a moment to run + await anyio.sleep(0.01) + + assert "Transport exception received" in caplog.text + + # The default message_handler logs the exception + assert "Transport exception received" in caplog.text + + @pytest.mark.anyio async def test_receive_loop_consumes_server_cancelled_without_reaching_message_handler(): """A server-sent notifications/cancelled is swallowed, matching the pre-swap contract.