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.
This commit is contained in:
Tin Chi Lo 2026-06-25 12:59:33 -07:00
parent 13b1dc18fb
commit 8e8a5d4f1e
3 changed files with 38 additions and 6 deletions

View file

@ -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] = {}
"""

View file

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

View file

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