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(