Fix ClientSession error handling: add transport_exception_handler callback (Fixes #1401)

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,

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")

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.