fix: Narrow retry exceptions and unwrap StreamableHTTPError at client boundary

Author: wukath

src/mcp/client/streamable_http.py

@@ -62,6 +62,21 @@ class ResumptionError(StreamableHTTPError):
     """Raised when resumption request is invalid."""
 
 
+def _unwrap_exception(exc: BaseException) -> BaseException:
+    """Recursively find and return the first StreamableHTTPError in an exception's tree/group."""
+    if isinstance(exc, StreamableHTTPError):
+        return exc
+
+    exceptions = getattr(exc, "exceptions", None)
+    if exceptions is not None:
+        for sub_exc in exceptions:
+            unwrapped = _unwrap_exception(sub_exc)
+            if isinstance(unwrapped, StreamableHTTPError):
+                return unwrapped
+    return exc
+
+
+
 @dataclass
 class RequestContext:
     """Context for a request operation."""
@@ -228,11 +243,12 @@ async def handle_get_stream(self, client: httpx2.AsyncClient, read_stream_writer
                     # Stream ended normally (server closed) - reset attempt counter
                     attempt = 0
 
-            except Exception as exc:
+            except httpx2.HTTPError as exc:
                 logger.debug("GET stream error", exc_info=True)
                 attempt += 1
                 last_exc = exc
 
+
             if attempt >= MAX_RECONNECTION_ATTEMPTS:
                 logger.debug(f"GET stream max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded")
                 raise StreamableHTTPError(
@@ -448,9 +464,10 @@ async def _handle_sse_response(
                 if is_complete:
                     await response.aclose()
                     return  # Normal completion, no reconnect needed
-        except Exception:
+        except httpx2.HTTPError:
             logger.debug("SSE stream ended", exc_info=True)  # pragma: lax no cover
 
+
         # Stream ended without response - reconnect if we received an event with ID
         if last_event_id is not None:
             logger.info("SSE stream disconnected, reconnecting...")
@@ -532,11 +549,12 @@ async def _handle_reconnection(
                 # Stream ended again without response - reconnect again (reset attempt counter)
                 logger.info("SSE stream disconnected, reconnecting...")
                 await self._handle_reconnection(ctx, reconnect_last_event_id, reconnect_retry_ms, 0)
-        except Exception as e:  # pragma: no cover
+        except httpx2.HTTPError as e:  # pragma: no cover
             logger.debug(f"Reconnection failed: {e}")
             # Try to reconnect again if we still have an event ID
             await self._handle_reconnection(ctx, last_event_id, retry_interval_ms, attempt + 1)
 
+
     async def post_writer(
         self,
         client: httpx2.AsyncClient,
@@ -684,39 +702,46 @@ async def streamable_http_client(
 
     logger.debug(f"Connecting to StreamableHTTP endpoint: {url}")
 
-    async with contextlib.AsyncExitStack() as stack:
-        # Only manage client lifecycle if we created it
-        if not client_provided:
-            await stack.enter_async_context(client)
-
-        read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
-        write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
+    try:
+        async with contextlib.AsyncExitStack() as stack:
+            # Only manage client lifecycle if we created it
+            if not client_provided:
+                await stack.enter_async_context(client)
 
-        async with (
-            read_stream_writer,
-            read_stream,
-            write_stream,
-            write_stream_reader,
-            anyio.create_task_group() as tg,
-        ):
+            read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
+            write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
 
-            def start_get_stream() -> None:
-                tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
-
-            tg.start_soon(
-                transport.post_writer,
-                client,
-                write_stream_reader,
+            async with (
                 read_stream_writer,
+                read_stream,
                 write_stream,
-                start_get_stream,
-                tg,
-            )
+                write_stream_reader,
+                anyio.create_task_group() as tg,
+            ):
+
+                def start_get_stream() -> None:
+                    tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
+
+                tg.start_soon(
+                    transport.post_writer,
+                    client,
+                    write_stream_reader,
+                    read_stream_writer,
+                    write_stream,
+                    start_get_stream,
+                    tg,
+                )
+
+                try:
+                    yield read_stream, write_stream
+                finally:
+                    if transport.session_id and terminate_on_close:
+                        await transport.terminate_session(client)
+                    tg.cancel_scope.cancel()
+            await resync_tracer()
+    except BaseException as exc:
+        unwrapped = _unwrap_exception(exc)
+        if isinstance(unwrapped, StreamableHTTPError):
+            raise unwrapped from exc
+        raise
 
-            try:
-                yield read_stream, write_stream
-            finally:
-                if transport.session_id and terminate_on_close:
-                    await transport.terminate_session(client)
-                tg.cancel_scope.cancel()
-        await resync_tracer()

tests/shared/test_streamable_http.py

@@ -2299,8 +2299,8 @@ async def test_reconnect_failure_propagates_error() -> None:
     # Create a context-aware stream writer (matches StreamWriter type alias)
     write_stream, read_stream = create_context_streams[SessionMessage | Exception](1)
 
-    # Mock client.sse to raise an exception
-    client.sse.side_effect = Exception("Connection refused")
+    # Mock client.sse to raise a connection error
+    client.sse.side_effect = httpx2.ConnectError("Connection refused")
 
     # Patch anyio.sleep to avoid waiting
     with patch("mcp.client.streamable_http.anyio.sleep", new_callable=AsyncMock) as mock_sleep:
@@ -2316,3 +2316,87 @@ async def test_reconnect_failure_propagates_error() -> None:
 
     await write_stream.aclose()
     await read_stream.aclose()
+
+
+@pytest.mark.anyio
+async def test_streamable_http_client_reconnect_failure_propagates_error() -> None:
+    """streamable_http_client context manager should propagate StreamableHTTPError
+    when the GET stream connection fails completely (max attempts exceeded).
+    """
+    client = AsyncMock(spec=httpx2.AsyncClient)
+
+    # Mock post_writer requests:
+    # 1. initialize request -> returns response with session ID
+    # 2. notifications/initialized -> returns 202 Accepted
+    mock_response = AsyncMock(spec=httpx2.Response)
+    mock_response.status_code = 200
+    mock_response.headers = {
+        "content-type": "application/json",
+        "mcp-session-id": "test-session",
+    }
+    mock_response.aread.return_value = json.dumps({
+        "jsonrpc": "2.0",
+        "id": 1,
+        "result": {
+            "protocolVersion": "2025-06-18",
+            "capabilities": {},
+            "serverInfo": {"name": "test-server", "version": "1.0"},
+        }
+    }).encode("utf-8")
+
+    mock_initialized_response = AsyncMock(spec=httpx2.Response)
+    mock_initialized_response.status_code = 202
+
+    # Use asynccontextmanager to mock client.stream
+    responses = [mock_response, mock_initialized_response]
+
+    @asynccontextmanager
+    async def mock_stream(*args: Any, **kwargs: Any):
+        yield responses.pop(0)
+
+    client.stream = mock_stream
+
+
+    # Mock client.sse to raise httpx2.HTTPError
+    client.sse.side_effect = httpx2.HTTPError("SSE connection refused")
+
+    # Patch anyio.sleep to avoid waiting during reconnect attempts
+    with patch("mcp.client.streamable_http.anyio.sleep", new_callable=AsyncMock) as mock_sleep:
+        with pytest.raises(StreamableHTTPError) as exc_info:
+            with anyio.fail_after(5):
+                async with streamable_http_client("http://localhost:8000/mcp", http_client=client) as (read_stream, write_stream):
+                    # Send initialize message
+                    await write_stream.send(SessionMessage(
+                        types.JSONRPCRequest(
+                            jsonrpc="2.0",
+                            id=1,
+                            method="initialize",
+                            params={
+                                "protocolVersion": "2025-06-18",
+                                "capabilities": {},
+                                "clientInfo": {"name": "test-client", "version": "1.0"},
+                            },
+                        )
+                    ))
+
+                    # Receive the response
+                    await read_stream.receive()
+
+                    # Send notifications/initialized (which will trigger start_get_stream)
+                    await write_stream.send(SessionMessage(
+                        types.JSONRPCNotification(
+                            jsonrpc="2.0",
+                            method="notifications/initialized",
+                        )
+                    ))
+
+                    # Wait for the task group to fail
+                    event = anyio.Event()
+                    await event.wait()
+
+
+        assert "Failed to connect to GET stream" in str(exc_info.value)
+        assert client.sse.call_count == 2
+        mock_sleep.assert_called_once_with(1.0)
+
+