From ec91056fcbfaae6c5752ed6e92bdcf212d0813db Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 18 Jun 2026 17:32:40 -0700 Subject: [PATCH] feat(mcp): thread inbound identity into the v2 resolver (Subject edge) Adds _to_subject(user_api_key_auth, subject_token) -> Subject, the single isolated mapping from v1's authenticated principal onto the v2 Subject (so it can later swap to auth_v2's Principal): subject_id <- user_id, tenant_id <- org_id (falling back to team_id), inbound_token <- subject_token; an unauthenticated caller yields empty ids. resolve_mcp_auth and resolve_v2_auth_value now accept user_api_key_auth and thread it (plus subject_token) through, and _create_mcp_client passes its auth context in. This is foundational for the per-user arms (BYOK api_key, token_exchange, authorization_code), which must reject an empty subject_id rather than share one credential slot across callers. The identity-free modes already grafted (none, api_key shared, client_credentials, aws_sigv4) ignore the subject, so threading it is additive and backward-compatible; the params default to None. Tests cover the mapping (org-over-team precedence, team fallback, missing user -> empty, anonymous when no auth) and that threading identity doesn't change the grafted static modes. 87 tests pass; the bridge typechecks clean and no new errors land on the manager or token cache. --- .../mcp_server/mcp_server_manager.py | 5 ++- .../mcp_server/oauth2_token_cache.py | 6 +++- .../mcp_server/v2_resolver_bridge.py | 28 +++++++++++++-- .../mcp_server/test_v2_resolver_bridge.py | 36 +++++++++++++++++++ 4 files changed, 71 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0eabc5051d1..ca042e2d348 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1940,7 +1940,10 @@ class MCPServerManager: Configured MCP client instance. """ auth_value = await resolve_mcp_auth( - server, mcp_auth_header, subject_token=subject_token + server, + mcp_auth_header, + subject_token=subject_token, + user_api_key_auth=user_api_key_auth, ) transport = server.transport or MCPTransport.sse diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index d7f098908a4..5961c8752c1 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -43,6 +43,7 @@ except ImportError as _e: resolve_v2_auth_value = None # type: ignore[assignment] if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -289,6 +290,7 @@ async def resolve_mcp_auth( server: "MCPServer", mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, subject_token: Optional[str] = None, + user_api_key_auth: Optional["UserAPIKeyAuth"] = None, ) -> Optional[Union[str, Dict[str, str]]]: """Resolve the auth value for an MCP server. @@ -301,7 +303,9 @@ async def resolve_mcp_auth( if mcp_auth_header: return mcp_auth_header if resolve_v2_auth_value is not None: - v2_auth_value = await resolve_v2_auth_value(server) + v2_auth_value = await resolve_v2_auth_value( + server, user_api_key_auth=user_api_key_auth, subject_token=subject_token + ) if v2_auth_value is not None: return v2_auth_value if server.has_token_exchange_config: diff --git a/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py b/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py index 1f34ed48aae..bc1c6c8eefe 100644 --- a/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py +++ b/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py @@ -57,6 +57,7 @@ from litellm.proxy.gateway.mcp.result import Error, Result from litellm.types.mcp import MCPAuth if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer _V2_ENV_FLAG = "LITELLM_USE_V2_MCP_RESOLVER" @@ -158,7 +159,30 @@ def _added_headers(auth: httpx.Auth) -> Dict[str, str]: } -async def resolve_v2_auth_value(server: MCPServer) -> Optional[Dict[str, str]]: +def _to_subject( + user_api_key_auth: Optional[UserAPIKeyAuth], subject_token: Optional[str] +) -> Subject: + """Map v1's authenticated principal onto the v2 Subject. + + Isolated so it can later swap to auth_v2's Principal. tenant_id/subject_id are empty for an + unauthenticated caller; the per-user arms (BYOK api_key, token_exchange, authorization_code) + must reject an empty subject_id rather than share one credential slot across callers. + """ + inbound = SecretStr(subject_token) if subject_token else None + if user_api_key_auth is None: + return Subject(tenant_id="", subject_id="", inbound_token=inbound) + return Subject( + tenant_id=user_api_key_auth.org_id or user_api_key_auth.team_id or "", + subject_id=user_api_key_auth.user_id or "", + inbound_token=inbound, + ) + + +async def resolve_v2_auth_value( + server: MCPServer, + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + subject_token: Optional[str] = None, +) -> 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 @@ -166,7 +190,7 @@ async def resolve_v2_auth_value(server: MCPServer) -> Optional[Dict[str, str]]: if spec is None: return None result = await _provider().resolve( - Subject(tenant_id="", subject_id="", inbound_token=None), spec + _to_subject(user_api_key_auth, subject_token), spec ) if isinstance(result, Error): verbose_logger.warning( 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 5231bab5573..31c20c556fb 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,9 +11,11 @@ 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, resolve_v2_auth_value, resolve_v2_aws_auth, ) +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp import MCPAuth, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -238,3 +240,37 @@ async def test_aws_sigv4_config_defaults_to_ambient(v2_on): cfg = _to_aws_sigv4_config(_aws_server()) assert cfg is not None assert isinstance(cfg.credentials, Ambient) + + +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") + assert subj.subject_id == "u1" + assert subj.tenant_id == "org1" # org preferred over team + assert subj.inbound_token is not None + assert subj.inbound_token.get_secret_value() == "inbound-jwt" + + +async def test_to_subject_anonymous_when_no_auth(): + subj = _to_subject(None, None) + assert subj.subject_id == "" + assert subj.tenant_id == "" + assert subj.inbound_token is None + + +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) + assert subj.tenant_id == "team-x" # no org -> team + assert ( + subj.subject_id == "" + ) # missing user -> empty; per-user arms fail closed on this + + +async def test_resolve_v2_auth_value_threads_identity_without_breaking_static(v2_on): + auth = UserAPIKeyAuth(token="sk-test", user_id="u1", org_id="org1") + server = _server(MCPAuth.api_key, "up-secret") + result = await resolve_v2_auth_value( + server, user_api_key_auth=auth, subject_token="jwt" + ) + assert result == {"X-API-Key": "up-secret"}