diff --git a/src/mcp/client/streamable_http.py b/src/mcp/client/streamable_http.py index c95cfcf50..9ce42e2a6 100644 --- a/src/mcp/client/streamable_http.py +++ b/src/mcp/client/streamable_http.py @@ -1,6 +1,6 @@ """Implements StreamableHTTP transport for MCP clients.""" -from __future__ import annotations as _annotations +from __future__ import annotations import contextlib import logging @@ -62,6 +62,20 @@ 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.""" @@ -200,6 +214,7 @@ async def handle_get_stream(self, client: httpx2.AsyncClient, read_stream_writer last_event_id: str | None = None retry_interval_ms: int | None = None attempt: int = 0 + last_exc: Exception | None = None while attempt < MAX_RECONNECTION_ATTEMPTS: # pragma: no branch try: @@ -227,13 +242,16 @@ async def handle_get_stream(self, client: httpx2.AsyncClient, read_stream_writer # Stream ended normally (server closed) - reset attempt counter attempt = 0 - except Exception: + except httpx2.HTTPError as exc: logger.debug("GET stream error", exc_info=True) attempt += 1 + last_exc = exc - if attempt >= MAX_RECONNECTION_ATTEMPTS: # pragma: no cover + if attempt >= MAX_RECONNECTION_ATTEMPTS: logger.debug(f"GET stream max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded") - return + raise StreamableHTTPError( + f"Failed to connect to GET stream after {MAX_RECONNECTION_ATTEMPTS} attempts" + ) from last_exc # Wait before reconnecting delay_ms = retry_interval_ms if retry_interval_ms is not None else DEFAULT_RECONNECTION_DELAY_MS @@ -444,7 +462,7 @@ 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 @@ -528,7 +546,7 @@ 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) @@ -680,39 +698,45 @@ 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) - - async with ( - read_stream_writer, - read_stream, - write_stream, - write_stream_reader, - anyio.create_task_group() as tg, - ): + 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) - def start_get_stream() -> None: - tg.start_soon(transport.handle_get_stream, client, read_stream_writer) + read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0) + write_stream, write_stream_reader = create_context_streams[SessionMessage](0) - 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() + 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 diff --git a/tests/shared/test_streamable_http.py b/tests/shared/test_streamable_http.py index d7eeccdfd..f907f25e2 100644 --- a/tests/shared/test_streamable_http.py +++ b/tests/shared/test_streamable_http.py @@ -12,7 +12,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass, field from typing import Any -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock, patch from urllib.parse import urlparse import anyio @@ -45,7 +45,11 @@ from mcp import MCPError from mcp.client import ClientRequestContext from mcp.client.session import ClientSession -from mcp.client.streamable_http import StreamableHTTPTransport, streamable_http_client +from mcp.client.streamable_http import ( + StreamableHTTPError, + StreamableHTTPTransport, + streamable_http_client, +) from mcp.server import Server, ServerRequestContext from mcp.server.streamable_http import ( GET_STREAM_KEY, @@ -2283,3 +2287,128 @@ async def asgi_receive() -> Message: assert body_chunks[-1] == {"type": "http.response.body", "body": b"", "more_body": False} assert "Error in standalone SSE writer" not in caplog.text assert "Error in standalone SSE response" not in caplog.text + + +@pytest.mark.anyio +async def test_reconnect_failure_propagates_error() -> None: + """Client should raise StreamableHTTPError when reconnection fails completely.""" + transport = StreamableHTTPTransport(url="http://localhost:8000/mcp") + transport.session_id = "test-session" + client = AsyncMock(spec=httpx2.AsyncClient) + + # 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 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: + with pytest.raises(StreamableHTTPError) as exc_info: + await transport.handle_get_stream(client, write_stream) + assert "Failed to connect to GET stream" in str(exc_info.value) + # Should have attempted MAX_RECONNECTION_ATTEMPTS times (default is 2) + assert client.sse.call_count == 2 + # Should have slept between attempts (attempts - 1 times) + assert mock_sleep.call_count == 1 + # Verify it slept with the default delay (1000ms / 1000.0 = 1.0s) + mock_sleep.assert_called_once_with(1.0) + + 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 client.delete for session termination + mock_delete_response = AsyncMock(spec=httpx2.Response) + mock_delete_response.status_code = 204 + client.delete.return_value = mock_delete_response + + # 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 ( # pragma: no branch + 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)