From 8e8a5d4f1eba1d74fdb50c9a230bf7645d3d1fb3 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 25 Jun 2026 12:59:33 -0700 Subject: [PATCH] feat(mcp): route oauth2 per-user (authorization_code) servers through the v2 resolver to_server_spec maps an oauth2 server to AuthorizationCodeConfig when it relies on per-user tokens (needs_user_oauth_token and not delegate_auth_to_upstream); client_credentials (M2M), delegated upstream OAuth, token exchange, and SigV4 still defer to v1. The manager injects V1PerUserTokenStore (resolving through v1's shared egress core) into the credential provider. The v2 path is live but still defers to a token v1 places in extra_headers; the cutover that makes v1 step aside lands next, alongside the unified challenge. --- .../mcp_server/mcp_server_manager.py | 7 ++++++- .../outbound_credentials/adapter.py | 19 +++++++++++++++---- .../outbound_credentials/test_adapter.py | 18 +++++++++++++++++- 3 files changed, 38 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b9e0379e445..2f8831be0ea 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -66,6 +66,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import to_server_spec, to_subject, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.v1_token_store import ( + V1PerUserTokenStore, +) from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, MCPMissingUserEnvVarsError, @@ -523,7 +526,9 @@ class MCPServerManager: return None def __init__(self, cred_provider: Optional[UpstreamCredentialProvider] = None): - self._cred_provider = cred_provider or UpstreamCredentialProvider() + self._cred_provider = cred_provider or UpstreamCredentialProvider( + oauth_token_store=V1PerUserTokenStore(self.get_mcp_server_by_id) + ) self.registry: Dict[str, MCPServer] = {} self.config_mcp_servers: Dict[str, MCPServer] = {} """ diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 91876df762f..f119cbd416b 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -20,6 +20,7 @@ from typing_extensions import assert_never from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, + AuthorizationCodeConfig, CredError, NoneConfig, ServerSpec, @@ -61,8 +62,9 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: Dispatches on the declared ``auth_type``. The match is exhaustive over ``MCPAuthType`` with an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is explicitly mapped or explicitly deferred, rather than silently falling through to v1. Live - modes: ``none`` and the static-header family (``api_key`` plus the Authorization schemes), - all shared-key; every other mode returns None and stays on v1. + modes: ``none``, the static-header family (``api_key`` plus the Authorization schemes, + all shared-key), and ``oauth2`` per-user tokens (``authorization_code``); client_credentials + (M2M), delegated/passthrough oauth2, token exchange, and SigV4 return None and stay on v1. """ if server.is_byok: return ( @@ -89,8 +91,17 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: return _shared_key_spec( server, resource, "Authorization", "Basic", encode=True ) - case MCPAuth.oauth2 | MCPAuth.oauth2_token_exchange | MCPAuth.aws_sigv4: - return None # OAuth grants and SigV4 are not migrated yet -> defer to v1 + case MCPAuth.oauth2: + if server.needs_user_oauth_token and not server.delegate_auth_to_upstream: + return ServerSpec( + server_id=server.server_id, + resource=resource, + config=AuthorizationCodeConfig(), + ) + # client_credentials (M2M) and delegate/passthrough oauth2 stay on v1 + return None + case MCPAuth.oauth2_token_exchange | MCPAuth.aws_sigv4: + return None # token exchange and SigV4 are not migrated yet -> defer to v1 assert_never(auth_type) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 594e9dcc969..e660462b9bc 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -18,6 +18,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ) from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, + AuthorizationCodeConfig, CredError, NoneConfig, SharedKey, @@ -71,12 +72,27 @@ def test_basic_scheme_base64_encodes_the_token(): assert spec.config.key_source.value.get_secret_value() == expected +@pytest.mark.parametrize( + "oauth2_flow", + [None, "authorization_code"], +) +def test_oauth2_user_token_maps_to_authorization_code(oauth2_flow): + # oauth2 without client_credentials is the per-user authorization_code mode. + spec = to_server_spec(_server(auth_type=MCPAuth.oauth2, oauth2_flow=oauth2_flow)) + assert spec is not None and isinstance(spec.config, AuthorizationCodeConfig) + + @pytest.mark.parametrize( "server", [ _server(auth_type=MCPAuth.api_key), # no token configured _server(auth_type=MCPAuth.bearer_token), # no token configured - _server(auth_type=MCPAuth.oauth2), + _server( + auth_type=MCPAuth.oauth2, oauth2_flow="client_credentials" + ), # M2M -> v1 + _server( + auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True + ), # delegated upstream OAuth -> v1 _server(auth_type=MCPAuth.oauth2_token_exchange), _server(auth_type=MCPAuth.aws_sigv4), _server(