From 153e23e1b419ab16ff15a44fd97b9f6e09febfbc Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 19 Jun 2026 18:05:39 -0700 Subject: [PATCH] feat(mcp/v2): MCPServerManagerV2 egress override for the list path Step 6b: override _get_tools_from_server to resolve auth via resolve() and list tools through the v2 UpstreamConnection (none / api_key / client_credentials / authorization_code). Inherits v1's namespacing (_create_prefixed_tools) and static-header resolution. super() is reserved for non-egress concerns that migrate as their own subsystems: OpenAPI tools (registry, S1.8), the per-request mcp_auth_header override path, the JWT-signer guardrail, and any not-yet-mapped mode. List errors map per the scheme: unauthorized -> MCPUpstreamAuthError (401 + WWW-Authenticate, the LIT-3795 behavior); anything else degrades to an empty list (logged). Also: - Make the bridge's v1->v2 adapters public (to_subject / to_server_spec / provider); the egress manager reuses them (was reportPrivateUsage). - Make global_mcp_server_manager a PEP 562 lazy singleton, fixing a v1<->v2 import cycle (eager module-load instantiation imported the v2 subclass mid-load; import-order dependent). Live-validated: the m2m mode through the override against the harness lists echo via v2 (real token fetch + Bearer-authed connection); the none mode via an in-process server integration test. --- .../mcp_server/mcp_server_manager.py | 32 ++++- .../mcp_server/mcp_server_manager_v2.py | 128 ++++++++++++++++-- .../mcp_server/v2_resolver_bridge.py | 14 +- .../mcp_server/test_mcp_server_manager_v2.py | 74 +++++++++- .../mcp_server/test_v2_resolver_bridge.py | 36 ++--- 5 files changed, 246 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 2cc14a84668..91209b36c90 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 ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Literal, + Optional, + Set, + Tuple, + Union, + cast, +) from urllib.parse import urlparse import anyio @@ -4523,4 +4535,20 @@ def _make_global_mcp_server_manager() -> MCPServerManager: return MCPServerManagerV2() -global_mcp_server_manager: MCPServerManager = _make_global_mcp_server_manager() +_global_mcp_server_manager: Optional[MCPServerManager] = None + +if TYPE_CHECKING: + global_mcp_server_manager: MCPServerManager + + +def __getattr__(name: str) -> MCPServerManager: + # PEP 562 lazy singleton: instantiating at import time would import the v2 subclass while this + # module is still loading (a v1<->v2 cycle, import-order dependent), so the manager is built on + # first access instead. Tests that reassign global_mcp_server_manager set a real attribute that + # shadows this. + if name == "global_mcp_server_manager": + global _global_mcp_server_manager + if _global_mcp_server_manager is None: + _global_mcp_server_manager = _make_global_mcp_server_manager() + return _global_mcp_server_manager + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py index 4290f5d1b99..eaf6a11c5fd 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py @@ -1,22 +1,130 @@ """v2-owned MCP egress manager (the cutover seam). -``MCPServerManagerV2`` subclasses v1's ``MCPServerManager`` and, in later steps, overrides the -per-server egress methods (``_get_tools_from_server``, ``call_tool``, the prompt/resource ops) to -route through the v2 ``UpstreamConnection`` + ``resolve()`` instead of ``_create_mcp_client``. -Registry, RBAC, cross-server aggregation, namespacing, and static-header resolution are inherited -from v1 unchanged. It is the egress manager, constructed at the composition root (see +``MCPServerManagerV2`` subclasses v1's ``MCPServerManager`` and overrides the per-server egress +methods to route through the v2 ``UpstreamConnection`` + ``resolve()`` instead of +``_create_mcp_client``. Registry, RBAC, cross-server aggregation, namespacing +(``_create_prefixed_tools``), and static-header resolution are inherited from v1 unchanged. It is +the egress manager, constructed at the composition root (see ``mcp_server_manager._make_global_mcp_server_manager``); there is no opt-in flag (v2 is the egress -implementation). v1's per-server methods remain reachable via ``super()`` for modes not yet -overridden, so the migration is the override progression, not a runtime toggle. +implementation). -Step 6a lands the skeleton only: no overrides, so behavior is identical to v1. The egress overrides -land in 6b/6c. +Migration is the override progression: each egress mode is wired through ``resolve()`` + +``UpstreamConnection``, live-validated, and committed one at a time. Modes that are not yet wired +(passthrough / token_exchange, which need the caller's inbound token) simply fail closed until their +commit. ``super()`` is reserved for non-egress concerns that migrate as their own subsystems: +OpenAPI tools (registry, S1.8), the per-request ``mcp_auth_header`` override/inbound-token path, and +the JWT-signer guardrail. """ from __future__ import annotations +from typing import TYPE_CHECKING, Dict, List, Optional, Union + +from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager +from litellm.proxy._types import MCPTransport + +if TYPE_CHECKING: + from mcp.types import Tool as MCPTool + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.gateway.mcp.outbound_credentials.types import CredError + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + from .v2_egress import ConnError class MCPServerManagerV2(MCPServerManager): - """v2 egress manager; see the module docstring. No overrides yet (step 6a).""" + """v2 egress manager; see the module docstring.""" + + @staticmethod + def _jwt_signer_configured() -> bool: + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + get_mcp_jwt_signer, + ) + + return get_mcp_jwt_signer() is not None + + async def _get_tools_from_server( + self, + server: MCPServer, + mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, + extra_headers: Optional[Dict[str, str]] = None, + add_prefix: bool = True, + raw_headers: Optional[Dict[str, str]] = None, + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + ) -> List[MCPTool]: + from litellm.proxy._experimental.mcp_server.v2_egress import UpstreamConnection + from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( + provider, + to_server_spec, + to_subject, + ) + from litellm.proxy.gateway.mcp.result import Error + + spec = to_server_spec(server) + # Stay on v1 for what hasn't migrated to the egress transport: unmapped/not-yet-wired modes + # (spec is None, e.g. aws_sigv4), OpenAPI tools (registry, S1.8), the per-request + # mcp_auth_header override/inbound-token path (migrated with passthrough/token_exchange), and + # the JWT-signer guardrail. + if ( + spec is None + or server.spec_path + or mcp_auth_header + or self._jwt_signer_configured() + ): + return await super()._get_tools_from_server( + server, + mcp_auth_header, + extra_headers, + add_prefix, + raw_headers, + user_api_key_auth, + ) + + auth = await provider().resolve(to_subject(user_api_key_auth, None), spec) + if isinstance(auth, Error): + return self._egress_list_failure(server, auth.error) + + resolved_static = await self._resolve_static_headers_with_env_vars( + server, user_api_key_auth, raise_on_missing=False + ) + headers = {**(extra_headers or {}), **(resolved_static or {})} or None + is_stdio = server.transport == MCPTransport.stdio + + result = await UpstreamConnection( + server.url, + transport=server.transport, + auth=auth.ok, + extra_headers=headers, + command=server.command, + args=server.args, + env=self._build_stdio_env(server, raw_headers) if is_stdio else None, + ).list_tools() + if isinstance(result, Error): + return self._egress_list_failure(server, result.error) + return self._create_prefixed_tools(result.ok, server, add_prefix=add_prefix) + + def _egress_list_failure( + self, server: MCPServer, error: "CredError | ConnError" + ) -> List[MCPTool]: + # List path: an upstream 401/403 (or a per-user mode with no usable credential) surfaces as + # MCPUpstreamAuthError so the client gets a 401 + WWW-Authenticate and starts the OAuth flow + # (this is the LIT-3795 behavior for non-delegated interactive oauth2). Any other failure + # degrades to an empty tool list (logged), so one bad server never collapses the federated + # catalog. The typed partial-failure marker is a later, separate surface. + if error.tag == "unauthorized": + from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPUpstreamAuthError, + ) + + raise MCPUpstreamAuthError( + status_code=401, www_authenticate=None, server_name=server.name + ) + verbose_logger.warning( + "v2 egress: tools unavailable for %s (%s): %s", + server.name, + server.server_id, + error.summary, + ) + return [] diff --git a/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py b/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py index c4733e05a84..1e02527f8bf 100644 --- a/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py +++ b/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py @@ -91,7 +91,7 @@ class _Unwired: @functools.lru_cache(maxsize=1) -def _provider() -> UpstreamCredentialProvider: +def provider() -> UpstreamCredentialProvider: # Real bodies are wired as their modes are grafted. token_refresher stays unwired: the bridge's # V1OAuthTokenStore returns currently-valid tokens (v1 refreshes on read), so the resolver's # proactive-refresh path is inert until v1 retires. @@ -108,7 +108,7 @@ def _provider() -> UpstreamCredentialProvider: ) -def _to_server_spec(server: MCPServer) -> Optional[ServerSpec]: +def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: resource = server.url or server.server_id if server.auth_type in (None, MCPAuth.none): # A none server opted into upstream OAuth passthrough forwards the caller's bearer; the @@ -219,7 +219,7 @@ def _added_headers(auth: httpx.Auth) -> Dict[str, str]: } -def _to_subject( +def to_subject( user_api_key_auth: Optional[UserAPIKeyAuth], subject_token: Optional[str] ) -> Subject: """Map v1's authenticated principal onto the v2 Subject. @@ -246,11 +246,11 @@ async def resolve_v2_auth_value( """Resolve `none`/`api_key` via the v2 resolver, or return None to defer to v1.""" if not v2_resolver_enabled(): return None - spec = _to_server_spec(server) + spec = to_server_spec(server) if spec is None: return None - result = await _provider().resolve( - _to_subject(user_api_key_auth, subject_token), spec + result = await provider().resolve( + to_subject(user_api_key_auth, subject_token), spec ) if isinstance(result, Error): verbose_logger.warning( @@ -315,7 +315,7 @@ async def resolve_v2_aws_auth(server: MCPServer) -> Optional[httpx.Auth]: config = _to_aws_sigv4_config(server) if config is None: return None - result = await _provider().resolve( + result = await provider().resolve( Subject(tenant_id="", subject_id="", inbound_token=None), ServerSpec( server_id=server.server_id, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py index 08b831fa4cf..0dd3dd6f5a7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py @@ -1,4 +1,13 @@ -"""Tests for the v2 egress manager factory + skeleton (step 6a).""" +"""Tests for the v2 egress manager: factory, skeleton, and the egress override (step 6).""" + +import contextlib +import socket +import threading +import time + +import httpx +import pytest +import uvicorn from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, @@ -7,6 +16,9 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( from litellm.proxy._experimental.mcp_server.mcp_server_manager_v2 import ( MCPServerManagerV2, ) +from litellm.proxy._types import MCPTransport +from litellm.types.mcp import MCPAuth +from litellm.types.mcp_server.mcp_server_manager import MCPServer def test_v2_is_a_manager_subclass(): @@ -16,3 +28,63 @@ def test_v2_is_a_manager_subclass(): def test_factory_constructs_the_v2_manager(): # v2 is the egress implementation; there is no opt-in flag. assert isinstance(_make_global_mcp_server_manager(), MCPServerManagerV2) + + +@contextlib.contextmanager +def _serve_echo(): + """A no-auth streamable-http FastMCP server with one `echo` tool, in a background thread.""" + from mcp.server.fastmcp import FastMCP + + mcp = FastMCP("mgr-echo-test", stateless_http=True) + + @mcp.tool() + def echo(text: str) -> str: + return f"echo: {text}" + + 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" + ) + ) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + for _ in range(100): + try: + httpx.get(url, timeout=0.3) + break + except httpx.ConnectError: + time.sleep(0.05) + except Exception: + break + try: + yield url + finally: + server.should_exit = True + thread.join(timeout=5) + + +@pytest.fixture +def echo_server_url(): + with _serve_echo() as url: + yield url + + +@pytest.mark.asyncio +async def test_v2_override_lists_tools_via_upstream_connection(echo_server_url): + # The `none`-mode list path goes through resolve() + UpstreamConnection (v2), and the tools come + # back namespaced via the inherited _create_prefixed_tools. + manager = MCPServerManagerV2() + server = MCPServer( + server_id="echo1", + name="echo1", + transport=MCPTransport.http, + url=echo_server_url, + auth_type=MCPAuth.none, + ) + tools = await manager._get_tools_from_server(server, add_prefix=True) + assert any(t.name.endswith("echo") for t in tools) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py index e717a3651a7..2b78b86af76 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py @@ -11,7 +11,7 @@ import pytest from litellm.experimental_mcp_client.client import MCPClient from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mcp_auth from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( - _to_subject, + to_subject, resolve_v2_auth_value, resolve_v2_aws_auth, ) @@ -108,13 +108,13 @@ def _m2m_server(): async def test_client_credentials_maps_to_config(): from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( - _to_server_spec, + to_server_spec, ) from litellm.proxy.gateway.mcp.outbound_credentials.types import ( ClientCredentialsConfig, ) - spec = _to_server_spec(_m2m_server()) + spec = to_server_spec(_m2m_server()) assert spec is not None assert isinstance(spec.config, ClientCredentialsConfig) assert spec.config.client_id == "cid" @@ -244,7 +244,7 @@ async def test_aws_sigv4_config_defaults_to_ambient(v2_on): async def test_to_subject_maps_v1_identity(): auth = UserAPIKeyAuth(token="sk-test", user_id="u1", org_id="org1", team_id="t1") - subj = _to_subject(auth, "inbound-jwt") + subj = to_subject(auth, "inbound-jwt") assert subj.subject_id == "u1" assert subj.tenant_id == "org1" # org preferred over team assert subj.inbound_token is not None @@ -252,7 +252,7 @@ async def test_to_subject_maps_v1_identity(): async def test_to_subject_anonymous_when_no_auth(): - subj = _to_subject(None, None) + subj = to_subject(None, None) assert subj.subject_id == "" assert subj.tenant_id == "" assert subj.inbound_token is None @@ -260,7 +260,7 @@ async def test_to_subject_anonymous_when_no_auth(): async def test_to_subject_falls_back_to_team_and_blanks_missing_user(): auth = UserAPIKeyAuth(token="sk-test", team_id="team-x") - subj = _to_subject(auth, None) + subj = to_subject(auth, None) assert subj.tenant_id == "team-x" # no org -> team assert ( subj.subject_id == "" @@ -278,7 +278,7 @@ async def test_resolve_v2_auth_value_threads_identity_without_breaking_static(v2 async def test_byok_server_maps_to_byok_key_source(): from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( - _to_server_spec, + to_server_spec, ) from litellm.proxy.gateway.mcp.outbound_credentials.types import ApiKeyConfig, Byok @@ -290,7 +290,7 @@ async def test_byok_server_maps_to_byok_key_source(): auth_type=MCPAuth.api_key, is_byok=True, ) - spec = _to_server_spec(server) + spec = to_server_spec(server) assert spec is not None assert isinstance(spec.config, ApiKeyConfig) assert isinstance(spec.config.key_source, Byok) @@ -299,7 +299,7 @@ async def test_byok_server_maps_to_byok_key_source(): async def test_token_exchange_server_maps_to_config(): from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( - _to_server_spec, + to_server_spec, ) from litellm.proxy.gateway.mcp.outbound_credentials.types import TokenExchangeConfig @@ -315,7 +315,7 @@ async def test_token_exchange_server_maps_to_config(): audience="https://aud.example", scopes=["a", "b"], ) - spec = _to_server_spec(server) + spec = to_server_spec(server) assert spec is not None assert isinstance(spec.config, TokenExchangeConfig) assert spec.config.token_exchange_endpoint == "https://idp/exchange" @@ -330,7 +330,7 @@ async def test_interactive_oauth2_no_delegate_maps_to_authorization_code(): # LIT-3795: a non-delegated interactive oauth2 server resolves to authorization_code # (gateway-stored per-user token), so the caller JWT is never forwarded upstream. from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( - _to_server_spec, + to_server_spec, ) from litellm.proxy.gateway.mcp.outbound_credentials.types import ( AuthorizationCodeConfig, @@ -350,7 +350,7 @@ async def test_interactive_oauth2_no_delegate_maps_to_authorization_code(): token_url="https://idp/token", scopes=["openid"], ) - spec = _to_server_spec(server) + spec = to_server_spec(server) assert spec is not None assert isinstance(spec.config, AuthorizationCodeConfig) assert spec.config.token_url == "https://idp/token" @@ -359,7 +359,7 @@ async def test_interactive_oauth2_no_delegate_maps_to_authorization_code(): async def test_interactive_oauth2_delegate_maps_to_passthrough(): from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( - _to_server_spec, + to_server_spec, ) from litellm.proxy.gateway.mcp.outbound_credentials.types import PassthroughConfig @@ -372,14 +372,14 @@ async def test_interactive_oauth2_delegate_maps_to_passthrough(): oauth2_flow="authorization_code", delegate_auth_to_upstream=True, ) - spec = _to_server_spec(server) + spec = to_server_spec(server) assert spec is not None assert isinstance(spec.config, PassthroughConfig) async def test_none_oauth_passthrough_maps_to_passthrough(): from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( - _to_server_spec, + to_server_spec, ) from litellm.proxy.gateway.mcp.outbound_credentials.types import PassthroughConfig @@ -392,14 +392,14 @@ async def test_none_oauth_passthrough_maps_to_passthrough(): oauth_passthrough=True, extra_headers=["Authorization"], ) - spec = _to_server_spec(server) + spec = to_server_spec(server) assert spec is not None assert isinstance(spec.config, PassthroughConfig) async def test_none_without_passthrough_maps_to_none(): from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( - _to_server_spec, + to_server_spec, ) from litellm.proxy.gateway.mcp.outbound_credentials.types import NoneConfig @@ -410,6 +410,6 @@ async def test_none_without_passthrough_maps_to_none(): url="https://up.example/mcp", auth_type=MCPAuth.none, ) - spec = _to_server_spec(server) + spec = to_server_spec(server) assert spec is not None assert isinstance(spec.config, NoneConfig)