Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py: 28%
192 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""The one credential resolver: dispatch on the declared mode, fail closed.
3`resolve_credentials` selects exactly one arm off the server's typed `config` and either
4produces an `httpx2.Auth` or returns a typed `CredError`. The `match` is over the `AuthConfig`
5variant, so each arm receives its own fully-typed config with no field-presence inference and
6no precedence cascade. It is wildcard-free with an `assert_never` tail, so adding a mode without
7an arm fails the type gate (basedpyright `reportMatchNotExhaustive`); a bypassed gate fails loudly
8at runtime instead of returning `None`.
10`none`, `api_key` (shared-key source), and `passthrough` (forwards the caller's own inbound token)
11are live, as is `authorization_code`, which reads the user's token from the injected
12`OAuthTokenStore`, `token_exchange`, which swaps the caller's inbound token through the injected
13`TokenExchanger`, `client_credentials`, which mints and caches the gateway's M2M token through the
14injected `ClientCredentialsTokenSource`, and `id_jag`, which runs the two-leg identity-assertion
15grant against a subject token taken from the request or from the injected `SSOAssertionStore`. The
16remaining arms are `not_implemented` stubs that each land in a follow-up PR with their seam. Pure
17v2: no imports from v1.
18"""
20from __future__ import annotations
22import hashlib
23from datetime import datetime, timezone
24from functools import partial
25from typing import Final
27import httpx
28import httpx2
29from typing_extensions import assert_never
31from litellm._logging import verbose_proxy_logger
32from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import (
33 ClientCredentialsBearerAuth,
34 ClientCredentialsTokenSource,
35)
36from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import (
37 NoOpAuth,
38 StaticHeaderAuth,
39)
40from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
41 OAuthToken,
42 OAuthTokenStore,
43 TokenStoreUnavailable,
44)
45from litellm.proxy._experimental.mcp_server.outbound_credentials.result import (
46 Error,
47 Ok,
48 Result,
49)
50from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import (
51 default_sso_assertion_store,
52)
53from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
54 AssertionStoreUnavailable,
55 SSOAssertionStore,
56 assertion_expired,
57)
58from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import (
59 ExchangedToken,
60 ExchangedTokenCache,
61 TokenEndpointClient,
62)
63from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import (
64 TokenExchanger,
65)
66from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
67 ApiKeyConfig,
68 AuthorizationCodeConfig,
69 AuthResolution,
70 AuthSpecKind,
71 AwsSigV4Config,
72 Byok,
73 ClientAuth,
74 ClientCredentialsConfig,
75 ClientSecretAuth,
76 CredError,
77 IdJagConfig,
78 NoneConfig,
79 PassthroughConfig,
80 PrivateKeyJwtAuth,
81 ResolvedCredential,
82 ServerSpec,
83 SharedKey,
84 Subject,
85 TokenExchangeConfig,
86)
88_TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"
89_JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer"
90_ID_JAG_REQUESTED_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:id-jag"
93class _NullOAuthTokenStore:
94 """Fail-closed default: with no token store wired, every user reads as not authorized."""
96 async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
97 return None
100class _NullTokenExchanger:
101 """Fail-closed default: with no exchanger wired, token_exchange cannot produce a credential."""
103 async def exchange(
104 self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig, *, tenant_id: str = ""
105 ) -> Result[OAuthToken, CredError]:
106 return Error(CredError.of_misconfigured("token exchange collaborator not wired"))
108 async def invalidate(
109 self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig, *, tenant_id: str = ""
110 ) -> None:
111 return None
114class UpstreamCredentialProvider:
115 """Produces the one `httpx.Auth` for a `(subject, upstream)` pair, per declared mode.
117 Collaborators (the per-mode credential stores and token fetchers) are injected as each arm is
118 built; the live `none` and `api_key`-shared arms read from the config and need none, while
119 `authorization_code` reads the user's token from the injected `OAuthTokenStore` and
120 `token_exchange` swaps the caller's token through the injected `TokenExchanger`.
121 """
123 def __init__(
124 self,
125 oauth_token_store: OAuthTokenStore | None = None,
126 token_exchanger: TokenExchanger | None = None,
127 token_endpoint: TokenEndpointClient | None = None,
128 exchanged_tokens: ExchangedTokenCache | None = None,
129 client_credentials_source: ClientCredentialsTokenSource | None = None,
130 sso_assertion_store: SSOAssertionStore | None = None,
131 ) -> None:
132 self._oauth_token_store: OAuthTokenStore = oauth_token_store or _NullOAuthTokenStore()
133 self._token_exchanger: TokenExchanger = token_exchanger or _NullTokenExchanger()
134 self._token_endpoint: TokenEndpointClient = token_endpoint or TokenEndpointClient()
135 self._exchanged_tokens: ExchangedTokenCache = exchanged_tokens or ExchangedTokenCache()
136 self._client_credentials_source = client_credentials_source or ClientCredentialsTokenSource()
137 self._sso_assertion_store: SSOAssertionStore = sso_assertion_store or default_sso_assertion_store()
139 async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx2.Auth, CredError]:
140 match server.config:
141 case NoneConfig():
142 return self._none(server)
143 case ApiKeyConfig() as config: 143 ↛ 144line 143 didn't jump to line 144 because the pattern on line 143 never matched
144 return self._api_key(config)
145 case PassthroughConfig(): 145 ↛ 147line 145 didn't jump to line 147 because the pattern on line 145 always matched
146 return self._passthrough(subject)
147 case ClientCredentialsConfig() as config:
148 return await self._client_credentials(server.server_id, config)
149 case TokenExchangeConfig() as config:
150 return await self._token_exchange(subject, server, config)
151 case IdJagConfig() as config:
152 return await self._id_jag(subject, server, config)
153 case AuthorizationCodeConfig() as config:
154 return await self._authorization_code(subject, server, config)
155 case AwsSigV4Config():
156 return _not_implemented(AuthSpecKind.aws_sigv4)
157 assert_never(server.config)
159 def _none(self, server: ServerSpec) -> Result[httpx2.Auth, CredError]:
160 try:
161 resource: Final = httpx.URL(server.resource)
162 except httpx.InvalidURL:
163 return Ok(NoOpAuth())
164 if resource.userinfo: 164 ↛ 165line 164 didn't jump to line 165 because the condition on line 164 was never true
165 return Error(CredError.of_url_credentials_not_allowed())
166 return Ok(NoOpAuth())
168 async def has_user_token(self, subject: Subject, server: ServerSpec) -> bool:
169 """Whether a usable per-user token exists for this server (the preemptive 401's check).
171 Reads from the same per-user store as the ``authorization_code`` arm, so the discovery
172 challenge and the egress agree on whether the user is authorized. Returns a typed ``bool``
173 (no ``httpx2.Auth``), unlike ``resolve_credentials``. A non-per-user mode has no token in the
174 store, so it reads as False without a per-mode branch here.
175 """
176 return await self._authz_token(subject, server) is not None
178 def _passthrough(self, subject: Subject) -> Result[httpx2.Auth, CredError]:
179 """Forward the caller's own upstream credential verbatim; the gateway mints nothing.
181 The inbound token is the caller's already-disambiguated ``Authorization`` (never the LiteLLM
182 admission credential; the edge adapter drops that before building the ``Subject``). When it is
183 absent the request is sent unauthenticated so the upstream's own 401 surfaces, rather than the
184 gateway challenging on the upstream's behalf.
185 """
186 if subject.inbound_token is None: 186 ↛ 188line 186 didn't jump to line 188 because the condition on line 186 was always true
187 return Ok(NoOpAuth())
188 return Ok(StaticHeaderAuth(subject.inbound_token.get_secret_value(), header_name="Authorization"))
190 def _api_key(self, config: ApiKeyConfig) -> Result[httpx2.Auth, CredError]:
191 match config.key_source:
192 case SharedKey() as source:
193 header_name, header_value = config.header(source.value.get_secret_value())
194 return Ok(StaticHeaderAuth(header_value, header_name=header_name))
195 case Byok():
196 # Per-user key pulled from the credential store; lands with that seam.
197 return Error(CredError.of_not_implemented("api_key BYOK source not implemented yet"))
198 assert_never(config.key_source)
200 async def _id_jag(
201 self, subject: Subject, server: ServerSpec, config: IdJagConfig
202 ) -> Result[httpx2.Auth, CredError]:
203 match await self._id_jag_subject_token(subject):
204 case Error(err):
205 return Error(err)
206 case Ok(subject_token):
207 return await self._id_jag_exchange(subject, subject_token, server, config)
209 async def _id_jag_subject_token(self, subject: Subject) -> Result[str, CredError]:
210 """The identity token ID-JAG leg 1 asserts, from the request or from the SSO login it was captured at.
212 A caller that presents its own IdP identity token wins: that is the strongest available
213 assertion of who is calling. Otherwise the subject is the assertion captured for this user
214 at LiteLLM SSO login, which is what lets an agent holding a brokered LiteLLM credential
215 reach an upstream as the user it was issued for. The user is always taken from the
216 authenticated principal, never from a caller-supplied field, so no caller can select whose
217 identity is asserted upstream.
219 Every miss is ``precondition_required`` (412) rather than a fall-through to a weaker
220 credential: ID-JAG exists to assert a specific user, so a missing subject has no safe
221 substitute. A store outage is the one exception: it is ``upstream_unavailable`` (503), not
222 412, because the user has nothing to fix by signing in again, and it is a value rather than
223 a raised error so a DB blip cannot 500 the egress or the upstream-401 retry.
224 """
225 if subject.inbound_token is not None:
226 return Ok(subject.inbound_token.get_secret_value())
227 if not subject.subject_id:
228 return Error(
229 CredError.of_precondition_required(
230 "ID-JAG requires an identified caller; this request carries neither an "
231 "identity token nor a resolved LiteLLM user."
232 )
233 )
234 try:
235 assertion: Final = await self._sso_assertion_store.fetch(subject.subject_id)
236 except AssertionStoreUnavailable as exc:
237 # The driver's message can name hosts, schemas or connection details, and this summary
238 # is returned to the caller verbatim as a 503 body. Operators get it from the log.
239 verbose_proxy_logger.warning(
240 "ID-JAG: the IdP identity assertion store is unreachable for user_id=%s: %s",
241 subject.subject_id,
242 exc,
243 )
244 return Error(
245 CredError.of_upstream_unavailable(
246 "The IdP identity assertion store is unreachable, so ID-JAG cannot resolve a subject."
247 )
248 )
249 if assertion is None:
250 return Error(
251 CredError.of_precondition_required(
252 "ID-JAG requires an IdP identity assertion for this user and none is stored. "
253 "Sign in through LiteLLM SSO so the gateway captures one."
254 )
255 )
256 if assertion_expired(assertion, datetime.now(timezone.utc)):
257 return Error(
258 CredError.of_precondition_required(
259 "The stored IdP identity assertion for this user has expired. Sign in through "
260 "LiteLLM SSO again to capture a current one."
261 )
262 )
263 return Ok(assertion.id_token.get_secret_value())
265 async def _id_jag_exchange(
266 self, subject: Subject, token: str, server: ServerSpec, config: IdJagConfig
267 ) -> Result[httpx2.Auth, CredError]:
268 slot: Final = _id_jag_slot_key(subject, server)
269 fingerprint: Final = _id_jag_fingerprint(token, server.server_id, config)
271 async def _exchange() -> Result[ExchangedToken, CredError]:
272 leg1_params: Final = {
273 "grant_type": _TOKEN_EXCHANGE_GRANT_TYPE,
274 "requested_token_type": _ID_JAG_REQUESTED_TOKEN_TYPE,
275 "subject_token": token,
276 "subject_token_type": config.subject_token_type,
277 **({"audience": config.audience} if config.audience else {}),
278 **({"resource": config.resource} if config.resource else {}),
279 **({"scope": " ".join(config.scopes)} if config.scopes else {}),
280 }
281 match await self._token_endpoint.fetch(
282 config.org_token_endpoint,
283 config.client_id,
284 leg1_params,
285 config.client_auth,
286 ):
287 case Error(err):
288 return Error(err)
289 case Ok(id_jag):
290 leg2_params: Final = {
291 "grant_type": _JWT_BEARER_GRANT_TYPE,
292 "assertion": id_jag.access_token,
293 }
294 return await self._token_endpoint.fetch(
295 config.resource_token_endpoint,
296 config.client_id,
297 leg2_params,
298 config.client_auth,
299 )
301 match await self._exchanged_tokens.get_or_compute(slot, _exchange, fingerprint=fingerprint):
302 case Ok(access_token):
303 header_name, header_value = config.header(access_token)
304 return Ok(StaticHeaderAuth(header_value, header_name=header_name))
305 case Error(err):
306 return Error(err)
308 async def _authorization_code(
309 self, subject: Subject, server: ServerSpec, config: AuthorizationCodeConfig
310 ) -> Result[StaticHeaderAuth, CredError]:
311 token: Final = await self._authz_token(subject, server)
312 if token is None:
313 return Error(CredError.of_unauthorized("Authorization required: complete the OAuth flow for this server."))
314 header_name, header_value = config.header(token.access_token)
315 return Ok(StaticHeaderAuth(header_value, header_name=header_name))
317 async def _client_credentials(
318 self, server_id: str, config: ClientCredentialsConfig
319 ) -> Result[httpx2.Auth, CredError]:
320 """The M2M arm: resolve a cached (or freshly minted) gateway token; no user context.
322 The token is resolved here, before any upstream request, so a misconfigured grant or an
323 unreachable IdP surfaces as a typed ``CredError``. The returned auth carries the source's
324 ``refetch``, so an upstream 401 is retried exactly once with a freshly minted token (the
325 contract's invalid-token recovery); a second 401 surfaces the upstream's own error.
326 """
327 match await self._client_credentials_source.get(server_id, config):
328 case Ok(token):
329 refetch: Final = partial(self._client_credentials_source.refetch, server_id, config)
330 return Ok(ClientCredentialsBearerAuth(token.access_token, refetch, config))
331 case Error(err):
332 return Error(err)
334 async def _token_exchange(
335 self, subject: Subject, server: ServerSpec, config: TokenExchangeConfig
336 ) -> Result[StaticHeaderAuth, CredError]:
337 """RFC 8693 OBO: exchange the caller's inbound token for an upstream-bound bearer.
339 No inbound token means there is nothing to exchange, so the arm fails closed with a 401 rather
340 than falling through to a weaker source (§1.5); the exchanger handles the IdP round-trip and
341 caching and returns the upstream token or a typed error.
342 """
343 inbound: Final = subject.inbound_token
344 if inbound is None:
345 return Error(
346 CredError.of_unauthorized(
347 "Token exchange requires a caller token to exchange (OBO).",
348 www_authenticate='Bearer error="invalid_request"',
349 )
350 )
351 match await self._token_exchanger.exchange(
352 inbound.get_secret_value(), server, config, tenant_id=subject.tenant_id
353 ):
354 case Ok(token):
355 header_name, header_value = config.header(token.access_token)
356 return Ok(StaticHeaderAuth(header_value, header_name=header_name))
357 case Error(err):
358 return Error(err)
360 async def invalidate_credentials(self, subject: Subject, server: ServerSpec) -> None:
361 """Drop any cached credential the resolver owns for this `(subject, server)`.
363 Used after an upstream rejects the injected credential, so the next resolve re-mints rather
364 than serving the same rejected token until TTL. `token_exchange` and `id_jag` hold a
365 re-mintable cached credential here; `client_credentials` recovers inside its own auth flow
366 (`ClientCredentialsBearerAuth` retries the 401'd request once with a fresh token), and
367 other modes are a no-op.
369 `id_jag` evicts by a slot key derived from the principal, so it needs no lookup against the
370 assertion store on this path; the fingerprint stored beside the entry is what keeps a slot
371 shared between callers safe.
372 """
373 if isinstance(server.config, IdJagConfig):
374 self._invalidate_id_jag(subject, server)
375 elif isinstance(server.config, TokenExchangeConfig) and subject.inbound_token is not None:
376 await self._token_exchanger.invalidate(
377 subject.inbound_token.get_secret_value(), server, server.config, tenant_id=subject.tenant_id
378 )
380 def _invalidate_id_jag(self, subject: Subject, server: ServerSpec) -> None:
381 """Evict the bearer this `(subject, server)` last resolved, without depending on the store.
383 The slot is addressed by the principal (plus the caller's own token when it presented one),
384 never by the credential material, so it stays computable when the assertion store is down.
385 The fingerprint stored with the entry is what keeps that safe: an entry minted for different
386 inputs reads as a miss rather than being served.
387 """
388 self._exchanged_tokens.invalidate(_id_jag_slot_key(subject, server))
390 async def _authz_token(self, subject: Subject, server: ServerSpec) -> OAuthToken | None:
391 """The user's authorization_code token, or None when absent or the store is unreachable.
393 A store outage is mapped to None (the OAuth challenge), not raised, so a transient outage
394 does not 500; it is the store, not this resolver, that declines to cache the failure.
395 """
396 try:
397 return await self._oauth_token_store.fetch(subject.subject_id, server.server_id)
398 except TokenStoreUnavailable:
399 return None
402def _id_jag_slot_key(subject: Subject, server: ServerSpec) -> str:
403 """Which cache slot this caller's bearer for this upstream lives in.
405 Addressed by the principal, plus the caller's own token when it presented one so two callers
406 sharing an empty principal do not contend for one slot. Deliberately free of the stored
407 assertion, which is what lets invalidation compute this while the assertion store is down. The
408 entry's fingerprint, not this key, is what guarantees a cached bearer matches current inputs.
409 """
410 inbound: Final = subject.inbound_token.get_secret_value() if subject.inbound_token is not None else ""
411 material: Final = "\x00".join((subject.tenant_id, subject.subject_id, server.server_id, inbound))
412 return hashlib.sha256(material.encode()).hexdigest()
415def _id_jag_fingerprint(subject_token: str, server_id: str, config: IdJagConfig) -> str:
416 """What the cached leg-2 bearer was minted from: the subject token, the server, and the config.
418 Stored beside the bearer and compared on every read, so a rotated assertion or an edited server
419 config reads as a miss and re-mints instead of serving a bearer authorized under the old policy.
421 Every exchange parameter derives from the config (endpoints, audience, resource, scopes, client
422 auth), so a server update that changes any of them must change the key; otherwise the old bearer,
423 authorized under the old policy, keeps being served until its TTL. Everything is hashed, so no
424 secret is held in the key.
425 """
426 material: Final = "\x00".join(
427 (
428 subject_token,
429 server_id,
430 config.org_token_endpoint,
431 config.resource_token_endpoint,
432 config.client_id,
433 _client_auth_fingerprint(config.client_auth),
434 config.subject_token_type,
435 config.audience or "",
436 config.resource or "",
437 " ".join(config.scopes),
438 )
439 )
440 return hashlib.sha256(material.encode()).hexdigest()
443def _client_auth_fingerprint(client_auth: ClientAuth) -> str:
444 match client_auth:
445 case PrivateKeyJwtAuth() as auth:
446 return "\x00".join(
447 ("private_key_jwt", auth.private_key.get_secret_value(), auth.key_id or "", auth.signing_alg)
448 )
449 case ClientSecretAuth() as auth:
450 return "\x00".join(("client_secret", auth.client_secret.get_secret_value()))
451 assert_never(client_auth)
454def _not_implemented(kind: AuthSpecKind) -> Result[httpx2.Auth, CredError]:
455 return Error(CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet"))
458async def resolve_credentials_with_source(
459 provider: UpstreamCredentialProvider, subject: Subject, server: ServerSpec
460) -> Result[ResolvedCredential, CredError]:
461 match await provider.resolve_credentials(subject, server):
462 case Error(err): 462 ↛ 463line 462 didn't jump to line 463 because the pattern on line 462 never matched
463 return Error(err)
464 case Ok(auth): 464 ↛ exitline 464 didn't return from function 'resolve_credentials_with_source' because the pattern on line 464 always matched
465 if isinstance(auth, NoOpAuth): 465 ↛ 467line 465 didn't jump to line 467 because the condition on line 465 was always true
466 return Ok(ResolvedCredential(auth, AuthResolution.no_auth))
467 match server.config:
468 case NoneConfig():
469 return Ok(ResolvedCredential(auth, AuthResolution.no_auth))
470 case ApiKeyConfig():
471 return Ok(ResolvedCredential(auth, AuthResolution.static_token))
472 case PassthroughConfig():
473 return Ok(ResolvedCredential(auth, AuthResolution.oauth2_passthrough))
474 case ClientCredentialsConfig():
475 return Ok(ResolvedCredential(auth, AuthResolution.client_credentials))
476 case TokenExchangeConfig():
477 return Ok(ResolvedCredential(auth, AuthResolution.token_exchange))
478 case IdJagConfig():
479 return Ok(ResolvedCredential(auth, AuthResolution.id_jag))
480 case AuthorizationCodeConfig():
481 return Ok(ResolvedCredential(auth, AuthResolution.stored_user_token))
482 case AwsSigV4Config():
483 return Ok(ResolvedCredential(auth, AuthResolution.aws_sigv4))
484 assert_never(server.config)