diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 849ce023109..a084bf51ded 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -135,6 +135,26 @@ class MCPClient: ) return transport_ctx, http_client + @staticmethod + def _extract_root_cause(exc: BaseException) -> Optional[BaseException]: + """Extract the first meaningful root-cause from an ExceptionGroup. + + When the MCP SDK's streamable-HTTP transport encounters an HTTP error + (e.g. 401 Unauthorized), the task-group cancels the session, so the + caller only sees ``CancelledError``. The *real* error surfaces later + inside the ``ExceptionGroup`` raised during transport cleanup. This + helper digs it out so we can re-raise something actionable. + """ + if isinstance(exc, BaseExceptionGroup): + for inner in exc.exceptions: + cause = MCPClient._extract_root_cause(inner) + if cause is not None: + return cause + return exc.exceptions[0] if exc.exceptions else None + if isinstance(exc, (asyncio.CancelledError, GeneratorExit)): + return None + return exc + async def _execute_session_operation( self, transport_ctx: Any, @@ -146,6 +166,7 @@ class MCPClient: Handles entering/exiting contexts and running the operation. """ transport = await transport_ctx.__aenter__() + original_error: Optional[BaseException] = None try: read_stream, write_stream = transport[0], transport[1] session_ctx = ClientSession(read_stream, write_stream) @@ -153,16 +174,28 @@ class MCPClient: try: await session.initialize() return await operation(session) + except asyncio.CancelledError: + original_error = None + raise finally: try: await session_ctx.__aexit__(None, None, None) except BaseException as e: verbose_logger.debug(f"Error during session context exit: {e}") + except asyncio.CancelledError: + raise finally: try: await transport_ctx.__aexit__(None, None, None) except BaseException as e: verbose_logger.debug(f"Error during transport context exit: {e}") + root = self._extract_root_cause(e) + if root is not None: + original_error = root + if original_error is not None: + raise ConnectionError( + f"MCP connection failed: {original_error}" + ) from original_error async def run_with_session( self, operation: Callable[[ClientSession], Awaitable[TSessionResult]] @@ -172,6 +205,18 @@ class MCPClient: try: transport_ctx, http_client = self._create_transport_context() return await self._execute_session_operation(transport_ctx, operation) + except asyncio.CancelledError as e: + verbose_logger.warning( + "MCP client run_with_session cancelled for %s: %s", + self.server_url or "stdio", + e, + ) + raise ConnectionError( + f"MCP session was cancelled while connecting to " + f"{self.server_url or 'stdio'}. The remote server may have " + f"rejected the request (e.g. authentication failure) or is " + f"unreachable." + ) from e except Exception: verbose_logger.warning( "MCP client run_with_session failed for %s", self.server_url or "stdio" diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 6082e9bd606..e5ae7d9387f 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -564,6 +564,41 @@ if MCP_AVAILABLE: NewMCPServerRequest, ) + def _build_mcp_error_message(exc: BaseException) -> str: + """Build a user-facing error message from an MCP connection failure. + + Walks the exception chain (__cause__, __context__) looking for + actionable details such as HTTP status codes or connection errors. + """ + import httpx + + parts: List[str] = [] + seen: set = set() + current: Optional[BaseException] = exc + while current is not None and id(current) not in seen: + seen.add(id(current)) + if isinstance(current, httpx.HTTPStatusError): + status = current.response.status_code + reason = current.response.reason_phrase or "" + parts.append( + f"Server returned HTTP {status} {reason}".strip() + ) + break + if isinstance(current, (httpx.ConnectError, httpx.TimeoutException)): + parts.append(str(current)) + break + if isinstance(current, ConnectionError) and not isinstance( + current, ConnectionResetError + ): + parts.append(str(current)) + current = current.__cause__ or current.__context__ + + if parts: + return "Failed to connect to MCP server: " + "; ".join(parts) + return ( + "Failed to connect to MCP server. Check proxy logs for details." + ) + def _extract_credentials( request: NewMCPServerRequest, ) -> tuple: @@ -655,10 +690,11 @@ if MCP_AVAILABLE: raise except BaseException as e: verbose_logger.error("Error in MCP operation: %s", e, exc_info=True) + user_message = _build_mcp_error_message(e) return { "status": "error", "error": True, - "message": "Failed to connect to MCP server. Check proxy logs for details.", + "message": user_message, } async def _preview_openapi_tools(spec_path: str) -> dict: @@ -715,6 +751,26 @@ if MCP_AVAILABLE: """ Test if we can connect to the provided MCP server before adding it """ + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + mcp_auth_header: Optional[str] = None + if new_mcp_server_request.auth_type in { + MCPAuth.api_key, + MCPAuth.bearer_token, + MCPAuth.basic, + MCPAuth.authorization, + }: + credentials = getattr(new_mcp_server_request, "credentials", None) + if isinstance(credentials, dict): + mcp_auth_header = credentials.get("auth_value") + + oauth2_headers: Optional[Dict[str, str]] = None + if new_mcp_server_request.auth_type == MCPAuth.oauth2: + oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers( + request.headers + ) async def _test_connection_operation(client): async def _noop(session): @@ -726,6 +782,8 @@ if MCP_AVAILABLE: return await _execute_with_mcp_client( new_mcp_server_request, _test_connection_operation, + mcp_auth_header=mcp_auth_header, + oauth2_headers=oauth2_headers, raw_headers=_safe_get_request_headers(request), ) diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index febc7c454bd..679227cf6d7 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -246,5 +246,153 @@ class TestMCPClient: await test_client.aclose() +class TestMCPClientErrorHandling: + """Test improved error handling for MCP connection failures.""" + + def test_extract_root_cause_http_status_error(self): + """Should extract HTTPStatusError from an ExceptionGroup.""" + import httpx + + http_err = httpx.HTTPStatusError( + "Client error '401 Unauthorized'", + request=httpx.Request("POST", "http://example.com/mcp"), + response=httpx.Response(401), + ) + group = ExceptionGroup("test", [http_err]) + result = MCPClient._extract_root_cause(group) + assert result is http_err + + def test_extract_root_cause_nested_group(self): + """Should extract root cause from nested ExceptionGroups.""" + import httpx + + http_err = httpx.HTTPStatusError( + "Server error '500'", + request=httpx.Request("POST", "http://example.com/mcp"), + response=httpx.Response(500), + ) + inner = ExceptionGroup("inner", [http_err]) + outer = ExceptionGroup("outer", [inner]) + result = MCPClient._extract_root_cause(outer) + assert result is http_err + + def test_extract_root_cause_skips_cancelled(self): + """Should skip CancelledError and find the real cause.""" + import asyncio + + import httpx + + cancelled = asyncio.CancelledError("cancel scope") + http_err = httpx.HTTPStatusError( + "Client error '401 Unauthorized'", + request=httpx.Request("POST", "http://example.com/mcp"), + response=httpx.Response(401), + ) + group = BaseExceptionGroup("test", [cancelled, http_err]) + result = MCPClient._extract_root_cause(group) + assert result is http_err + + def test_extract_root_cause_connection_error(self): + """Should extract connection errors from ExceptionGroup.""" + conn_err = ConnectionRefusedError("Connection refused") + group = ExceptionGroup("test", [conn_err]) + result = MCPClient._extract_root_cause(group) + assert result is conn_err + + def test_extract_root_cause_returns_none_for_only_cancelled(self): + """Should return None when group only contains CancelledError.""" + import asyncio + + cancelled = asyncio.CancelledError("cancel scope") + result = MCPClient._extract_root_cause(cancelled) + assert result is None + + def test_extract_root_cause_plain_exception(self): + """Should return a non-cancelled plain exception directly.""" + err = RuntimeError("something broke") + result = MCPClient._extract_root_cause(err) + assert result is err + + @pytest.mark.asyncio + @patch.object(mcp_client_module, "streamable_http_client") + @patch("litellm.experimental_mcp_client.client.ClientSession") + async def test_cancelled_error_converted_to_connection_error( + self, mock_session, mock_streamable_http_client + ): + """When session.initialize() raises CancelledError and transport + cleanup yields an ExceptionGroup with the real cause, the client + should raise ConnectionError with the original HTTP error info.""" + import asyncio + + import httpx + + http_err = httpx.HTTPStatusError( + "Client error '401 Unauthorized'", + request=httpx.Request("POST", "http://example.com/mcp"), + response=httpx.Response(401), + ) + + mock_transport = (MagicMock(), MagicMock()) + mock_http_ctx = AsyncMock() + mock_http_ctx.__aenter__.return_value = mock_transport + mock_http_ctx.__aexit__.side_effect = ExceptionGroup("tasks", [http_err]) + mock_streamable_http_client.return_value = mock_http_ctx + + mock_session_instance = AsyncMock() + mock_session_instance.initialize = AsyncMock( + side_effect=asyncio.CancelledError("Cancelled via cancel scope") + ) + mock_session_ctx = AsyncMock() + mock_session_ctx.__aenter__.return_value = mock_session_instance + mock_session_ctx.__aexit__.return_value = None + mock_session.return_value = mock_session_ctx + + client = MCPClient( + server_url="http://example.com/mcp", + transport_type=MCPTransport.http, + ) + + with pytest.raises(ConnectionError, match="401 Unauthorized"): + async def _op(session): + return await session.list_tools() + await client.run_with_session(_op) + + @pytest.mark.asyncio + @patch.object(mcp_client_module, "streamable_http_client") + @patch("litellm.experimental_mcp_client.client.ClientSession") + async def test_cancelled_error_fallback_message( + self, mock_session, mock_streamable_http_client + ): + """When CancelledError occurs but no root cause can be extracted, + run_with_session should convert it to a ConnectionError with a + descriptive fallback message.""" + import asyncio + + mock_transport = (MagicMock(), MagicMock()) + mock_http_ctx = AsyncMock() + mock_http_ctx.__aenter__.return_value = mock_transport + mock_http_ctx.__aexit__.return_value = None + mock_streamable_http_client.return_value = mock_http_ctx + + mock_session_instance = AsyncMock() + mock_session_instance.initialize = AsyncMock( + side_effect=asyncio.CancelledError("Cancelled via cancel scope") + ) + mock_session_ctx = AsyncMock() + mock_session_ctx.__aenter__.return_value = mock_session_instance + mock_session_ctx.__aexit__.return_value = None + mock_session.return_value = mock_session_ctx + + client = MCPClient( + server_url="http://example.com/mcp", + transport_type=MCPTransport.http, + ) + + with pytest.raises(ConnectionError, match="rejected the request"): + async def _op(session): + return await session.list_tools() + await client.run_with_session(_op) + + if __name__ == "__main__": pytest.main([__file__])