From ac3d072561d995c8c7a3419ddfd7b4e324af742d Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 4 Jun 2026 21:34:01 +0200 Subject: [PATCH] fix(mcp): surface upstream tool-listing failures instead of empty success --- .../_experimental/mcp_server/exceptions.py | 29 ++++ .../mcp_server/mcp_server_manager.py | 147 ++++++++++++++---- .../test_mcp_oauth_passthrough_tools.py | 85 +++++++++- .../mcp_server/test_rest_endpoints.py | 132 ++++++++++++++++ 4 files changed, 364 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/exceptions.py b/litellm/proxy/_experimental/mcp_server/exceptions.py index fd8fc3d5e58..870c1b8b5be 100644 --- a/litellm/proxy/_experimental/mcp_server/exceptions.py +++ b/litellm/proxy/_experimental/mcp_server/exceptions.py @@ -78,3 +78,32 @@ class MCPUpstreamAuthError(Exception): detail=detail, headers={"www-authenticate": challenge} if challenge else None, ) + + +class MCPUpstreamError(Exception): + """Raised when an upstream MCP server fails during tool listing with a + non-auth failure (e.g. HTTP 429/5xx, or an interrupted/cancelled listing). + + Unlike :class:`MCPUpstreamAuthError` (which drives the upstream OAuth flow + for pass-through servers on 401/403), this exists purely so the gateway + surfaces the failure to the caller instead of laundering it into an + empty-but-successful ``tools/list`` response. Multi-server aggregators + absorb it per-server; single-server REST routes report it as a server + error. + """ + + def __init__( + self, + status_code: Optional[int], + server_name: str, + message: Optional[str] = None, + ) -> None: + self.status_code = status_code + self.server_name = server_name + super().__init__( + message + or ( + f"Upstream MCP server {server_name!r} failed during tool listing" + + (f" with HTTP {status_code}" if status_code is not None else "") + ) + ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0d2008cdade..1db1c9d5dbe 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,7 +13,19 @@ import json import os import re import time -from typing import Any, Callable, Dict, List, Literal, Optional, Set, Tuple, Union, cast +from typing import ( + Any, + Callable, + Dict, + Iterator, + List, + Literal, + Optional, + Set, + Tuple, + Union, + cast, +) from urllib.parse import urlparse import anyio @@ -48,7 +60,10 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) -from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError +from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPUpstreamAuthError, + MCPUpstreamError, +) from litellm.proxy._experimental.mcp_server.elicitation_handler import ( MCP_ELICITATION_AVAILABLE, ) @@ -171,19 +186,13 @@ def _should_strip_caller_authorization( ) -def _extract_upstream_auth_failure( - exc: BaseException, -) -> Optional[Tuple[int, Optional[str]]]: - """Walk the exception tree looking for an HTTP 401/403 response from the - upstream MCP server. +def _iter_exception_responses(exc: BaseException) -> Iterator[Any]: + """Yield every ``response`` object found by walking an exception tree. - The MCP SDK wraps transport errors in anyio ``ExceptionGroup`` objects and - may chain through ``__cause__`` / ``__context__``. We inspect all of those - layers for an ``httpx.Response``-bearing exception (typically - ``httpx.HTTPStatusError``) and extract the status code and any upstream - ``WWW-Authenticate`` header. - - Returns ``(status_code, www_authenticate)`` on match, else ``None``. + The MCP SDK wraps transport errors in anyio ``ExceptionGroup`` objects + (PEP 654) and may chain through ``__cause__`` / ``__context__``. This + iterator visits all of those layers and yields each ``response`` attribute + it finds (typically the ``httpx.Response`` on an ``httpx.HTTPStatusError``). """ seen: Set[int] = set() stack: List[BaseException] = [exc] @@ -195,16 +204,7 @@ def _extract_upstream_auth_failure( response = getattr(current, "response", None) if response is not None: - status_code = getattr(response, "status_code", None) - if isinstance(status_code, int) and status_code in (401, 403): - www_authenticate: Optional[str] = None - headers = getattr(response, "headers", None) - if headers is not None: - try: - www_authenticate = headers.get("www-authenticate") - except Exception: - www_authenticate = None - return status_code, www_authenticate + yield response # anyio / PEP 654 ExceptionGroup sub_exceptions = getattr(current, "exceptions", None) @@ -219,6 +219,58 @@ def _extract_upstream_auth_failure( ): stack.append(current.__context__) + +def _extract_upstream_http_status( + exc: BaseException, +) -> Optional[Tuple[int, Optional[str]]]: + """Return ``(status_code, www_authenticate)`` for the first upstream + ``httpx.Response``-bearing exception in the tree, for ANY status code. + + Returns ``None`` if no HTTP status is present. + """ + for response in _iter_exception_responses(exc): + status_code = getattr(response, "status_code", None) + if isinstance(status_code, int): + www_authenticate: Optional[str] = None + headers = getattr(response, "headers", None) + if headers is not None: + try: + www_authenticate = headers.get("www-authenticate") + except Exception: + www_authenticate = None + return status_code, www_authenticate + return None + + +def _extract_upstream_auth_failure( + exc: BaseException, +) -> Optional[Tuple[int, Optional[str]]]: + """Walk the exception tree looking for an HTTP 401/403 response from the + upstream MCP server. + + The MCP SDK wraps transport errors in anyio ``ExceptionGroup`` objects and + may chain through ``__cause__`` / ``__context__``. We inspect all of those + layers for an ``httpx.Response``-bearing exception (typically + ``httpx.HTTPStatusError``) and extract the status code and any upstream + ``WWW-Authenticate`` header. + + Returns ``(status_code, www_authenticate)`` on match, else ``None``. + (Contract unchanged; now shares the tree-walk with + ``_extract_upstream_http_status``. We iterate independently rather than + reusing its "first response" so a non-401/403 response earlier in the tree + does not stop us from finding a 401/403 deeper in it.) + """ + for response in _iter_exception_responses(exc): + status_code = getattr(response, "status_code", None) + if isinstance(status_code, int) and status_code in (401, 403): + www_authenticate: Optional[str] = None + headers = getattr(response, "headers", None) + if headers is not None: + try: + www_authenticate = headers.get("www-authenticate") + except Exception: + www_authenticate = None + return status_code, www_authenticate return None @@ -1826,6 +1878,12 @@ class MCPServerManager: # client triggers the upstream OAuth flow. The multi-server # aggregator catches this explicitly to keep absorbing. raise + except MCPUpstreamError: + # Surface non-auth upstream failures the same way: single-server + # routes report a server error; multi-server aggregators absorb + # it per-server. Without this re-raise the generic handler below + # would swallow it back into [] and re-launder it as success. + raise except Exception as e: verbose_logger.warning( f"Failed to get tools from server {server.name}: {str(e)}" @@ -2465,11 +2523,44 @@ class MCPServerManager: except TimeoutError: verbose_logger.warning(f"Timeout while listing tools from {server_name}") return [] - except asyncio.CancelledError: - verbose_logger.warning( - f"Task cancelled while listing tools from {server_name}" + except asyncio.CancelledError as e: + # An upstream HTTP failure (e.g. 429/5xx) on a non-pass-through + # server surfaces through the MCP SDK's anyio TaskGroup as a + # cancel-scope cancellation. Do not launder that into an empty + # "success" — surface it as a per-server error so the REST routes + # report it instead of "Successfully retrieved tools". + upstream = _extract_upstream_http_status(e) + if upstream is not None: + status_code, _ = upstream + verbose_logger.info( + f"Upstream MCP server {server_name} failed with HTTP " + f"{status_code} (surfaced via cancellation)" + ) + raise MCPUpstreamError( + status_code=status_code, server_name=server_name + ) from e + + # No recoverable HTTP status. Respect a genuine *outer* cancellation + # (shutdown / client disconnect) so cooperative cancellation still + # works; only treat an internally-originated cancel as a failure. + task = asyncio.current_task() + is_outer_cancel = bool( + task is not None + and getattr(task, "cancelling", None) is not None + and task.cancelling() > 0 ) - return [] + if is_outer_cancel: + raise + verbose_logger.warning( + f"Tool listing from {server_name} was interrupted before " + f"completion; surfacing as a server error (was previously " + f"reported as 0 tools / success)." + ) + raise MCPUpstreamError( + status_code=None, + server_name=server_name, + message=f"Tool listing from {server_name} was interrupted", + ) from e except ConnectionError as e: verbose_logger.warning( f"Connection error while listing tools from {server_name}: {str(e)}" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index d900f690c57..497301d447e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -1,5 +1,6 @@ """Unit tests for MCP OAuth passthrough tool-fetch behavior.""" +import asyncio import sys from unittest.mock import AsyncMock, MagicMock @@ -8,10 +9,14 @@ import pytest sys.path.insert(0, "../../../../../") -from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError +from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPUpstreamAuthError, + MCPUpstreamError, +) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, _extract_upstream_auth_failure, + _extract_upstream_http_status, ) from litellm.proxy._types import MCPTransport from litellm.types.mcp import MCPAuth @@ -194,4 +199,82 @@ async def test_fetch_tools_from_gateway_managed_swallows_errors(): mock_client, oauth2_server.name, server=oauth2_server ) assert tools == [] + + +def test_extract_upstream_http_status_finds_429(): + response = httpx.Response( + status_code=429, + headers={}, + request=httpx.Request("GET", "https://upstream/mcp"), + ) + exc = httpx.HTTPStatusError("429", request=response.request, response=response) + assert _extract_upstream_http_status(exc) == (429, None) + + +def test_extract_upstream_http_status_walks_exception_group(): + response = httpx.Response( + status_code=503, + headers={}, + request=httpx.Request("GET", "https://upstream/mcp"), + ) + inner = httpx.HTTPStatusError("503", request=response.request, response=response) + try: + raise ExceptionGroup("wrapped", [inner]) # noqa: F821 (PEP 654, py3.11+) + except Exception as group: + assert _extract_upstream_http_status(group) == (503, None) + + +def test_extract_upstream_http_status_none_for_non_http(): + assert _extract_upstream_http_status(RuntimeError("boom")) is None + + +@pytest.mark.asyncio +async def test_fetch_tools_surfaces_upstream_error_on_cancellation_with_status(): + """A 429 that arrives as a cancel-scope cancellation must surface, not return [].""" + manager = MCPServerManager() + server = MCPServer( + server_id="o1", + name="mock_databricks", + url="https://upstream/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + response = httpx.Response( + status_code=429, + headers={}, + request=httpx.Request("GET", "https://upstream/mcp"), + ) + upstream_429 = httpx.HTTPStatusError( + "429", request=response.request, response=response + ) + + async def _raise_cancel(*args, **kwargs): + raise asyncio.CancelledError() from upstream_429 + + mock_client = MagicMock() + mock_client.list_tools = AsyncMock(side_effect=_raise_cancel) + + with pytest.raises(MCPUpstreamError) as exc_info: + await manager._fetch_tools_with_timeout(mock_client, server.name, server=server) + assert exc_info.value.status_code == 429 + assert exc_info.value.server_name == "mock_databricks" + + +@pytest.mark.asyncio +async def test_fetch_tools_surfaces_error_on_bare_cancellation(): + """An interrupted listing with no recoverable status is a failure, not success-empty.""" + manager = MCPServerManager() + server = MCPServer( + server_id="o1", + name="mock_databricks", + url="https://upstream/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + mock_client = MagicMock() + mock_client.list_tools = AsyncMock(side_effect=asyncio.CancelledError()) + + with pytest.raises(MCPUpstreamError) as exc_info: + await manager._fetch_tools_with_timeout(mock_client, server.name, server=server) + assert exc_info.value.status_code is None mock_client.list_tools.assert_awaited_with(raise_on_error=False) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 6433e0f6360..b9a2d097b59 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -544,6 +544,138 @@ class TestListToolsRestAPI: assert result["error"] is None assert result["message"] == "Successfully retrieved tools" + async def test_single_server_upstream_error_surfaces_server_error( + self, monkeypatch + ): + """A non-auth upstream failure (e.g. 429) on the single-server route must + surface as ``server_error`` — not laundered into a 200 success-empty body.""" + from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPUpstreamError, + ) + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + class StubServer: + alias = "server-1" + server_name = "server-1" + name = "mock_databricks" + allowed_tools = None + mcp_info = {"server_name": "mock_databricks"} + available_on_public_internet = True + + stub_server = StubServer() + + async def fake_get_tools(*args, **kwargs): + raise MCPUpstreamError(status_code=429, server_name="mock_databricks") + + monkeypatch.setattr( + rest_endpoints, + "build_effective_auth_contexts", + fake_contexts, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints, + "_get_tools_for_single_server", + fake_get_tools, + raising=False, + ) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + result = await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert result["tools"] == [] + assert result["error"] == "server_error" + assert "mock_databricks" in result["message"] + assert result["message"] != "Successfully retrieved tools" + + async def test_multi_server_upstream_error_surfaces_partial_failure( + self, monkeypatch + ): + """When every queried server fails with an upstream error, the multi-server + route must report ``partial_failure`` — not a 200 ``Successfully retrieved + tools`` with an empty list.""" + from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPUpstreamError, + ) + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + class StubServer: + server_id = "server-1" + alias = "server-1" + server_name = "server-1" + name = "mock_databricks" + allowed_tools = None + mcp_info = {"server_name": "mock_databricks"} + available_on_public_internet = True + + stub_server = StubServer() + + async def fake_get_tools(*args, **kwargs): + raise MCPUpstreamError(status_code=429, server_name="mock_databricks") + + monkeypatch.setattr( + rest_endpoints, + "build_effective_auth_contexts", + fake_contexts, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints, + "_get_tools_for_single_server", + fake_get_tools, + raising=False, + ) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + result = await rest_endpoints.list_tool_rest_api( + request, + server_id=None, + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert result["tools"] == [] + assert result["error"] == "partial_failure" + assert "mock_databricks" in result["message"] + assert result["message"] != "Successfully retrieved tools" + @pytest.mark.parametrize("upstream_status", [401, 403]) async def test_upstream_auth_failure_surfaces_status_and_challenge( self, monkeypatch, upstream_status