From 9fafa0966675f7e8a71795451f723ba9c72ea051 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 19 Jun 2026 09:28:15 -0700 Subject: [PATCH] fix(mcp/v2): classify upstream 401 in UpstreamConnection via v1's helper Live-testing UpstreamConnection against the Bearer-protected harness surfaced that an upstream 401 mapped to upstream_unavailable (503) instead of unauthorized: the SDK wraps the httpx 401 in an anyio ExceptionGroup, so the top-level .response check never saw it. Reuse v1's extract_upstream_auth_failure (renamed from the private _extract_upstream_auth_failure) instead of a hand-rolled flattener: it already walks the ExceptionGroup and __cause__/__context__ chains portably (BaseExceptionGroup is not a builtin on the supported 3.10) and handles 401/403. ConnError drops the now-unused protocol_error tag. Adds a regression test (Bearer-protected in-process server: no auth -> unauthorized; correct bearer -> Ok) that fails with the old classifier. --- .../mcp_server/mcp_server_manager.py | 4 +- .../_experimental/mcp_server/v2_egress.py | 40 ++++------ .../test_mcp_oauth_passthrough_tools.py | 8 +- .../mcp_server/test_v2_egress.py | 80 +++++++++++++++---- 4 files changed, 85 insertions(+), 47 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index ca042e2d348..8a6759e1185 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -214,7 +214,7 @@ def _should_strip_caller_authorization( ) -def _extract_upstream_auth_failure( +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 @@ -2847,7 +2847,7 @@ class MCPServerManager: return [] except Exception as e: if should_surface_upstream_auth: - auth_info = _extract_upstream_auth_failure(e) + auth_info = extract_upstream_auth_failure(e) if auth_info is not None: status_code, www_authenticate = auth_info verbose_logger.info( diff --git a/litellm/proxy/_experimental/mcp_server/v2_egress.py b/litellm/proxy/_experimental/mcp_server/v2_egress.py index a77c34799da..850b1ee79b1 100644 --- a/litellm/proxy/_experimental/mcp_server/v2_egress.py +++ b/litellm/proxy/_experimental/mcp_server/v2_egress.py @@ -123,12 +123,12 @@ class MCPEgressManager(Protocol): class ConnError(BaseModel): """A connection/transport failure to an upstream MCP server, modeled as a value. - Discriminated on ``tag`` so callers can route on it (a 401 is surfaced to the client to - trigger the upstream OAuth flow; transient/transport failures map to a 503). + Discriminated on ``tag``: an upstream 401/403 is surfaced to the client to trigger the + upstream OAuth flow; any other transport failure maps to a 503. """ model_config = ConfigDict(frozen=True) - tag: Literal["unauthorized", "upstream_unavailable", "protocol_error"] + tag: Literal["unauthorized", "upstream_unavailable"] summary: str @classmethod @@ -139,31 +139,19 @@ class ConnError(BaseModel): def of_upstream_unavailable(cls, summary: str) -> "ConnError": return cls(tag="upstream_unavailable", summary=summary) - @classmethod - def of_protocol_error(cls, summary: str) -> "ConnError": - return cls(tag="protocol_error", summary=summary) - - -_TRANSIENT_ERROR_NAMES = ( - "ConnectError", - "ConnectTimeout", - "ReadTimeout", - "WriteTimeout", - "PoolTimeout", - "TimeoutException", - "ConnectionError", - "RemoteProtocolError", -) - def _classify_conn_error(error: Exception) -> ConnError: - response = getattr(error, "response", None) - if getattr(response, "status_code", None) == 401: - return ConnError.of_unauthorized(f"upstream returned 401: {error}") - if type(error).__name__ == "McpError": - return ConnError.of_protocol_error(f"MCP protocol error: {error}") - if type(error).__name__ in _TRANSIENT_ERROR_NAMES: - return ConnError.of_upstream_unavailable(f"upstream unreachable: {error}") + # Reuse v1's exception-tree walk (it handles the SDK's anyio ExceptionGroup and the + # __cause__/__context__ chains portably) to spot an upstream 401/403; anything else is a + # transport failure. Imported lazily to avoid an import cycle at the manager swap. + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + extract_upstream_auth_failure, + ) + + auth_failure = extract_upstream_auth_failure(error) + if auth_failure is not None: + status_code, _ = auth_failure + return ConnError.of_unauthorized(f"upstream returned {status_code}") return ConnError.of_upstream_unavailable(f"upstream connection failed: {error}") 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 d51cf8c5b72..3ca07a594f2 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 @@ -11,7 +11,7 @@ sys.path.insert(0, "../../../../../") from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, - _extract_upstream_auth_failure, + extract_upstream_auth_failure, ) from litellm.proxy._types import MCPTransport from litellm.types.mcp import MCPAuth @@ -26,7 +26,7 @@ def test_extract_upstream_auth_failure_finds_401_in_http_status_error(): ) exc = httpx.HTTPStatusError("401", request=response.request, response=response) - result = _extract_upstream_auth_failure(exc) + result = extract_upstream_auth_failure(exc) assert result == (401, 'Bearer resource_metadata="https://x"') @@ -41,13 +41,13 @@ def test_extract_upstream_auth_failure_walks_exception_group(): try: raise ExceptionGroup("wrapped", [inner]) # noqa: F821 (PEP 654, py3.11+) except Exception as group: - result = _extract_upstream_auth_failure(group) + result = extract_upstream_auth_failure(group) assert result == (401, "Bearer") def test_extract_upstream_auth_failure_returns_none_for_non_auth(): - assert _extract_upstream_auth_failure(RuntimeError("boom")) is None + assert extract_upstream_auth_failure(RuntimeError("boom")) is None @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_egress.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_egress.py index 6ec5317e3a8..4b21ceec69d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_egress.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_egress.py @@ -1,5 +1,6 @@ """Tests for the v2 MCP egress transport: the flag and the UpstreamConnection.""" +import contextlib import socket import threading import time @@ -30,27 +31,16 @@ def test_egress_flag_falsey_values(monkeypatch, value): assert v2_egress_enabled() is False -@pytest.fixture -def echo_server_url(): - """A no-auth streamable-http FastMCP server with one `echo` tool, in a background thread.""" - from mcp.server.fastmcp import FastMCP - - mcp = FastMCP("egress-echo-test", stateless_http=True) - - @mcp.tool() - def echo(text: str) -> str: - return f"echo: {text}" - +@contextlib.contextmanager +def _serve(app): + """Serve an ASGI app on a free port in a background thread; yield its /mcp url.""" sock = socket.socket() sock.bind(("127.0.0.1", 0)) port = sock.getsockname()[1] sock.close() url = f"http://127.0.0.1:{port}/mcp" - server = uvicorn.Server( - uvicorn.Config( - mcp.streamable_http_app(), host="127.0.0.1", port=port, log_level="error" - ) + uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error") ) thread = threading.Thread(target=server.run, daemon=True) thread.start() @@ -69,6 +59,44 @@ def echo_server_url(): thread.join(timeout=5) +def _echo_app(token=None): + """A stateless streamable-http FastMCP app with one `echo` tool, optionally Bearer-gated.""" + from mcp.server.fastmcp import FastMCP + + mcp = FastMCP("egress-test", stateless_http=True) + + @mcp.tool() + def echo(text: str) -> str: + return f"echo: {text}" + + app = mcp.streamable_http_app() + if token is not None: + from starlette.middleware.base import BaseHTTPMiddleware + from starlette.responses import JSONResponse + + class _BearerCheck(BaseHTTPMiddleware): + async def dispatch(self, request, call_next): + authorized = request.headers.get("authorization") == f"Bearer {token}" + if request.url.path.startswith("/mcp") and not authorized: + return JSONResponse({"error": "unauthorized"}, status_code=401) + return await call_next(request) + + app.add_middleware(_BearerCheck) + return app + + +@pytest.fixture +def echo_server_url(): + with _serve(_echo_app()) as url: + yield url + + +@pytest.fixture +def protected_server_url(): + with _serve(_echo_app(token="secret-token")) as url: + yield url, "secret-token" + + @pytest.mark.asyncio async def test_upstream_connection_lists_and_calls(echo_server_url): from litellm.proxy._experimental.mcp_server.v2_egress import UpstreamConnection @@ -95,3 +123,25 @@ async def test_upstream_connection_unreachable_is_upstream_unavailable(): result = await conn.list_tools() assert isinstance(result, Error) assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_upstream_connection_401_is_unauthorized(protected_server_url): + from litellm.proxy._experimental.mcp_server.v2_egress import UpstreamConnection + from litellm.proxy.gateway.mcp.outbound_credentials.httpx_auth import ( + NoOpAuth, + StaticHeaderAuth, + ) + from litellm.proxy.gateway.mcp.result import Error, Ok + + url, token = protected_server_url + + rejected = await UpstreamConnection(url, auth=NoOpAuth()).list_tools() + assert isinstance(rejected, Error) + assert rejected.error.tag == "unauthorized" + + authed = await UpstreamConnection( + url, auth=StaticHeaderAuth(f"Bearer {token}") + ).list_tools() + assert isinstance(authed, Ok) + assert any(t.name == "echo" for t in authed.ok)