mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
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:
parent
13b1dc18fb
commit
8e8a5d4f1e
3 changed files with 38 additions and 6 deletions
|
|
@ -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] = {}
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue