Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchanger.py: 27%
126 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""v2-native OBO token exchange: swap the caller's token for an upstream-bound one.
3The pure core of the ``token_exchange`` mode. Given the caller's inbound token and the server's
4``TokenExchangeConfig``, ``OboTokenExchanger.exchange`` POSTs the grant selected by ``config.profile``
5to the configured endpoint and returns the upstream-bound ``access_token`` as a typed ``OAuthToken``,
6or a typed ``CredError`` - never a raise (the HTTP edge is the injected ``ExchangeHttpPost``, whose
7adapter contains the I/O). Two profiles share this one engine: ``rfc8693`` (the RFC 8693 token-exchange
8grant) and ``entra_obo`` (Microsoft Entra On-Behalf-Of, which is the RFC 7523 ``jwt-bearer`` grant);
9only the request form differs, so the cache, single-flight, and TTL machinery are dialect-agnostic. The
10exchanged token is cached and single-flighted per ``(subject_token, tenant, config, server)`` so a
11repeated caller token skips the IdP round-trip and concurrent calls collapse to one exchange, reusing
12the shared in-process cache + coordinator foundation. A rotated caller token hashes to a new key and
13re-exchanges. Pure v2 apart from the shared RFC 6749 client-auth helper, which carries no v1 state.
15A missing/expired exchange is an error, never a fall-through to a weaker source (§1.5): the caller
16presenting no token is the resolver arm's 401, and an IdP that does not return a usable token is an
17``upstream_unavailable`` here.
18"""
20from __future__ import annotations
22import hashlib
23import time
24from collections.abc import Awaitable, Callable
25from typing import Final, Literal, Protocol
27from typing_extensions import assert_never
29from litellm._logging import verbose_logger
30from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import (
31 build_token_endpoint_client_auth,
32)
33from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
34 InMemoryTokenCacheBackend,
35 InProcessRefreshCoordinator,
36 OAuthToken,
37 RefreshCoordinator,
38 TokenCacheBackend,
39)
40from litellm.proxy._experimental.mcp_server.outbound_credentials.result import (
41 Error,
42 Ok,
43 Result,
44)
45from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
46 CredError,
47 ServerSpec,
48 TokenExchangeConfig,
49)
51# A token with no declared expiry is cached for this long; one with an expiry is cached until then
52# minus the skew buffer, floored at the minimum. Values mirror v1's MCP_OAUTH2_* constants; the
53# composition root injects the configured ones.
54_DEFAULT_TTL_SECONDS: Final = 3600.0
55_MIN_TTL_SECONDS: Final = 10.0
56_EXPIRY_BUFFER_SECONDS: Final = 60.0
58_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"
59# Microsoft Entra On-Behalf-Of speaks the RFC 7523 jwt-bearer grant, not RFC 8693, and gates delegation
60# behind ``requested_token_use=on_behalf_of`` (a Microsoft extension present in neither RFC).
61_JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer"
62_REQUESTED_TOKEN_USE_OBO: Final = "on_behalf_of"
64# RFC 8693 3 token-type URNs that are not usable as an upstream Bearer access token. token_type
65# already rejects the common non-access case (N_A); this catches a malformed STS that mints one of
66# these but still labels it Bearer. An access_token / jwt / absent / unknown type is accepted (lenient).
67_NON_ACCESS_ISSUED_TOKEN_TYPES: Final = frozenset(
68 {
69 "urn:ietf:params:oauth:token-type:refresh_token",
70 "urn:ietf:params:oauth:token-type:id_token",
71 "urn:ietf:params:oauth:token-type:saml1",
72 "urn:ietf:params:oauth:token-type:saml2",
73 }
74)
76# The IdP returns an opaque JSON object; the post adapter hands it over untyped and the exchanger
77# validates each field, so no Any leaks past this seam (None == any transport/HTTP failure). The
78# second dict is the form body; the third is the client-auth headers (HTTP Basic for
79# client_secret_basic, empty for client_secret_post).
80ExchangeHttpPost = Callable[[str, "dict[str, str]", "dict[str, str]"], Awaitable["dict[str, object] | None"]]
83class SubjectTokenRejected(Exception):
84 """The IdP refused to exchange the subject token (an RFC 8693 4xx, e.g. ``invalid_grant``).
86 Distinct from a transport / IdP-availability failure, which the post adapter maps to ``None`` ->
87 ``upstream_unavailable`` -> 503 (retryable). A rejected subject is the caller's problem, not the
88 gateway's, so the arm surfaces it as a non-retryable 401 (the OBO challenge) instead.
89 ``claims`` is the IdP's step-up challenge blob (Entra Conditional Access / CAE) from the
90 rejection body; it threads into the 401 challenge so the client can satisfy the step-up and
91 retry. The ``error_description`` is never carried (it can leak IdP internals).
92 """
94 def __init__(self, detail: str, *, claims: str | None = None) -> None:
95 super().__init__(detail)
96 self.claims = claims
99class TokenExchangeClientError(Exception):
100 """The IdP rejected the exchange for a reason that is the gateway's fault, not the caller's.
102 RFC 6749 5.2 codes such as ``invalid_client`` (the gateway's own STS credentials are wrong),
103 ``unauthorized_client`` / ``unsupported_grant_type`` (the gateway is not permitted to exchange),
104 ``invalid_target`` / ``invalid_scope`` (the gateway's audience/scope config for this server is
105 wrong). The caller cannot fix these by re-authenticating, so the arm surfaces them as a 500
106 (``misconfigured``), not the 401 OBO challenge. The IdP ``error_description`` is never carried.
107 """
110class TokenExchanger(Protocol):
111 """Exchanges a caller token for an upstream-bound one, per the server's token_exchange config."""
113 async def exchange( 113 ↛ exitline 113 didn't return from function 'exchange' because
114 self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig, *, tenant_id: str = ""
115 ) -> Result[OAuthToken, CredError]: ...
117 async def invalidate( 117 ↛ exitline 117 didn't return from function 'invalidate' because
118 self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig, *, tenant_id: str = ""
119 ) -> None: ...
122def _cache_key(subject_token: str, tenant_id: str, config: TokenExchangeConfig) -> str:
123 """Bind the cache entry to the caller token, the tenant, AND the exchange config that minted it.
125 A rotated caller token, a different tenant, profile, endpoint, audience, scope, client_id, secret,
126 auth method, or subject_token_type all change the key, so two tenants behind the same opaque token
127 never share an entry and a config change (including a profile flip that alters the wire form)
128 forces a fresh exchange instead of serving a token minted for the old config until TTL. Everything
129 is hashed, so no secret is held in the key.
130 """
131 secret: Final = config.client_secret.get_secret_value() if config.client_secret else ""
132 material: Final = "\x00".join(
133 (
134 subject_token,
135 tenant_id,
136 config.profile,
137 config.token_exchange_endpoint or "",
138 config.audience or "",
139 config.subject_token_type,
140 config.client_id or "",
141 secret,
142 config.token_endpoint_auth_method or "",
143 " ".join(config.scopes),
144 )
145 )
146 return hashlib.sha256(material.encode()).hexdigest()
149def _parse_expires_in(raw: object) -> int | None:
150 if isinstance(raw, bool):
151 return None
152 if isinstance(raw, (int, float)):
153 return int(raw)
154 if isinstance(raw, str):
155 try:
156 return int(float(raw))
157 except ValueError:
158 return None
159 return None
162def _rfc8693_form(
163 *,
164 subject_token: str,
165 subject_token_type: str,
166 audience: str | None,
167 scopes: tuple[str, ...],
168) -> dict[str, str]:
169 return {
170 "grant_type": _GRANT_TYPE,
171 "subject_token": subject_token,
172 "subject_token_type": subject_token_type,
173 **({"audience": audience} if audience else {}),
174 **({"scope": " ".join(scopes)} if scopes else {}),
175 }
178def _entra_obo_form(
179 *,
180 subject_token: str,
181 scopes: tuple[str, ...],
182) -> dict[str, str]:
183 # Microsoft Entra On-Behalf-Of (RFC 7523 jwt-bearer, not RFC 8693): the caller's inbound access
184 # token rides as ``assertion`` (its ``aud`` must be this gateway's ``client_id``); the target
185 # resource is carried in ``scope`` (e.g. api://<app-id>/.default), since Entra has no audience
186 # parameter and ignores subject_token_type; ``requested_token_use=on_behalf_of`` is the Microsoft
187 # extension that turns the jwt-bearer grant into a delegation. ``scope`` is required, and the
188 # exchange precondition rejects an empty one, so it is always present here. Client authentication
189 # (client_id/client_secret via post, or Basic) is layered on by the caller through
190 # build_token_endpoint_client_auth, so it is not built into the form here.
191 return {
192 "grant_type": _JWT_BEARER_GRANT_TYPE,
193 "assertion": subject_token,
194 "scope": " ".join(scopes),
195 "requested_token_use": _REQUESTED_TOKEN_USE_OBO,
196 }
199def _build_exchange_form(
200 *,
201 profile: Literal["rfc8693", "entra_obo"],
202 subject_token: str,
203 subject_token_type: str,
204 audience: str | None,
205 scopes: tuple[str, ...],
206) -> dict[str, str]:
207 match profile:
208 case "rfc8693":
209 return _rfc8693_form(
210 subject_token=subject_token,
211 subject_token_type=subject_token_type,
212 audience=audience,
213 scopes=scopes,
214 )
215 case "entra_obo":
216 return _entra_obo_form(
217 subject_token=subject_token,
218 scopes=scopes,
219 )
220 assert_never(profile)
223class OboTokenExchanger:
224 """``TokenExchanger`` that runs the profile's OBO grant once per caller token, then caches the result.
226 The HTTP post is injected (``None`` on any IdP failure, mirroring v1: a failed exchange is a miss,
227 not a 500). The cache and single-flight coordinator default to the in-process foundation; a
228 deployment with no shared state needs nothing more (v1's exchanged-token cache is per-process too).
229 The clock is injected so TTL/expiry is deterministic in tests.
230 """
232 def __init__(
233 self,
234 http_post: ExchangeHttpPost,
235 *,
236 cache: TokenCacheBackend | None = None,
237 coordinator: RefreshCoordinator | None = None,
238 clock: Callable[[], float] = time.time,
239 default_ttl_seconds: float = _DEFAULT_TTL_SECONDS,
240 min_ttl_seconds: float = _MIN_TTL_SECONDS,
241 expiry_buffer_seconds: float = _EXPIRY_BUFFER_SECONDS,
242 ) -> None:
243 self._http_post = http_post
244 self._cache: TokenCacheBackend = cache or InMemoryTokenCacheBackend(clock=clock)
245 self._coordinator: RefreshCoordinator = coordinator or InProcessRefreshCoordinator()
246 self._clock = clock
247 self._default_ttl_seconds = default_ttl_seconds
248 self._min_ttl_seconds = min_ttl_seconds
249 self._expiry_buffer_seconds = expiry_buffer_seconds
251 async def exchange(
252 self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig, *, tenant_id: str = ""
253 ) -> Result[OAuthToken, CredError]:
254 endpoint: Final = config.token_exchange_endpoint
255 client_id: Final = config.client_id
256 client_secret: Final = config.client_secret
257 if not endpoint:
258 # No endpoint configured and none discoverable: fail closed (412) rather than guess an IdP
259 # or fall back to a weaker source. The caller's token is never sent anywhere.
260 return Error(
261 CredError.of_precondition_required("token exchange endpoint is not configured for this server")
262 )
263 if not client_id or client_secret is None:
264 return Error(CredError.of_misconfigured("token_exchange requires client_id and client_secret"))
265 if config.profile == "entra_obo" and not config.scopes:
266 # Entra carries the target resource in ``scope`` (api://<app-id>/.default); with no scope the
267 # IdP cannot resolve an audience, so fail closed as misconfigured rather than POST a form the
268 # IdP will reject.
269 return Error(
270 CredError.of_misconfigured("entra_obo token exchange requires a scope (e.g. api://<app-id>/.default)")
271 )
273 cache_key: Final = _cache_key(subject_token, tenant_id, config)
274 server_id: Final = server.server_id
275 cached: Final = await self._cache.get(cache_key, server_id)
276 if cached is not None:
277 verbose_logger.debug("MCP token exchange cache hit for server %s", server_id)
278 return Ok(cached)
280 client_auth: Final = build_token_endpoint_client_auth(
281 auth_method=config.token_endpoint_auth_method,
282 client_id=client_id,
283 client_secret=client_secret.get_secret_value(),
284 )
285 form: Final = {
286 **_build_exchange_form(
287 profile=config.profile,
288 subject_token=subject_token,
289 subject_token_type=config.subject_token_type,
290 audience=config.audience,
291 scopes=config.scopes,
292 ),
293 **client_auth.body,
294 }
296 async def run_exchange() -> OAuthToken | None:
297 fresh: Final = await self._cache.get(cache_key, server_id)
298 if fresh is not None:
299 return fresh
300 verbose_logger.debug(
301 "Exchanging token for MCP server %s at %s (audience=%s)", server_id, endpoint, config.audience
302 )
303 body: Final = await self._http_post(endpoint, form, client_auth.headers)
304 if body is None:
305 return None
306 token: Final = self._token_from_body(body)
307 if token is None:
308 return None
309 await self._cache.set(cache_key, server_id, token, self._ttl_seconds(token))
310 verbose_logger.info("Token exchange succeeded for MCP server %s", server_id)
311 return token
313 async def reread() -> OAuthToken | None:
314 return await self._cache.get(cache_key, server_id)
316 try:
317 token: Final = await self._coordinator.run(cache_key, server_id, refresh=run_exchange, reread=reread)
318 except SubjectTokenRejected as rejected:
319 # The IdP rejected the subject token (4xx). This is non-retryable: the caller must
320 # re-authenticate with the IdP, so it surfaces as a 401 (the OBO challenge), not a 503.
321 # A step-up rejection (Entra Conditional Access) carries the claims blob through so the
322 # edge's challenge tells the client how to satisfy it.
323 return Error(
324 CredError.of_unauthorized(
325 str(rejected) or "subject token rejected by the IdP",
326 claims=rejected.claims,
327 )
328 )
329 except TokenExchangeClientError:
330 # RFC 6749 5.2 gateway-fault code (invalid_client / invalid_target / ...): the caller can't
331 # fix it by re-authenticating, so surface a 500 rather than the OBO 401 challenge. The
332 # specific code is logged at the edge; the user-facing summary stays generic.
333 return Error(
334 CredError.of_misconfigured(
335 "token exchange configuration error: the gateway's credentials, audience, or scope "
336 "for this server were not accepted by the IdP"
337 )
338 )
339 if token is None:
340 return Error(CredError.of_upstream_unavailable("token exchange did not return a usable access token"))
341 return Ok(token)
343 async def invalidate(
344 self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig, *, tenant_id: str = ""
345 ) -> None:
346 """Drop the cached exchanged token so the next call re-exchanges (e.g. after an upstream 401)."""
347 await self._cache.delete(_cache_key(subject_token, tenant_id, config), server.server_id)
349 def _token_from_body(self, body: dict[str, object]) -> OAuthToken | None:
350 access_token: Final = body.get("access_token")
351 if not isinstance(access_token, str) or not access_token:
352 return None
353 # token_type is forwarded downstream as Bearer, so a present-but-non-Bearer type (e.g. N_A)
354 # must fail closed rather than be minted as a bogus Bearer; an absent type defaults to Bearer.
355 token_type: Final = body.get("token_type")
356 if isinstance(token_type, str) and token_type.strip().lower() != "bearer":
357 verbose_logger.warning(
358 "MCP token exchange returned unusable token_type %r; refusing to forward it as Bearer", token_type
359 )
360 return None
361 # issued_token_type says what representation was minted; reject a clearly-non-access type
362 # (refresh/id/saml) even if token_type claimed Bearer. access_token / jwt / absent / unknown pass.
363 issued_token_type: Final = body.get("issued_token_type")
364 if isinstance(issued_token_type, str) and issued_token_type in _NON_ACCESS_ISSUED_TOKEN_TYPES:
365 return None
366 expires_in: Final = _parse_expires_in(body.get("expires_in"))
367 expires_at: Final = self._clock() + expires_in if expires_in is not None else None
368 return OAuthToken(access_token=access_token, expires_at=expires_at)
370 def _ttl_seconds(self, token: OAuthToken) -> float:
371 if token.expires_at is None:
372 return self._default_ttl_seconds
373 lifetime: Final = max(0.0, token.expires_at - self._clock())
374 # Floor at min_ttl, but never cache past the token's own expiry: a token whose remaining
375 # lifetime is below the buffer (or even below min_ttl) must not be served stale upstream.
376 return min(max(lifetime - self._expiry_buffer_seconds, self._min_ttl_seconds), lifetime)