scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 227

feat(mcp): source the ID-JAG subject token from the stored SSO assertion

Browse source 
The oauth2_id_jag egress arm spends the user's IdP id_token as its RFC 8693
subject_token, but a front-door client presents only a gateway-minted
llm_session_ bearer, so hop A had nothing usable and EMA was unreachable end
to end for exactly the clients it exists for — while a gateway-issued
credential presented via Authorization was forwarded verbatim to the IdP.

One rule per concern, each with a single owner:

- Provenance (credential_provenance.py, new): the inbound bearer is
  classified once at the edge (absent / gateway_credential / external_jwt /
  external_opaque) by the one property every gateway mint shares and no
  external token has: it verifies under this gateway's master key. A new
  mint format is refused as exchange material by construction, with no
  denylist to grow stale. id_jag exchanges the inbound bearer only when it
  is an external JWT; token_exchange (OBO) keeps exchange-what-was-presented
  but refuses gateway credentials.

- Header selection (_subject_bearer_token): every egress surface (tools
  call/list, prompts, resources, templates, OpenAPI) selects the caller's
  bearer through one manager rule; for id_jag the bearer must also be the
  credential that admitted the caller, so a free-riding Authorization value
  bound to nobody is never exchanged.

- Sourcing (_id_jag_subject_token): the caller's presented id_token when
  the inbound bearer is one, else the stored SSO assertion, else fail
  closed — shared verbatim by resolve and invalidate.

- Storage (sso_assertion_store.py): the generic-SSO callback persists the
  IdP id_token + refresh token as one encrypted row per user (latest login
  wins), write-through to the DB so every pod sees it, gated on an id_jag
  server existing per the config and DB authorities (never the registry
  snapshot). Reads compose the shared per-user credential stack —
  Cached(Refreshing(db)) — for positive-only caching, expiry-skewed renewal
  and per-user single-flight; the TokenRefresher protocol now documents the
  raise-through contract the EMA refresher relies on.

- Renewal lifecycle: expiry is judged by readers (30s buffer). A dead grant
  requires proof about the GRANT, and RFC 6749 has exactly one code that
  states it: invalid_grant. Every other outcome (invalid_client and friends
  are the gateway's own client config failing — ops-fixable), a 429, a 4xx
  without an error object, a 5xx, or a transport failure reads retryable
  and never touches the row. A proven verdict is recorded ON the row by
  stripping its refresh token — never in a cache beside it — so it holds on
  every pod, costs one POST total, and a re-login lifts it by construction.
  The verdict write re-reads the row and strips only a row still carrying
  the exact rejected token, preserving the row's other material, so a
  losing pod in a cross-pod rotation race cannot clobber the winner's
  renewed assertion.

- Failure semantics, same story on every surface: absence -> 412 naming the
  gateway-SSO remedy; dead grant -> 401 + re-authenticate challenge; outage
  -> retryable 503, never a false re-login and never absence. Single-server
  routes surface these at the transport edge via preflight_id_jag (sibling
  of the OBO preflight), where an HTTP status and WWW-Authenticate still
  reach the client — in-session list handlers can only serialize a raise
  into a JSON-RPC error. The aggregate absorbs per-server faults, and
  tools/list outcomes now preserve the retryable-vs-terminal class
  (unavailable -> 503, precondition -> 412) instead of reporting a
  dependency outage as the gateway's own 500.

Tests pin the rules, not the mechanisms: provenance totality over every
current mint plus adversarial input, the one-POST-per-dead-grant bound, the
re-signed-in caller never refused, absence-vs-outage by remedy, the
cross-pod clobber no-op, the invalid_grant-only proof predicate, scope
identity between login and refresh, winner-and-waiter verdict sharing in
the shared stack, and the preflight wiring gated to single-server routes.

Part of LIT-3856 (PR 2 of the EMA stack, on the #34072 assertion store).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Tin Chi Lo 2026-07-24 18:27:55 -07:00
parent 3eaf7b1c0a
commit f04398e3fd
24 changed files with 2695 additions and 208 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -1208,6 +1208,29 @@ class MCPRequestHandler:
)
return MCP_CLIENT_SIDE_AUTH_HEADER_NAME
@staticmethod
def authorization_is_free_rider(raw_headers: dict[str, str] | None) -> bool:
"""Whether the request's Authorization bearer was NOT the admission credential.
This is the inverse view of ``get_litellm_api_key_from_headers`` and must encode the
SAME header-preference rule: admission authenticates with the explicit primary header
only when its value is truthy (a missing or empty primary falls back to Authorization),
so the Authorization bearer was the admission credential exactly when the primary is
absent or empty. It lives in this module beside that accessor so one owner holds the
rule, and the agreement test pins the two against a shared header matrix.
"""
if not raw_headers:
return False
primary = next(
(
value
for key, value in raw_headers.items()
if str(key).lower() == MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY.lower()
),
None,
)
return bool(primary)
@staticmethod
def get_litellm_api_key_from_headers(headers: Headers) -> Optional[str]:
"""

24
litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py
View file

@ -30,6 +30,8 @@ ListFaultCategory: TypeAlias = Literal[
"timeout",
"unreachable",
"upstream_error",
"unavailable",
"precondition",
"internal",
]
@ -43,8 +45,11 @@ class ServerListOk(BaseModel):
class ServerListFault(BaseModel):
"""Why a server contributed nothing to a listing: the caller must authenticate upstream
(``auth_required``/``forbidden``), the upstream did not answer (``timeout``/``unreachable``),
the upstream answered outside its contract (``upstream_error``), or the gateway itself failed
(``internal``). ``status_code`` is the upstream HTTP status when one exists."""
the upstream answered outside its contract (``upstream_error``), a dependency the gateway
needs to authorize egress did not answer so the caller should retry (``unavailable``), a
per-caller requirement is unmet and only the caller can meet it, e.g. a gateway SSO sign-in
(``precondition``), or the gateway itself failed (``internal``). ``status_code`` is the
upstream HTTP status when one exists."""
model_config = ConfigDict(frozen=True)
tag: ListFaultCategory
@ -149,7 +154,16 @@ def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]:
match outcome.tag:
case "ok":
return {"status": "ok", "tool_count": outcome.tool_count}
case "auth_required" | "forbidden" | "timeout" | "unreachable" | "upstream_error" | "internal":
case (
"auth_required"
| "forbidden"
| "timeout"
| "unreachable"
| "upstream_error"
| "unavailable"
| "precondition"
| "internal"
):
return {
"status": outcome.tag,
**({"http_status": outcome.status_code} if outcome.status_code is not None else {}),
@ -171,6 +185,10 @@ def list_fault_http_status(fault: ServerListFault) -> int:
return 504
case "unreachable" | "upstream_error":
return 502
case "unavailable":
return 503
case "precondition":
return 412
case "internal":
return 500
case _:

155
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -90,6 +90,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto
from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import (
LazyPerUserOAuthTokenStore,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
LiveSsoAssertionSource,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_provider import (
build_token_exchanger,
)
@ -809,6 +812,25 @@ def _caller_authorization_fans_out(
)
def _obo_retry_covers_caller(
auth_type: MCPAuthType | None,
subject_token: str | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> bool:
"""Whether the invalidate-and-retry path has a subject to re-mint with after an upstream 401.
``token_exchange`` re-mints only from the caller's inbound token. ``id_jag`` also re-mints
from the stored SSO assertion, which the resolver looks up by the caller's user id, so a
caller with a user identity but no inbound assertion (an admission-key or session caller)
must route through the retry path too or its rejected cached bearer is never evicted.
"""
if auth_type not in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag):
return False
if subject_token:
return True
return auth_type == MCPAuth.oauth2_id_jag and user_api_key_auth is not None and bool(user_api_key_auth.user_id)
def _extract_upstream_auth_failure(
exc: BaseException,
) -> Optional[tuple[int, Optional[str]]]:
@ -1153,6 +1175,7 @@ class MCPServerManager:
self._cred_provider = cred_provider or UpstreamCredentialProvider(
oauth_token_store=self._per_user_oauth_token_store,
token_exchanger=build_token_exchanger(),
sso_assertions=LiveSsoAssertionSource(),
)
self.registry: dict[str, MCPServer] = {}
self.config_mcp_servers: dict[str, MCPServer] = {}
@ -1251,6 +1274,7 @@ class MCPServerManager:
mcp_auth_header=None,
extra_headers=extra_headers,
stdio_env=None,
user_api_key_auth=None, # internal path (cache warm-up / health check): no caller identity
)
async def _noop(_session):
@ -2739,20 +2763,35 @@ class MCPServerManager:
return auth_value
return None
def _obo_subject_token(
def _subject_bearer_token(
self,
server: MCPServer,
raw_headers: Optional[dict[str, str]],
oauth2_headers: dict[str, str] | None = None,
) -> Optional[str]:
"""The caller's bearer as the token_exchange (OBO) subject token, for that mode only.
"""The caller's bearer as the exchange subject material, resolved by the ONE rule every
egress surface shares: only the two exchange modes (token_exchange OBO and id_jag) read
the inbound bearer, and every surface (tools call and list, prompts, resources) resolves
it here so the RULE cannot diverge per surface. The INPUTS still can: the prompt/resource
helpers receive no ``oauth2_headers``, so a client-side oauth2 header participates only on
the tool surfaces until someone threads it through those signatures. Other modes return
None to avoid forwarding it.
Prompts/resources discovery and reads on a token_exchange server must exchange the caller's
token like the tools paths do, not connect with no credential. Other modes never read the
inbound bearer, so return None to avoid forwarding it.
For id_jag the bearer must additionally BE the credential that authenticated the caller:
admission prefers the explicit litellm key header, so when that header is present the
Authorization bearer is an unvalidated free rider bound to nobody, and exchanging it
would let a caller act upstream under any identity whose token they hold. Binding is
admission's job, so the rule is structural (was this the admission credential), never a
claims comparison re-deriving what admission already decided. Such callers resolve
through the stored assertion, which is bound to the authenticated user by construction.
The OBO mode keeps its documented exchange-what-was-presented semantics; its identical
free-rider shape predates this seam and is tracked as a follow-up.
"""
if server.auth_type != MCPAuth.oauth2_token_exchange:
if server.auth_type not in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag):
return None
return self._extract_bearer_token(None, raw_headers)
if server.auth_type == MCPAuth.oauth2_id_jag and MCPRequestHandler.authorization_is_free_rider(raw_headers):
return None
return self._extract_bearer_token(oauth2_headers, raw_headers)
def _build_stdio_env(
self,
@ -3036,6 +3075,39 @@ class MCPServerManager:
)
raise_public(err)
async def preflight_id_jag(
self,
server: MCPServer,
oauth2_headers: dict[str, str] | None,
raw_headers: dict[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> None:
"""Run the ID-JAG exchange at the transport edge, where an HTTP status and
``WWW-Authenticate`` still reach the client.
The list handlers cannot surface auth failures (the MCP session manager serializes a
raise into a JSON-RPC error, losing the header), so without this preflight a dead stored
assertion reads as an empty tool/prompt/resource catalog on single-server routes. Unlike
the OBO preflight there is no subject-not-yet-supplied case to wave through: id_jag
sources its subject from the caller's presented id_token or the stored SSO assertion, so
every resolution outcome is already the caller's answer — the re-login 401 (with the
resolver's challenge header), the retryable 503 of a store or IdP outage, or the 412 of a
user who never signed in. The subject bearer is selected by the same one rule egress uses,
so the preflight and the session it admits cannot disagree. A successful exchange is
cached by the resolver, so the session's list/call reuses it.
"""
if server.auth_type != MCPAuth.oauth2_id_jag:
return
spec = to_server_spec(server)
if spec is None or not isinstance(spec.config, IdJagConfig):
return
subject_token = self._subject_bearer_token(server, raw_headers, oauth2_headers)
match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec):
case Ok(_):
return
case Error(err):
raise_public(err)
async def _create_mcp_client(
self,
server: MCPServer,
@ -3280,14 +3352,7 @@ class MCPServerManager:
stdio_env = self._build_stdio_env(server, raw_headers)
# token_exchange (OBO) discovery needs the caller's token too: list it with the user's own
# token (mirrors the call path), not v1's deleted client_credentials fallback. Other modes
# never read the inbound bearer, so leave subject_token None to avoid forwarding it.
subject_token = (
self._extract_bearer_token(oauth2_headers, raw_headers)
if server.auth_type == MCPAuth.oauth2_token_exchange
else None
)
subject_token = self._subject_bearer_token(server, raw_headers, oauth2_headers=oauth2_headers)
client = await self._create_mcp_client(
server=server,
@ -3354,6 +3419,16 @@ class MCPServerManager:
server_name=server.name,
) from e
verbose_logger.warning(f"Failed to get tools from server {server.name}: {str(e)}")
if e.status_code == 503:
# The resolver's retryable verdict (credential store or IdP did not answer). Keep
# it retryable on the listing edge: "internal" would report the gateway itself
# failed and turn a transient outage into a 500.
raise MCPServerListError(ServerListFault(tag="unavailable", status_code=503), server.name) from e
if e.status_code == 412:
# The resolver's per-caller precondition (e.g. no stored SSO identity assertion,
# remedied only by the caller signing in to the gateway). Neither an upstream auth
# challenge nor a gateway fault.
raise MCPServerListError(ServerListFault(tag="precondition", status_code=412), server.name) from e
raise MCPServerListError(ServerListFault(tag="internal", status_code=e.status_code), server.name) from e
except MCPServerListError:
raise
@ -3368,6 +3443,7 @@ class MCPServerManager:
extra_headers: Optional[dict[str, str]] = None,
add_prefix: bool = True,
raw_headers: Optional[dict[str, str]] = None,
user_api_key_auth: UserAPIKeyAuth | None = None,
) -> list[Prompt]:
"""
Helper method to get prompts from a single MCP server with prefixed names.
@ -3392,7 +3468,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env = self._build_stdio_env(server, raw_headers)
subject_token = self._obo_subject_token(server, raw_headers)
subject_token = self._subject_bearer_token(server, raw_headers)
client = await self._create_mcp_client(
server=server,
@ -3400,6 +3476,7 @@ class MCPServerManager:
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
)
prompts = await client.list_prompts()
@ -3419,6 +3496,7 @@ class MCPServerManager:
extra_headers: Optional[dict[str, str]] = None,
add_prefix: bool = True,
raw_headers: Optional[dict[str, str]] = None,
user_api_key_auth: UserAPIKeyAuth | None = None,
) -> list[Resource]:
"""Fetch available resources from a single MCP server."""
@ -3434,7 +3512,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env = self._build_stdio_env(server, raw_headers)
subject_token = self._obo_subject_token(server, raw_headers)
subject_token = self._subject_bearer_token(server, raw_headers)
client = await self._create_mcp_client(
server=server,
@ -3442,6 +3520,7 @@ class MCPServerManager:
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
)
resources = await client.list_resources()
@ -3461,6 +3540,7 @@ class MCPServerManager:
extra_headers: Optional[dict[str, str]] = None,
add_prefix: bool = True,
raw_headers: Optional[dict[str, str]] = None,
user_api_key_auth: UserAPIKeyAuth | None = None,
) -> list[ResourceTemplate]:
"""Fetch available resource templates from a single MCP server."""
@ -3476,7 +3556,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env = self._build_stdio_env(server, raw_headers)
subject_token = self._obo_subject_token(server, raw_headers)
subject_token = self._subject_bearer_token(server, raw_headers)
client = await self._create_mcp_client(
server=server,
@ -3484,6 +3564,7 @@ class MCPServerManager:
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
)
resource_templates = await client.list_resource_templates()
@ -3505,6 +3586,7 @@ class MCPServerManager:
mcp_auth_header: Optional[Union[str, dict[str, str]]] = None,
extra_headers: Optional[dict[str, str]] = None,
raw_headers: Optional[dict[str, str]] = None,
user_api_key_auth: UserAPIKeyAuth | None = None,
) -> ReadResourceResult:
"""Read resource contents from a specific MCP server."""
@ -3517,7 +3599,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env = self._build_stdio_env(server, raw_headers)
subject_token = self._obo_subject_token(server, raw_headers)
subject_token = self._subject_bearer_token(server, raw_headers)
client = await self._create_mcp_client(
server=server,
@ -3525,6 +3607,7 @@ class MCPServerManager:
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
)
return await client.read_resource(url)
@ -3537,6 +3620,7 @@ class MCPServerManager:
mcp_auth_header: Optional[Union[str, dict[str, str]]] = None,
extra_headers: Optional[dict[str, str]] = None,
raw_headers: Optional[dict[str, str]] = None,
user_api_key_auth: UserAPIKeyAuth | None = None,
) -> GetPromptResult:
"""Fetch a specific prompt definition from a single MCP server."""
@ -3549,7 +3633,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env = self._build_stdio_env(server, raw_headers)
subject_token = self._obo_subject_token(server, raw_headers)
subject_token = self._subject_bearer_token(server, raw_headers)
client = await self._create_mcp_client(
server=server,
@ -3557,6 +3641,7 @@ class MCPServerManager:
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
)
get_prompt_request_params = GetPromptRequestParams(
@ -4593,12 +4678,14 @@ class MCPServerManager:
subject_token: Optional[str],
user_api_key_auth: Optional[UserAPIKeyAuth],
) -> CallToolResult:
"""Call a token_exchange (OBO) tool; on an upstream 401/403 re-mint the token once and retry.
"""Call an exchange-mode (token_exchange OBO or id_jag) tool; on an upstream 401/403
re-mint the token once and retry.
The exchanged token is baked into the client at build time, so the retry invalidates the
cached exchange and rebuilds the client (which re-exchanges). One retry only: a non-auth
failure or a second auth failure degrades to the normal ``isError`` result, and a re-exchange
that now fails surfaces its own 401 challenge from ``_create_mcp_client``.
cached exchange and rebuilds the client (which re-exchanges — for id_jag stored-assertion
callers that includes re-reading the assertion). One retry only: a non-auth failure or a
second auth failure degrades to the normal ``isError`` result, and a re-exchange that now
fails surfaces its own 401 challenge from ``_create_mcp_client``.
"""
try:
return await client.call_tool(
@ -4633,7 +4720,7 @@ class MCPServerManager:
proxy_logging_obj: Optional[ProxyLogging],
host_progress_callback: Optional[Callable] = None,
hook_extra_headers: Optional[dict[str, str]] = None,
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
user_api_key_auth: UserAPIKeyAuth | None = None,
) -> CallToolResult:
"""
Call a regular MCP tool using the MCP client.
@ -4680,15 +4767,10 @@ class MCPServerManager:
if server_auth_header is None:
server_auth_header = mcp_auth_header
# Extract subject token for OAuth2 Token Exchange (OBO) and ID-JAG flows
subject_token: Optional[str] = None
subject_token: str | None = self._subject_bearer_token(mcp_server, raw_headers, oauth2_headers=oauth2_headers)
# The exchange modes leave extra_headers None (the resolver injects the minted bearer).
extra_headers: Optional[dict[str, str]] = None
if mcp_server.auth_type in (
MCPAuth.oauth2_token_exchange,
MCPAuth.oauth2_id_jag,
):
subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
elif mcp_server.auth_type == MCPAuth.oauth2:
if mcp_server.auth_type == MCPAuth.oauth2:
if mcp_server.has_client_credentials:
# For M2M OAuth servers, Authorization must come from token fetch.
extra_headers = None
@ -4788,7 +4870,7 @@ class MCPServerManager:
arguments=arguments,
)
if mcp_server.auth_type in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag) and subject_token:
if _obo_retry_covers_caller(mcp_server.auth_type, subject_token, user_api_key_auth):
# OBO / ID-JAG: the exchanged token may have been revoked/rotated upstream since it was
# cached, so an upstream 401 gets one invalidate + re-mint + retry. Gated to these modes;
# all others keep the plain single call below.
@ -5041,7 +5123,7 @@ class MCPServerManager:
subject_token: str | None = None
if isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)):
subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
subject_token = self._subject_bearer_token(mcp_server, raw_headers, oauth2_headers=oauth2_headers)
elif isinstance(spec.config, PassthroughConfig):
inbound_token, forwarded_headers = _take_forwarded_authorization(forwarded_headers)
per_server_token = _passthrough_token_from_mcp_auth_header(mcp_auth_header)
@ -5703,6 +5785,7 @@ class MCPServerManager:
mcp_auth_header=None,
extra_headers=extra_headers,
stdio_env=None,
user_api_key_auth=None, # internal path (cache warm-up / health check): no caller identity
)
try:
@ -5795,7 +5878,7 @@ class MCPServerManager:
async def get_all_allowed_mcp_servers(
self,
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
user_api_key_auth: UserAPIKeyAuth | None = None,
) -> list[LiteLLM_MCPServerTable]:
"""
Get all MCP servers that the user has access to.

11
litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py
View file

@ -18,6 +18,9 @@ from fastapi import HTTPException
from pydantic import SecretStr
from typing_extensions import assert_never
from litellm.proxy._experimental.mcp_server.outbound_credentials.credential_provenance import (
classify_inbound_provenance,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
ApiKeyConfig,
AuthorizationCodeConfig,
@ -48,15 +51,19 @@ def to_subject(user_api_key_auth: Optional[UserAPIKeyAuth], subject_token: Optio
"""Map v1's authenticated principal onto the resolver's Subject.
tenant_id / subject_id are empty for an unauthenticated caller; the per-user arms must reject
an empty subject rather than share one credential slot across callers.
an empty subject rather than share one credential slot across callers. The inbound token's
provenance (gateway-issued vs the caller's own external credential) is classified here, at the
edge, so the resolver core never has to recognize gateway credentials itself.
"""
inbound = SecretStr(subject_token) if subject_token else None
provenance = classify_inbound_provenance(subject_token)
if user_api_key_auth is None:
return Subject(tenant_id="", subject_id="", inbound_token=inbound)
return Subject(tenant_id="", subject_id="", inbound_token=inbound, inbound_provenance=provenance)
return Subject(
tenant_id=user_api_key_auth.org_id or user_api_key_auth.team_id or "",
subject_id=user_api_key_auth.user_id or "",
inbound_token=inbound,
inbound_provenance=provenance,
)

180
litellm/proxy/_experimental/mcp_server/outbound_credentials/credential_provenance.py Normal file
View file

@ -0,0 +1,180 @@
"""Provenance of the inbound subject token, decided at the resolver's v1 edge.
The resolver must never forward a gateway-issued or gateway-honored credential to an external
token endpoint as exchange subject material. Recognizing "is this a credential this gateway
issued" by an enumerated denylist of prefixes is a leak generator: the set of minted formats
grows, and every format the list has not caught yet is disclosed. So this module recognizes a
gateway credential by the one property every one of them shares and no external token shares:
it is cryptographically ours, signed or encrypted under this gateway's master key. A new mint
format is caught by construction, with no edit here.
``classify_inbound_provenance`` is the single owner both exchange arms read through
(``Subject.inbound_provenance``). id_jag's subject is an id_token (a JWT), so it exchanges the
inbound bearer only when it is ``external_jwt`` and otherwise sources the stored assertion (an
``external_opaque`` value is not an id_token and is never forwarded); token_exchange (OBO) keeps
its exchange-what-was-presented contract for both ``external_jwt`` and ``external_opaque`` tokens
but refuses a ``gateway_credential``.
"""
from __future__ import annotations
import hashlib
import secrets
import jwt
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
ENVELOPE_ISSUER,
ENVELOPE_PREFIX,
REFRESH_ENVELOPE_PREFIX,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
SESSION_ISSUER,
SESSION_REFRESH_PREFIX,
SESSION_TOKEN_PREFIX,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import InboundTokenProvenance
_VIRTUAL_KEY_PREFIX = "sk-"
_GATEWAY_MINT_PREFIXES = (
_VIRTUAL_KEY_PREFIX,
SESSION_TOKEN_PREFIX,
SESSION_REFRESH_PREFIX,
ENVELOPE_PREFIX,
REFRESH_ENVELOPE_PREFIX,
)
_RESERVED_GATEWAY_ISSUERS = frozenset({SESSION_ISSUER, ENVELOPE_ISSUER})
# Recognition asks whether a token is THIS gateway's, not whether it is still usable, so every
# decode below disables claim validation (expiry, not-before, issued-at, audience, issuer). An
# expired or not-yet-valid gateway credential is still a gateway credential; letting a claim check
# reject it would drop it back to the exchangeable-external path and disclose it upstream. Only the
# signature (for the master-key check) speaks to provenance.
_RECOGNITION_CLAIM_CHECKS_DISABLED = {
"verify_exp": False,
"verify_nbf": False,
"verify_iat": False,
"verify_aud": False,
"verify_iss": False,
}
def _master_key() -> str | None:
from litellm.proxy.proxy_server import master_key
return master_key
def _equals_master_key(token: str, master_key: str) -> bool:
"""Constant-time equality with the master key, TOTAL over any input string.
``secrets.compare_digest`` raises ``TypeError`` on a ``str`` carrying non-ASCII characters, and
recognition must never raise into egress (an inbound bearer is fully attacker-controlled). Both
sides are hashed to fixed-length digests first, so any string compares without raising and the
master key's length is not leaked; a non-ASCII bearer reads "not the master key" rather than
500-ing the request.
"""
token_digest = hashlib.sha256(token.encode("utf-8", "surrogatepass")).digest()
master_key_digest = hashlib.sha256(master_key.encode("utf-8", "surrogatepass")).digest()
return secrets.compare_digest(token_digest, master_key_digest)
def _verifies_under_master_key(token: str, master_key: str) -> bool:
"""True when ``token`` is a JWT this gateway signed with the master key (HS256).
Covers every master-key-signed mint (UI login session, onboarding, byok_session) and any
future one, since the check is the SIGNATURE, not the format and not the claims. Claim
validation is disabled: an expired or not-yet-valid gateway token is still this gateway's and
must be recognized, or it would fall through to the exchangeable-external path and leak.
"""
try:
jwt.decode(
token,
master_key,
algorithms=["HS256"],
options={"verify_signature": True, **_RECOGNITION_CLAIM_CHECKS_DISABLED},
)
return True
except Exception: # noqa: BLE001 # signature (not claims) failed -> not signed by us; other checks still run
return False
def _claims_a_reserved_issuer(token: str) -> bool:
"""True when ``token`` is a JWT whose ``iss`` is one this gateway reserves (session/envelope).
These are signed with master-key-DERIVED keys, so they do not verify under the master key
directly; they are already caught by their mint prefix, and this is the belt-and-suspenders
for an inner token presented without its prefix. The issuer is refused on the unverified
claim alone because the value is reserved to this gateway: no external token legitimately
carries it, so treating it as ours fails closed. Claim validation stays off for the same
reason as the master-key check: an expired reserved-issuer token is still ours.
"""
try:
claims = jwt.decode(token, options={"verify_signature": False, **_RECOGNITION_CLAIM_CHECKS_DISABLED})
except Exception: # noqa: BLE001 # not a JWT
return False
return claims.get("iss") in _RESERVED_GATEWAY_ISSUERS
def _decrypts_under_gateway_key(token: str) -> bool:
"""True when ``token`` decrypts under this gateway's salt key (the CLI / experimental UI
login blobs, which are encrypted rather than JWT-signed). Decryption succeeding is proof the
value is ours; an external token decrypts to nothing."""
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
try:
decrypted = decrypt_value_helper(
token, key="mcp_subject_token_provenance", exception_type="debug", return_original_value=False
)
except Exception: # noqa: BLE001 # defensive: never let recognition raise into egress
return False
return decrypted is not None
def is_gateway_issued_credential(token: str) -> bool:
"""Whether ``token`` is a credential this gateway issued or honors as its own.
Complete by construction rather than by enumeration: a gateway credential is caught if it is
the master key, carries a gateway mint prefix, verifies under the master key, claims a
reserved gateway issuer, or decrypts under the gateway salt key. Every current mint (virtual
key, master key, MCP session/refresh, bridge envelope, UI login session, onboarding,
byok_session, CLI/experimental login) matches at least one, and a future mint signed or
encrypted with the master key matches without a change here.
"""
master_key = _master_key()
if master_key and _equals_master_key(token, master_key):
return True
if token.startswith(_GATEWAY_MINT_PREFIXES):
return True
if master_key and _verifies_under_master_key(token, master_key):
return True
if _claims_a_reserved_issuer(token):
return True
return _decrypts_under_gateway_key(token)
def _is_decodable_jwt(token: str) -> bool:
"""Whether ``token`` is structurally a JWT (an id_token is a JWT). Signature and every
registered-claim check are disabled: this is a shape test to tell an id_token candidate from
an opaque bearer, not a validation. The gateway check runs first, so this only ever runs on a
non-gateway token; the org authorization server remains the authority on the id_token itself.
"""
try:
jwt.decode(token, options={"verify_signature": False, **_RECOGNITION_CLAIM_CHECKS_DISABLED})
return True
except Exception: # noqa: BLE001 # not a JWT -> an opaque bearer, not an id_token candidate
return False
def classify_inbound_provenance(subject_token: str | None) -> InboundTokenProvenance:
"""The one classification both exchange arms read. A gateway-issued credential is never
exchange subject material (any mode). Among the caller's own credentials, a JWT is an id_token
candidate (``external_jwt``) that id_jag can exchange, while an opaque value (``external_opaque``)
is not an id_token and id_jag falls to the stored assertion; OBO may forward either."""
if not subject_token:
return "absent"
if is_gateway_issued_credential(subject_token):
return "gateway_credential"
if _is_decodable_jwt(subject_token):
return "external_jwt"
return "external_opaque"

26
litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py
View file

@ -91,6 +91,13 @@ class TokenRefresher(Protocol):
``server_id`` selects the upstream's config (token endpoint, client credentials, scopes) the
grant runs against; ``(user_id, server_id)`` is the key the new token is persisted under. They
are not derivable from ``token``, so the seam threads them alongside it.
CONTRACT: an implementation MAY raise a domain exception instead of returning, to carry an
outcome richer than ``OAuthToken | None`` (the EMA assertion refresher raises its dead-grant
and store-outage verdicts, each mapping to a different caller remedy). ``RefreshingTokenStore``
and the refresh coordinator MUST propagate such exceptions to the refresh winner and every
single-flight waiter alike — wrapping a refresher call in ``try/except -> None`` here would
silently collapse those verdicts into "no token" and change caller-visible statuses.
"""
async def refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None: ...
@ -159,6 +166,7 @@ class CachedOAuthTokenStore:
*,
default_ttl_seconds: float,
expiry_skew_seconds: float = 60.0,
max_ttl_seconds: float | None = None,
max_size: int = 4096,
backend: TokenCacheBackend | None = None,
clock: Callable[[], float] = time.time,
@ -166,13 +174,25 @@ class CachedOAuthTokenStore:
self._inner = inner
self._default_ttl_seconds = default_ttl_seconds
self._expiry_skew_seconds = expiry_skew_seconds
self._max_ttl_seconds = max_ttl_seconds
self._clock = clock
self._backend: TokenCacheBackend = backend or InMemoryTokenCacheBackend(max_size=max_size, clock=clock)
def _ttl(self, token: OAuthToken) -> float:
if token.expires_at is not None:
return max(0.0, token.expires_at - self._expiry_skew_seconds - self._clock())
return self._default_ttl_seconds
"""How long this token may be served from cache.
A token that declares an expiry is held until then (minus the skew); one that does not is
held for the default. ``max_ttl_seconds`` caps both, for a credential whose holder must
converge on a replacement sooner than the token's own life: an identity assertion is
re-issued by a re-login that may REDUCE the user's claims, so serving the superseded one
for its full life would keep the old claims usable.
"""
ttl = (
max(0.0, token.expires_at - self._expiry_skew_seconds - self._clock())
if token.expires_at is not None
else self._default_ttl_seconds
)
return ttl if self._max_ttl_seconds is None else min(ttl, self._max_ttl_seconds)
async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
hit = await self._backend.get(user_id, server_id)

111
litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py
View file

@ -19,6 +19,7 @@ from __future__ import annotations
import hashlib
from functools import partial
from typing import Protocol
import httpx
from typing_extensions import assert_never
@ -41,6 +42,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.result import (
Ok,
Result,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
SsoAssertionUnrenewable,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import (
ExchangedToken,
ExchangedTokenCache,
@ -74,6 +78,29 @@ _JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
_ID_JAG_REQUESTED_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:id-jag"
def _inbound_id_token(subject: Subject) -> str | None:
"""The inbound bearer when it is the caller's own presented id_token, else None so id_jag
sourcing falls to the stored assertion. The ID-JAG exchange takes an id_token (a JWT), so the
inbound is used only when the edge classified it ``external_jwt``: a gateway-issued credential
and an opaque (non-JWT) bearer are both rejected here and resolve through the store instead."""
if subject.inbound_token is None or subject.inbound_provenance != "external_jwt":
return None
return subject.inbound_token.get_secret_value()
class SsoAssertionSource(Protocol):
"""Where the ID-JAG arm finds the caller's stored SSO identity assertion by user id."""
async def fetch_usable(self, user_id: str) -> OAuthToken | None: ...
class _NullSsoAssertionSource:
"""Fail-closed default: with no assertion source wired, no stored assertion is ever found."""
async def fetch_usable(self, user_id: str) -> OAuthToken | None:
return None
class _NullOAuthTokenStore:
"""Fail-closed default: with no token store wired, every user reads as not authorized."""
@ -111,12 +138,14 @@ class UpstreamCredentialProvider:
token_endpoint: TokenEndpointClient | None = None,
exchanged_tokens: ExchangedTokenCache | None = None,
client_credentials_source: ClientCredentialsTokenSource | None = None,
sso_assertions: SsoAssertionSource | None = None,
) -> None:
self._oauth_token_store: OAuthTokenStore = oauth_token_store or _NullOAuthTokenStore()
self._token_exchanger: TokenExchanger = token_exchanger or _NullTokenExchanger()
self._token_endpoint: TokenEndpointClient = token_endpoint or TokenEndpointClient()
self._exchanged_tokens: ExchangedTokenCache = exchanged_tokens or ExchangedTokenCache()
self._client_credentials_source = client_credentials_source or ClientCredentialsTokenSource()
self._sso_assertions: SsoAssertionSource = sso_assertions or _NullSsoAssertionSource()
async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]:
match server.config:
@ -170,15 +199,55 @@ class UpstreamCredentialProvider:
return Error(CredError.of_not_implemented("api_key BYOK source not implemented yet"))
assert_never(config.key_source)
async def _id_jag(self, subject: Subject, server: ServerSpec, config: IdJagConfig) -> Result[httpx.Auth, CredError]:
if subject.inbound_token is None:
async def _id_jag_subject_token(self, subject: Subject) -> Result[str, CredError]:
"""The ONE sourcing rule for the ID-JAG subject token, shared by resolve and invalidate.
The ID-JAG exchange takes an id_token, so the subject is the caller's presented id_token
when the inbound bearer is one (the live, freshest identity evidence), else the stored SSO
assertion (itself an id_token, renewed when expired), else fail closed. Anything that is
not an id_token, a gateway-issued credential or an opaque bearer, is never forwarded: it
resolves through the store or not at all. An expired-and-unrenewable stored assertion is a
401 re-login challenge, never a silent fall-through: the user HAS connected and the only
fix is theirs to perform."""
inbound = _inbound_id_token(subject)
if inbound is not None:
return Ok(inbound)
try:
assertion = await self._sso_assertions.fetch_usable(subject.subject_id)
except SsoAssertionUnrenewable:
return Error(
CredError.of_precondition_required(
"ID-JAG requires a caller identity token; it asserts the calling "
"user's identity upstream and cannot use a static credential."
CredError.of_unauthorized(
"The stored identity assertion for this user has expired and could not "
"be renewed; sign in to the gateway again to re-establish it.",
www_authenticate=(
'Bearer error="invalid_token", error_description="identity assertion expired; re-authenticate"'
),
)
)
token = subject.inbound_token.get_secret_value()
except TokenStoreUnavailable:
# Covers both causes with one honest, cause-neutral message: the assertion store (DB)
# not answering on the read and the IdP not answering on the renewal. Naming either
# one would misdiagnose the other.
return Error(
CredError.of_upstream_unavailable(
"the stored identity assertion could not be read or renewed; retry shortly"
)
)
if assertion is None:
return Error(
CredError.of_precondition_required(
"no stored identity assertion exists for this user; ID-JAG asserts the "
"calling user's identity upstream, so sign in to the gateway via SSO to "
"establish one (or present an IdP id_token as the bearer)."
)
)
return Ok(assertion.access_token)
async def _id_jag(self, subject: Subject, server: ServerSpec, config: IdJagConfig) -> Result[httpx.Auth, CredError]:
match await self._id_jag_subject_token(subject):
case Error(err):
return Error(err)
case Ok(resolved):
token = resolved
cache_key = _id_jag_cache_key(token, server.server_id, config)
async def _exchange() -> Result[ExchangedToken, CredError]:
@ -248,6 +317,12 @@ class UpstreamCredentialProvider:
No inbound token means there is nothing to exchange, so the arm fails closed with a 401 rather
than falling through to a weaker source (§1.5); the exchanger handles the IdP round-trip and
caching and returns the upstream token or a typed error.
The arm keeps its exchange-what-was-presented contract for the caller's own external token,
but a gateway-issued credential (an admission key, session token, or gateway-signed JWT that
arrived on Authorization) is never disclosed to the external token endpoint; it fails closed
with the same 401 as a missing token so a misdirected gateway credential is refused, not
forwarded.
"""
inbound = subject.inbound_token
if inbound is None:
@ -257,6 +332,14 @@ class UpstreamCredentialProvider:
www_authenticate='Bearer error="invalid_request"',
)
)
if subject.inbound_provenance == "gateway_credential":
return Error(
CredError.of_unauthorized(
"Token exchange refuses a gateway-issued credential as the subject token; "
"present the caller's own identity token to exchange (OBO).",
www_authenticate='Bearer error="invalid_request"',
)
)
match await self._token_exchanger.exchange(
inbound.get_secret_value(), server, config, tenant_id=subject.tenant_id
):
@ -272,18 +355,22 @@ class UpstreamCredentialProvider:
than serving the same rejected token until TTL. `token_exchange` and `id_jag` hold a
re-mintable cached credential here; `client_credentials` recovers inside its own auth flow
(`ClientCredentialsBearerAuth` retries the 401'd request once with a fresh token), and
other modes are a no-op.
other modes are a no-op. The `id_jag` eviction resolves its subject token through the same
rule that minted the cache entry, so a stored-assertion caller (no inbound token) evicts
the entry its resolve created; a failed resolution means no entry could exist, a no-op.
"""
if subject.inbound_token is None:
return
if isinstance(server.config, TokenExchangeConfig):
if subject.inbound_token is None:
return
await self._token_exchanger.invalidate(
subject.inbound_token.get_secret_value(), server, server.config, tenant_id=subject.tenant_id
)
if isinstance(server.config, IdJagConfig):
self._exchanged_tokens.invalidate(
_id_jag_cache_key(subject.inbound_token.get_secret_value(), server.server_id, server.config)
)
match await self._id_jag_subject_token(subject):
case Ok(resolved):
self._exchanged_tokens.invalidate(_id_jag_cache_key(resolved, server.server_id, server.config))
case Error(_):
return
async def _authz_token(self, subject: Subject, server: ServerSpec) -> OAuthToken | None:
"""The user's authorization_code token, or None when absent or the store is unreachable.

385
litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py
View file

@ -5,8 +5,12 @@ The ``oauth2_id_jag`` egress arm needs the user's IdP ``id_token`` as its RFC 86
IdP assertion, so the assertion captured at the one SSO login is the only usable subject
source for it. This module owns both sides of that state: the SSO callback persists here
(write-through to the DB so a login on one pod is visible to every pod) and the resolver
seam reads back by ``user_id``. Retention is gated on an ``oauth2_id_jag`` server actually
being registered, so a gateway with no EMA upstream never stores bearer material.
seam reads back by ``user_id``. Capture currently hooks the generic SSO connector only (the
one whose ``GENERIC_*`` client also renews assertions below); Microsoft/Google SSO logins do
not feed this store, so EMA on those connectors answers 412 until they do. Retention is gated
on an ``oauth2_id_jag`` server actually being registered, so a gateway with no EMA upstream
never BEGINS storing bearer material; rows persisted while one existed are not purged by its
removal and remain until the user's next login replaces them.
The row is one encrypted payload per user, latest login wins. ``expires_at`` mirrors the
id_token ``exp`` claim and is judged by the reader, never enforced by deletion here: an
@ -17,6 +21,10 @@ truth, the same contract as the per-user OAuth credential store.
from __future__ import annotations
import json
import os
import time
import weakref
from collections.abc import Awaitable, Callable
from datetime import datetime, timezone
from typing import TYPE_CHECKING
@ -24,6 +32,12 @@ import jwt
from pydantic import BaseModel, ConfigDict, SecretStr, TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
CachedOAuthTokenStore,
OAuthToken,
RefreshingTokenStore,
TokenStoreUnavailable,
)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
@ -32,6 +46,20 @@ _ASSERTION_DECRYPT_LOG_KEY = "sso_identity_assertion"
_STR_ADAPTER: TypeAdapter[str] = TypeAdapter(str)
_MAYBE_STR_ADAPTER: TypeAdapter[str | None] = TypeAdapter(str | None)
# An assertion this close to expiry is treated as expired, so a subject token is never handed
# to the exchange with less lifetime than the two token-endpoint legs need to complete.
_ASSERTION_EXPIRY_BUFFER_SECONDS = 30
# The ceiling on how long any assertion is served from this pod's cache before the DB row is
# re-read: also the recovery bound for the one race the write chokepoint's invalidate cannot
# close (a read in flight across a re-login re-caching the superseded row). An assertion
# declaring an earlier expiry leaves the cache at that expiry (minus the buffer above) instead.
_ASSERTION_CACHE_TTL_SECONDS = 60.0
# The assertion is the user's SSO identity rather than a per-upstream credential, so it occupies a
# single per-user slot in the shared (user_id, server_id) keyed stack.
_ASSERTION_SERVER_KEY = ""
class SSOIdentityAssertion(BaseModel):
"""The IdP material an EMA exchange needs: ``id_token`` is the RFC 8693 subject token,
@ -127,31 +155,42 @@ async def persist_sso_identity_assertion(user_id: str, assertion: SSOIdentityAss
"update": {"assertion_b64": encoded},
},
)
# This is the one place the row is replaced, so it is the one place that must drop a cached
# predecessor: otherwise a re-login (notably one that reduces the user's IdP claims) would keep
# serving the superseded assertion until the old id_token expired.
await _drop_cached_assertion(user_id)
async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | None:
"""The stored assertion for ``user_id``, or ``None`` when absent, undecryptable (salt-key
rotation), or unparseable. Expiry is not judged here; the reader owns that policy."""
"""The stored assertion for ``user_id``, or ``None`` only when no row was ever written.
``None`` is reserved for that one determinate fact, so a fault can never be reported as "this
user has not signed in". A store that cannot be reached raises ``TokenStoreUnavailable``, which
is the contract ``OAuthTokenStore`` already defines for exactly this, and a row that cannot be
decrypted (salt-key rotation) or parsed raises ``SsoAssertionUnrenewable`` because the material
is unusable and only a fresh sign-in replaces it. Expiry is not judged here; the reader owns
that policy.
"""
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper # noqa: PLC0415 # runtime global
from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime global
if prisma_client is None:
return None
row = await prisma_client.db.litellm_ssoidentityassertion.find_unique(where={"user_id": user_id})
raise TokenStoreUnavailable("the assertion store is not connected")
try:
row = await prisma_client.db.litellm_ssoidentityassertion.find_unique(where={"user_id": user_id})
except Exception as exc: # noqa: BLE001 # an unreadable store is indeterminate, never an absence
raise TokenStoreUnavailable("the assertion store could not be read") from exc
if row is None:
return None
raw = _MAYBE_STR_ADAPTER.validate_python(
decrypt_value_helper(row.assertion_b64, _ASSERTION_DECRYPT_LOG_KEY, exception_type="debug")
)
if raw is None:
return None
raise SsoAssertionUnrenewable("the stored identity assertion could not be decrypted")
try:
payload = _StoredAssertionPayload.model_validate_json(raw)
except ValidationError:
verbose_proxy_logger.warning(
"Stored SSO identity assertion for user_id=%s could not be parsed; treating as absent.", user_id
)
return None
except ValidationError as exc:
raise SsoAssertionUnrenewable("the stored identity assertion could not be parsed") from exc
return SSOIdentityAssertion(
id_token=SecretStr(payload.id_token),
refresh_token=SecretStr(payload.refresh_token) if payload.refresh_token else None,
@ -198,6 +237,328 @@ async def rotate_sso_identity_assertions_master_key(prisma_client: PrismaClient,
)
class SsoAssertionUnrenewable(Exception):
"""The stored assertion is expired and its grant is definitively dead; only a fresh sign-in
produces a new one. Raised rather than returned so the verdict survives the shared refresh
stack, whose ``OAuthToken | None`` would otherwise collapse "the grant is dead" (re-login)
into "this user never connected" (a different remedy, and a different status)."""
# The renewal POST answers with the token body or raises the verdict its failure proves, so the
# outcome needs no intermediate representation: an unrenewable grant and an unreachable IdP are
# exactly the two terminal states the shared stack already propagates.
SsoTokenEndpointPost = Callable[[str, dict[str, str]], Awaitable[dict[str, object]]]
class _SsoRefreshClient(BaseModel):
"""The gateway's own SSO client registration at the IdP. The stored refresh token was
issued to THIS client (the generic SSO app), never to an ``id_jag`` server's client, so
renewing the assertion must authenticate as it or the IdP answers ``invalid_grant``."""
model_config = ConfigDict(frozen=True)
client_id: str
client_secret: SecretStr | None
token_endpoint: str
def generic_sso_scopes(getenv: Callable[[str], str | None] = os.getenv) -> list[str]:
"""The scopes the generic SSO client requests, shared by the login authorize request and the
EMA assertion refresh so the two can never diverge. ``openid`` must be among them or the IdP
returns no ``id_token`` on refresh, which would strand a renewable assertion as expired and
force a needless re-login; requesting exactly what login was granted keeps the refresh within
RFC 6749 scope while guaranteeing the id_token comes back. The refresh passes the source's
injected ``getenv`` so it reads the same environment as the rest of its client config."""
raw = getenv("GENERIC_SCOPE")
return (raw if raw is not None else "openid email profile").split(" ")
def _sso_refresh_client_from_env(getenv: Callable[[str], str | None]) -> _SsoRefreshClient | None:
client_id = getenv("GENERIC_CLIENT_ID")
token_endpoint = getenv("GENERIC_TOKEN_ENDPOINT")
if not client_id or not token_endpoint:
return None
secret = getenv("GENERIC_CLIENT_SECRET")
return _SsoRefreshClient(
client_id=client_id,
client_secret=SecretStr(secret) if secret else None,
token_endpoint=token_endpoint,
)
def _oauth_error_code(response: object) -> str | None:
"""The RFC 6749 section 5.2 ``error`` code from a token-endpoint response body, or None
when the body is not a JSON object carrying one."""
try:
body = response.json() # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType,reportUnknownVariableType] # httpx response is partially typed
except Exception: # noqa: BLE001 # an unparseable body simply carries no code
return None
if not isinstance(body, dict):
return None
error = body.get("error") # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # narrowed below
return error if isinstance(error, str) and error else None
async def _post_sso_token_endpoint(url: str, form: dict[str, str]) -> dict[str, object]:
"""The renewal grant, answering with the token body or raising what its failure proves.
Unrenewable requires PROOF of a verdict on the GRANT, and RFC 6749 section 5.2 has exactly
one code that states that fact: ``invalid_grant`` ("the provided ... refresh token is
invalid, expired, revoked"). Every other code describes the CLIENT or the REQUEST —
``invalid_client`` is this gateway's own credentials failing, ``invalid_scope`` its scope
config — which an operator fixes with no help from the user, so recording them on the row
would permanently destroy renewable assertions org-wide over an ops mistake. Those, plus a
429, a 4xx without the error object (an intermediary answering, not the token endpoint), a
5xx, a transport failure, or a 2xx whose body is not a JSON object, all read unavailable, so
the caller retries instead of being told to sign in again.
"""
import httpx # noqa: PLC0415
from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # http_handler is untyped
)
from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415
try:
client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) # pyright: ignore[reportUnknownVariableType] # http_handler is untyped
response = await client.post(url, headers={"Accept": "application/json"}, data=form) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # httpx handler partially typed
response.raise_for_status() # pyright: ignore[reportUnknownMemberType] # httpx handler partially typed
body: object = response.json() # pyright: ignore[reportUnknownMemberType] # shape-validated below
except httpx.HTTPStatusError as exc:
status_code = exc.response.status_code
error_code = _oauth_error_code(exc.response)
if 400 <= status_code < 500 and status_code != 429 and error_code == "invalid_grant":
raise SsoAssertionUnrenewable(f"the identity provider rejected the refresh grant ({status_code})") from exc
raise TokenStoreUnavailable(
f"the identity provider answered {status_code}" + (f" ({error_code})" if error_code else "")
) from exc
except (SsoAssertionUnrenewable, TokenStoreUnavailable):
raise
except Exception as exc: # noqa: BLE001 # transport/parse failure carries no verdict on the grant
raise TokenStoreUnavailable(f"the identity provider could not be reached ({type(exc).__name__})") from exc
if not isinstance(body, dict):
raise TokenStoreUnavailable("the identity provider returned a non-object JSON body")
return body
def _to_oauth_token(assertion: SSOIdentityAssertion) -> OAuthToken:
"""The assertion as the shared stack's credential: the id_token is the value EMA spends (the
RFC 8693 subject token), so it rides in ``access_token``. ``issuer`` has no reader."""
return OAuthToken(
access_token=assertion.id_token.get_secret_value(),
expires_at=assertion.expires_at.timestamp() if assertion.expires_at is not None else None,
refresh_token=assertion.refresh_token.get_secret_value() if assertion.refresh_token else None,
)
class _SsoAssertionDbStore:
"""``OAuthTokenStore`` over the EMA assertion row. The assertion is the user's SSO identity,
not a per-upstream credential, so ``server_id`` is ignored: every id_jag server spends it."""
def __init__(self, fetch: Callable[[str], Awaitable[SSOIdentityAssertion | None]]) -> None:
self._fetch = fetch
async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
assertion = await self._fetch(user_id)
return None if assertion is None else _to_oauth_token(assertion)
class _SsoAssertionRefresher:
"""``TokenRefresher`` that renews the assertion as the gateway's OWN generic SSO client, whose
registration issued the stored refresh token (an id_jag server's client would get
``invalid_grant``). A renewed token is returned and persisted best-effort, so a rotated refresh
token is not lost while an in-hand one is never discarded; a dead grant raises
``SsoAssertionUnrenewable`` (re-login) and anything proving nothing about the grant raises
``TokenStoreUnavailable`` (retry), so a down IdP never tells a user to sign in again.
A PROVEN dead grant is recorded where the assertion lives, by persisting the row without the
refresh token the IdP rejected, never in a cache beside it. Later reads then answer from the
row with no token-endpoint call, on every pod and for as long as the row stands, and a re-login
lifts it by replacing that row. A verdict in a side cache guarantees neither: it is per-pod, it
expires on its own clock, and it can refuse a caller who has just signed in again.
"""
def __init__(
self,
fetch: Callable[[str], Awaitable[SSOIdentityAssertion | None]],
persist: Callable[[str, SSOIdentityAssertion], Awaitable[None]],
post: SsoTokenEndpointPost,
getenv: Callable[[str], str | None],
now: Callable[[], float] = time.time,
) -> None:
self._fetch = fetch
self._persist = persist
self._post = post
self._getenv = getenv
self._now = now
async def _record_dead_grant(self, user_id: str, rejected_refresh_token: str) -> None:
"""Strip the rejected refresh token from the row — but only while the row still carries it.
The verdict is a fact about ONE grant, so it is written by re-reading the row and rewriting
its own current material minus that grant, never by fabricating a row from this pod's stale
snapshot. When the re-read shows a different refresh token, another pod's renewal (the IdP
rotated the token this pod's POST lost the race to) or a fresh login already replaced the
grant this verdict describes, and writing it would clobber a live assertion into a forced
re-login. Best effort throughout: failing to record the verdict costs one more POST next
time, never a wrong answer."""
try:
current = await self._fetch(user_id)
if (
current is None
or current.refresh_token is None
or current.refresh_token.get_secret_value() != rejected_refresh_token
):
return
await self._persist(
user_id,
SSOIdentityAssertion(
id_token=current.id_token,
refresh_token=None,
issuer=current.issuer,
expires_at=current.expires_at,
),
)
except Exception as exc: # noqa: BLE001 # recording the verdict must not mask it
verbose_proxy_logger.warning(
"Could not record the rejected SSO refresh grant for user_id=%s: %s", user_id, exc
)
async def refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None:
spent_refresh_token = token.refresh_token
if spent_refresh_token is None:
# No grant to spend, because the login never returned one or a proven rejection stripped
# it. Either way only a re-login helps, and answering costs no token-endpoint call.
raise SsoAssertionUnrenewable("the stored assertion carries no refresh token")
client = _sso_refresh_client_from_env(self._getenv)
if client is None:
verbose_proxy_logger.warning(
"Stored SSO assertion for user_id=%s is expired and the SSO client env "
"(GENERIC_CLIENT_ID/GENERIC_TOKEN_ENDPOINT) is not set; renewal is unavailable "
"until the deployment restores it.",
user_id,
)
raise TokenStoreUnavailable("the SSO client environment is not configured")
scope = " ".join(generic_sso_scopes(self._getenv))
form = {
"grant_type": "refresh_token",
"refresh_token": spent_refresh_token,
"client_id": client.client_id,
**({"client_secret": client.client_secret.get_secret_value()} if client.client_secret else {}),
**({"scope": scope} if scope else {}),
}
try:
body = await self._post(client.token_endpoint, form)
except SsoAssertionUnrenewable:
await self._record_dead_grant(user_id, spent_refresh_token)
raise
rotated = body.get("refresh_token")
carried_refresh = rotated if isinstance(rotated, str) and rotated else spent_refresh_token
refreshed = assertion_from_sso_login(body.get("id_token"), carried_refresh)
renewed = _to_oauth_token(refreshed) if refreshed is not None else None
if renewed is not None and (
renewed.expires_at is None or self._now() < renewed.expires_at - _ASSERTION_EXPIRY_BUFFER_SECONDS
):
try:
await self._persist(user_id, refreshed)
except Exception as exc: # noqa: BLE001 # a write-back failure must not discard the in-hand token
verbose_proxy_logger.warning(
"SSO assertion refresh for user_id=%s succeeded but the write-back failed; serving the "
"in-hand assertion (a rotated refresh token was not stored): %s",
user_id,
exc,
)
return renewed
# A 2xx that produced no usable id_token, because the IdP omitted one or returned one already
# past expiry, spent the grant for nothing. It is recorded dead on the row exactly like a 4xx
# rejection, so the next read answers re-login without another POST; keeping the refresh token
# to retry would re-spend a one-time grant on every egress and never converge, and a fresh
# sign-in (which requests openid at the authorize leg) is what actually restores a usable one.
verbose_proxy_logger.warning(
"SSO assertion refresh for user_id=%s produced no usable id_token; treating as expired.", user_id
)
await self._record_dead_grant(user_id, spent_refresh_token)
raise SsoAssertionUnrenewable("the refresh produced no usable id_token")
# Live sources, so the one write chokepoint can drop a superseded assertion from their caches the
# moment the row is replaced. Weak, so a discarded source never keeps an instance alive.
_LIVE_ASSERTION_SOURCES: weakref.WeakSet[LiveSsoAssertionSource] = weakref.WeakSet()
async def _drop_cached_assertion(user_id: str) -> None:
for source in tuple(_LIVE_ASSERTION_SOURCES):
try:
await source.invalidate(user_id)
except Exception as exc: # noqa: BLE001 # a cache drop must never fail the login write
verbose_proxy_logger.warning(
"Could not drop the cached SSO assertion for user_id=%s after its row was replaced: %s", user_id, exc
)
class LiveSsoAssertionSource:
"""The resolver's view of the assertion store, built on the shared per-user credential stack.
``Cached(Refreshing(db))`` is the composition ``authorization_code`` uses, so the read-through
cache (positive-only, never caching an absence, so a fresh login is seen at once), the
expiry-skewed renewal and the per-user single-flight all come from ``oauth_token_store``
instead of being rebuilt here; passing a Redis cache and coordinator would extend that
single-flight across replicas without touching this class.
That stack answers ``OAuthToken | None`` while EMA needs four outcomes, each with its own
remedy, so the refresher raises its two terminal verdicts and this adapter maps them with no
second read: a token is usable, ``SsoAssertionUnrenewable`` is re-login (401),
``TokenStoreUnavailable`` is retry (503), and ``None`` then means no row was ever stored (412).
"""
def __init__(
self,
fetch: Callable[[str], Awaitable[SSOIdentityAssertion | None]] = fetch_sso_identity_assertion,
persist: Callable[[str, SSOIdentityAssertion], Awaitable[None]] = persist_sso_identity_assertion,
post: SsoTokenEndpointPost = _post_sso_token_endpoint,
getenv: Callable[[str], str | None] = os.getenv,
cache_ttl_seconds: float = _ASSERTION_CACHE_TTL_SECONDS,
now: Callable[[], float] = time.time,
) -> None:
self._store = CachedOAuthTokenStore(
RefreshingTokenStore(
_SsoAssertionDbStore(fetch),
_SsoAssertionRefresher(fetch, persist, post, getenv, now=now),
expiry_skew_seconds=_ASSERTION_EXPIRY_BUFFER_SECONDS,
clock=now,
),
default_ttl_seconds=cache_ttl_seconds,
expiry_skew_seconds=_ASSERTION_EXPIRY_BUFFER_SECONDS,
max_ttl_seconds=cache_ttl_seconds,
clock=now,
)
_LIVE_ASSERTION_SOURCES.add(self)
async def invalidate(self, user_id: str) -> None:
"""Drop this pod's cached assertion for ``user_id`` so the next read sees the replaced row.
Called from the write chokepoint, so a re-login (in particular one that reduces the user's
claims) is honored immediately rather than when the superseded id_token finally expires.
"""
await self._store.invalidate(user_id, _ASSERTION_SERVER_KEY)
async def fetch_usable(self, user_id: str) -> OAuthToken | None:
"""The caller's usable assertion, ``None`` only when no row was ever stored.
``SsoAssertionUnrenewable`` (re-login) and ``TokenStoreUnavailable`` (retry) propagate, so
the arm maps three remedies without a parallel result union. ``None`` carries exactly one
meaning: anything indeterminate reads as unavailable, so a store outage is a retry rather
than a false "this user has never signed in", and never a 500.
"""
if not user_id:
return None
try:
return await self._store.fetch(user_id, _ASSERTION_SERVER_KEY)
except (SsoAssertionUnrenewable, TokenStoreUnavailable):
raise
except Exception as exc: # noqa: BLE001 # indeterminate, so never absence and never a 500
raise TokenStoreUnavailable("the stored identity assertion could not be read") from exc
async def retain_sso_identity_assertion_for_ema(user_id: str, assertion: SSOIdentityAssertion | None) -> None:
"""The SSO-callback hook: a no-op unless there is material AND an EMA server is registered.
A store failure is logged and swallowed because the login itself must not fail on an

11
litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py
View file

@ -383,6 +383,15 @@ AuthConfig = Annotated[
]
# What the inbound bearer is, classified once at the edge so each exchange arm gates on the
# exact shape its mode requires. `gateway_credential` is one this gateway issued (never exchange
# material, any mode). `external_jwt` is a non-gateway JWT, i.e. an id_token candidate id_jag can
# exchange. `external_opaque` is a non-gateway non-JWT: id_jag rejects it (an id_token is a JWT)
# and falls to the stored assertion, while token_exchange (OBO) may still forward it per its
# exchange-what-was-presented contract.
InboundTokenProvenance = Literal["absent", "gateway_credential", "external_jwt", "external_opaque"]
class Subject(BaseModel):
"""The validated inbound principal. NOT the v1 request object and NOT the LiteLLM key."""
@ -392,6 +401,8 @@ class Subject(BaseModel):
subject_id: str
# Opaque, already-validated inbound identity. Only `token_exchange` / `passthrough` read it.
inbound_token: SecretStr | None = None
# What `inbound_token` is, classified once at the edge (`to_subject`); each arm gates on it.
inbound_provenance: InboundTokenProvenance = "absent"
class ServerSpec(BaseModel):

23
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -2164,6 +2164,7 @@ if MCP_AVAILABLE:
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
raw_headers=raw_headers,
user_api_key_auth=user_api_key_auth,
)
all_prompts.extend(prompts)
@ -2217,6 +2218,7 @@ if MCP_AVAILABLE:
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
raw_headers=raw_headers,
user_api_key_auth=user_api_key_auth,
)
all_resources.extend(resources)
@ -2268,6 +2270,7 @@ if MCP_AVAILABLE:
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
raw_headers=raw_headers,
user_api_key_auth=user_api_key_auth,
)
all_resource_templates.extend(resource_templates)
verbose_logger.debug(
@ -3113,6 +3116,7 @@ if MCP_AVAILABLE:
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
raw_headers=raw_headers,
user_api_key_auth=user_api_key_auth,
)
async def mcp_read_resource(
@ -3162,6 +3166,7 @@ if MCP_AVAILABLE:
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
raw_headers=raw_headers,
user_api_key_auth=user_api_key_auth,
)
def _get_standard_logging_mcp_tool_call(
@ -3634,6 +3639,7 @@ if MCP_AVAILABLE:
user_api_key_auth: Optional[UserAPIKeyAuth],
client_ip: Optional[str],
allowed_server_ids: Optional[Set[str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
) -> None:
"""Fail fast with HTTP 401 for MCP servers that need user auth but
didn't receive it on this request. Covers both gateway-managed OAuth2
@ -3748,6 +3754,21 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
)
# id_jag (EMA): resolve the exchange here at the transport edge for the same reason as
# the OBO preflight above — the in-session list handlers can only serialize a raise
# into a JSON-RPC error, so a dead stored SSO assertion would read as an empty catalog
# instead of the re-login 401 the resolver minted for it. There is no subject-header
# gate: id_jag sources its subject from the caller's id_token or the stored assertion,
# so resolution always has an answer. Gated to single-server routes; the multi-server
# aggregate keeps absorbing per-server auth failures.
if server and server.auth_type == MCPAuth.oauth2_id_jag and len(mcp_servers or []) == 1:
await global_mcp_server_manager.preflight_id_jag(
server=server,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
user_api_key_auth=user_api_key_auth,
)
# Pass-through OAuth: when the admin has opted a server into
# forwarding the client's bearer token (is_oauth_passthrough) and
# the client hasn't supplied one, fail fast with 401 and point
@ -4071,6 +4092,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
client_ip=_client_ip,
allowed_server_ids=toolset_allowed_server_ids,
raw_headers=raw_headers,
)
# Pre-flight auth check for pass-through servers. Must run after
@ -4393,6 +4415,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
client_ip=_sse_client_ip,
allowed_server_ids=toolset_allowed_server_ids,
raw_headers=raw_headers,
)
# Pre-flight auth check for pass-through servers: surface upstream

5
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -66,6 +66,7 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
SSOIdentityAssertion,
assertion_from_sso_login,
generic_sso_scopes,
retain_sso_identity_assertion_for_ema,
)
from litellm.proxy._types import (
@ -1130,7 +1131,7 @@ def _setup_generic_sso_env_vars(
) -> Tuple[str, List[str], str, str, str, bool]:
"""Setup and validate Generic SSO environment variables."""
generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None)
generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ")
generic_scope = generic_sso_scopes()
generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None)
generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None)
generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None)
@ -2708,7 +2709,7 @@ class SSOAuthenticationHandler:
from fastapi_sso.sso.generic import create_provider
generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None)
generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ")
generic_scope = generic_sso_scopes()
generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None)
generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None)
generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None)

2
tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py
View file

@ -94,6 +94,8 @@ def test_wire_value_carries_no_prose():
("timeout", None, 504),
("unreachable", None, 502),
("upstream_error", 500, 502),
("unavailable", 503, 503),
("precondition", 412, 412),
("internal", None, 500),
],
)

36
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py
View file

@ -358,6 +358,34 @@ def test_to_subject_maps_principal_fields():
assert subject.tenant_id == "org1"
assert subject.subject_id == "user1"
assert subject.inbound_token is None
assert subject.inbound_provenance == "absent"
def test_to_subject_classifies_an_external_jwt_as_id_token_candidate():
import jwt as pyjwt
external_jwt = pyjwt.encode({"iss": "https://idp.example.com", "sub": "alice"}, "idp-key", algorithm="HS256")
subject = to_subject(SimpleNamespace(org_id="", team_id="", user_id="alice"), external_jwt)
assert subject.inbound_provenance == "external_jwt"
def test_to_subject_classifies_an_opaque_inbound_token_as_external_opaque():
subject = to_subject(SimpleNamespace(org_id="", team_id="", user_id="alice"), "opaque-not-a-jwt-token")
assert subject.inbound_provenance == "external_opaque"
def test_to_subject_classifies_a_gateway_credential_inbound_token(monkeypatch):
"""The edge is where a gateway-issued bearer that rode Authorization is recognized, so the
resolver core never has to; a virtual key is stamped `gateway_credential`."""
import litellm.proxy.proxy_server as proxy_server
monkeypatch.setattr(proxy_server, "master_key", "sk-master-adapter-test")
subject = to_subject(SimpleNamespace(org_id="", team_id="", user_id="alice"), "sk-a-caller-virtual-key")
assert subject.inbound_provenance == "gateway_credential"
def test_to_subject_absent_token_classifies_absent():
assert to_subject(SimpleNamespace(org_id="", team_id="", user_id="alice"), None).inbound_provenance == "absent"
@pytest.mark.parametrize(
@ -512,9 +540,7 @@ def test_id_jag_client_secret_maps_to_config():
# ID-JAG asserts the user's id_token; the access_token default maps to id_token.
assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:id_token"
assert isinstance(spec.config.client_auth, ClientSecretAuth)
assert spec.config.client_auth.client_secret.get_secret_value() == (
"litellm-client-secret"
)
assert spec.config.client_auth.client_secret.get_secret_value() == ("litellm-client-secret")
def test_id_jag_private_key_maps_to_private_key_jwt_auth():
@ -540,9 +566,7 @@ def test_id_jag_private_key_wins_over_client_secret():
def test_id_jag_honors_explicit_subject_token_type():
spec = to_server_spec(
_id_jag_server(subject_token_type="urn:ietf:params:oauth:token-type:saml2")
)
spec = to_server_spec(_id_jag_server(subject_token_type="urn:ietf:params:oauth:token-type:saml2"))
assert spec is not None and isinstance(spec.config, IdJagConfig)
assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:saml2"

66
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py
View file

@ -38,19 +38,13 @@ def _endpoint(body, sink=None):
def _recording_persist(sink):
async def persist(
user_id, server_id, access_token, refresh_token, expires_in, scopes
):
sink.append(
(user_id, server_id, access_token, refresh_token, expires_in, scopes)
)
async def persist(user_id, server_id, access_token, refresh_token, expires_in, scopes):
sink.append((user_id, server_id, access_token, refresh_token, expires_in, scopes))
return persist
def _refresher(
server=None, body=None, *, post_sink=None, persist_sink=None, clock=lambda: 1000.0
):
def _refresher(server=None, body=None, *, post_sink=None, persist_sink=None, clock=lambda: 1000.0):
return AuthorizationCodeRefresher(
_lookup(server if server is not None else _Server()),
_endpoint(body, post_sink),
@ -73,9 +67,7 @@ async def test_refreshes_persists_and_returns_typed_token():
post_sink=posted,
persist_sink=persisted,
)
token = await refresher.refresh(
"alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")
)
token = await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt"))
assert token is not None
assert token.access_token == "new-at"
@ -108,9 +100,7 @@ async def test_client_secret_basic_sends_authorization_header_not_body():
body={"access_token": "new-at"},
post_sink=posted,
)
token = await refresher.refresh(
"alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")
)
token = await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt"))
assert token is not None
_url, form, headers = posted[0]
@ -136,38 +126,22 @@ async def test_client_secret_basic_without_secret_is_a_failed_refresh():
async def test_no_refresh_token_is_not_refreshable():
posted = []
refresher = _refresher(body={"access_token": "x"}, post_sink=posted)
assert (
await refresher.refresh("alice", "srv", OAuthToken(access_token="old")) is None
)
assert await refresher.refresh("alice", "srv", OAuthToken(access_token="old")) is None
assert posted == [] # never hit the IdP
@pytest.mark.asyncio
async def test_unknown_server_or_no_token_url_yields_none():
assert (
await _refresher(server=None).refresh(
"a", "s", OAuthToken("old", refresh_token="rt")
)
is None
)
assert await _refresher(server=None).refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None
no_url = _Server(token_url=None)
assert (
await _refresher(server=no_url).refresh(
"a", "s", OAuthToken("old", refresh_token="rt")
)
is None
)
assert await _refresher(server=no_url).refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None
@pytest.mark.asyncio
async def test_grant_failure_does_not_persist():
persisted = []
refresher = _refresher(
body=None, persist_sink=persisted
) # token_endpoint signals failure
assert (
await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None
)
refresher = _refresher(body=None, persist_sink=persisted) # token_endpoint signals failure
assert await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None
assert persisted == []
@ -175,9 +149,7 @@ async def test_grant_failure_does_not_persist():
async def test_response_without_access_token_does_not_persist():
persisted = []
refresher = _refresher(body={"expires_in": 60}, persist_sink=persisted)
assert (
await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None
)
assert await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None
assert persisted == []
@ -185,13 +157,9 @@ async def test_response_without_access_token_does_not_persist():
async def test_unrotated_refresh_token_is_carried_forward():
persisted = []
refresher = _refresher(body={"access_token": "new-at"}, persist_sink=persisted)
token = await refresher.refresh(
"a", "s", OAuthToken("old", refresh_token="keep-rt")
)
token = await refresher.refresh("a", "s", OAuthToken("old", refresh_token="keep-rt"))
assert token is not None
assert (
token.refresh_token == "keep-rt"
) # response omitted refresh_token -> reuse the old one
assert token.refresh_token == "keep-rt" # response omitted refresh_token -> reuse the old one
assert token.expires_at is None # no expires_in -> no known expiry
assert persisted[0][3] == "keep-rt"
@ -200,9 +168,7 @@ async def test_unrotated_refresh_token_is_carried_forward():
async def test_unrecorded_scope_is_carried_forward():
persisted = []
refresher = _refresher(body={"access_token": "new-at"}, persist_sink=persisted)
token = await refresher.refresh(
"a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write"))
)
token = await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write")))
assert token is not None
# response omitted "scope" -> the user's recorded grant is preserved, not dropped
assert token.scopes == ("read", "write")
@ -216,9 +182,7 @@ async def test_returned_scope_overrides_prior_when_present():
body={"access_token": "new-at", "scope": "read"},
persist_sink=persisted,
)
token = await refresher.refresh(
"a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write"))
)
token = await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write")))
assert token is not None
assert token.scopes == ("read",) # a present scope replaces the prior grant
assert persisted[0][5] == ("read",)

183
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_credential_provenance.py Normal file
View file

@ -0,0 +1,183 @@
"""The gateway-credential recognizer must be complete by construction: every credential this
gateway mints classifies as `gateway_credential` (so it is never exchange subject material), and
a real external token classifies as `external`. These tests mint the real gateway formats so a
newly added mint that this recognizer fails to catch fails a test here, not in production."""
from datetime import datetime, timezone
import jwt as pyjwt
import pytest
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._experimental.mcp_server.outbound_credentials.credential_provenance import (
classify_inbound_provenance,
is_gateway_issued_credential,
)
_MASTER = "sk-master-key-for-provenance-tests"
@pytest.fixture(autouse=True)
def _master_key(monkeypatch):
monkeypatch.setattr(proxy_server, "master_key", _MASTER)
def _hs256(claims: dict, key: str) -> str:
return pyjwt.encode(claims, key, algorithm="HS256")
def _real_session_token() -> str:
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
session_keys_from_master_key,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
SessionPrincipal,
mint_session_token,
)
minted = mint_session_token(
SessionPrincipal(user_id="u", client_id="c"),
session_keys_from_master_key(_MASTER),
datetime.now(timezone.utc),
)
return minted.token.get_secret_value()
def _gateway_credentials() -> dict[str, str]:
now = 2_000_000_000
return {
"master_key": _MASTER,
"virtual_key": "sk-" + "a" * 40,
# The three the prefix denylist missed: master-key HS256 JWTs with no prefix and no iss.
"ui_login_session": _hs256({"user_id": "u", "key": "sk-embedded", "user_role": "proxy_admin"}, _MASTER),
"onboarding": _hs256({"token_type": "litellm_onboarding", "user_id": "u", "exp": now}, _MASTER),
"byok_session": _hs256({"user_id": "u", "server_id": "s", "type": "byok_session", "exp": now}, _MASTER),
# Prefixed gateway JWTs.
"mcp_session_real": _real_session_token(),
"envelope_prefixed": "llm_env_" + _hs256({"iss": "litellm-mcp-bridge"}, "derived-key"),
# Reserved-issuer inner tokens presented without their prefix (derived-key signed).
"session_inner_no_prefix": _hs256({"iss": "litellm-mcp-gateway", "exp": now}, "some-derived-key"),
"envelope_inner_no_prefix": _hs256({"iss": "litellm-mcp-bridge", "exp": now}, "some-derived-key"),
}
def _external_jwt_tokens() -> dict[str, str]:
"""Non-gateway JWTs: id_token candidates id_jag can exchange (`external_jwt`)."""
now = 2_000_000_000
return {
"external_idp_jwt": _hs256({"iss": "https://okta.example.com", "sub": "u", "aud": "gw", "exp": now}, "idp-key"),
"external_jwt_no_exp": _hs256({"iss": "https://idp.example.com", "sub": "u"}, "another-idp-key"),
}
def _external_opaque_tokens() -> dict[str, str]:
"""Non-gateway, non-JWT bearers: not id_tokens, so id_jag rejects them (`external_opaque`)."""
return {
"opaque_access_token": "opaque-access-token-from-an-idp-1234567890",
# Three dot-segments but not a decodable JWT; the finding's example.
"jwt_shaped_but_undecodable": "abc.def.ghi",
}
# Evaluated once, and the case NAME is the test id. A minted session token embeds the current time
# and a random jti, so deriving the id from the token value instead would make it differ on every
# collection; under xdist the workers then disagree about which tests exist and the run errors out.
_GATEWAY_CREDENTIALS = _gateway_credentials()
_EXTERNAL_JWT_TOKENS = _external_jwt_tokens()
_EXTERNAL_OPAQUE_TOKENS = _external_opaque_tokens()
@pytest.mark.parametrize("name,token", list(_GATEWAY_CREDENTIALS.items()), ids=list(_GATEWAY_CREDENTIALS))
def test_every_gateway_credential_is_recognized(name, token):
assert is_gateway_issued_credential(token) is True
assert classify_inbound_provenance(token) == "gateway_credential"
@pytest.mark.parametrize("name,token", list(_EXTERNAL_JWT_TOKENS.items()), ids=list(_EXTERNAL_JWT_TOKENS))
def test_external_jwt_is_an_id_token_candidate(name, token):
assert is_gateway_issued_credential(token) is False
assert classify_inbound_provenance(token) == "external_jwt"
@pytest.mark.parametrize("name,token", list(_EXTERNAL_OPAQUE_TOKENS.items()), ids=list(_EXTERNAL_OPAQUE_TOKENS))
def test_external_opaque_token_is_not_an_id_token(name, token):
"""A non-gateway non-JWT bearer is not an id_token; id_jag must not forward it (the finding)."""
assert is_gateway_issued_credential(token) is False
assert classify_inbound_provenance(token) == "external_opaque"
@pytest.mark.parametrize("empty", [None, ""])
def test_absent_inbound_token_classifies_absent(empty):
assert classify_inbound_provenance(empty) == "absent"
@pytest.mark.parametrize(
"temporal_claims",
[
{"exp": 1}, # long expired
{"nbf": 9_999_999_999}, # not yet valid
{"iat": 9_999_999_999}, # issued in the future
],
)
def test_temporally_invalid_gateway_jwt_is_still_recognized(temporal_claims):
"""Recognition determines provenance, not usability: an expired or not-yet-valid gateway JWT
is still this gateway's and must classify as gateway_credential, or it drops to the
exchangeable-external path and is disclosed upstream."""
token = _hs256({"user_id": "u", "key": "sk-embedded", **temporal_claims}, _MASTER)
assert is_gateway_issued_credential(token) is True
assert classify_inbound_provenance(token) == "gateway_credential"
def test_expired_external_jwt_stays_an_id_token_candidate():
"""An expired token that is NOT ours stays external_jwt; the org authorization server, not
this recognizer, is the authority on whether the caller's own token is still valid."""
token = _hs256({"iss": "https://idp.example.com", "sub": "u", "exp": 1}, "external-idp-key")
assert is_gateway_issued_credential(token) is False
assert classify_inbound_provenance(token) == "external_jwt"
def test_encrypted_login_blob_is_recognized_by_decryption():
"""CLI / experimental UI login tokens are encrypted, not JWT-signed; decryption succeeding is
proof they are this gateway's, so they are recognized without a prefix."""
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
blob = encrypt_value_helper('{"token": "cli-session-abc", "is_session_token": true}')
assert is_gateway_issued_credential(blob) is True
def test_recognizer_does_not_crash_without_a_master_key(monkeypatch):
"""With no master key configured the crypto checks are skipped; an external token is still
classified external rather than raising into egress."""
monkeypatch.setattr(proxy_server, "master_key", None)
assert is_gateway_issued_credential("some.external.jwt") is False
# A prefix-bearing gateway credential is still caught with no master key.
assert is_gateway_issued_credential("sk-still-caught-by-prefix") is True
@pytest.mark.parametrize(
"adversarial",
[
"café-tökèn-with-non-ascii", # the reported crash: non-ASCII bytes hit the master-key compare
"bearer-\U0001f600-emoji", # supplementary-plane characters
"\x00\x01control-bytes", # control characters
"aaa.é.ccc", # non-ASCII inside a jwt-shaped value
],
)
def test_non_ascii_inbound_bearer_never_raises_into_egress(adversarial):
"""The inbound bearer is fully attacker-controlled, so recognition must be TOTAL: a non-ASCII
bearer used to reach ``secrets.compare_digest(token, master_key)`` and raise ``TypeError``,
500-ing every egress. It must now classify as a non-gateway external token without raising.
Reverting the digest-based master-key compare makes this test raise."""
assert is_gateway_issued_credential(adversarial) is False
assert classify_inbound_provenance(adversarial) in ("external_jwt", "external_opaque")
def test_master_key_compare_matches_only_the_exact_key(monkeypatch):
"""The digest-based compare must still be exact equality, not a length match or near miss. A
master key with no gateway prefix isolates the compare from the other checks: the exact key is
recognized, a value that only shares its length or a prefix of it is not."""
plain_master = "plain-master-key-no-gateway-prefix"
monkeypatch.setattr(proxy_server, "master_key", plain_master)
assert is_gateway_issued_credential(plain_master) is True
assert is_gateway_issued_credential(plain_master + "x") is False
assert is_gateway_issued_credential(plain_master[:-1] + "Z") is False

6
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_oauth_token_store.py
View file

@ -269,6 +269,12 @@ async def test_stale_read_after_refresh_rereads_before_starting_new_refresh():
async def test_refresh_failure_is_shared_by_joiners_not_re_run():
"""Pins the ``TokenRefresher`` contract: a refresher's raised exception MUST propagate to the
refresh winner and every single-flight waiter alike, and the failed attempt runs once. Domain
refreshers (the EMA assertion refresher) raise verdicts richer than ``OAuthToken | None``;
catching them here and returning None would silently collapse "the grant is dead" (re-login)
into "no token" (a different remedy and status)."""
class _FailingRefresher:
def __init__(self) -> None:
self.calls = 0

302
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py
View file

@ -68,7 +68,21 @@ class _FakeTokenEndpoint:
def _with_inbound(token: str) -> Subject:
return Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr(token))
"""A caller presenting their own id_token (the edge classified it `external_jwt`)."""
return Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr(token), inbound_provenance="external_jwt")
def _idp_jwt(sub: str = "user") -> str:
import jwt as pyjwt
return pyjwt.encode(
{"sub": sub, "iss": "https://idp.example.com"}, "test-signing-key-32-bytes-long-xx", algorithm="HS256"
)
_USER_IDP_JWT = _idp_jwt("user")
_ALICE_IDP_JWT = _idp_jwt("alice")
_BOB_IDP_JWT = _idp_jwt("bob")
def _spec(config):
@ -247,9 +261,17 @@ _OBO = TokenExchangeConfig(
@pytest.mark.asyncio
async def test_token_exchange_emits_the_exchanged_bearer():
@pytest.mark.parametrize("external_provenance", ["external_jwt", "external_opaque"])
async def test_token_exchange_emits_the_exchanged_bearer(external_provenance):
"""OBO exchanges what was presented for the caller's own external token, JWT or opaque; the
id_token-only restriction is id_jag's, not OBO's."""
exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at")))
subject = Subject(tenant_id="acme", subject_id="alice", inbound_token=SecretStr("caller-jwt"))
subject = Subject(
tenant_id="acme",
subject_id="alice",
inbound_token=SecretStr("caller-jwt"),
inbound_provenance=external_provenance,
)
result = await UpstreamCredentialProvider(token_exchanger=exchanger).resolve_credentials(subject, _spec(_OBO))
assert isinstance(result, Ok)
assert _emitted(result.ok)["Authorization"] == "Bearer exchanged-at"
@ -257,6 +279,25 @@ async def test_token_exchange_emits_the_exchanged_bearer():
assert exchanger.calls == [("caller-jwt", "acme", "s")]
@pytest.mark.asyncio
async def test_token_exchange_refuses_a_gateway_credential_subject_token():
"""OBO keeps exchange-what-was-presented for the caller's own external token, but a
gateway-issued credential (the edge classified it `gateway_credential`) is refused with a
401 and never disclosed to the external token endpoint."""
exchanger = _FakeExchanger(Ok(OAuthToken(access_token="never")))
subject = Subject(
tenant_id="acme",
subject_id="alice",
inbound_token=SecretStr("sk-a-gateway-admission-key"),
inbound_provenance="gateway_credential",
)
result = await UpstreamCredentialProvider(token_exchanger=exchanger).resolve_credentials(subject, _spec(_OBO))
assert isinstance(result, Error)
assert result.error.tag == "unauthorized"
# The gateway credential is never handed to the exchanger.
assert exchanger.calls == []
@pytest.mark.asyncio
async def test_invalidate_credentials_drops_the_exchanged_token_for_the_subject_and_tenant():
exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at")))
@ -446,22 +487,16 @@ async def test_id_jag_runs_both_legs_and_returns_the_leg2_bearer():
]
)
provider = UpstreamCredentialProvider(token_endpoint=endpoint)
result = await provider.resolve_credentials(
_with_inbound("user-id-token"), _spec(_id_jag_config())
)
result = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(_id_jag_config()))
assert isinstance(result, Ok)
assert _emitted(result.ok)["Authorization"] == "Bearer final-access"
leg1_endpoint, _, leg1_params = endpoint.calls[0]
assert leg1_endpoint == "https://idp.example.com/token"
assert (
leg1_params["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange"
)
assert (
leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag"
)
assert leg1_params["subject_token"] == "user-id-token"
assert leg1_params["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange"
assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag"
assert leg1_params["subject_token"] == _USER_IDP_JWT
leg2_endpoint, _, leg2_params = endpoint.calls[1]
assert leg2_endpoint == "https://mcp-as.example.com/token"
@ -474,9 +509,7 @@ async def test_id_jag_runs_both_legs_and_returns_the_leg2_bearer():
async def test_id_jag_without_inbound_token_is_precondition_required_no_http():
endpoint = _FakeTokenEndpoint([])
provider = UpstreamCredentialProvider(token_endpoint=endpoint)
result = await provider.resolve_credentials(
Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())
)
result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()))
assert isinstance(result, Error)
assert result.error.tag == "precondition_required"
@ -485,13 +518,9 @@ async def test_id_jag_without_inbound_token_is_precondition_required_no_http():
@pytest.mark.asyncio
async def test_id_jag_propagates_a_leg1_error_without_calling_leg2():
endpoint = _FakeTokenEndpoint(
[Error(CredError.of_upstream_unavailable("leg1 down"))]
)
endpoint = _FakeTokenEndpoint([Error(CredError.of_upstream_unavailable("leg1 down"))])
provider = UpstreamCredentialProvider(token_endpoint=endpoint)
result = await provider.resolve_credentials(
_with_inbound("user-id-token"), _spec(_id_jag_config())
)
result = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(_id_jag_config()))
assert isinstance(result, Error)
assert result.error.tag == "upstream_unavailable"
@ -508,9 +537,7 @@ async def test_id_jag_propagates_a_leg2_error():
]
)
provider = UpstreamCredentialProvider(token_endpoint=endpoint)
result = await provider.resolve_credentials(
_with_inbound("user-id-token"), _spec(_id_jag_config())
)
result = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(_id_jag_config()))
assert isinstance(result, Error)
assert result.error.tag == "upstream_unavailable"
@ -530,8 +557,8 @@ async def test_id_jag_reuses_the_cached_bearer_for_an_unchanged_config():
endpoint = _FakeTokenEndpoint(_two_leg_ok("first-bearer"))
provider = UpstreamCredentialProvider(token_endpoint=endpoint)
first = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config()))
second = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config()))
first = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(_id_jag_config()))
second = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(_id_jag_config()))
assert isinstance(first, Ok) and isinstance(second, Ok)
assert _emitted(second.ok)["Authorization"] == "Bearer first-bearer"
@ -566,8 +593,8 @@ async def test_id_jag_config_change_forces_a_fresh_exchange(changed):
endpoint = _FakeTokenEndpoint(_two_leg_ok("old-policy-bearer") + _two_leg_ok("new-policy-bearer"))
provider = UpstreamCredentialProvider(token_endpoint=endpoint)
before = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config()))
after = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(changed))
before = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(_id_jag_config()))
after = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(changed))
assert isinstance(before, Ok) and isinstance(after, Ok)
assert _emitted(after.ok)["Authorization"] == "Bearer new-policy-bearer"
@ -579,8 +606,8 @@ async def test_id_jag_does_not_share_the_cached_bearer_across_caller_tokens():
endpoint = _FakeTokenEndpoint(_two_leg_ok("alice-bearer") + _two_leg_ok("bob-bearer"))
provider = UpstreamCredentialProvider(token_endpoint=endpoint)
alice = await provider.resolve_credentials(_with_inbound("alice-id-token"), _spec(_id_jag_config()))
bob = await provider.resolve_credentials(_with_inbound("bob-id-token"), _spec(_id_jag_config()))
alice = await provider.resolve_credentials(_with_inbound(_ALICE_IDP_JWT), _spec(_id_jag_config()))
bob = await provider.resolve_credentials(_with_inbound(_BOB_IDP_JWT), _spec(_id_jag_config()))
assert isinstance(alice, Ok) and isinstance(bob, Ok)
assert _emitted(bob.ok)["Authorization"] == "Bearer bob-bearer"
@ -591,7 +618,7 @@ async def test_id_jag_does_not_share_the_cached_bearer_across_caller_tokens():
async def test_invalidate_credentials_evicts_the_id_jag_bearer_so_the_next_resolve_re_exchanges():
endpoint = _FakeTokenEndpoint(_two_leg_ok("rejected-bearer") + _two_leg_ok("fresh-bearer"))
provider = UpstreamCredentialProvider(token_endpoint=endpoint)
subject = _with_inbound("user-id-token")
subject = _with_inbound(_USER_IDP_JWT)
first = await provider.resolve_credentials(subject, _spec(_id_jag_config()))
await provider.invalidate_credentials(subject, _spec(_id_jag_config()))
@ -606,7 +633,7 @@ async def test_invalidate_credentials_evicts_the_id_jag_bearer_so_the_next_resol
async def test_invalidate_credentials_for_id_jag_is_a_noop_without_a_caller_token():
endpoint = _FakeTokenEndpoint(_two_leg_ok("cached-bearer"))
provider = UpstreamCredentialProvider(token_endpoint=endpoint)
subject = _with_inbound("user-id-token")
subject = _with_inbound(_USER_IDP_JWT)
first = await provider.resolve_credentials(subject, _spec(_id_jag_config()))
await provider.invalidate_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()))
@ -615,3 +642,212 @@ async def test_invalidate_credentials_for_id_jag_is_a_noop_without_a_caller_toke
assert isinstance(first, Ok) and isinstance(second, Ok)
assert _emitted(second.ok)["Authorization"] == "Bearer cached-bearer"
assert len(endpoint.calls) == 2
# -- ID-JAG subject sourcing: stored SSO assertion seam (EMA) --------------------------------
class _FakeAssertionSource:
"""Returns the canned lookup and records which user ids were asked for."""
def __init__(self, lookup):
self._lookup = lookup
self.asked: list[str] = []
async def fetch_usable(self, user_id: str):
self.asked.append(user_id)
if isinstance(self._lookup, Exception):
raise self._lookup
return self._lookup
def _usable_assertion(id_token: str = "hdr.stored-id-token.sig"):
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken
return OAuthToken(access_token=id_token)
def _no_assertion():
return None
def _expired_assertion():
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
SsoAssertionUnrenewable,
)
return SsoAssertionUnrenewable("dead grant")
def _no_inbound(subject_id: str = "alice") -> Subject:
return Subject(tenant_id="", subject_id=subject_id)
@pytest.mark.asyncio
async def test_id_jag_sources_the_stored_assertion_when_no_inbound_token():
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
source = _FakeAssertionSource(_usable_assertion())
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertions=source)
result = await provider.resolve_credentials(_no_inbound(), _spec(_id_jag_config()))
assert isinstance(result, Ok)
assert source.asked == ["alice"]
_, _, leg1_params = endpoint.calls[0]
assert leg1_params["subject_token"] == "hdr.stored-id-token.sig"
@pytest.mark.asyncio
async def test_id_jag_prefers_the_presented_idp_jwt_over_the_stored_assertion():
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
source = _FakeAssertionSource(_usable_assertion())
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertions=source)
result = await provider.resolve_credentials(_with_inbound(_idp_jwt("live")), _spec(_id_jag_config()))
assert isinstance(result, Ok)
assert source.asked == []
_, _, leg1_params = endpoint.calls[0]
assert leg1_params["subject_token"] == _idp_jwt("live")
@pytest.mark.asyncio
async def test_id_jag_never_exchanges_a_gateway_credential_inbound_token():
"""A gateway-issued bearer that rode the Authorization header (the edge classified it
`gateway_credential`) must fall through to the stored assertion, never reach an external
IdP, no matter that a token happens to be present."""
gateway_bearer = "sk-litellm-admission-key"
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
source = _FakeAssertionSource(_usable_assertion())
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertions=source)
subject = Subject(
tenant_id="",
subject_id="alice",
inbound_token=SecretStr(gateway_bearer),
inbound_provenance="gateway_credential",
)
result = await provider.resolve_credentials(subject, _spec(_id_jag_config()))
assert isinstance(result, Ok)
assert source.asked == ["alice"]
_, _, leg1_params = endpoint.calls[0]
assert leg1_params["subject_token"] == "hdr.stored-id-token.sig"
assert all(gateway_bearer not in str(call) for call in endpoint.calls)
@pytest.mark.asyncio
async def test_id_jag_never_forwards_an_opaque_inbound_token():
"""The ID-JAG exchange takes an id_token (a JWT). An opaque non-gateway bearer is not an
id_token, so id_jag sources the stored assertion and never forwards the opaque value."""
opaque_bearer = "abc.def.ghi"
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
source = _FakeAssertionSource(_usable_assertion())
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertions=source)
subject = Subject(
tenant_id="", subject_id="alice", inbound_token=SecretStr(opaque_bearer), inbound_provenance="external_opaque"
)
result = await provider.resolve_credentials(subject, _spec(_id_jag_config()))
assert isinstance(result, Ok)
assert source.asked == ["alice"]
_, _, leg1_params = endpoint.calls[0]
assert leg1_params["subject_token"] == "hdr.stored-id-token.sig"
assert all(opaque_bearer not in str(call) for call in endpoint.calls)
@pytest.mark.asyncio
async def test_id_jag_expired_stored_assertion_is_a_401_challenge_not_a_412():
endpoint = _FakeTokenEndpoint([])
provider = UpstreamCredentialProvider(
token_endpoint=endpoint, sso_assertions=_FakeAssertionSource(_expired_assertion())
)
result = await provider.resolve_credentials(_no_inbound(), _spec(_id_jag_config()))
assert isinstance(result, Error)
assert result.error.tag == "unauthorized"
unauthorized = result.error.unauthorized
assert unauthorized is not None and unauthorized.www_authenticate is not None
assert "invalid_token" in unauthorized.www_authenticate
assert endpoint.calls == []
@pytest.mark.asyncio
async def test_id_jag_absent_stored_assertion_and_no_inbound_is_precondition_required():
endpoint = _FakeTokenEndpoint([])
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertions=_FakeAssertionSource(_no_assertion()))
result = await provider.resolve_credentials(_no_inbound(), _spec(_id_jag_config()))
assert isinstance(result, Error)
assert result.error.tag == "precondition_required"
assert endpoint.calls == []
@pytest.mark.asyncio
async def test_id_jag_cache_key_is_the_resolved_token_so_a_new_login_re_exchanges():
endpoint = _FakeTokenEndpoint(_two_leg_ok("first-bearer") + _two_leg_ok("second-bearer"))
provider = UpstreamCredentialProvider(
token_endpoint=endpoint, sso_assertions=_FakeAssertionSource(_usable_assertion())
)
first = await provider.resolve_credentials(_no_inbound(), _spec(_id_jag_config()))
relogged = UpstreamCredentialProvider(
token_endpoint=endpoint,
exchanged_tokens=provider._exchanged_tokens,
sso_assertions=_FakeAssertionSource(_usable_assertion("hdr.new-login-id-token.sig")),
)
second = await relogged.resolve_credentials(_no_inbound(), _spec(_id_jag_config()))
assert isinstance(first, Ok) and isinstance(second, Ok)
assert _emitted(first.ok)["Authorization"] == "Bearer first-bearer"
assert _emitted(second.ok)["Authorization"] == "Bearer second-bearer"
assert len(endpoint.calls) == 4
@pytest.mark.asyncio
async def test_invalidate_for_id_jag_evicts_the_stored_assertion_caller_entry():
endpoint = _FakeTokenEndpoint(_two_leg_ok("rejected-bearer") + _two_leg_ok("fresh-bearer"))
source = _FakeAssertionSource(_usable_assertion())
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertions=source)
subject = _no_inbound()
first = await provider.resolve_credentials(subject, _spec(_id_jag_config()))
await provider.invalidate_credentials(subject, _spec(_id_jag_config()))
second = await provider.resolve_credentials(subject, _spec(_id_jag_config()))
assert isinstance(first, Ok) and isinstance(second, Ok)
assert _emitted(second.ok)["Authorization"] == "Bearer fresh-bearer"
assert len(endpoint.calls) == 4
@pytest.mark.asyncio
async def test_id_jag_empty_subject_id_with_null_source_stays_precondition_required():
provider = UpstreamCredentialProvider(token_endpoint=_FakeTokenEndpoint([]))
result = await provider.resolve_credentials(Subject(tenant_id="", subject_id=""), _spec(_id_jag_config()))
assert isinstance(result, Error)
assert result.error.tag == "precondition_required"
@pytest.mark.asyncio
async def test_id_jag_unavailable_renewal_is_503_not_a_relogin_challenge():
"""An unreachable IdP during renewal proves nothing about the stored assertion, so the
caller gets a retryable upstream_unavailable, never the re-login challenge."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
TokenStoreUnavailable,
)
endpoint = _FakeTokenEndpoint([])
provider = UpstreamCredentialProvider(
token_endpoint=endpoint, sso_assertions=_FakeAssertionSource(TokenStoreUnavailable("idp down"))
)
result = await provider.resolve_credentials(_no_inbound(), _spec(_id_jag_config()))
assert isinstance(result, Error)
assert result.error.tag == "upstream_unavailable"
assert endpoint.calls == []

812
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py
View file

@ -1,10 +1,10 @@
"""Tests for the SSO identity assertion store (EMA subject-token capture).
Pins the contract of the store that PR 2's ``_id_jag`` subject-sourcing seam will read:
the carrier validates untyped IdP token-response values at the boundary, retention is
gated on an ``oauth2_id_jag`` server being registered, the row is encrypted at rest and
round-trips exactly, a store failure never escapes into the login path, and a salt-key
rotation re-encrypts stored rows like the sibling per-user credential tables.
Pins the contract the resolver's ``_id_jag`` subject-sourcing seam reads: the carrier
validates untyped IdP token-response values at the boundary, retention is gated on an
``oauth2_id_jag`` server being registered, the row is encrypted at rest and round-trips
exactly, a store failure never escapes into the login path, and a salt-key rotation
re-encrypts stored rows like the sibling per-user credential tables.
"""
import json
@ -14,7 +14,11 @@ from unittest.mock import AsyncMock, MagicMock, patch
import jwt as pyjwt
import pytest
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
TokenStoreUnavailable,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
SsoAssertionUnrenewable,
assertion_from_sso_login,
ema_assertion_retention_enabled,
fetch_sso_identity_assertion,
@ -49,9 +53,7 @@ def _make_prisma(stored: dict, db_has_id_jag_server: bool = False):
``db_has_id_jag_server`` drives the retention gate's authoritative DB fallback;
it is wired explicitly so the gate never reads a truthy bare MagicMock."""
prisma = MagicMock()
prisma.db.litellm_mcpservertable.find_first = AsyncMock(
return_value=MagicMock() if db_has_id_jag_server else None
)
prisma.db.litellm_mcpservertable.find_first = AsyncMock(return_value=MagicMock() if db_has_id_jag_server else None)
async def _upsert(where, data):
stored[where["user_id"]] = data["update"]["assertion_b64"]
@ -240,19 +242,35 @@ async def test_fetch_missing_row_returns_none():
@pytest.mark.asyncio
async def test_fetch_undecryptable_row_returns_none():
async def test_fetch_undecryptable_row_is_unrenewable_not_absent():
"""A row that will not decrypt (salt-key rotation) is unusable material, not a user who never
signed in; the remedy is a fresh sign-in, so it must not read as absence."""
prisma = _make_prisma({"user-a": "not-an-encrypted-blob"})
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
assert await fetch_sso_identity_assertion("user-a") is None
with patch("litellm.proxy.proxy_server.prisma_client", prisma), pytest.raises(SsoAssertionUnrenewable):
await fetch_sso_identity_assertion("user-a")
@pytest.mark.asyncio
async def test_fetch_unparseable_payload_returns_none():
async def test_fetch_unparseable_payload_is_unrenewable_not_absent():
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
prisma = _make_prisma({"user-a": encrypt_value_helper("]]not json")})
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
assert await fetch_sso_identity_assertion("user-a") is None
with patch("litellm.proxy.proxy_server.prisma_client", prisma), pytest.raises(SsoAssertionUnrenewable):
await fetch_sso_identity_assertion("user-a")
@pytest.mark.asyncio
async def test_fetch_reports_an_unreachable_store_as_unavailable_never_absent():
"""The reported defect: a transient store outage for a user who DOES have an assertion must be
retryable (503), never the precondition-required shape that means "you have never signed in".
``None`` is reserved for the one determinate fact that no row was ever written."""
with patch("litellm.proxy.proxy_server.prisma_client", None), pytest.raises(TokenStoreUnavailable):
await fetch_sso_identity_assertion("user-a")
exploding = _make_prisma({})
exploding.db.litellm_ssoidentityassertion.find_unique = AsyncMock(side_effect=RuntimeError("db down"))
with patch("litellm.proxy.proxy_server.prisma_client", exploding), pytest.raises(TokenStoreUnavailable):
await fetch_sso_identity_assertion("user-a")
@pytest.mark.asyncio
@ -341,3 +359,769 @@ async def test_rotation_skips_unreadable_rows_but_rotates_readable_ones():
await rotate_sso_identity_assertions_master_key(prisma_client=prisma, new_master_key="another-new-salt-key-0000")
assert stored["bad"] == "garbage-blob"
assert stored["good"] != good_blob_before
# -- LiveSsoAssertionSource: the resolver-facing usable-assertion lookup + refresh -------------
from datetime import datetime, timedelta, timezone # noqa: E402
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( # noqa: E402
LiveSsoAssertionSource,
SSOIdentityAssertion,
SsoAssertionUnrenewable,
)
from pydantic import SecretStr # noqa: E402
_SSO_ENV = {
"GENERIC_CLIENT_ID": "sso-client",
"GENERIC_CLIENT_SECRET": "sso-secret",
"GENERIC_TOKEN_ENDPOINT": "https://idp.example.com/token",
}
def _stored(expires_in_seconds: int | None, refresh_token: str | None = "rt_stored") -> SSOIdentityAssertion:
return SSOIdentityAssertion(
id_token=SecretStr(_make_id_token()),
refresh_token=SecretStr(refresh_token) if refresh_token else None,
expires_at=(
datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds)
if expires_in_seconds is not None
else None
),
)
def _source(stored: SSOIdentityAssertion | None, post_bodies: list, env: dict | None = None):
fetches: list[str] = []
persisted: list[tuple[str, SSOIdentityAssertion]] = []
posts: list[tuple[str, dict]] = []
state = {"stored": stored}
async def fetch(user_id: str):
fetches.append(user_id)
return state["stored"]
async def persist(user_id: str, assertion: SSOIdentityAssertion):
persisted.append((user_id, assertion))
state["stored"] = assertion
async def post(url: str, form: dict):
posts.append((url, dict(form)))
body = post_bodies.pop(0) if post_bodies else None
if body is None:
raise TokenStoreUnavailable("idp unreachable")
return body
environment = _SSO_ENV if env is None else env
source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=environment.get)
return source, fetches, persisted, posts
@pytest.mark.asyncio
async def test_source_empty_user_id_is_absent_without_a_db_read():
source, fetches, _, _ = _source(_stored(3600), [])
assert await source.fetch_usable("") is None
assert fetches == []
@pytest.mark.asyncio
async def test_source_missing_row_is_absent():
source, _, _, posts = _source(None, [])
assert await source.fetch_usable("alice") is None
assert posts == []
@pytest.mark.asyncio
async def test_source_unexpired_assertion_is_usable_without_refresh():
stored = _stored(3600)
source, _, _, posts = _source(stored, [])
lookup = await source.fetch_usable("alice")
assert lookup is not None
assert lookup.access_token == stored.id_token.get_secret_value()
assert posts == []
@pytest.mark.asyncio
async def test_source_near_expiry_assertion_counts_as_expired():
"""A token inside the buffer would die mid-exchange; it must refresh, not be served."""
source, _, _, posts = _source(_stored(10, refresh_token=None), [])
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
assert posts == []
@pytest.mark.asyncio
async def test_source_expired_with_refresh_renews_persists_and_returns_usable():
new_id_token = _make_id_token(exp_offset=7200)
source, _, persisted, posts = _source(
_stored(-100), [{"id_token": new_id_token, "refresh_token": "rt_rotated", "access_token": "at"}]
)
lookup = await source.fetch_usable("alice")
assert lookup is not None
assert lookup.access_token == new_id_token
assert lookup.refresh_token is not None
assert lookup.refresh_token == "rt_rotated"
assert len(persisted) == 1 and persisted[0][0] == "alice"
url, form = posts[0]
assert url == "https://idp.example.com/token"
assert form["grant_type"] == "refresh_token"
assert form["refresh_token"] == "rt_stored"
assert form["client_id"] == "sso-client"
assert form["client_secret"] == "sso-secret"
@pytest.mark.asyncio
async def test_source_refresh_requests_openid_scope_so_the_idp_returns_an_id_token():
"""The refresh grant must carry a scope containing ``openid`` or the IdP returns no id_token,
stranding a renewable assertion as expired and forcing a needless re-login (the round-9
finding). With no GENERIC_SCOPE configured it falls to the same default the login uses, which
contains ``openid``. Dropping the scope from the refresh form makes this fail."""
source, _, _, posts = _source(_stored(-100), [{"id_token": _make_id_token(exp_offset=7200)}])
await source.fetch_usable("alice")
_, form = posts[0]
assert "openid" in form["scope"].split(" ")
@pytest.mark.asyncio
async def test_source_refresh_scope_matches_the_configured_login_scope():
"""Refresh requests exactly the scopes the login was granted, read through the same shared
owner off the injected env, so the two can never diverge and the refresh stays within RFC 6749
granted scope. Sourcing the scope from the global environment instead of the injected getenv
makes this fail."""
env = {**_SSO_ENV, "GENERIC_SCOPE": "openid email offline_access"}
source, _, _, posts = _source(_stored(-100), [{"id_token": _make_id_token(exp_offset=7200)}], env=env)
await source.fetch_usable("alice")
_, form = posts[0]
assert form["scope"] == "openid email offline_access"
@pytest.mark.asyncio
async def test_source_refresh_without_rotated_token_carries_the_stored_one_forward():
source, _, _, _ = _source(_stored(-100), [{"id_token": _make_id_token(exp_offset=7200)}])
lookup = await source.fetch_usable("alice")
assert lookup is not None
assert lookup.refresh_token is not None
assert lookup.refresh_token == "rt_stored"
@pytest.mark.asyncio
async def test_source_expired_without_refresh_token_is_expired():
source, _, _, posts = _source(_stored(-100, refresh_token=None), [])
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
assert posts == []
@pytest.mark.asyncio
@pytest.mark.parametrize("body", [{}, {"id_token": ""}, {"id_token": 123}, {"access_token": "at-only"}])
async def test_source_refresh_granting_no_usable_id_token_records_a_dead_grant(body):
"""A 2xx that carries no usable id_token spent the grant for nothing, so it is recorded dead on
the row (refresh token dropped) rather than left to re-POST. A second egress must not spend the
grant again: one-time refresh tokens and IdPs that omit id_token on refresh would otherwise
thrash the token endpoint on every call."""
source, _, persisted, posts = _source(_stored(-100), [body, body])
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
assert len(persisted) == 1
assert persisted[0][1].refresh_token is None
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
assert len(posts) == 1 # the dead-grant row short-circuits the second attempt with no POST
@pytest.mark.asyncio
async def test_source_refresh_rejection_is_expired_and_unreachable_is_unavailable():
"""A 4xx is a verdict on the grant (re-login fixes it); an unreachable IdP proves nothing
about the grant, so the user must NOT be told to re-login for a transient outage."""
async def fetch(user_id: str):
return _stored(-100)
async def persist(user_id, assertion):
return None
async def rejected_post(url, form):
raise SsoAssertionUnrenewable("rejected")
async def unreachable_post(url, form):
raise TokenStoreUnavailable("unreachable")
rejected_source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=rejected_post, getenv=_SSO_ENV.get)
with pytest.raises(SsoAssertionUnrenewable):
await rejected_source.fetch_usable("alice")
unreachable_source = LiveSsoAssertionSource(
fetch=fetch, persist=persist, post=unreachable_post, getenv=_SSO_ENV.get
)
with pytest.raises(TokenStoreUnavailable):
await unreachable_source.fetch_usable("alice")
@pytest.mark.asyncio
async def test_source_missing_sso_env_is_unavailable_not_expired():
"""Env absence is a deployment problem; a re-login cannot fix it (SSO needs the same env),
so the state must be Unavailable, not Expired."""
source, _, _, posts = _source(_stored(-100), [], env={})
with pytest.raises(TokenStoreUnavailable):
await source.fetch_usable("alice")
assert posts == []
@pytest.mark.asyncio
async def test_a_dead_grant_costs_one_post_because_the_verdict_lands_on_the_row():
"""A dead refresh grant costs one token-endpoint POST, not one per egress call, because the
rejection strips the refresh token from the row itself; later reads answer from the row. The
bound therefore holds on every pod and for as long as the row stands, and a re-login lifts it
by replacing the row rather than by waiting out a timer."""
clock = {"now": time.time()}
posts: list[dict] = []
state = {"stored": _stored(-100)}
async def fetch(user_id: str):
return state["stored"]
async def persist(user_id, assertion):
state["stored"] = assertion
async def post(url, form):
posts.append(dict(form))
raise SsoAssertionUnrenewable("rejected")
source = LiveSsoAssertionSource(
fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get, now=lambda: clock["now"]
)
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
assert len(posts) == 1
clock["now"] = clock["now"] + 31.0
state["stored"] = _stored(3600)
lookup = await source.fetch_usable("alice")
assert lookup is not None
assert len(posts) == 1
@pytest.mark.asyncio
async def test_a_late_dead_grant_verdict_never_clobbers_anothers_renewal():
"""The verdict is a fact about the ONE grant the IdP rejected. When another pod's renewal
already replaced that grant (rotating IdPs reject the loser of a cross-pod race with
``invalid_grant``), the loser's verdict describes a token the row no longer holds, so
recording it must be a no-op — not an overwrite that destroys the winner's live assertion
and forces an org-visible re-login."""
winner_row = _stored(7200, refresh_token="rt_winner")
state = {"stored": _stored(-100, refresh_token="rt_loser")}
persisted: list[tuple[str, SSOIdentityAssertion]] = []
async def fetch(user_id: str):
return state["stored"]
async def persist(user_id, assertion):
persisted.append((user_id, assertion))
state["stored"] = assertion
async def post(url, form):
# The winner's write lands while this pod's POST is in flight.
state["stored"] = winner_row
raise SsoAssertionUnrenewable("invalid_grant: rt_loser was already spent")
source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get)
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
assert persisted == []
assert state["stored"] is winner_row
@pytest.mark.asyncio
async def test_dead_grant_recording_preserves_the_rows_current_material():
"""The verdict strips exactly the rejected grant; every other fact on the row (id_token,
issuer, expiry) is the row's own current material, re-read at recording time, never this
pod's stale snapshot of it."""
stored = SSOIdentityAssertion(
id_token=SecretStr(_make_id_token(exp_offset=-100)),
refresh_token=SecretStr("rt_stored"),
issuer=ISSUER,
expires_at=datetime.now(timezone.utc) - timedelta(seconds=100),
)
state = {"stored": stored}
persisted: list[tuple[str, SSOIdentityAssertion]] = []
async def fetch(user_id: str):
return state["stored"]
async def persist(user_id, assertion):
persisted.append((user_id, assertion))
state["stored"] = assertion
async def post(url, form):
raise SsoAssertionUnrenewable("rejected")
source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get)
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
assert len(persisted) == 1
recorded = persisted[0][1]
assert recorded.refresh_token is None
assert recorded.id_token.get_secret_value() == stored.id_token.get_secret_value()
assert recorded.issuer == stored.issuer
assert recorded.expires_at == stored.expires_at
@pytest.mark.asyncio
async def test_recording_a_dead_grant_is_best_effort_when_the_store_is_unreachable():
"""Failing to record the verdict costs one more POST next time, never a wrong answer: the
caller still gets the re-login challenge (that verdict is proven), and nothing is written."""
outage = {"armed": False}
persisted: list[tuple[str, SSOIdentityAssertion]] = []
async def fetch(user_id: str):
# The verdict recording's re-read is the first read after the rejected POST; killing
# exactly that one leaves every other layer's reads (initial, single-flight double-check,
# authority fall-through) working.
if outage["armed"]:
outage["armed"] = False
raise TokenStoreUnavailable("store went away")
return _stored(-100)
async def persist(user_id, assertion):
persisted.append((user_id, assertion))
async def post(url, form):
outage["armed"] = True
raise SsoAssertionUnrenewable("rejected")
source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get)
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
assert persisted == []
@pytest.mark.asyncio
async def test_source_unavailable_refresh_keeps_the_grant_for_the_next_attempt():
"""A refresh failure that proves nothing about the grant (the gateway's own client config
rejected, the IdP down) must leave the stored refresh token in place, so the next attempt
after the operator's fix renews with no re-login."""
new_id_token = _make_id_token(exp_offset=7200)
source, _, persisted, posts = _source(
_stored(-100), [None, {"id_token": new_id_token, "refresh_token": "rt_rotated"}]
)
with pytest.raises(TokenStoreUnavailable):
await source.fetch_usable("alice")
lookup = await source.fetch_usable("alice")
assert lookup is not None
assert lookup.access_token == new_id_token
assert [form["refresh_token"] for _, form in posts] == ["rt_stored", "rt_stored"]
assert len(persisted) == 1
@pytest.mark.asyncio
async def test_source_relogin_immediately_after_a_rejection_is_honored():
"""The 401 challenge tells the user to re-login; the fresh row must be served the moment it
lands (no clock advance, same pod). Nothing beside the row may hold the verdict: a rejection
is recorded ON the row, so a re-login lifts it by construction, and no per-pod state may
mask the authoritative read, or the recovery the challenge asks for silently fails."""
posts: list[dict] = []
state = {"stored": _stored(-100)}
async def fetch(user_id: str):
return state["stored"]
async def persist(user_id, assertion):
state["stored"] = assertion
async def post(url, form):
posts.append(dict(form))
raise SsoAssertionUnrenewable("rejected")
source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get)
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
assert len(posts) == 1
state["stored"] = _stored(3600) # re-login lands right after the rejection, no clock advance
lookup = await source.fetch_usable("alice")
assert lookup is not None
assert len(posts) == 1 # the fresh usable row short-circuits before any refresh POST
@pytest.mark.asyncio
async def test_source_waiter_of_an_expired_winner_does_not_respend_the_grant():
"""The single-flight hand-off carries negative verdicts too: a waiter behind a winner whose
refresh was rejected observes the winner's raised verdict instead of re-POSTing the spent
token."""
import asyncio
posts: list[dict] = []
async def fetch(user_id: str):
await asyncio.sleep(0)
return _stored(-100)
async def persist(user_id, assertion):
return None
async def post(url, form):
posts.append(dict(form))
await asyncio.sleep(0)
raise SsoAssertionUnrenewable("rejected")
source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get)
results = await asyncio.gather(source.fetch_usable("alice"), source.fetch_usable("alice"), return_exceptions=True)
assert all(isinstance(r, SsoAssertionUnrenewable) for r in results)
assert len(posts) == 1
@pytest.mark.asyncio
async def test_source_concurrent_expired_fetches_refresh_once():
"""Two flights hit the same expired assertion; the single-flight lock makes the second
re-read the row the first refreshed instead of spending the refresh token again. The fake
post yields to the loop so the flights genuinely interleave."""
import asyncio
fetches: list[str] = []
posts: list[tuple[str, dict]] = []
state = {"stored": _stored(-100)}
bodies = [{"id_token": _make_id_token(exp_offset=7200)}, {"id_token": _make_id_token(exp_offset=7200)}]
async def fetch(user_id: str):
await asyncio.sleep(0)
fetches.append(user_id)
return state["stored"]
async def persist(user_id: str, assertion: SSOIdentityAssertion):
await asyncio.sleep(0)
state["stored"] = assertion
async def post(url: str, form: dict):
posts.append((url, dict(form)))
await asyncio.sleep(0)
return bodies.pop(0)
source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get)
results = await asyncio.gather(source.fetch_usable("alice"), source.fetch_usable("alice"))
assert all(r is not None for r in results)
assert len(posts) == 1
@pytest.mark.asyncio
async def test_source_memoizes_usable_lookups_within_ttl():
"""A usable assertion is served from the in-process memo, so repeated egress calls do not
pay a DB read per call; the memo expires with the TTL and the row is re-read."""
clock = {"now": time.time()}
fetches: list[str] = []
stored = _stored(3600)
async def fetch(user_id: str):
fetches.append(user_id)
return stored
async def persist(user_id, assertion):
return None
async def post(url, form):
return None
source = LiveSsoAssertionSource(
fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get, now=lambda: clock["now"], cache_ttl_seconds=60.0
)
first = await source.fetch_usable("alice")
second = await source.fetch_usable("alice")
assert first is not None and second is not None
assert fetches == ["alice"]
clock["now"] = clock["now"] + 61.0
third = await source.fetch_usable("alice")
assert third is not None
assert fetches == ["alice", "alice"]
@pytest.mark.asyncio
async def test_a_dead_grant_verdict_never_refuses_a_caller_who_signed_in_again():
"""A rejected grant must not refuse the NEXT assertion, even one that is itself near expiry.
The verdict is recorded on the row (its refresh token is stripped), so replacing the row at
re-login clears it by construction. A verdict parked in a side cache keyed by user would still
be inside its window here and would raise a false 401 at the moment the fresh assertion first
needs renewing, despite a perfectly good refresh token.
"""
state = {"stored": _stored(-100, refresh_token="rt_dead")}
outcomes = {"reject": True}
async def fetch(user_id: str):
return state["stored"]
async def persist(user_id, assertion):
state["stored"] = assertion
async def post(url, form):
if outcomes["reject"]:
raise SsoAssertionUnrenewable("grant is dead")
return {"id_token": _make_id_token(exp_offset=3600), "refresh_token": "rt_fresh"}
source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get)
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
# The verdict lives on the row now: the dead refresh token is gone, so a repeat costs no POST.
assert state["stored"].refresh_token is None
# The user signs in again. The new assertion is ALSO near expiry, so renewal runs immediately;
# it must use the new grant rather than a remembered verdict.
outcomes["reject"] = False
state["stored"] = _stored(-5, refresh_token="rt_after_relogin")
renewed = await source.fetch_usable("alice")
assert renewed is not None
assert renewed.refresh_token == "rt_fresh"
@pytest.mark.asyncio
async def test_replacing_the_row_drops_a_cached_superseded_assertion():
"""A re-login must take effect at once, not when the superseded id_token finally expires.
The assertion is identity material, so a user who re-authenticates after their IdP claims are
REDUCED must stop being able to spend the old one. ``persist_sso_identity_assertion`` is the
single place the row is replaced, so it drops the cached predecessor; without that hook the
warm entry below keeps serving the old id_token for its full remaining life.
"""
superseded = _make_id_token(exp_offset=3600, iss="https://old.example.com")
replacement = _make_id_token(exp_offset=3600, iss="https://new.example.com")
state = {"stored": SSOIdentityAssertion(id_token=SecretStr(superseded))}
async def fetch(user_id: str):
return state["stored"]
async def persist(user_id, assertion):
state["stored"] = assertion
async def post(url, form):
raise TokenStoreUnavailable("unreachable")
source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get)
warm = await source.fetch_usable("alice")
assert warm is not None
assert warm.access_token == superseded
# The user signs in again; the row is replaced through the real write chokepoint.
state["stored"] = SSOIdentityAssertion(id_token=SecretStr(replacement))
with patch("litellm.proxy.proxy_server.prisma_client", _make_prisma({})):
await persist_sso_identity_assertion("alice", state["stored"])
after = await source.fetch_usable("alice")
assert after is not None
assert after.access_token == replacement
@pytest.mark.asyncio
async def test_source_never_memoizes_absence_so_a_fresh_login_is_seen_immediately():
fetches: list[str] = []
state = {"stored": None}
async def fetch(user_id: str):
fetches.append(user_id)
return state["stored"]
async def persist(user_id, assertion):
return None
async def post(url, form):
return None
source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get)
assert await source.fetch_usable("alice") is None
state["stored"] = _stored(3600)
assert await source.fetch_usable("alice") is not None
assert fetches == ["alice", "alice"]
@pytest.mark.asyncio
async def test_source_store_outage_is_retryable_and_a_warm_cache_survives_it():
"""A warm cache keeps serving through an outage, and once cold the outage is RETRYABLE.
It must never read as absence: this user has an assertion, so answering "you have never signed
in" would send them to a sign-in that fixes nothing while the real remedy is to retry. ``None``
is reserved for the one case where no row was ever written.
"""
calls = {"n": 0}
stored = _stored(3600)
async def flaky_fetch(user_id: str):
calls["n"] += 1
if calls["n"] > 1:
raise RuntimeError("db down")
return stored
async def persist(user_id, assertion):
return None
async def post(url, form):
return None
clock = {"now": time.time()}
source = LiveSsoAssertionSource(
fetch=flaky_fetch,
persist=persist,
post=post,
getenv=_SSO_ENV.get,
now=lambda: clock["now"],
cache_ttl_seconds=60.0,
)
assert await source.fetch_usable("alice") is not None
assert await source.fetch_usable("alice") is not None
clock["now"] = clock["now"] + 61.0
with pytest.raises(TokenStoreUnavailable):
await source.fetch_usable("alice")
@pytest.mark.asyncio
async def test_source_persist_failure_after_refresh_still_serves_the_in_hand_token():
"""A write-back failure is not a read failure: the freshly refreshed assertion is valid
and in hand, and discarding it would also strand a one-time rotated refresh token."""
new_id_token = _make_id_token(exp_offset=7200)
posts: list[tuple[str, dict]] = []
async def fetch(user_id: str):
return _stored(-100)
async def failing_persist(user_id, assertion):
raise RuntimeError("db write down")
async def post(url, form):
posts.append((url, dict(form)))
return {"id_token": new_id_token, "refresh_token": "rt_rotated"}
source = LiveSsoAssertionSource(fetch=fetch, persist=failing_persist, post=post, getenv=_SSO_ENV.get)
lookup = await source.fetch_usable("alice")
assert lookup is not None
assert lookup.access_token == new_id_token
assert len(posts) == 1
@pytest.mark.asyncio
async def test_source_refresh_returning_an_already_expired_token_records_a_dead_grant():
"""Usable has ONE construction gate: a refreshed assertion passes the same expiry predicate a
stored one does, so an IdP handing back an already-dead token raises the re-login verdict
rather than being sent to an exchange. A short-lived token that lands expired every time would
thrash if its rotated refresh token were kept, so this is the same dead-grant outcome as a
missing id_token: the grant is recorded dead and re-login is the remedy, not another spend of
the rotation."""
expired_new = _make_id_token(exp_offset=-10)
persisted: list[tuple[str, SSOIdentityAssertion]] = []
async def fetch(user_id: str):
return _stored(-100)
async def persist(user_id, assertion):
persisted.append((user_id, assertion))
async def post(url, form):
return {"id_token": expired_new, "refresh_token": "rt_rotated"}
source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get)
with pytest.raises(SsoAssertionUnrenewable):
await source.fetch_usable("alice")
assert len(persisted) == 1
assert persisted[0][1].refresh_token is None
@pytest.mark.asyncio
async def test_source_waiter_gets_the_winners_token_when_persist_failed():
"""The memo is the single-flight hand-off: with the write-back failing and a one-time
refresh token, the waiter must receive the winner's in-hand assertion from the memo under
the lock instead of re-reading the stale row and re-spending the refresh grant."""
import asyncio
new_id_token = _make_id_token(exp_offset=7200)
posts: list[dict] = []
async def fetch(user_id: str):
await asyncio.sleep(0)
return _stored(-100)
async def failing_persist(user_id, assertion):
raise RuntimeError("db write down")
async def post(url, form):
posts.append(dict(form))
await asyncio.sleep(0)
if form["refresh_token"] != "rt_stored" or len(posts) > 1:
raise SsoAssertionUnrenewable("rejected")
return {"id_token": new_id_token, "refresh_token": "rt_rotated"}
source = LiveSsoAssertionSource(fetch=fetch, persist=failing_persist, post=post, getenv=_SSO_ENV.get)
results = await asyncio.gather(source.fetch_usable("alice"), source.fetch_usable("alice"))
assert all(r is not None for r in results)
assert len(posts) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize(
"status_code,body,expected_type",
[
(400, {"error": "invalid_grant"}, SsoAssertionUnrenewable),
(401, {"error": "invalid_grant"}, SsoAssertionUnrenewable),
(401, {"error": "invalid_client"}, TokenStoreUnavailable),
(400, {"error": "invalid_scope"}, TokenStoreUnavailable),
(400, {"error": "invalid_request"}, TokenStoreUnavailable),
(429, {"error": "rate_limited"}, TokenStoreUnavailable),
(400, "not-json-object", TokenStoreUnavailable),
(400, None, TokenStoreUnavailable),
(502, {"error": "invalid_grant"}, TokenStoreUnavailable),
],
)
async def test_post_helper_rejection_requires_a_proven_oauth_verdict(status_code, body, expected_type, monkeypatch):
"""Rejected needs proof of a verdict ON THE GRANT, and only ``invalid_grant`` states one.
``invalid_client``/``invalid_scope`` are the gateway's own client config failing — an ops
problem a re-login cannot fix, so treating them as grant death would destroy every renewable
assertion over a rotated GENERIC_CLIENT_SECRET. A rate limit, a 4xx without the error object
(an intermediary, not the token endpoint), or any 5xx likewise proves nothing about the
grant and must not send the user to re-login."""
import json as _json
import httpx
from litellm.proxy._experimental.mcp_server.outbound_credentials import sso_assertion_store
response = MagicMock()
response.status_code = status_code
if body is None:
response.json.side_effect = _json.JSONDecodeError("x", "y", 0)
else:
response.json.return_value = body
response.raise_for_status.side_effect = httpx.HTTPStatusError("err", request=MagicMock(), response=response)
client = MagicMock()
client.post = AsyncMock(return_value=response)
monkeypatch.setattr("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", lambda llm_provider: client)
with pytest.raises(expected_type):
await sso_assertion_store._post_sso_token_endpoint("https://idp.example.com/token", {"a": "b"})
@pytest.mark.asyncio
@pytest.mark.parametrize("body", [None, [], "string", 42])
async def test_post_helper_2xx_non_object_body_is_unreachable_not_a_crash(body, monkeypatch):
"""A 2xx whose body is not a JSON object must read Unreachable; letting it construct a
Granted outcome detonates downstream validation into the broad catch, which misreads the
state as Absent and misdirects the user to a 412 instead of a retryable failure."""
client = MagicMock()
response = MagicMock()
response.status_code = 200
response.raise_for_status = MagicMock()
response.json.return_value = body
client.post = AsyncMock(return_value=response)
monkeypatch.setattr("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", lambda llm_provider: client)
from litellm.proxy._experimental.mcp_server.outbound_credentials import sso_assertion_store
with pytest.raises(TokenStoreUnavailable):
await sso_assertion_store._post_sso_token_endpoint("https://idp.example.com/token", {"a": "b"})
@pytest.mark.asyncio
async def test_post_helper_transport_failure_is_unreachable(monkeypatch):
client = MagicMock()
client.post = AsyncMock(side_effect=ConnectionError("down"))
monkeypatch.setattr("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", lambda llm_provider: client)
from litellm.proxy._experimental.mcp_server.outbound_credentials import sso_assertion_store
with pytest.raises(TokenStoreUnavailable):
await sso_assertion_store._post_sso_token_endpoint("https://idp.example.com/token", {"a": "b"})

8
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py
View file

@ -19,9 +19,7 @@ def _wrapping_codec():
def test_round_trips_the_access_token():
codec = _wrapping_codec()
token = codec.decode(
codec.encode(OAuthToken(access_token="at-123", expires_at=1234.5))
)
token = codec.decode(codec.encode(OAuthToken(access_token="at-123", expires_at=1234.5)))
assert token is not None
assert token.access_token == "at-123"
@ -30,9 +28,7 @@ def test_encode_encrypts_and_omits_the_refresh_token():
codec = _wrapping_codec()
blob = codec.encode(OAuthToken(access_token="at", refresh_token="super-secret-rt"))
assert blob.startswith("enc:") # encryption was applied
assert (
"super-secret-rt" not in blob
) # the long-lived secret never reaches the cache
assert "super-secret-rt" not in blob # the long-lived secret never reaches the cache
decoded = codec.decode(blob)
assert decoded is not None and decoded.refresh_token is None

5
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py
View file

@ -35,10 +35,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
)
from pydantic import SecretStr
_PATCH_TARGET = (
"litellm.proxy._experimental.mcp_server.outbound_credentials."
"token_endpoint.get_async_httpx_client"
)
_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint.get_async_httpx_client"
_ENDPOINT = "https://idp.example.com/oauth2/token"
_CLIENT_ID = "litellm-client-id"

8
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py
View file

@ -98,9 +98,7 @@ def test_aws_sigv4_defaults_to_the_ambient_credential_chain():
def test_authconfig_discriminates_on_kind():
api_key = _AUTH_CONFIG.validate_python(
{"kind": "api_key", "key_source": {"source": "shared", "value": "k"}}
)
api_key = _AUTH_CONFIG.validate_python({"kind": "api_key", "key_source": {"source": "shared", "value": "k"}})
assert isinstance(api_key, ApiKeyConfig)
assert isinstance(api_key.key_source, SharedKey)
@ -197,9 +195,7 @@ def test_id_jag_client_auth_discriminates_on_source():
def test_id_jag_client_auth_rejects_unknown_source():
with pytest.raises(ValidationError):
_AUTH_CONFIG.validate_python(
{**_ID_JAG_MINIMAL, "client_auth": {"source": "mystery"}}
)
_AUTH_CONFIG.validate_python({**_ID_JAG_MINIMAL, "client_auth": {"source": "mystery"}})
def test_id_jag_config_defaults_id_token_subject_and_empty_optionals():

4
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_v2_token_store.py
View file

@ -43,9 +43,7 @@ async def test_optional_fields_absent_yield_none():
@pytest.mark.asyncio
async def test_unparseable_expiry_is_dropped_not_raised():
store = V2PerUserTokenStore(
_reader({"access_token": "at", "expires_at": "not-a-date"})
)
store = V2PerUserTokenStore(_reader({"access_token": "at", "expires_at": "not-a-date"}))
token = await store.fetch("alice", "srv")
assert token is not None and token.expires_at is None

61
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -805,6 +805,9 @@ async def test_mcp_get_prompt_success():
mcp_auth_header={"Authorization": "token"},
extra_headers={"X-Test": "1"},
raw_headers=None,
# An id_jag server resolves its stored SSO assertion by user id, so dropping the caller
# identity here would make prompts fail closed while tools/call still succeeded.
user_api_key_auth=user_api_key_auth,
)
assert result is prompt_result
@ -866,6 +869,7 @@ async def test_mcp_read_resource_success():
mcp_auth_header={"Authorization": "token"},
extra_headers={"X-Test": "1"},
raw_headers=None,
user_api_key_auth=user_api_key_auth,
)
assert result is read_result
@ -7726,3 +7730,60 @@ class TestPreemptive401ModeAware:
await self._run(delegate, None, has_stored_token=False)
assert exc.value.status_code == 401
await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False)
class TestPreemptiveIdJagPreflight:
"""The transport-edge chokepoint runs the id_jag preflight on single-server routes only,
with the caller's raw headers, so a dead stored assertion surfaces its re-login 401 where a
WWW-Authenticate still reaches the client — and a multi-server aggregate keeps absorbing
per-server failures instead of letting one server 401 the whole connect."""
def _id_jag_server(self, alias: str):
return MCPServer(
server_id=f"id-{alias}",
name=alias,
alias=alias,
server_name=alias,
url=f"https://{alias}.test/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_id_jag,
client_id="gw-client",
client_secret="gw-secret",
token_exchange_endpoint="https://org-idp.test/token",
id_jag_resource_token_endpoint="https://res-as.test/token",
mcp_info={"server_name": alias},
)
async def _run(self, mcp_servers, server, raw_headers=None):
from litellm.proxy._experimental.mcp_server import server as server_module
preflight = AsyncMock(return_value=None)
with (
patch.object(server_module.global_mcp_server_manager, "get_mcp_server_by_name", return_value=server),
patch.object(server_module.global_mcp_server_manager, "preflight_id_jag", preflight),
):
await server_module._raise_preemptive_401_for_unauthenticated_servers(
scope={"type": "http", "method": "POST", "path": f"/mcp/{server.alias}", "headers": []},
mcp_servers=mcp_servers,
oauth2_headers=None,
mcp_server_auth_headers=None,
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key"),
client_ip=None,
raw_headers=raw_headers,
)
return preflight
@pytest.mark.asyncio
async def test_single_server_route_runs_the_preflight_with_the_raw_headers(self):
server = self._id_jag_server("emasrv")
raw = {"x-litellm-api-key": "sk-litellm-virtual-key"}
preflight = await self._run([server.alias], server, raw_headers=raw)
preflight.assert_awaited_once()
assert preflight.await_args.kwargs["raw_headers"] == raw
assert preflight.await_args.kwargs["server"] is server
@pytest.mark.asyncio
async def test_multi_server_aggregate_skips_the_preflight(self):
server = self._id_jag_server("emasrv")
preflight = await self._run([server.alias, "other-server"], server)
preflight.assert_not_awaited()

456
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -1020,15 +1020,26 @@ class TestMCPServerManager:
assert exc_info.value.www_authenticate == challenge
@pytest.mark.asyncio
async def test_list_surfaces_non_auth_httpexception_as_internal_fault(self):
"""A non-auth HTTP error (e.g. 412 no endpoint, 503 IdP down) now raises MCPServerListError
with an "internal" fault carrying the status code instead of absorbing to []: the silent
empty list made a misconfigured/unavailable server indistinguishable from a healthy server
with no tools. The aggregate absorbs it into that server's outcome, so one broken server
still does not blank the whole aggregate listing."""
@pytest.mark.parametrize(
"status_code,expected_tag",
[
(503, "unavailable"), # resolver's retryable verdict (store/IdP outage) stays retryable
(412, "precondition"), # resolver's per-caller precondition (e.g. gateway SSO sign-in)
(500, "internal"), # only the gateway's own failure reads as the gateway's fault
],
)
async def test_list_preserves_the_resolvers_verdict_class_on_non_auth_httpexceptions(
self, status_code, expected_tag
):
"""A non-auth HTTP error raises MCPServerListError with a fault that preserves the
resolver's retryable-vs-terminal verdict instead of absorbing to [] (the silent empty list
made a broken server indistinguishable from a healthy empty one) and instead of collapsing
everything to "internal" (which reported an IdP outage as the gateway's own 500). The
aggregate absorbs the fault into that server's outcome, so one broken server still does
not blank the whole aggregate listing."""
server = MCPServer(
server_id="te-412",
name="te-412-server",
server_id=f"te-{status_code}",
name=f"te-{status_code}-server",
url="https://up.example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
@ -1037,13 +1048,11 @@ class TestMCPServerManager:
client_secret="csec",
)
manager = MCPServerManager()
manager._create_mcp_client = AsyncMock(
side_effect=HTTPException(status_code=412, detail="token exchange endpoint is not configured")
)
manager._create_mcp_client = AsyncMock(side_effect=HTTPException(status_code=status_code, detail="resolver"))
with pytest.raises(MCPServerListError) as exc_info:
await manager._get_tools_from_server(server=server, oauth2_headers={"Authorization": "Bearer subj-jwt"})
assert exc_info.value.fault == ServerListFault(tag="internal", status_code=412)
assert exc_info.value.server_name == "te-412-server"
assert exc_info.value.fault == ServerListFault(tag=expected_tag, status_code=status_code)
assert exc_info.value.server_name == f"te-{status_code}-server"
def _upstream_status_error(self, status_code: int, www_authenticate: Optional[str] = None) -> httpx.HTTPStatusError:
"""Build an httpx.HTTPStatusError shaped like the one the MCP SDK surfaces for an upstream
@ -1875,6 +1884,133 @@ class TestMCPServerManager:
await manager.preflight_token_exchange(server=server, oauth2_headers=None, user_api_key_auth=None)
assert resolved == ["good-subject"]
@staticmethod
def _id_jag_server(server_id: str) -> MCPServer:
return MCPServer(
server_id=server_id,
name=f"{server_id}-name",
url="https://upstream.example/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_id_jag,
client_id="gateway-client",
client_secret="gateway-secret",
token_exchange_endpoint="https://org-idp.example/oauth2/token",
id_jag_resource_token_endpoint="https://resource-as.example/oauth2/token",
)
@pytest.mark.asyncio
async def test_preflight_id_jag_surfaces_the_relogin_challenge_at_the_transport_edge(self):
"""A dead stored assertion must raise the resolver's re-login 401 (challenge header
intact) from the preflight, so a single-server front-door client sees the remedy instead
of the session opening onto an empty tool/prompt/resource catalog."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError
class _FakeProvider:
async def resolve_credentials(self, subject, server):
return Error(
CredError.of_unauthorized(
"the stored identity assertion has expired; sign in to the gateway again",
www_authenticate='Bearer error="invalid_token", error_description="re-authenticate"',
)
)
manager = MCPServerManager(cred_provider=_FakeProvider())
with pytest.raises(HTTPException) as exc_info:
await manager.preflight_id_jag(
server=self._id_jag_server("idjag-preflight-401"),
oauth2_headers=None,
raw_headers={"x-litellm-api-key": "Bearer sk-admission"},
user_api_key_auth=None,
)
assert exc_info.value.status_code == 401
headers = exc_info.value.headers or {}
www_authenticate = headers.get("WWW-Authenticate") or headers.get("www-authenticate") or ""
assert "invalid_token" in www_authenticate
@pytest.mark.asyncio
@pytest.mark.parametrize(
"error_factory,expected_status",
[
(lambda CredError: CredError.of_upstream_unavailable("the assertion store could not be read"), 503),
(lambda CredError: CredError.of_precondition_required("no stored identity assertion"), 412),
],
)
async def test_preflight_id_jag_preserves_retryable_vs_terminal(self, error_factory, expected_status):
"""A store/IdP outage must surface as the retryable 503 and a never-signed-in user as the
412 — at the transport edge, with the same statuses the call surface uses, so the two
surfaces never tell the caller different stories."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError
class _FakeProvider:
async def resolve_credentials(self, subject, server):
return Error(error_factory(CredError))
manager = MCPServerManager(cred_provider=_FakeProvider())
with pytest.raises(HTTPException) as exc_info:
await manager.preflight_id_jag(
server=self._id_jag_server("idjag-preflight-status"),
oauth2_headers=None,
raw_headers=None,
user_api_key_auth=None,
)
assert exc_info.value.status_code == expected_status
@pytest.mark.asyncio
async def test_preflight_id_jag_subject_follows_the_shared_bearer_rule(self):
"""The preflight selects its subject with the same one rule egress uses
(`_subject_bearer_token`), including the free-rider guard: an Authorization bearer that
did not admit the caller is never handed to the resolver as exchange subject material."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import (
StaticHeaderAuth,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok
subjects = []
class _FakeProvider:
async def resolve_credentials(self, subject, server):
subjects.append(subject.inbound_token.get_secret_value() if subject.inbound_token else None)
return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization"))
manager = MCPServerManager(cred_provider=_FakeProvider())
server = self._id_jag_server("idjag-preflight-subject")
# Admission consumed the explicit key header, so the Authorization bearer is a free
# rider: the preflight must resolve through the stored assertion (no inbound subject).
await manager.preflight_id_jag(
server=server,
oauth2_headers=None,
raw_headers={"x-litellm-api-key": "sk-admission", "authorization": "Bearer free-rider"},
user_api_key_auth=None,
)
# The bearer WAS the admission credential, so it is the subject.
await manager.preflight_id_jag(
server=server,
oauth2_headers=None,
raw_headers={"authorization": "Bearer caller-id-token"},
user_api_key_auth=None,
)
assert subjects == [None, "caller-id-token"]
@pytest.mark.asyncio
async def test_preflight_id_jag_noop_for_other_auth_types(self):
"""The preflight is mode-terminal: a non-id_jag server never reaches the resolver here
(its own preflight or challenge path owns it)."""
class _FakeProvider:
async def resolve_credentials(self, subject, server):
raise AssertionError("resolver must not run for non-id_jag servers")
manager = MCPServerManager(cred_provider=_FakeProvider())
await manager.preflight_id_jag(
server=self._token_exchange_server("te-not-id-jag"),
oauth2_headers={"Authorization": "Bearer subject"},
raw_headers=None,
user_api_key_auth=None,
)
@pytest.mark.asyncio
async def test_call_regular_mcp_tool_passthrough_strips_authorization_when_admission_consumed_litellm_key(
self,
@ -2670,6 +2806,7 @@ class TestMCPServerManager:
extra_headers=None,
stdio_env=None,
subject_token=None,
user_api_key_auth=None,
)
mock_client.list_resource_templates.assert_awaited_once()
mock_prefix.assert_called_once_with(mock_templates, server, add_prefix=False)
@ -3443,7 +3580,7 @@ class TestMCPServerManager:
# Capture the extra_headers passed to _create_mcp_client
captured_extra_headers = None
async def capture_create_mcp_client(server, mcp_auth_header, extra_headers, stdio_env):
async def capture_create_mcp_client(server, mcp_auth_header, extra_headers, stdio_env, **kwargs):
nonlocal captured_extra_headers
captured_extra_headers = extra_headers
return mock_client
@ -7573,7 +7710,7 @@ class TestHealthCheckInterpolatesGlobalEnvVars:
mock_client = AsyncMock()
mock_client.run_with_session = AsyncMock(return_value="ok")
async def _create(server, mcp_auth_header, extra_headers, stdio_env):
async def _create(server, mcp_auth_header, extra_headers, stdio_env, **kwargs):
captured["extra_headers"] = extra_headers
return mock_client
@ -9240,3 +9377,292 @@ class TestDiscoveryFailureLogging:
assert "typo_row" in caplog.text
assert "authorization_url, token_url" in caplog.text
assert "unresolved" in caplog.text
@pytest.mark.parametrize(
"auth_type,subject_token,user_id,expected",
[
(MCPAuth.oauth2_token_exchange, "hdr.jwt.sig", None, True),
(MCPAuth.oauth2_token_exchange, None, "alice", False),
(MCPAuth.oauth2_id_jag, "hdr.jwt.sig", None, True),
(MCPAuth.oauth2_id_jag, None, "alice", True),
(MCPAuth.oauth2_id_jag, None, None, False),
(MCPAuth.oauth2_id_jag, None, "", False),
(MCPAuth.oauth2, "hdr.jwt.sig", "alice", False),
(None, "hdr.jwt.sig", "alice", False),
],
)
def test_obo_retry_covers_caller_matrix(auth_type, subject_token, user_id, expected):
"""token_exchange re-mints only from the inbound token; id_jag also re-mints from the
stored SSO assertion keyed by user id, so an id_jag caller with a user identity but no
inbound assertion must still route through the invalidate-and-retry path."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import _obo_retry_covers_caller
from litellm.proxy._types import UserAPIKeyAuth
auth = UserAPIKeyAuth(api_key="sk-x", user_id=user_id) if user_id is not None else None
assert _obo_retry_covers_caller(auth_type, subject_token, auth) is expected
def test_obo_retry_requires_an_authenticated_caller_for_stored_assertion_retry():
from litellm.proxy._experimental.mcp_server.mcp_server_manager import _obo_retry_covers_caller
assert _obo_retry_covers_caller(MCPAuth.oauth2_id_jag, None, None) is False
class TestSubjectBearerTokenSymmetry:
"""One rule decides whether the caller's bearer becomes exchange subject material, and every
egress surface (tools call, tools list, prompts, resources) resolves through it; an id_jag
caller presenting a fresh IdP JWT must see identical sourcing on list and call."""
def _server(self, auth_type):
server = MagicMock()
server.auth_type = auth_type
return server
@pytest.mark.parametrize(
"auth_type,expected",
[
(MCPAuth.oauth2_token_exchange, "hdr.jwt.sig"),
(MCPAuth.oauth2_id_jag, "hdr.jwt.sig"),
(MCPAuth.oauth2, None),
(MCPAuth.api_key, None),
(None, None),
],
)
def test_subject_bearer_mode_matrix(self, auth_type, expected):
manager = MCPServerManager()
raw_headers = {"authorization": "Bearer hdr.jwt.sig"}
assert manager._subject_bearer_token(self._server(auth_type), raw_headers) == expected
def test_oauth2_headers_win_over_raw(self):
manager = MCPServerManager()
token = manager._subject_bearer_token(
self._server(MCPAuth.oauth2_id_jag),
{"authorization": "Bearer raw.jwt.sig"},
oauth2_headers={"Authorization": "Bearer oauth2.jwt.sig"},
)
assert token == "oauth2.jwt.sig"
@pytest.mark.parametrize("key_header", ["x-litellm-api-key", "X-Litellm-Api-Key"])
def test_id_jag_ignores_a_free_rider_jwt_when_admission_used_the_explicit_key(self, key_header):
"""A caller admitted with the explicit litellm key can ride any valid IdP JWT in
Authorization; it was never validated as THIS caller's identity, so exchanging it would
let the caller act upstream as whoever the token belongs to. The bearer is subject
material only when it was the admission credential; such callers resolve through the
stored assertion, which is bound to the authenticated user."""
manager = MCPServerManager()
raw_headers = {key_header: "sk-caller-a", "authorization": "Bearer stolen.user-b.jwt"}
assert manager._subject_bearer_token(self._server(MCPAuth.oauth2_id_jag), raw_headers) is None
assert (
manager._subject_bearer_token(
self._server(MCPAuth.oauth2_id_jag),
raw_headers,
oauth2_headers={"Authorization": "Bearer stolen.user-b.jwt"},
)
is None
)
def test_id_jag_uses_the_bearer_when_it_was_the_admission_credential(self):
manager = MCPServerManager()
raw_headers = {"authorization": "Bearer admitted.idp.jwt"}
assert (
manager._subject_bearer_token(self._server(MCPAuth.oauth2_id_jag), raw_headers) == "admitted.idp.jwt"
)
@pytest.mark.parametrize(
"raw_headers",
[
{"x-litellm-api-key": "", "authorization": "Bearer admitted.idp.jwt"},
{"X-Litellm-Api-Key": "", "authorization": "Bearer admitted.idp.jwt"},
],
)
def test_id_jag_empty_explicit_key_is_not_a_free_rider(self, raw_headers):
"""Admission skips a falsy primary header and authenticates with Authorization, so an
empty explicit key means the bearer IS the admission credential and stays usable."""
manager = MCPServerManager()
assert (
manager._subject_bearer_token(self._server(MCPAuth.oauth2_id_jag), raw_headers) == "admitted.idp.jwt"
)
@pytest.mark.parametrize(
"raw_headers",
[
None,
{},
{"authorization": "Bearer j.w.t"},
{"x-litellm-api-key": "", "authorization": "Bearer j.w.t"},
{"x-litellm-api-key": "sk-a", "authorization": "Bearer j.w.t"},
{"X-LiteLLM-Api-Key": "sk-a", "authorization": "Bearer j.w.t"},
{"x-litellm-api-key": " ", "authorization": "Bearer j.w.t"},
{"x-litellm-api-key": "sk-a"},
],
)
def test_free_rider_predicate_agrees_with_the_admission_accessor(self, raw_headers):
"""The predicate and get_litellm_api_key_from_headers must encode ONE header-preference
rule: the bearer is a free rider exactly when admission's chosen credential is NOT the
Authorization value. A mutation to either side breaks this agreement."""
from starlette.datastructures import Headers
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
headers = Headers(headers=dict(raw_headers or {}))
admission_credential = MCPRequestHandler.get_litellm_api_key_from_headers(headers)
authorization_value = headers.get("authorization")
expected_free_rider = (
admission_credential != authorization_value if authorization_value is not None else bool(admission_credential)
)
assert MCPRequestHandler.authorization_is_free_rider(raw_headers) == expected_free_rider
@pytest.mark.asyncio
async def test_openapi_egress_uses_the_same_subject_bearer_rule(self):
"""spec_path servers egress outside _create_mcp_client, so their extraction is the
eighth surface; it must resolve through the identical rule or the binding invariant
holds on MCP paths while the OpenAPI path exchanges a free-rider token."""
manager = MCPServerManager()
server = MagicMock()
server.auth_type = MCPAuth.oauth2_id_jag
server.server_id = "idjag-openapi"
server.url = None
server.spec_path = "https://spec.example.com/openapi.yaml"
captured: dict[str, object] = {}
def fake_rule(mcp_server, raw_headers, oauth2_headers=None):
captured["called_with"] = (raw_headers, oauth2_headers)
return None
raw = {"x-litellm-api-key": "sk-a", "authorization": "Bearer free.rider.jwt"}
with (
patch.object(manager, "_subject_bearer_token", side_effect=fake_rule),
patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.to_server_spec"
) as spec_mock,
):
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
ClientSecretAuth,
IdJagConfig,
ServerSpec,
)
from pydantic import SecretStr
spec_mock.return_value = ServerSpec(
server_id="idjag-openapi",
resource="https://up.example.com",
config=IdJagConfig(
org_token_endpoint="https://idp.example.com/token",
resource_token_endpoint="https://ras.example.com/token",
client_id="gw",
client_auth=ClientSecretAuth(client_secret=SecretStr("s")),
),
)
with patch.object(
manager._cred_provider, "resolve_credentials", new=AsyncMock(side_effect=RuntimeError("stop"))
):
try:
await manager.resolve_openapi_upstream_auth(
mcp_server=server,
oauth2_headers={"Authorization": "Bearer free.rider.jwt"},
raw_headers=raw,
mcp_auth_header=None,
user_api_key_auth=None,
forwarded_headers=None,
)
except Exception:
pass
assert captured.get("called_with") == (raw, {"Authorization": "Bearer free.rider.jwt"})
def test_obo_keeps_exchange_what_was_presented_semantics(self):
manager = MCPServerManager()
raw_headers = {"x-litellm-api-key": "sk-caller-a", "authorization": "Bearer presented.obo.jwt"}
assert (
manager._subject_bearer_token(self._server(MCPAuth.oauth2_token_exchange), raw_headers)
== "presented.obo.jwt"
)
@pytest.mark.asyncio
async def test_list_path_threads_the_id_jag_subject_token(self):
manager = MCPServerManager()
server = MagicMock()
server.auth_type = MCPAuth.oauth2_id_jag
server.server_id = "idjag-1"
server.name = "idjag_srv"
server.alias = "idjag_srv"
server.transport = MCPTransport.http
server.spec_path = None
captured: dict[str, object] = {}
async def fake_create_client(**kwargs):
captured.update(kwargs)
raise RuntimeError("stop after capture")
from litellm.proxy._experimental.mcp_server.exceptions import MCPServerListError
manager._create_mcp_client = fake_create_client
with patch.object(manager, "_build_stdio_env", return_value=None):
with pytest.raises(MCPServerListError):
await manager._get_tools_from_server(
server,
mcp_auth_header=None,
oauth2_headers={"Authorization": "Bearer caller.idp.jwt"},
raw_headers={"authorization": "Bearer caller.idp.jwt"},
user_api_key_auth=None,
)
assert captured.get("subject_token") == "caller.idp.jwt"
class TestEgressSurfacesThreadCallerIdentity:
"""Every id_jag/OBO egress surface must forward the caller's ``user_api_key_auth`` into
``_create_mcp_client`` so the stored SSO assertion resolves to the right user. A surface that
drops it silently breaks id_jag for virtual-key callers; prompts and resources did exactly
that. This pins the whole prompt/resource family so a surface cannot regress the threading."""
def _server(self):
server = MagicMock()
server.static_headers = None
server.transport = MCPTransport.http
server.url = "https://up.example.com/mcp"
server.name = "srv"
server.alias = "srv"
server.server_name = "srv"
server.server_id = "srv-1"
return server
@pytest.mark.asyncio
@pytest.mark.parametrize(
"invoke",
[
lambda mgr, srv, uk: mgr.get_prompts_from_server(server=srv, raw_headers={}, user_api_key_auth=uk),
lambda mgr, srv, uk: mgr.get_resources_from_server(server=srv, raw_headers={}, user_api_key_auth=uk),
lambda mgr, srv, uk: mgr.get_resource_templates_from_server(server=srv, raw_headers={}, user_api_key_auth=uk),
lambda mgr, srv, uk: mgr.get_prompt_from_server(
server=srv, prompt_name="p", raw_headers={}, user_api_key_auth=uk
),
lambda mgr, srv, uk: mgr.read_resource_from_server(
server=srv, url="https://up.example.com/r", raw_headers={}, user_api_key_auth=uk
),
],
ids=["list_prompts", "list_resources", "list_resource_templates", "get_prompt", "read_resource"],
)
async def test_prompt_and_resource_surfaces_thread_user_api_key_auth(self, invoke):
from litellm.proxy._types import UserAPIKeyAuth
manager = MCPServerManager()
caller = UserAPIKeyAuth(api_key="sk-caller", user_id="alice")
captured: dict = {}
async def capture(**kwargs):
captured.update(kwargs)
return AsyncMock()
with (
patch.object(manager, "_create_mcp_client", side_effect=capture),
patch.object(manager, "_build_stdio_env", return_value=None),
patch.object(manager, "_subject_bearer_token", return_value=None),
):
try:
await invoke(manager, self._server(), caller)
except Exception:
pass
assert captured.get("user_api_key_auth") is caller