Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py: 35%
97 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""The identity-provider side of the RFC 8693 token exchange on ``POST /token``: a native
2client that already holds a JWT from the customer's IdP trades it for the same proxy-API
3credential ``lite login`` stores, proven by the proxy's own JWT auth (signature, claims,
4and the user and team sync it performs), so no browser round trip is needed."""
6from __future__ import annotations
8from collections.abc import Awaitable, Callable, Mapping
9from dataclasses import dataclass
10from typing import Final, Literal, Protocol
12from fastapi import HTTPException, Request
13from typing_extensions import assert_never
15from litellm._logging import verbose_proxy_logger
16from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal
17from litellm.proxy._types import JWTAuthBuilderResult, ProxyException
18from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
19from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
21EXCHANGE_ROUTE: Final = "/token"
22REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth"
23SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = (
24 "the gateway could not verify subject_token because its identity provider or database is unavailable; retry"
25)
26SUBJECT_TOKEN_CHECK_FAULTED: Final = (
27 "the gateway could not verify subject_token because its database reported a fault that is not a transient "
28 "outage; retrying will not help until the gateway deployment is repaired"
29)
30GatewayOutage = Literal["retryable", "faulted"]
33@dataclass(frozen=True, slots=True)
34class TokenExchangePrerequisites:
35 """The deployment-level gates ``user_api_key_auth`` applies before it verifies any JWT
36 bearer, plus the JWT-to-virtual-key mapping it consults first: a gateway that maps
37 tokens authenticates a JWT as its mapped key, with that key's models and budget, or
38 refuses an unmapped one, and the exchange proves the token through ``auth_builder``
39 alone, so it would mint the user's own credential past that policy. Discovery and
40 registration advertise the exchange grant only when every gate holds, and an exchange
41 attempt is refused naming the first one that does not."""
43 jwt_auth_enabled: bool
44 has_database: bool
45 licensed: bool
46 maps_jwts_to_virtual_keys: bool
48 @property
49 def available(self) -> bool:
50 return self.jwt_auth_enabled and self.has_database and self.licensed and not self.maps_jwts_to_virtual_keys
52 def refusal(self) -> SubjectTokenRefusal | None:
53 if not self.jwt_auth_enabled:
54 return SubjectTokenRefusal(
55 error="unsupported_grant_type",
56 description="JWT auth is not enabled on this gateway, so it cannot exchange IdP tokens",
57 )
58 if not self.has_database:
59 return SubjectTokenRefusal(
60 error="unsupported_grant_type",
61 description="this gateway has no database, so it cannot exchange IdP tokens",
62 )
63 if not self.licensed:
64 return SubjectTokenRefusal(
65 error="unsupported_grant_type",
66 description="JWT auth is an enterprise only feature; no license is set",
67 )
68 if self.maps_jwts_to_virtual_keys:
69 return SubjectTokenRefusal(
70 error="unsupported_grant_type",
71 description="this gateway maps IdP tokens to virtual keys, which the exchange does not serve",
72 )
73 return None
76def read_token_exchange_prerequisites() -> TokenExchangePrerequisites:
77 from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call
78 general_settings,
79 jwt_handler,
80 premium_user,
81 prisma_client,
82 )
84 return TokenExchangePrerequisites(
85 jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True,
86 has_database=prisma_client is not None,
87 licensed=premium_user is True,
88 maps_jwts_to_virtual_keys=_maps_jwts_to_virtual_keys(jwt_handler),
89 )
92def _maps_jwts_to_virtual_keys(jwt_handler: JWTHandler) -> bool:
93 if not hasattr(jwt_handler, "litellm_jwtauth"): 93 ↛ 94line 93 didn't jump to line 94 because the condition on line 93 was never true
94 return False
95 return jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured()
98def token_exchange_available() -> bool:
99 return read_token_exchange_prerequisites().available
102class AuthorizeSubjectToken(Protocol):
103 """Injected JWT authorization ``(subject_token, request_headers)``: the proxy's
104 ``JWTAuthManager.auth_builder`` in production, which raises when the token is not
105 acceptable and otherwise names the user and team it resolved."""
107 def __call__( 107 ↛ exitline 107 didn't return from function '__call__' because
108 self, subject_token: str, request_headers: Mapping[str, str], /
109 ) -> Awaitable[JWTAuthBuilderResult]: ...
112async def exchange_idp_subject_token(subject_token: str, request: Request) -> SubjectIdentity | SubjectTokenRefusal:
113 from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call
114 general_settings,
115 jwt_handler,
116 prisma_client,
117 proxy_logging_obj,
118 user_api_key_cache,
119 )
121 async def authorize(token: str, request_headers: Mapping[str, str]) -> JWTAuthBuilderResult:
122 return await JWTAuthManager.auth_builder(
123 api_key=token,
124 jwt_handler=jwt_handler,
125 request_data={},
126 general_settings=general_settings,
127 route=EXCHANGE_ROUTE,
128 prisma_client=prisma_client,
129 user_api_key_cache=user_api_key_cache,
130 parent_otel_span=None,
131 proxy_logging_obj=proxy_logging_obj,
132 request_headers=request_headers,
133 request_method="POST",
134 )
136 return await identity_from_subject_token(
137 subject_token,
138 request_headers=request.headers,
139 prerequisites=read_token_exchange_prerequisites(),
140 is_jwt=jwt_handler.is_jwt,
141 authorize=authorize,
142 )
145async def identity_from_subject_token(
146 subject_token: str,
147 request_headers: Mapping[str, str],
148 prerequisites: TokenExchangePrerequisites,
149 is_jwt: Callable[[str], bool],
150 authorize: AuthorizeSubjectToken,
151) -> SubjectIdentity | SubjectTokenRefusal:
152 """Apply the same gates ``user_api_key_auth`` applies to a JWT bearer, then let the
153 proxy's JWT auth prove the token. A rejection comes back as ``invalid_request``, which
154 RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token, and a
155 check the gateway could not complete (the IdP's JWKS unreachable with no cached copy,
156 the auth database down) as ``temporarily_unavailable``, so the client retries instead
157 of treating a valid token as bad, worded by whether retrying can help. The reason stays
158 in the proxy log: this endpoint is public and JWT auth's own wording can name the JWKS
159 URL it fetched or quote the IdP's response."""
160 unmet: Final = prerequisites.refusal()
161 if unmet is not None:
162 return unmet
163 if not is_jwt(subject_token):
164 return SubjectTokenRefusal(error="invalid_request", description="subject_token is not a JWT")
165 try:
166 result: Final = await authorize(subject_token, request_headers)
167 except HTTPException as denied:
168 return _refusal_for(denied, denied.detail)
169 except ProxyException as denied:
170 return _refusal_for(denied, denied.message)
171 except Exception as denied: # noqa: BLE001 # auth_jwt raises a plain Exception on signature and claim failures
172 return _refusal_for(denied, denied)
173 user_id: Final = result["user_id"]
174 if user_id is None:
175 return SubjectTokenRefusal(error="invalid_request", description="subject_token names no user the gateway knows")
176 return SubjectIdentity(user_id=user_id, team_id=result["team_id"])
179def _refusal_for(denied: Exception, reason: object) -> SubjectTokenRefusal:
180 outage: Final = _gateway_could_not_verify(denied)
181 if outage is None:
182 verbose_proxy_logger.warning("token exchange refused a subject_token: %s", reason)
183 return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN)
184 verbose_proxy_logger.error("token exchange could not verify a subject_token, %s: %s", outage, reason)
185 return SubjectTokenRefusal(error="temporarily_unavailable", description=_check_unavailable_description(outage))
188def _check_unavailable_description(outage: GatewayOutage) -> str:
189 match outage:
190 case "retryable":
191 return SUBJECT_TOKEN_CHECK_UNAVAILABLE
192 case "faulted":
193 return SUBJECT_TOKEN_CHECK_FAULTED
194 case _:
195 assert_never(outage)
198def _gateway_could_not_verify(denied: Exception) -> GatewayOutage | None:
199 """A database fault anywhere in the chain or a 5xx from JWT auth (the IdP's JWKS
200 unreachable with no cached copy) is the gateway failing, not the token. A fault retrying cannot clear (a missing or
201 version-skewed query engine) is named as such, the way the mint path words it, so the
202 client is not told to wait on a deployment that needs repair."""
203 fault: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(denied)
204 if fault is not None:
205 return "faulted" if PrismaDBExceptionHandler.is_permanent_database_fault(fault) else "retryable"
206 return "retryable" if _is_server_error(denied) else None
209def _is_server_error(denied: Exception) -> bool:
210 match denied:
211 case HTTPException(status_code=status_code):
212 return status_code >= 500
213 case ProxyException(code=code):
214 return code.isdigit() and int(code) >= 500
215 case _:
216 return False