fix(proxy): keep the token exchange off gateways that map JWTs to virtual keys

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
 

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 []

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