diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py index cdefaf76d49..7a453e85cce 100644 --- a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -14,7 +14,7 @@ from fastapi import HTTPException, Request from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy._types import JWTAuthBuilderResult, ProxyException -from litellm.proxy.auth.handle_jwt import JWTAuthManager +from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler EXCHANGE_ROUTE: Final = "/token" REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth" @@ -23,16 +23,21 @@ REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT @dataclass(frozen=True, slots=True) class TokenExchangePrerequisites: """The deployment-level gates ``user_api_key_auth`` applies before it verifies any JWT - bearer. Discovery and registration advertise the exchange grant only when every one of - them holds, and an exchange attempt is refused naming the first one that does not.""" + bearer, plus the JWT-to-virtual-key mapping it consults first: a gateway that maps + tokens authenticates a JWT as its mapped key, with that key's models and budget, or + refuses an unmapped one, and the exchange proves the token through ``auth_builder`` + alone, so it would mint the user's own credential past that policy. Discovery and + registration advertise the exchange grant only when every gate holds, and an exchange + attempt is refused naming the first one that does not.""" jwt_auth_enabled: bool has_database: bool licensed: bool + maps_jwts_to_virtual_keys: bool @property def available(self) -> bool: - return self.jwt_auth_enabled and self.has_database and self.licensed + return self.jwt_auth_enabled and self.has_database and self.licensed and not self.maps_jwts_to_virtual_keys def refusal(self) -> SubjectTokenRefusal | None: if not self.jwt_auth_enabled: @@ -50,12 +55,18 @@ class TokenExchangePrerequisites: error="unsupported_grant_type", description="JWT auth is an enterprise only feature; no license is set", ) + if self.maps_jwts_to_virtual_keys: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway maps IdP tokens to virtual keys, which the exchange does not serve", + ) return None def read_token_exchange_prerequisites() -> TokenExchangePrerequisites: from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call general_settings, + jwt_handler, premium_user, prisma_client, ) @@ -64,9 +75,16 @@ def read_token_exchange_prerequisites() -> TokenExchangePrerequisites: jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True, has_database=prisma_client is not None, licensed=premium_user is True, + maps_jwts_to_virtual_keys=_maps_jwts_to_virtual_keys(jwt_handler), ) +def _maps_jwts_to_virtual_keys(jwt_handler: JWTHandler) -> bool: + if not hasattr(jwt_handler, "litellm_jwtauth"): + return False + return jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured() + + def token_exchange_available() -> bool: return read_token_exchange_prerequisites().available diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 6965b3b4ebe..d7666f5e694 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11111,13 +11111,31 @@ def test_native_client_login_walks_discovery_consent_token_refresh_and_revoke(mo assert stranger.json()["error"] == "invalid_client" -@pytest.mark.parametrize("exchange_servable", [True, False]) -def test_discovery_advertises_the_exchange_grant_only_where_the_gateway_can_serve_it(monkeypatch, exchange_servable): +@pytest.mark.parametrize( + "jwt_auth_enabled, virtual_key_claim_field, exchange_servable", + [(True, None, True), (False, None, False), (True, "client_id", False)], + ids=["jwt auth on", "jwt auth off", "jwts mapped to virtual keys"], +) +def test_discovery_advertises_the_exchange_grant_only_where_the_gateway_can_serve_it( + monkeypatch, jwt_auth_enabled, virtual_key_claim_field, exchange_servable +): """Every document a native client reads before it picks a grant (the versioned contract, the aggregate authorization-server metadata, and the registration response) lists the RFC 8693 - exchange exactly when the running proxy can serve it: JWT auth on, a database, and a license.""" + exchange exactly when the running proxy can serve it: JWT auth on, a database, a license, and + no JWT-to-virtual-key mapping, since the exchange would mint past the mapped key's policy.""" + from litellm.caching.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.handle_jwt import JWTHandler + client, _session_cookie, _minted = _native_client_app(monkeypatch) - monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": exchange_servable}) + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(virtual_key_claim_field=virtual_key_claim_field), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", handler) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": jwt_auth_enabled}) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", object()) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) exchange_grant = ["urn:ietf:params:oauth:grant-type:token-exchange"] if exchange_servable else [] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py index d1b049dddd5..e12c8823f99 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py @@ -3,6 +3,7 @@ import logging import pytest from fastapi import HTTPException +from litellm.caching.caching import DualCache from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( REJECTED_SUBJECT_TOKEN, @@ -10,12 +11,17 @@ from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( identity_from_subject_token, token_exchange_available, ) -from litellm.proxy._types import ProxyException +from litellm.proxy._types import JWTIssuerConfig, LiteLLM_JWTAuth, ProxyException from litellm.proxy.auth.handle_jwt import JWTHandler IDP_JWT = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature" REQUEST_HEADERS = {"x-litellm-team-id": "team-b", "user-agent": "lite/0.1"} -EVERY_GATE_HOLDS = {"jwt_auth_enabled": True, "has_database": True, "licensed": True} +EVERY_GATE_HOLDS = { + "jwt_auth_enabled": True, + "has_database": True, + "licensed": True, + "maps_jwts_to_virtual_keys": False, +} JWKS_URL = "https://idp.example.com/.well-known/jwks.json" @@ -82,6 +88,7 @@ async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity(): ({"jwt_auth_enabled": False}, IDP_JWT, "unsupported_grant_type", "JWT auth is not enabled"), ({"has_database": False}, IDP_JWT, "unsupported_grant_type", "no database"), ({"licensed": False}, IDP_JWT, "unsupported_grant_type", "enterprise"), + ({"maps_jwts_to_virtual_keys": True}, IDP_JWT, "unsupported_grant_type", "virtual keys"), ({}, "sk-litellm-virtual-key", "invalid_request", "not a JWT"), ], ) @@ -96,28 +103,53 @@ async def test_the_gates_user_api_key_auth_applies_refuse_before_any_verificatio assert authorizer.calls == [] -@pytest.mark.parametrize("unmet", [{}, {"jwt_auth_enabled": False}, {"has_database": False}, {"licensed": False}]) +@pytest.mark.parametrize( + "unmet", + [ + {}, + {"jwt_auth_enabled": False}, + {"has_database": False}, + {"licensed": False}, + {"maps_jwts_to_virtual_keys": True}, + ], +) def test_the_grant_is_available_exactly_when_every_gate_holds(unmet): prerequisites = TokenExchangePrerequisites(**{**EVERY_GATE_HOLDS, **unmet}) assert prerequisites.available is (unmet == {}) assert (prerequisites.refusal() is None) is prerequisites.available +MAPPED_ISSUER = JWTIssuerConfig( + issuer="https://idp.example.test", audience="litellm-gateway", virtual_key_claim_field="client_id" +) + + +def _running_jwt_handler(litellm_jwtauth): + handler = JWTHandler() + if litellm_jwtauth is not None: + handler.update_environment(prisma_client=None, user_api_key_cache=DualCache(), litellm_jwtauth=litellm_jwtauth) + return handler + + @pytest.mark.parametrize( - "general_settings, prisma_client, premium_user, expected", + "general_settings, prisma_client, premium_user, litellm_jwtauth, expected", [ - ({"enable_jwt_auth": True}, object(), True, True), - ({}, object(), True, False), - ({"enable_jwt_auth": True}, None, True, False), - ({"enable_jwt_auth": True}, object(), False, False), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(), True), + ({"enable_jwt_auth": True}, object(), True, None, True), + ({}, object(), True, LiteLLM_JWTAuth(), False), + ({"enable_jwt_auth": True}, None, True, LiteLLM_JWTAuth(), False), + ({"enable_jwt_auth": True}, object(), False, LiteLLM_JWTAuth(), False), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(virtual_key_claim_field="client_id"), False), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(issuers=[MAPPED_ISSUER]), False), ], ) def test_availability_is_read_from_the_running_proxy( - monkeypatch, general_settings, prisma_client, premium_user, expected + monkeypatch, general_settings, prisma_client, premium_user, litellm_jwtauth, expected ): monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", premium_user) + monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", _running_jwt_handler(litellm_jwtauth)) assert token_exchange_available() is expected