Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py: 79%
224 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""The upstream-credential vocabulary — the typed seam the resolver dispatches on.
3This module ships the data types only; the resolver lands in a later PR. It is the contract
4the credential build implements and the spec tests assert against.
6Design invariants encoded here:
8- **Mode is the single source of truth.** A server declares exactly one per-mode `config`
9 (the `AuthConfig` discriminated union); `auth_spec_kind` is *derived* from it, never a
10 second field that can drift. The resolver dispatches on the config variant, one arm per
11 mode. No field-presence inference, no precedence cascade.
12- **Illegal states unrepresentable.** Each mode's config is its own frozen model holding
13 only that mode's fields — an `aws_sigv4` server cannot hold OAuth fields, and a config
14 missing a required field is rejected at construction, not at call time.
15- **Fail-closed at the boundary.** A raw mode string can only enter through
16 `parse_auth_spec_kind()`, which returns a typed `CredError`.
17- **Errors as values.** Every seam returns `Result[_, CredError]`; only edge adapters raise.
18- **No v1 imports.** This vocabulary stays free of `MCPServer` and the rest of v1; the
19 v1 -> v2 adapter maps onto these types in a later PR.
21Sum types are Expression `@tagged_union`s discriminated on a `Literal` `tag`, matched via
22`self.tag` with an `assert_never` tail; `Result` is this package's vendored `Ok | Error`
23union (see `result.py`), not `expression.Result`.
24"""
26from __future__ import annotations
28from collections.abc import Mapping
29from dataclasses import dataclass, field
30from enum import Enum
31from typing import Annotated, Final, Literal
33import httpx2
34from expression import case, tag, tagged_union
35from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator
36from typing_extensions import assert_never
38from litellm.proxy._experimental.mcp_server.outbound_credentials.result import (
39 Error,
40 Ok,
41 Result,
42)
43from litellm.types.mcp import (
44 DEFAULT_CREDENTIAL_HEADER,
45 DEFAULT_SUBJECT_TOKEN_TYPE,
46 normalize_upstream_header_name,
47)
50class AuthResolution(str, Enum):
51 no_auth = "no-auth"
52 stored_user_token = "stored-user-token"
53 static_token = "static-token"
54 per_request_header = "per-request-header"
55 oauth2_passthrough = "oauth2-passthrough"
56 client_credentials = "m2m-client-credentials"
57 token_exchange = "token-exchange"
58 id_jag = "id-jag"
59 aws_sigv4 = "aws-sigv4"
60 extra_headers = "extra-headers"
61 not_applicable = "not-applicable"
62 unresolved = "unresolved"
63 failed = "resolution-failed"
64 multiple = "multiple"
67@dataclass(frozen=True, slots=True)
68class ResolvedCredential:
69 auth: httpx2.Auth = field(repr=False)
70 source: AuthResolution
73class AuthSpecKind(str, Enum):
74 """The server's statically-declared upstream-auth mode — derived from its `config`.
76 Covers v1's full `MCPAuth` surface, not only OAuth grants: the three grant modes, the
77 collapsed static-header family, client passthrough, no-auth, and AWS request signing.
78 BYOK is *not* a member: it is the `api_key` mode seeded per-user, a source selector
79 inside that arm. The static-header schemes v1 splits into separate `MCPAuth` values
80 (`bearer_token`/`api_key`/`basic`/`token`/`authorization`) collapse into `api_key`; the
81 scheme is a parameter the arm carries, not its own mode.
82 """
84 authorization_code = "authorization_code" # per-user 3LO; gateway-stored token
85 client_credentials = "client_credentials" # gateway service account (M2M)
86 token_exchange = "token_exchange" # RFC 8693: token endpoint + subject_token (OBO)
87 id_jag = "id_jag" # draft-ietf-oauth-identity-assertion-authz-grant: two-leg exchange then jwt-bearer
88 api_key = "api_key" # static header, any scheme (BYOK = per-user-seeded source)
89 passthrough = "passthrough" # client forwards an upstream-audience token
90 none = "none" # no upstream credential; resolve yields a no-op auth, never an error
91 aws_sigv4 = "aws_sigv4" # AWS SigV4 per-request signing (e.g. Bedrock AgentCore)
94@dataclass(frozen=True, slots=True)
95class Unauthorized:
96 """A 401 plus the optional challenge a client needs to recover.
98 ``detail`` is the human message; ``www_authenticate`` and ``body`` carry a scheme-specific
99 challenge (e.g. BYOK's provisioning prompt) so the edge can reproduce it verbatim.
100 ``claims`` carries an IdP step-up challenge (e.g. Entra Conditional Access) so the edge can
101 fold it into the ``WWW-Authenticate`` it builds; the client replays the claims to the IdP to
102 satisfy the step-up, then retries with the fresh token.
103 """
105 detail: str
106 www_authenticate: str | None = None
107 body: Mapping[str, str] | None = None
108 claims: str | None = None
111@tagged_union(frozen=True)
112class CredError:
113 """Why a credential could not be produced. Fail-closed: an arm yields this or an `httpx2.Auth`.
115 Discriminated on the `Literal` `tag`; consumers `match self.tag` (see `summary`) so the
116 type checker can prove exhaustiveness. Construct via the `of_*` factories.
117 """
119 tag: Literal[
120 "unauthorized",
121 "misconfigured",
122 "url_credentials_not_allowed",
123 "upstream_unavailable",
124 "unsupported_mode",
125 "precondition_required",
126 "not_implemented",
127 ] = tag()
129 unauthorized: Unauthorized = case() # no usable credential for this (subject, server) -> 401 challenge
130 misconfigured: str = case() # the declared mode is missing required config -> 5xx (operator)
131 url_credentials_not_allowed: None = case()
132 upstream_unavailable: str = case() # the IdP / token endpoint could not be reached -> 503
133 unsupported_mode: str = case() # a raw mode string did not parse into AuthSpecKind (boundary)
134 precondition_required: str = case() # a required per-user value (e.g. an env var) has not been provided -> 412
135 not_implemented: str = case() # the declared mode's resolver arm is not built yet -> 501 (not operator error)
137 @staticmethod
138 def of_unauthorized(
139 detail: str,
140 *,
141 www_authenticate: str | None = None,
142 body: Mapping[str, str] | None = None,
143 claims: str | None = None,
144 ) -> CredError:
145 return CredError(
146 unauthorized=Unauthorized(
147 detail=detail,
148 www_authenticate=www_authenticate,
149 body=body,
150 claims=claims,
151 )
152 )
154 @staticmethod
155 def of_misconfigured(detail: str) -> CredError:
156 return CredError(misconfigured=detail)
158 @staticmethod
159 def of_url_credentials_not_allowed() -> CredError:
160 return CredError(url_credentials_not_allowed=None)
162 @staticmethod
163 def of_upstream_unavailable(detail: str) -> CredError:
164 return CredError(upstream_unavailable=detail)
166 @staticmethod
167 def of_unsupported_mode(detail: str) -> CredError:
168 return CredError(unsupported_mode=detail)
170 @staticmethod
171 def of_precondition_required(detail: str) -> CredError:
172 return CredError(precondition_required=detail)
174 @staticmethod
175 def of_not_implemented(detail: str) -> CredError:
176 return CredError(not_implemented=detail)
178 @property
179 def summary(self) -> str:
180 # Exhaustiveness: every Literal tag has an arm; the trailing assert_never typechecks
181 # only while that stays true (a `case _` would defeat reportMatchNotExhaustive).
182 match self.tag:
183 case "unauthorized": 183 ↛ 184line 183 didn't jump to line 184 because the pattern on line 183 never matched
184 return f"unauthorized: {self.unauthorized.detail}"
185 case "misconfigured": 185 ↛ 187line 185 didn't jump to line 187 because the pattern on line 185 always matched
186 return f"misconfigured: {self.misconfigured}"
187 case "url_credentials_not_allowed":
188 return (
189 "misconfigured: auth_type none cannot be used with credentials embedded in the upstream URL; "
190 "remove them from the URL and configure Basic Auth with auth_type: basic and "
191 "auth_value: username:password"
192 )
193 case "upstream_unavailable":
194 return f"upstream unavailable: {self.upstream_unavailable}"
195 case "unsupported_mode":
196 return self.unsupported_mode
197 case "precondition_required":
198 return f"precondition required: {self.precondition_required}"
199 case "not_implemented":
200 return f"not implemented: {self.not_implemented}"
201 assert_never(self.tag)
204def validate_header_name(raw: str) -> Result[str, CredError]:
205 """``normalize_upstream_header_name`` with this package's error-as-value policy.
207 The grammar itself lives in ``litellm.types.mcp`` so the v1 model, the management endpoint and
208 this vocabulary all judge a header name the same way while each keeps its own failure shape.
209 """
210 normalized: Final = normalize_upstream_header_name(raw)
211 if normalized is None:
212 return Error(CredError.of_misconfigured(f"invalid upstream header name: {raw!r}"))
213 return Ok(normalized)
216class HeaderCarrier(BaseModel):
217 """Where a resolved credential is written upstream, and how its value is formatted.
219 ``Authorization: Bearer`` is only OAuth's *default* conveyance (RFC 6750 section 2.1), not its
220 only one: an ESB or API gateway commonly terminates its own credential in a private header while
221 a second credential passes through to the origin, so a credential has to be able to say which
222 slot it owns. Modeled like OpenAPI's apiKey scheme, so any upstream convention is expressible
223 (Authorization + Bearer, a raw value on X-API-Key, Ocp-Apim-Subscription-Key, esb-oauth, ...).
225 Every config whose credential the gateway mints or holds inherits this, so no resolver arm names
226 a header itself and the conflict rule in ``_resolve_v2_auth`` can always ask the auth object
227 which slot it is about to occupy. ``passthrough`` deliberately does not: it forwards the
228 caller's own credential into the slot the caller used, and mints nothing to place.
229 """
231 model_config = ConfigDict(frozen=True)
232 header_name: str = DEFAULT_CREDENTIAL_HEADER
233 value_prefix: str = "Bearer"
235 @field_validator("header_name")
236 @classmethod
237 def _check_header_name(cls, value: str) -> str:
238 match validate_header_name(value):
239 case Ok(name):
240 return name
241 case Error(err):
242 raise ValueError(err.summary)
244 def header(self, value: str) -> tuple[str, str]:
245 formatted: Final = f"{self.value_prefix} {value}" if self.value_prefix else value
246 return self.header_name, formatted
249class AuthorizationCodeConfig(HeaderCarrier):
250 """Per-user 3LO; the gateway is the OAuth client and stores the user's token.
252 Endpoints are discovered (RFC 9728 -> RFC 8414) and the client is registered via DCR
253 (RFC 7591), so the common case carries none of the fields below; they are optional manual
254 overrides for IdPs without discovery / DCR. The per-user token is read from the token store
255 at resolve time, not held here.
256 """
258 model_config = ConfigDict(frozen=True)
259 kind: Literal[AuthSpecKind.authorization_code] = AuthSpecKind.authorization_code
260 scopes: tuple[str, ...] = ()
261 client_id: str | None = None
262 client_secret: SecretStr | None = None
263 authorization_url: str | None = None
264 token_url: str | None = None
267class ClientCredentialsConfig(HeaderCarrier):
268 """M2M service account; one upstream identity for every user.
270 Fields are optional so the config can be built incomplete: a value may be supplied at
271 runtime (`token_url` via RFC 8414 discovery, `client_id`/`secret` via DCR), and the
272 resolver arm returns `CredError.misconfigured` when a needed field is still absent.
274 `audience` is the IdP-specific audience parameter some authorization servers require on
275 the client_credentials grant (sent as `audience` in the token request when set).
276 `token_endpoint_auth_method` selects how the client authenticates to the token endpoint
277 (RFC 6749 section 2.3.1); `None` defaults to `client_secret_post`.
278 """
280 model_config = ConfigDict(frozen=True)
281 kind: Literal[AuthSpecKind.client_credentials] = AuthSpecKind.client_credentials
282 client_id: str | None = None
283 client_secret: SecretStr | None = None
284 token_url: str | None = None
285 scopes: tuple[str, ...] = ()
286 audience: str | None = None
287 upstream_resource: str | None = None
288 token_endpoint_auth_method: Literal["client_secret_post", "client_secret_basic"] | None = None
291class TokenExchangeConfig(HeaderCarrier):
292 """OBO: swap the caller's live inbound token for a token bound to the upstream's audience. The
293 gateway authenticates to the exchange endpoint as an OAuth client (`client_id`/`client_secret`);
294 the inbound token is sent only to that endpoint, never to the upstream.
296 `profile` selects the wire dialect, since not every IdP speaks RFC 8693:
297 - `rfc8693` (default) is the standard token-exchange grant: the inbound token is the
298 `subject_token` (typed by `subject_token_type`), the target is the optional `audience`.
299 - `entra_obo` is Microsoft Entra On-Behalf-Of, which is the RFC 7523 `jwt-bearer` grant rather
300 than 8693: the inbound token rides as `assertion`, the target resource is carried in `scopes`
301 (`api://<app-id>/.default`, since Entra has no audience parameter), and the Microsoft-only
302 `requested_token_use=on_behalf_of` extension makes the jwt-bearer grant a delegation.
303 `subject_token_type` and `audience` are unused in this profile.
305 `audience` (rfc8693 only) is optional and sent only when the operator configured one, since both
306 `audience` and `resource` are optional in RFC 8693 and the authorization server applies its own
307 default when neither is sent (fabricating one risks `invalid_target`).
308 """
310 model_config = ConfigDict(frozen=True)
311 kind: Literal[AuthSpecKind.token_exchange] = AuthSpecKind.token_exchange
312 profile: Literal["rfc8693", "entra_obo"] = "rfc8693"
313 subject_token_type: str = DEFAULT_SUBJECT_TOKEN_TYPE
314 token_exchange_endpoint: str | None = None
315 audience: str | None = None
316 client_id: str | None = None
317 client_secret: SecretStr | None = None
318 token_endpoint_auth_method: Literal["client_secret_basic", "client_secret_post"] | None = None
319 scopes: tuple[str, ...] = ()
322class PrivateKeyJwtAuth(BaseModel):
323 """RFC 7523 private-key-JWT client authentication: the gateway signs a `client_assertion`."""
325 model_config = ConfigDict(frozen=True)
326 source: Literal["private_key_jwt"] = "private_key_jwt"
327 private_key: SecretStr
328 key_id: str | None = None
329 signing_alg: str = "RS256"
332class ClientSecretAuth(BaseModel):
333 """`client_secret_post` client authentication: the gateway posts `client_id` + `client_secret`."""
335 model_config = ConfigDict(frozen=True)
336 source: Literal["client_secret"] = "client_secret"
337 client_secret: SecretStr
340ClientAuth = Annotated[PrivateKeyJwtAuth | ClientSecretAuth, Field(discriminator="source")]
343class IdJagConfig(HeaderCarrier):
344 """draft-ietf-oauth-identity-assertion-authz-grant (Okta "AI agent token exchange").
346 Two legs: leg 1 is an RFC 8693 token exchange at the IdP org AS (`org_token_endpoint`) that
347 swaps the caller's identity token for an ID-JAG assertion; leg 2 is an RFC 7523 jwt-bearer at
348 the upstream resource AS (`resource_token_endpoint`) that swaps the assertion for the access
349 token. The gateway authenticates to both endpoints as `client_id` via `client_auth`. Required
350 fields are enforced at construction so a half-configured server cannot reach the arm.
351 """
353 model_config = ConfigDict(frozen=True)
354 kind: Literal[AuthSpecKind.id_jag] = AuthSpecKind.id_jag
355 org_token_endpoint: str
356 resource_token_endpoint: str
357 client_id: str
358 client_auth: ClientAuth
359 subject_token_type: str = "urn:ietf:params:oauth:token-type:id_token"
360 audience: str | None = None
361 resource: str | None = None
362 scopes: tuple[str, ...] = ()
365class SharedKey(BaseModel):
366 """A fixed key configured on the server, identical for every caller."""
368 model_config = ConfigDict(frozen=True)
369 source: Literal["shared"] = "shared"
370 value: SecretStr
373class Byok(BaseModel):
374 """A key the user brings via the entry flow, stored per-user and pulled from the credential
375 store at resolve time. Missing means the user must provide it, a 401 + WWW-Authenticate
376 challenge."""
378 model_config = ConfigDict(frozen=True)
379 source: Literal["byok"] = "byok"
382ApiKeySource = Annotated[SharedKey | Byok, Field(discriminator="source")]
385class ApiKeyConfig(HeaderCarrier):
386 """A fixed credential injected as a header. The value is shared (in config) or seeded
387 per-user (pulled from the store); the inherited `header_name` and `value_prefix` say where
388 and how it is written.
389 """
391 model_config = ConfigDict(frozen=True)
392 kind: Literal[AuthSpecKind.api_key] = AuthSpecKind.api_key
393 key_source: ApiKeySource
396class PassthroughConfig(BaseModel):
397 """Client-driven upstream OAuth; the gateway forwards the client's upstream token."""
399 model_config = ConfigDict(frozen=True)
400 kind: Literal[AuthSpecKind.passthrough] = AuthSpecKind.passthrough
403class NoneConfig(BaseModel):
404 """No upstream credential; the request is sent unauthenticated."""
406 model_config = ConfigDict(frozen=True)
407 kind: Literal[AuthSpecKind.none] = AuthSpecKind.none
410class StaticKeys(BaseModel):
411 """Long-lived AWS access keys configured on the server."""
413 model_config = ConfigDict(frozen=True)
414 source: Literal["static_keys"] = "static_keys"
415 access_key_id: str
416 secret_access_key: SecretStr
417 session_token: SecretStr | None = None
420class AssumeRole(BaseModel):
421 """An IAM role the gateway assumes via STS for short-lived, auto-refreshed credentials."""
423 model_config = ConfigDict(frozen=True)
424 source: Literal["assume_role"] = "assume_role"
425 role_arn: str
426 session_name: str | None = None
427 external_id: str | None = None
430class Ambient(BaseModel):
431 """The environment's default AWS credential chain (instance profile, IRSA, env vars)."""
433 model_config = ConfigDict(frozen=True)
434 source: Literal["ambient"] = "ambient"
437AwsCredentialSource = Annotated[StaticKeys | AssumeRole | Ambient, Field(discriminator="source")]
440class AwsSigV4Config(BaseModel):
441 """AWS SigV4 per-request signing for an AWS-hosted upstream (e.g. Bedrock AgentCore). The
442 gateway signs with its own AWS identity, never the caller's; `credentials` selects how that
443 identity is obtained, defaulting to the ambient credential chain."""
445 model_config = ConfigDict(frozen=True)
446 kind: Literal[AuthSpecKind.aws_sigv4] = AuthSpecKind.aws_sigv4
447 region: str
448 service: str = "bedrock-agentcore"
449 credentials: AwsCredentialSource = Ambient()
452AuthConfig = Annotated[
453 AuthorizationCodeConfig
454 | ClientCredentialsConfig
455 | TokenExchangeConfig
456 | IdJagConfig
457 | ApiKeyConfig
458 | PassthroughConfig
459 | NoneConfig
460 | AwsSigV4Config,
461 Field(discriminator="kind"),
462]
465class Subject(BaseModel):
466 """The validated inbound principal. NOT the v1 request object and NOT the LiteLLM key."""
468 model_config = ConfigDict(frozen=True)
470 tenant_id: str
471 subject_id: str
472 # Opaque, already-validated inbound identity. Read by `token_exchange`, `passthrough`, and
473 # `id_jag` (which falls back to the user's stored SSO assertion when it is absent).
474 inbound_token: SecretStr | None = None
477class ServerSpec(BaseModel):
478 """The declared upstream. A v2-native type; the v1 -> v2 adapter maps onto this."""
480 model_config = ConfigDict(frozen=True)
482 server_id: str
483 resource: str # RFC 8707 audience URI this upstream's tokens are bound to
484 config: AuthConfig
486 @property
487 def auth_spec_kind(self) -> AuthSpecKind:
488 return self.config.kind
491def parse_auth_spec_kind(raw: str) -> Result[AuthSpecKind, CredError]:
492 """Boundary parser — the *only* place an unknown mode is handled, and it fails closed.
494 Inside the core the mode is always a valid `AuthSpecKind`, so the resolver never needs a
495 wildcard arm and basedpyright can prove its `match` exhaustive.
496 """
497 try:
498 return Ok(AuthSpecKind(raw))
499 except ValueError:
500 return Error(CredError.of_unsupported_mode(f"unknown auth_spec_kind: {raw!r}"))