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.
This commit is contained in:
Tin Chi Lo 2026-06-19 18:05:39 -07:00
parent 86c8ab5739
commit 153e23e1b4
5 changed files with 246 additions and 38 deletions

View file

@ -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}")

View file

@ -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 []

View file

@ -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,

View file

@ -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)

View file

@ -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)