From 1337ebeb72afb78fd371f2e78a653e0a43ea7ba0 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 18 Jun 2026 09:24:45 -0700 Subject: [PATCH] feat(mcp): graft the v2 resolver into v1 for none + api_key (flag-gated) First strangler-fig graft of the v2 UpstreamCredentialProvider into the live v1 MCP request path, off by default. When LITELLM_USE_V2_MCP_RESOLVER is set (the --use_v2_migration_resolver CLI flag exports it), resolve_mcp_auth routes the none and api_key modes through the clean-room v2 resolver and returns the resolved credential as a header dict, which MCPClient merges verbatim; every other mode, a missing api_key, or any v2 error falls back to v1 unchanged. The hook sits after the per-request override check so that precedence is preserved. The bridge lives on the v1 side (v2 core keeps its no-v1-imports invariant). It adapts a v1 MCPServer to a v2 ServerSpec (none -> NoneConfig, api_key -> ApiKeyConfig on the X-API-Key header), builds the provider with inert in-memory/unwired ports (none and api_key resolve from config alone, so no real bodies are needed yet), runs resolve(), and extracts the produced headers via httpx .raw to preserve casing (X-API-Key, not httpx's lowercased x-api-key). Parity is asserted as the integration method: the v2-grafted upstream headers are byte-identical to v1's for none and api_key, verified through the real MCPClient._get_auth_headers; flag off, non-grafted modes, and tokenless api_key all defer to v1. v1 auth-priority tests still pass. 7 new tests; bridge typechecks clean and gates green. --- .../mcp_server/oauth2_token_cache.py | 6 + .../mcp_server/v2_resolver_bridge.py | 155 ++++++++++++++++++ litellm/proxy/proxy_cli.py | 2 + .../mcp_server/test_v2_resolver_bridge.py | 87 ++++++++++ 4 files changed, 250 insertions(+) create mode 100644 litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index 92ef57d8cd5..0592f083c40 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -27,6 +27,9 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy._experimental.mcp_server.auth import token_exchange +from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( + resolve_v2_auth_value, +) from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: @@ -287,6 +290,9 @@ async def resolve_mcp_auth( """ if mcp_auth_header: return mcp_auth_header + v2_auth_value = await resolve_v2_auth_value(server) + if v2_auth_value is not None: + return v2_auth_value if server.has_token_exchange_config: if subject_token: return await token_exchange.mcp_token_exchange_handler.exchange_token( diff --git a/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py b/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py new file mode 100644 index 00000000000..fa65ddb53c2 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py @@ -0,0 +1,155 @@ +"""Bridge: route v1 MCP auth resolution through the v2 UpstreamCredentialProvider. + +Strangler-fig graft. When ``LITELLM_USE_V2_MCP_RESOLVER`` is enabled (set by the +``--use_v2_migration_resolver`` CLI flag), ``resolve_mcp_auth()`` routes the ``none`` and +``api_key`` modes through the clean-room v2 resolver instead of v1's logic, returning the +resolved credential as a header dict (which ``MCPClient`` merges verbatim). Every other mode, +and any v2 error, falls back to v1. The v2 output is header-for-header identical to v1 for these +two modes; parity is the integration method, so the graft is observable but behavior-preserving. + +This lives on the v1 side so the v2 core keeps its no-v1-imports invariant. +""" + +from __future__ import annotations + +import functools +import os +from typing import TYPE_CHECKING, Dict, Optional + +import httpx +from pydantic import SecretStr + +from litellm._logging import verbose_logger +from litellm.proxy.gateway.mcp.outbound_credentials.clock import SystemClock +from litellm.proxy.gateway.mcp.outbound_credentials.credential_store import ( + InMemoryCredentialStore, +) +from litellm.proxy.gateway.mcp.outbound_credentials.resolver import ( + UpstreamCredentialProvider, +) +from litellm.proxy.gateway.mcp.outbound_credentials.service_token_store import ( + InMemoryServiceTokenStore, +) +from litellm.proxy.gateway.mcp.outbound_credentials.token_store import ( + InMemoryTokenStore, + StoredToken, +) +from litellm.proxy.gateway.mcp.outbound_credentials.types import ( + ApiKeyConfig, + AuthorizationCodeConfig, + AwsSigV4Config, + ClientCredentialsConfig, + CredError, + NoneConfig, + ServerSpec, + SharedKey, + Subject, + TokenExchangeConfig, +) +from litellm.proxy.gateway.mcp.result import Error, Result +from litellm.types.mcp import MCPAuth + +if TYPE_CHECKING: + from litellm.types.mcp_server.mcp_server_manager import MCPServer + +_V2_ENV_FLAG = "LITELLM_USE_V2_MCP_RESOLVER" + + +def v2_resolver_enabled() -> bool: + return os.getenv(_V2_ENV_FLAG, "").strip().lower() in ("1", "true", "yes", "on") + + +class _Unwired: + """Fail-closed stand-ins for the ports the grafted modes (none, api_key) never touch.""" + + async def refresh( + self, config: AuthorizationCodeConfig, refresh_token: SecretStr + ) -> Result[StoredToken, CredError]: + return Error(CredError.of_not_implemented("token_refresher not wired")) + + async def fetch( + self, config: ClientCredentialsConfig + ) -> Result[StoredToken, CredError]: + return Error( + CredError.of_not_implemented("client_credentials_fetcher not wired") + ) + + async def exchange( + self, config: TokenExchangeConfig, subject_token: SecretStr, resource: str + ) -> Result[StoredToken, CredError]: + return Error(CredError.of_not_implemented("token_exchanger not wired")) + + async def build(self, config: AwsSigV4Config) -> Result[httpx.Auth, CredError]: + return Error(CredError.of_not_implemented("signer_factory not wired")) + + +@functools.lru_cache(maxsize=1) +def _provider() -> UpstreamCredentialProvider: + # none + api_key resolve from config alone, so the stores/ports below are inert placeholders; + # real bodies get wired in as their modes are grafted. + unwired = _Unwired() + return UpstreamCredentialProvider( + credential_store=InMemoryCredentialStore(), + token_store=InMemoryTokenStore(), + token_refresher=unwired, + clock=SystemClock(), + service_token_store=InMemoryServiceTokenStore(), + client_credentials_fetcher=unwired, + token_exchanger=unwired, + signer_factory=unwired, + ) + + +def _to_server_spec(server: MCPServer) -> Optional[ServerSpec]: + resource = server.url or server.server_id + if server.auth_type in (None, MCPAuth.none): + return ServerSpec( + server_id=server.server_id, resource=resource, config=NoneConfig() + ) + if server.auth_type == MCPAuth.api_key: + token = server.authentication_token + if not token: + return None # api_key with no key: let v1 handle it (parity-safe) + return ServerSpec( + server_id=server.server_id, + resource=resource, + config=ApiKeyConfig( + header_name="X-API-Key", + value_prefix="", + key_source=SharedKey(value=SecretStr(token)), + ), + ) + return None # other modes are not grafted yet + + +def _added_headers(auth: httpx.Auth) -> Dict[str, str]: + # Use .raw (not .items()) so the auth's original header casing survives, e.g. X-API-Key + # rather than httpx's lowercased x-api-key, keeping byte-parity with v1. + request = httpx.Request("GET", "https://placeholder.invalid") + base = {(name.lower(), value) for name, value in request.headers.raw} + signed = next(auth.auth_flow(request)) + return { + name.decode(): value.decode() + for name, value in signed.headers.raw + if (name.lower(), value) not in base + } + + +async def resolve_v2_auth_value(server: MCPServer) -> Optional[Dict[str, str]]: + """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) + if spec is None: + return None + result = await _provider().resolve( + Subject(tenant_id="", subject_id="", inbound_token=None), spec + ) + if isinstance(result, Error): + verbose_logger.warning( + "v2 MCP resolver failed for server %s: %s; falling back to v1", + server.server_id, + result.error.summary, + ) + return None + return _added_headers(result.ok) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 9c4d7b1bb5d..72a0446004f 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -917,6 +917,8 @@ def run_server( use_v2_migration_resolver: bool, reload: bool, ): + if use_v2_migration_resolver: + os.environ["LITELLM_USE_V2_MCP_RESOLVER"] = "true" if cli_args: if cli_args == ("xai-oauth", "login"): from litellm.llms.xai.oauth import XAIOAuthAuthenticator 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 new file mode 100644 index 00000000000..1c30de9a8b8 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py @@ -0,0 +1,87 @@ +"""Parity tests for the v1->v2 MCP resolver graft (none + api_key). + +When the flag is on, resolve_mcp_auth routes none/api_key through the v2 resolver; the upstream +headers must be byte-identical to what v1 produces, and every other mode (or any v2 error) must +fall back to v1 unchanged. +""" + +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 ( + resolve_v2_auth_value, +) +from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +pytestmark = pytest.mark.asyncio + +FLAG = "LITELLM_USE_V2_MCP_RESOLVER" + + +def _server(auth_type, token=None): + return MCPServer( + server_id="s1", + name="s1", + transport=MCPTransport.http, + url="https://up.example/mcp", + auth_type=auth_type, + authentication_token=token, + ) + + +def _v1_headers(auth_type, auth_value): + """The final upstream headers v1's MCPClient would send for this resolved value.""" + return MCPClient(auth_type=auth_type, auth_value=auth_value)._get_auth_headers() + + +@pytest.fixture +def v2_on(monkeypatch): + monkeypatch.setenv(FLAG, "true") + + +@pytest.fixture +def v2_off(monkeypatch): + monkeypatch.delenv(FLAG, raising=False) + + +async def test_flag_off_defers_to_v1(v2_off): + assert await resolve_v2_auth_value(_server(MCPAuth.api_key, "k")) is None + + +async def test_api_key_parity(v2_on): + token = "up-secret" + server = _server(MCPAuth.api_key, token) + v2_value = await resolve_v2_auth_value(server) + assert v2_value == {"X-API-Key": token} + # byte-identical to v1's final upstream headers + assert _v1_headers(MCPAuth.api_key, token) == _v1_headers(MCPAuth.api_key, v2_value) + + +async def test_none_attaches_no_auth(v2_on): + server = _server(MCPAuth.none) + v2_value = await resolve_v2_auth_value(server) + assert v2_value == {} + # v1 none -> no auth header; v2 -> empty dict merged -> no auth header + assert _v1_headers(MCPAuth.none, None) == _v1_headers(MCPAuth.none, v2_value) == {} + + +async def test_non_grafted_mode_defers_to_v1(v2_on): + # bearer_token is not grafted yet -> v2 returns None so v1 handles it + assert await resolve_v2_auth_value(_server(MCPAuth.bearer_token, "k")) is None + + +async def test_api_key_without_token_defers_to_v1(v2_on): + assert await resolve_v2_auth_value(_server(MCPAuth.api_key, None)) is None + + +async def test_resolve_mcp_auth_hook_routes_api_key_when_on(v2_on): + server = _server(MCPAuth.api_key, "up-secret") + assert await resolve_mcp_auth(server) == {"X-API-Key": "up-secret"} + + +async def test_resolve_mcp_auth_hook_uses_v1_when_off(v2_off): + # v1 returns the raw token string; MCPClient then maps it to the X-API-Key header + server = _server(MCPAuth.api_key, "up-secret") + assert await resolve_mcp_auth(server) == "up-secret"