scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-28 06:23:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 94

feat(mcp): backport the ID-JAG stack to rc/1.107.0 (#46748, #46747, #46745, #47004) (#47058)

Browse source 
* feat(mcp): accept an ID-JAG as an RFC 7523 jwt-bearer grant on the gateway token endpoint (#46748)

Backport of #46748 to rc/1.107.0.
Cherry-picked from bca7650e8f (main).

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(sso): capture the IdP id_token on Microsoft and Google logins for oauth2_id_jag (#46747)

Backport of #46747 to rc/1.107.0.
Cherry-picked from 3501e36fa9 (main).

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(sso): retain the SAML assertion at login as the saml2 subject token for oauth2_id_jag (#46745)

Backport of #46745 to rc/1.107.0.
Cherry-picked from f02449ad38 (main).

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): keep the OIDC login refresh token as the ID-JAG subject when the SSO client is the ID-JAG client (#47004)

Backport of #47004 to rc/1.107.0.
Cherry-picked from fc484be77a (main).

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam 2026-10-17 09:21:26 -07:00 • committed by GitHub
parent 4bf1633452
commit 1f8bdaa657
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
44 changed files with 4936 additions and 500 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1345,7 +1345,9 @@ class MCPRequestHandler:
if auth is None:
return None
refreshed = auth.model_copy()
if auth.mcp_admitted_user_subject and auth.user_id:
if (auth.mcp_admitted_user_subject or auth.is_session_token) and auth.user_id:
# A session credential (``lite login`` or a redeemed ID-JAG) has no key row, its authority is
# sealed in the token: refresh the live user policy and keep the token's team, routes and models.
current = await MCPRequestHandler.reload_admitted_user(auth.user_id, requires_fresh_policy=True)
refreshed.object_permission = current.object_permission
refreshed.object_permission_id = current.object_permission_id

56
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -56,6 +56,7 @@ from litellm.proxy._experimental.mcp_server.faults import (
render_token_fault,
)
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
JWT_BEARER_GRANT_TYPE,
VendorCredentialState,
aggregate_authorize,
aggregate_token,
@ -71,6 +72,11 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
revoke_refresh_token,
supported_grant_types,
)
from litellm.proxy._experimental.mcp_server.id_jag_grant import (
basic_client_id,
jwt_bearer_available,
redeem_id_jag_assertion,
)
from litellm.proxy._experimental.mcp_server.idp_token_exchange import (
exchange_idp_subject_token,
token_exchange_available,
@ -2396,7 +2402,7 @@ async def token_endpoint(
grant_type: str = Form(...),
code: str = Form(None),
redirect_uri: str = Form(None),
client_id: str = Form(...),
client_id: str | None = Form(None),
client_secret: str | None = Form(None),
code_verifier: str = Form(None),
refresh_token: str | None = Form(None),
@ -2405,6 +2411,7 @@ async def token_endpoint(
subject_token: str | None = Form(None),
subject_token_type: str | None = Form(None),
requested_token_type: str | None = Form(None),
assertion: str | None = Form(None),
mcp_server_name: str | None = None,
):
"""
@ -2416,7 +2423,14 @@ async def token_endpoint(
3. Return the token
4. Return a virtual key in this response
"""
if mcp_server_name is None and is_gateway_dcr_client_id(client_id):
presented_client_id: Final = client_id or basic_client_id(request.headers.get("authorization"))
if presented_client_id is None:
raise HTTPException(
status_code=400, detail={"error": "invalid_request", "error_description": "client_id is required"}
)
if mcp_server_name is None and (
is_gateway_dcr_client_id(presented_client_id) or grant_type == JWT_BEARER_GRANT_TYPE
):
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # circular import at module load
master_key,
user_api_key_cache,
@ -2427,7 +2441,7 @@ async def token_endpoint(
grant_type=grant_type,
code=code,
redirect_uri=redirect_uri,
client_id=client_id,
client_id=presented_client_id,
code_verifier=code_verifier,
refresh_token=refresh_token,
master_key=master_key,
@ -2439,9 +2453,12 @@ async def token_endpoint(
subject_token_type=subject_token_type,
requested_token_type=requested_token_type,
exchange_subject_token=exchange_idp_subject_token,
assertion=assertion,
client_secret=client_secret,
redeem_assertion=redeem_id_jag_assertion,
)
lookup_name: Final = mcp_server_name or client_id
lookup_name: Final = mcp_server_name or presented_client_id
client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
mcp_server = _resolve_mcp_server_by_name_or_id(lookup_name, client_ip)
if mcp_server is None and mcp_server_name is None:
@ -2454,7 +2471,7 @@ async def token_endpoint(
grant_type=grant_type,
code=code,
redirect_uri=redirect_uri,
client_id=client_id,
client_id=presented_client_id,
client_secret=client_secret,
code_verifier=code_verifier,
refresh_token=refresh_token,
@ -2568,7 +2585,8 @@ async def native_client_auth_discovery(request: Request) -> JSONResponse:
"""The versioned contract a native client (``lite login --pkce``, or a CLI in any other
language) reads to sign a user in through the browser and obtain a proxy credential."""
return JSONResponse(
native_client_auth_contract(request, token_exchange_available()), headers=TOKEN_NO_CACHE_HEADERS
native_client_auth_contract(request, token_exchange_available(), jwt_bearer_available()),
headers=TOKEN_NO_CACHE_HEADERS,
)
@ -3135,7 +3153,7 @@ def _aggregate_header_scopes(request: Request) -> tuple[str, ...]:
def _build_aggregate_authorization_server_response(
request: Request, token_exchange_available: bool
request: Request, token_exchange_available: bool, jwt_bearer_available: bool = False
) -> dict[str, object]:
"""RFC 8414 metadata for the gateway as the aggregate authorization server.
@ -3155,9 +3173,13 @@ def _build_aggregate_authorization_server_response(
"registration_endpoint": f"{request_base_url}/register",
"response_types_supported": ["code"],
"scopes_supported": [],
"grant_types_supported": supported_grant_types(token_exchange_available),
"grant_types_supported": supported_grant_types(token_exchange_available, jwt_bearer_available),
"code_challenge_methods_supported": ["S256"],
"token_endpoint_auth_methods_supported": ["none", "client_secret_post"],
"token_endpoint_auth_methods_supported": (
["none", "client_secret_post", "client_secret_basic"]
if jwt_bearer_available
else ["none", "client_secret_post"]
),
}
@ -3193,7 +3215,7 @@ async def oauth_authorization_server_aggregate(request: Request):
per-server row win here instead would serve an issuer of {base} against a resource that
advertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door.
"""
return _build_aggregate_authorization_server_response(request, token_exchange_available())
return _build_aggregate_authorization_server_response(request, token_exchange_available(), jwt_bearer_available())
# Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name}
@ -3273,6 +3295,10 @@ def _build_oauth_authorization_server_response(
if explicitly_named
else request_base_url
)
# With no upstream resolved the document is the gateway's own: its issuer is the bare origin
# and its /token serves every grant the gateway does, including the ID-JAG jwt-bearer grant.
gateway_owned: Final = mcp_server_name is None and issuer_path is None
jwt_bearer: Final = gateway_owned and jwt_bearer_available()
return {
"issuer": issuer,
@ -3280,9 +3306,15 @@ def _build_oauth_authorization_server_response(
"token_endpoint": token_endpoint,
"response_types_supported": ["code"],
"scopes_supported": (mcp_server.scopes if mcp_server and mcp_server.scopes else []),
"grant_types_supported": ["authorization_code", "refresh_token"],
"grant_types_supported": (
list(supported_grant_types(token_exchange_available(), jwt_bearer))
if gateway_owned
else ["authorization_code", "refresh_token"]
),
"code_challenge_methods_supported": ["S256"],
"token_endpoint_auth_methods_supported": ["client_secret_post"],
"token_endpoint_auth_methods_supported": (
["none", "client_secret_post", "client_secret_basic"] if jwt_bearer else ["client_secret_post"]
),
# Claude expects a registration endpoint, even if we just fake it
"registration_endpoint": (
f"{request_base_url}/{mcp_server_name}/register" if mcp_server_name else f"{request_base_url}/register"

162
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
View file

@ -198,20 +198,26 @@ class MintProxyCredential(Protocol):
and mints the same credential ``lite login`` mints."""
def __call__(
self, user_id: str, team_id: str | None, /
self, user_id: str, team_id: str | None, /, *, allowed_routes: Sequence[str] | None = None
) -> Awaitable[MintedProxyCredential | ProxyCredentialMintFailure]: ...
TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"
JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer"
def supported_grant_types(token_exchange_available: bool) -> tuple[str, ...]:
def supported_grant_types(token_exchange_available: bool, jwt_bearer_available: bool = False) -> tuple[str, ...]:
"""The grants ``/token`` can serve on this deployment. The RFC 8693 exchange is listed
only where the JWT auth that proves a subject token is on, backed by a database, and
licensed, so a client never selects a grant the gateway would then refuse."""
if token_exchange_available:
return ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE)
return ("authorization_code", "refresh_token")
licensed, so a client never selects a grant the gateway would then refuse. The RFC 7523
jwt-bearer grant (an ID-JAG from the identity provider) is listed only where, on top of
that, the operator has allowlisted at least one requesting app."""
return (
"authorization_code",
"refresh_token",
*((TOKEN_EXCHANGE_GRANT_TYPE,) if token_exchange_available else ()),
*((JWT_BEARER_GRANT_TYPE,) if jwt_bearer_available else ()),
)
"""RFC 8693: a native client that already holds a token from the customer's identity
@ -236,10 +242,29 @@ class SubjectIdentity(LiteLLMBaseModel):
class SubjectTokenRefusal(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
error: Literal["unsupported_grant_type", "invalid_request", "temporarily_unavailable"]
error: Literal[
"unsupported_grant_type",
"invalid_request",
"invalid_grant",
"invalid_client",
"invalid_scope",
"temporarily_unavailable",
]
description: str = Field(min_length=1)
class AssertionIdentity(SubjectIdentity):
"""The litellm user an RFC 7523 assertion (an ID-JAG) stands for, plus what the grant
tail needs from the assertion itself: the single-use key for its ``jti``, how long that
marker must outlive it, the ``scope`` the credential is issued for, and the route groups that
scope narrows the credential to (``None`` leaves it the full ``lite login`` credential)."""
replay_key: str = Field(min_length=1)
replay_ttl_seconds: int = Field(gt=0)
scope: str | None = None
allowed_routes: tuple[str, ...] | None = None
class ExchangeSubjectToken(Protocol):
"""Injected RFC 8693 subject-token verifier ``(subject_token, request)``: proves the
IdP token the way the proxy's own JWT auth does and names the litellm user and team it
@ -272,7 +297,9 @@ class LookupConsentTeams(Protocol):
def __call__(self, user_id: str, /) -> Awaitable[tuple[ConsentTeam, ...] | ReloadUserFailure]: ...
async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCredentialMintFailure:
async def _refuse_proxy_credential(
user_id: str, team_id: str | None, *, allowed_routes: Sequence[str] | None = None
) -> ProxyCredentialMintFailure:
return "unresolvable"
@ -282,6 +309,25 @@ async def _refuse_subject_token(subject_token: str, request: Request) -> Subject
)
class RedeemAssertion(Protocol):
"""Injected RFC 7523 assertion verifier ``(assertion, client_id, client_secret, request)``:
proves an ID-JAG the identity provider minted for the requesting app ``client_id``,
authenticates that app, and names the litellm user the assertion stands for, or says why
this gateway will not take it."""
def __call__(
self, assertion: str, client_id: str, client_secret: str | None, request: Request, /
) -> Awaitable[AssertionIdentity | SubjectTokenRefusal]: ...
async def _refuse_assertion(
assertion: str, client_id: str, client_secret: str | None, request: Request
) -> SubjectTokenRefusal:
return SubjectTokenRefusal(
error="unsupported_grant_type", description="this gateway is not configured to redeem identity assertions"
)
async def _unavailable_vendor_credential(
user_id: str, server_id: str, requested_scopes: tuple[str, ...] = ()
) -> VendorCredentialState:
@ -797,7 +843,9 @@ class NativeClientAuthContract(TypedDict):
revocation_endpoint_auth_methods_supported: ReadOnly[tuple[str, ...]]
def native_client_auth_contract(request: Request, token_exchange_available: bool) -> NativeClientAuthContract:
def native_client_auth_contract(
request: Request, token_exchange_available: bool, jwt_bearer_available: bool = False
) -> NativeClientAuthContract:
"""The versioned discovery document at ``/.well-known/litellm-cli-auth``: everything a
native client (in any language) needs to run the sign-in without reading LiteLLM
source. ``resource`` is the exact value to send as the RFC 8707 ``resource`` parameter
@ -812,7 +860,7 @@ def native_client_auth_contract(request: Request, token_exchange_available: bool
"revocation_endpoint": f"{base_url}/revoke",
"resource": base_url,
"response_types_supported": ("code",),
"grant_types_supported": supported_grant_types(token_exchange_available),
"grant_types_supported": supported_grant_types(token_exchange_available, jwt_bearer_available),
"code_challenge_methods_supported": ("S256",),
"token_endpoint_auth_methods_supported": ("none",),
"revocation_endpoint_auth_methods_supported": ("none",),
@ -1403,10 +1451,11 @@ class _ProxyCredentialTokenResponse(TypedDict):
access_token: ReadOnly[str]
token_type: ReadOnly[Literal["Bearer"]]
expires_in: ReadOnly[int]
refresh_token: ReadOnly[str]
refresh_token: NotRequired[ReadOnly[str]]
user_id: ReadOnly[str]
team_id: ReadOnly[str | None]
issued_token_type: NotRequired[ReadOnly[_IssuedTokenType]]
scope: NotRequired[ReadOnly[str]]
def proxy_credential_response(
@ -1415,6 +1464,7 @@ def proxy_credential_response(
signing: SessionSigning,
issued_token_type: _IssuedTokenType | None = None,
family: str | None = None,
scope: str | None = None,
) -> Response:
"""The proxy-API token response: the access token is the very credential ``lite
login`` stores (accepted on every proxy route with user and team attribution), and
@ -1426,7 +1476,7 @@ def proxy_credential_response(
refresh: Final = mint_session_refresh_token(bound_principal, signing.keys, signing.now, family=family)
if not isinstance(refresh, MintedSessionToken):
return _oauth_error(500, "server_error", "failed to mint the session credential")
credential: Final[_ProxyCredentialTokenResponse] = {
renewable: Final[_ProxyCredentialTokenResponse] = {
"access_token": minted.key,
"token_type": "Bearer",
"expires_in": minted.expires_in,
@ -1434,12 +1484,27 @@ def proxy_credential_response(
"user_id": minted.user_id,
"team_id": minted.team_id,
}
credential: Final[_ProxyCredentialTokenResponse] = renewable if scope is None else {**renewable, "scope": scope}
if issued_token_type is None:
return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS)
exchanged: Final[_ProxyCredentialTokenResponse] = {**credential, "issued_token_type": issued_token_type}
return JSONResponse(status_code=200, content=exchanged, headers=TOKEN_NO_CACHE_HEADERS)
def assertion_credential_response(minted: MintedProxyCredential, scope: str | None) -> Response:
"""The proxy-API token response for an assertion grant: the same access token, with the
scope the assertion carried, and no refresh token."""
unscoped: Final[_ProxyCredentialTokenResponse] = {
"access_token": minted.key,
"token_type": "Bearer",
"expires_in": minted.expires_in,
"user_id": minted.user_id,
"team_id": minted.team_id,
}
credential: Final[_ProxyCredentialTokenResponse] = unscoped if scope is None else {**unscoped, "scope": scope}
return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS)
def _reload_failure_response(failure: ReloadUserFailure) -> Response:
"""Map the live-user revalidation failure onto its OAuth error, exhaustively, so a new
``ReloadUserFailure`` member is a type error here rather than silently 400ing."""
@ -1460,7 +1525,9 @@ def _subject_token_refusal_response(refusal: SubjectTokenRefusal) -> Response:
match refusal.error:
case "temporarily_unavailable":
return _oauth_error(503, refusal.error, refusal.description)
case "unsupported_grant_type" | "invalid_request":
case "invalid_client":
return _oauth_error(401, refusal.error, refusal.description)
case "unsupported_grant_type" | "invalid_request" | "invalid_grant" | "invalid_scope":
return _oauth_error(400, refusal.error, refusal.description)
case _:
assert_never(refusal.error)
@ -1515,6 +1582,9 @@ async def aggregate_token(
subject_token_type: str | None = None,
requested_token_type: str | None = None,
exchange_subject_token: ExchangeSubjectToken = _refuse_subject_token,
assertion: str | None = None,
client_secret: str | None = None,
redeem_assertion: RedeemAssertion = _refuse_assertion,
) -> Response:
"""The aggregate token verb: authorization_code and refresh_token grants for the
identity-only session pair, or for the proxy-API credential when the grant was issued
@ -1561,10 +1631,19 @@ async def aggregate_token(
exchange_subject_token=exchange_subject_token,
issue=issue,
)
if grant_type == JWT_BEARER_GRANT_TYPE:
return await _jwt_bearer_grant(
request=request,
assertion=assertion,
client_id=client_id,
client_secret=client_secret,
redeem_assertion=redeem_assertion,
issue=issue,
)
return _oauth_error(
400,
"unsupported_grant_type",
f"grant_type must be authorization_code, refresh_token, or {TOKEN_EXCHANGE_GRANT_TYPE}",
f"grant_type must be authorization_code, refresh_token, {TOKEN_EXCHANGE_GRANT_TYPE}, or {JWT_BEARER_GRANT_TYPE}",
)
@ -1690,6 +1769,48 @@ class _GrantIssuer:
return _mint_failure_response(minted)
return proxy_credential_response(minted, principal, self._signing, issued_token_type=ACCESS_TOKEN_TOKEN_TYPE)
async def redeem_assertion(
self, assertion: str, client_id: str, client_secret: str | None, redeem_assertion: RedeemAssertion
) -> Response:
"""The RFC 7523 tail: prove the assertion and the requesting app, claim the assertion's
``jti`` before anything is minted (RFC 7523 section 3 makes a replayed ``jti`` a
rejection, and the identity provider mints a fresh assertion cheaply, so a mint that
fails after the claim costs the client one more round trip, never a replay window),
then mint the proxy credential the user is revalidated for. No refresh token comes
with it (RFC 7521 section 4.1): the app renews by asking its identity provider for a
fresh assertion, which keeps the user's standing there in the loop."""
target_refusal: Final = self._assertion_target_refusal()
if target_refusal is not None:
return target_refusal
identity: Final = await redeem_assertion(assertion, client_id, client_secret, self._request)
if isinstance(identity, SubjectTokenRefusal):
return _subject_token_refusal_response(identity)
refused: Final = _claim_refusal(
await self._guard.claim(identity.replay_key, identity.replay_ttl_seconds),
_oauth_error(400, "invalid_grant", "the assertion has already been redeemed"),
)
if refused is not None:
return refused
minted: Final = await self._mint_proxy_credential(
identity.user_id, identity.team_id, allowed_routes=identity.allowed_routes
)
if not isinstance(minted, MintedProxyCredential):
return _mint_failure_response(minted)
return assertion_credential_response(minted, identity.scope)
def _assertion_target_refusal(self) -> Response | None:
"""An ID-JAG names the proxy origin or its MCP resource (``{base}/mcp``) as ``aud``, and an RFC
8707 client sends that same value as ``resource``, so both are this grant's targets."""
if self._resource is None or is_proxy_api_resource(self._request, self._resource):
return None
if self._resource.rstrip("/") == f"{get_request_base_url(self._request)}/mcp":
return None
return _oauth_error(
400,
"invalid_target",
"resource does not match the proxy API or MCP resource this gateway issues credentials for",
)
def _proxy_api_target_refusal(self) -> Response | None:
if self._resource is None or is_proxy_api_resource(self._request, self._resource):
return None
@ -1807,6 +1928,19 @@ async def refresh_proxy_credential(
)
async def _jwt_bearer_grant(
request: Request,
assertion: str | None,
client_id: str,
client_secret: str | None,
redeem_assertion: RedeemAssertion,
issue: _GrantIssuer,
) -> Response:
if not assertion:
return _oauth_error(400, "invalid_request", "assertion is required")
return await issue.redeem_assertion(assertion, client_id, client_secret, redeem_assertion)
async def _token_exchange_grant(
subject_token: str | None,
subject_token_type: str | None,

540
litellm/proxy/_experimental/mcp_server/id_jag_grant.py Normal file
View file

@ -0,0 +1,540 @@
"""The RFC 7523 jwt-bearer grant on ``POST /token``: a requesting app that obtained an
ID-JAG (Identity Assertion Authorization Grant, the cross-app access profile) from the
customer's identity provider redeems it for the same proxy-API credential ``lite login``
stores. The assertion is proven with the JWT auth issuer and JWKS plumbing the proxy
already trusts, its claims are checked the way the ID-JAG draft prescribes, the requesting
app must be one the operator allowlisted, and the ``jti`` is single-use until the
assertion expires."""
from __future__ import annotations
import base64
import binascii
import hashlib
import hmac
import math
from collections.abc import Awaitable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import partial
from types import MappingProxyType
from typing import Final, Protocol
from urllib.parse import unquote_plus
import jwt
from fastapi import Request
from litellm._logging import verbose_proxy_logger
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
AssertionIdentity,
ConsentTeam,
SubjectIdentity,
SubjectTokenRefusal,
)
from litellm.proxy._experimental.mcp_server.idp_token_exchange import (
gateway_outage_refusal,
log_line,
read_token_exchange_prerequisites,
)
from litellm.proxy._experimental.mcp_server.oauth_utils import proxy_base_url_from_env
from litellm.proxy._types import IdJagRequestingApp
from litellm.proxy.auth.handle_jwt import ID_JAG_TYP, AssertionRejection, JWTHandler
ID_JAG_REQUIRED_CLAIMS: Final = ("iss", "sub", "aud", "exp", "iat", "jti", "client_id")
REJECTED_ASSERTION: Final = "assertion was rejected"
TEAM_HEADER: Final = "x-litellm-team-id"
_INACTIVE_SUBJECT: Final = frozenset(("no_active_key", "unresolvable"))
_REPLAY_KEY_PREFIX: Final = "mcp_gateway_dcr:id_jag:"
SCOPE_MCP: Final = "litellm:mcp"
SCOPE_LLM: Final = "litellm:llm"
SCOPE_ROUTE_GROUPS: Final = MappingProxyType(
{SCOPE_MCP: ("mcp_inference_routes",), SCOPE_LLM: ("llm_api_routes_without_mcp",)}
)
@dataclass(frozen=True, slots=True)
class GrantedScope:
"""What an assertion's ``scope`` claim buys: the scope echoed in the token response, exactly the
tokens enforced, and the route groups the credential is narrowed to. No claim means the full
credential, so both are ``None``."""
scope: str | None
allowed_routes: tuple[str, ...] | None
def granted_scope(claim: object, *, mcp_resource: bool) -> GrantedScope | SubjectTokenRefusal:
"""The routes the assertion buys. With ``aud`` the bare origin the ``scope`` claim decides and
no claim is the full credential. With ``aud`` the MCP resource (``{base}/mcp``) the identity
provider named that resource, so the credential is ceilinged to the MCP routes whatever the
claim says: ``litellm:llm`` cannot be granted and is refused."""
if claim is not None and not isinstance(claim, str):
return SubjectTokenRefusal(error="invalid_scope", description=_SCOPE_DESCRIPTION)
requested: Final = frozenset(claim.split()) if claim else frozenset()
if not requested <= SCOPE_ROUTE_GROUPS.keys():
return SubjectTokenRefusal(error="invalid_scope", description=_SCOPE_DESCRIPTION)
if mcp_resource:
if SCOPE_LLM in requested:
return SubjectTokenRefusal(error="invalid_scope", description=_MCP_RESOURCE_DESCRIPTION)
return GrantedScope(scope=SCOPE_MCP, allowed_routes=SCOPE_ROUTE_GROUPS[SCOPE_MCP])
if not requested:
return GrantedScope(scope=None, allowed_routes=None)
granted: Final = tuple(scope for scope in SCOPE_ROUTE_GROUPS if scope in requested)
return GrantedScope(scope=" ".join(granted), allowed_routes=tuple(_route_groups(granted)))
def names_mcp_resource(aud: object, base_url: str | None) -> bool:
"""Whether the verified ``aud`` names only the MCP resource, ``{base}/mcp``, and not the
authorization server itself (the bare origin)."""
audiences: Final = frozenset((aud,) if isinstance(aud, str) else aud if isinstance(aud, list) else ())
return base_url is not None and f"{base_url}/mcp" in audiences and base_url not in audiences
def _route_groups(scopes: Sequence[str]) -> Iterator[str]:
for scope in scopes:
yield from SCOPE_ROUTE_GROUPS[scope]
_SCOPE_DESCRIPTION: Final = (
f"scope must be {SCOPE_MCP} (MCP gateway routes), {SCOPE_LLM} (LLM API routes), or both, space separated"
)
_MCP_RESOURCE_DESCRIPTION: Final = (
f"the /mcp audience only grants MCP routes ({SCOPE_MCP}); use the bare origin as audience for LLM routes"
)
@dataclass(frozen=True, slots=True)
class IdJagPrerequisites:
"""The deployment-level gates for redeeming an ID-JAG: the JWT auth whose issuer plumbing
proves the assertion must be on, backed by a database, and licensed, not mapping IdP tokens to
virtual keys (whose blocks and budgets a minted credential would sidestep), with a shared Redis
so an assertion is redeemed once across every worker and replica (the in-memory marker is per
process, so a second worker would mint again), ``PROXY_BASE_URL``
must fix the gateway URL the assertion's ``aud`` has to name (a request's ``Host`` is the
caller's to choose), and the operator must have allowlisted the requesting apps. Discovery advertises the grant only when every gate
holds, and a redemption is refused naming the first one that does not."""
jwt_auth_enabled: bool
has_database: bool
licensed: bool
maps_jwts_to_virtual_keys: bool
shared_replay_store: bool
base_url: str | None
requesting_apps: tuple[IdJagRequestingApp, ...]
@property
def available(self) -> bool:
return (
self.jwt_auth_enabled
and self.has_database
and self.licensed
and not self.maps_jwts_to_virtual_keys
and self.shared_replay_store
and self.base_url is not None
and bool(self.requesting_apps)
)
def refusal(self) -> SubjectTokenRefusal | None:
if not self.jwt_auth_enabled:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="JWT auth is not enabled on this gateway, so it cannot redeem identity assertions",
)
if not self.has_database:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="this gateway has no database, so it cannot redeem identity assertions",
)
if not self.licensed:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="JWT auth is an enterprise only feature; no license is set",
)
if self.maps_jwts_to_virtual_keys:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="this gateway maps IdP tokens to virtual keys, which the jwt-bearer grant does not serve",
)
if not self.shared_replay_store:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="configure Redis (cache_params or coordination_redis) so this gateway can redeem an identity assertion exactly once across its workers",
)
if self.base_url is None:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="set PROXY_BASE_URL to this gateway's public URL, the audience an identity assertion must name",
)
if not self.requesting_apps:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="this gateway has no requesting apps allowlisted for identity assertions",
)
return None
def read_id_jag_prerequisites() -> IdJagPrerequisites:
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call
jwt_handler,
redis_usage_cache,
)
shared: Final = read_token_exchange_prerequisites()
return IdJagPrerequisites(
jwt_auth_enabled=shared.jwt_auth_enabled,
has_database=shared.has_database,
licensed=shared.licensed,
maps_jwts_to_virtual_keys=shared.maps_jwts_to_virtual_keys,
shared_replay_store=redis_usage_cache is not None,
base_url=proxy_base_url_from_env(),
requesting_apps=_requesting_apps(jwt_handler),
)
def _requesting_apps(jwt_handler: JWTHandler) -> tuple[IdJagRequestingApp, ...]:
if not hasattr(jwt_handler, "litellm_jwtauth"):
return ()
return tuple(jwt_handler.litellm_jwtauth.id_jag_requesting_apps or ())
def jwt_bearer_available() -> bool:
return read_id_jag_prerequisites().available
def gateway_audiences(base_url: str | None) -> tuple[str, ...]:
"""The issuer values this gateway answers to, which an ID-JAG's ``aud`` must name: its MCP
resource and its authorization server, both under the fixed ``PROXY_BASE_URL``."""
return () if base_url is None else (f"{base_url}/mcp", base_url)
class DecodeAssertion(Protocol):
"""Injected signature and claim verification ``(assertion, audience, required_claims)``:
``JWTHandler.decode_trusted_assertion`` in production, which returns the verified claims,
a rejection, or raises when the gateway could not reach the issuer's keys."""
def __call__(
self, assertion: str, audience: Sequence[str], required_claims: Sequence[str], /
) -> Awaitable[Mapping[str, object] | AssertionRejection]: ...
class ClaimReader(Protocol):
"""The JWT auth user mapping as ``JWTHandler`` exposes it: the litellm user id and email the
verified claims name through ``user_id_jwt_field`` and ``user_email_jwt_field``."""
def get_user_id(self, token: dict[str, object], default_value: str | None) -> str | None: ...
def get_user_email(self, token: dict[str, object], default_value: str | None) -> str | None: ...
def get_all_jwt_team_ids(self, token: dict[str, object]) -> list[str]: ...
def is_allowed_domain(self, user_email: str) -> bool: ...
class ResolveSubject(Protocol):
"""Injected subject lookup ``(claims, team_header)``: the litellm user the verified claims
name through the JWT auth user mapping (``user_id_jwt_field`` or ``user_email_jwt_field``,
else ``sub``) and the team the credential is minted for (see ``select_team``), ``None`` when
the gateway knows no such user, or a refusal when the team cannot be settled. Raises on a
database outage."""
def __call__(
self, claims: Mapping[str, object], team_header: str | None, /
) -> Awaitable[SubjectIdentity | SubjectTokenRefusal | None]: ...
class LookupUser(Protocol):
"""Injected ``(user_id, user_email)`` lookup: the stored id of the existing litellm user the
claims name, ``None`` when the gateway knows no such user. Raises on a database outage."""
def __call__(self, user_id: str, user_email: str | None, /) -> Awaitable[str | None]: ...
class LookupTeams(Protocol):
"""Injected lookup of the teams a user is on, or the failure text when they cannot be read."""
def __call__(self, user_id: str, /) -> Awaitable[Sequence[ConsentTeam] | str]: ...
class GetUser(Protocol):
def __call__(
self, *, user_id: str, user_email: str | None, sso_user_id: str
) -> Awaitable[LiteLLM_UserTable | None]: ...
def proxy_user_lookup(get_user: GetUser) -> LookupUser:
"""``LookupUser`` over the proxy's user table. The assertion's subject is matched as the
local ``user_id`` and as the ``sso_user_id`` an SSO login linked, the way the JWT and SSO
login paths match it, so a linked account needs no email claim."""
async def lookup_user(user_id: str, user_email: str | None) -> str | None:
user: Final = await get_user(user_id=user_id, user_email=user_email, sso_user_id=user_id)
return None if user is None else user.user_id
return lookup_user
def subject_resolver(reader: ClaimReader, lookup_user: LookupUser, lookup_teams: LookupTeams) -> ResolveSubject:
"""``ResolveSubject`` over the JWT auth claim mapping and the proxy's user and team lookups."""
async def resolve_subject(
claims: Mapping[str, object], team_header: str | None
) -> SubjectIdentity | SubjectTokenRefusal | None:
token: Final = dict(claims)
subject: Final = claims.get("sub")
claimed_user_id: Final = reader.get_user_id(token=token, default_value=None) or (
subject if isinstance(subject, str) else None
)
if claimed_user_id is None:
return None
user_email: Final = reader.get_user_email(token=token, default_value=None)
if not reader.is_allowed_domain(user_email or ""):
return SubjectTokenRefusal(
error="invalid_grant",
description="the assertion's subject email is missing or outside the allowed domain",
)
user_id: Final = await lookup_user(claimed_user_id, user_email)
if user_id is None:
return None
teams: Final = await lookup_teams(user_id)
if teams in _INACTIVE_SUBJECT:
return SubjectTokenRefusal(
error="invalid_grant", description="the assertion's subject is not an active user"
)
if isinstance(teams, str):
return SubjectTokenRefusal(
error="temporarily_unavailable", description="the gateway could not read this user's teams; retry"
)
team: Final = select_team(team_header, tuple(reader.get_all_jwt_team_ids(token=token)), teams)
if isinstance(team, SubjectTokenRefusal):
return team
return SubjectIdentity(user_id=user_id, team_id=team)
return resolve_subject
def select_team(
team_header: str | None, claim_teams: Sequence[str], teams: Sequence[ConsentTeam]
) -> str | None | SubjectTokenRefusal:
"""The team an ID-JAG redemption mints for, settled before the assertion is consumed since
the minter refuses a teamless credential for a user who has a team. The ``x-litellm-team-id``
header (a team id or alias) wins and must name one of the user's teams, then the assertion's
mapped team claim, then the user's only team. A user on several teams must name one."""
if team_header is not None:
named: Final = next((team for team in teams if team_header in (team.team_id, team.team_alias)), None)
if named is None:
return SubjectTokenRefusal(
error="invalid_request", description=f"{TEAM_HEADER} does not name one of this user's teams"
)
return named.team_id
asserted: Final = tuple(team.team_id for team in teams if team.team_id in claim_teams)
if len(asserted) == 1:
return asserted[0]
if len(asserted) > 1:
return SubjectTokenRefusal(
error="invalid_request",
description=f"the assertion names several of this user's teams; send the {TEAM_HEADER} header naming the team for this credential",
)
if len(teams) == 1:
return teams[0].team_id
if teams:
return SubjectTokenRefusal(
error="invalid_request",
description=f"this user belongs to several teams; send the {TEAM_HEADER} header naming the team for this credential",
)
return None
@dataclass(frozen=True, slots=True)
class AssertionRedemption:
"""One ``jwt-bearer`` request as the grant sees it: the assertion, how the requesting app
identified and authenticated itself, the team the credential is asked for, the issuer
values this gateway answers to (what the assertion's ``aud`` must name), the clock, and the
leeway the verifier grants ``exp``, so the replay marker outlives every moment the assertion
would still verify."""
assertion: str
client_id: str
client_secret: str | None
authorization_header: str | None
team_id: str | None
audiences: tuple[str, ...]
now: float
leeway_seconds: float
async def redeem_id_jag_assertion(
assertion: str, client_id: str, client_secret: str | None, request: Request
) -> AssertionIdentity | SubjectTokenRefusal:
from litellm.proxy._experimental.mcp_server.proxy_api_credentials import ( # noqa: PLC0415 # circular import at module load
lookup_consent_teams_from_database,
)
from litellm.proxy.auth.auth_checks import get_user_object # noqa: PLC0415 # circular import at module load
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call
jwt_handler,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
resolve_subject: Final = subject_resolver(
_claim_reader(jwt_handler),
proxy_user_lookup(
partial(
get_user_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
proxy_logging_obj=proxy_logging_obj,
check_db_only=True,
)
),
lookup_consent_teams_from_database,
)
prerequisites: Final = read_id_jag_prerequisites()
return await identity_from_id_jag(
AssertionRedemption(
assertion=assertion,
client_id=client_id,
client_secret=client_secret,
authorization_header=request.headers.get("authorization"),
team_id=request.headers.get(TEAM_HEADER),
audiences=gateway_audiences(prerequisites.base_url),
now=datetime.now(timezone.utc).timestamp(),
leeway_seconds=jwt_handler.leeway,
),
prerequisites=prerequisites,
decode=jwt_handler.decode_trusted_assertion,
resolve_subject=resolve_subject,
)
def _claim_reader(handler: ClaimReader) -> ClaimReader:
return handler
async def identity_from_id_jag(
redemption: AssertionRedemption,
prerequisites: IdJagPrerequisites,
decode: DecodeAssertion,
resolve_subject: ResolveSubject,
) -> AssertionIdentity | SubjectTokenRefusal:
"""Prove an ID-JAG for the grant tail. The requesting app is checked first, the way a token
endpoint authenticates its client before it looks at the grant, so an app the operator never
allowlisted learns nothing about the assertion. Everything wrong with the assertion itself is
``invalid_grant`` (RFC 7523 section 3.1), a check the gateway could not complete is
``temporarily_unavailable``, and the reason stays in the proxy log because this endpoint is
public and the verifier's wording can quote the assertion."""
unmet: Final = prerequisites.refusal()
if unmet is not None:
return unmet
app: Final = next((app for app in prerequisites.requesting_apps if app.client_id == redemption.client_id), None)
if app is None:
return SubjectTokenRefusal(error="invalid_client", description="client_id is not an allowlisted requesting app")
if not _client_authenticated(app, redemption.client_secret, redemption.authorization_header):
return SubjectTokenRefusal(error="invalid_client", description="client authentication failed")
try:
header: Final = jwt.get_unverified_header(redemption.assertion)
except jwt.PyJWTError:
return SubjectTokenRefusal(error="invalid_grant", description="assertion is not a JWT")
if header.get("typ") != ID_JAG_TYP:
return SubjectTokenRefusal(error="invalid_grant", description=f"assertion typ must be {ID_JAG_TYP}")
try:
verified: Final = await decode(redemption.assertion, redemption.audiences, ID_JAG_REQUIRED_CLAIMS)
except Exception as denied: # noqa: BLE001 # the verifier raises its 503 for a JWKS outage, and anything else is a bad assertion
return gateway_outage_refusal(denied, denied, what="an assertion", parameter="assertion") or _rejected(
str(denied)
)
if isinstance(verified, AssertionRejection):
return _rejected(verified.reason)
marker: Final = _replay_marker(verified, redemption)
if isinstance(marker, SubjectTokenRefusal):
return marker
replay_key, replay_ttl_seconds = marker
scope: Final = granted_scope(
verified.get("scope"), mcp_resource=names_mcp_resource(verified.get("aud"), prerequisites.base_url)
)
if isinstance(scope, SubjectTokenRefusal):
return scope
try:
subject: Final = await resolve_subject(verified, redemption.team_id)
except Exception as failed: # noqa: BLE001 # the user lookup raises for a missing user and for a database outage alike
return gateway_outage_refusal(
failed, failed, what="an assertion's subject", parameter="assertion"
) or _rejected(f"subject lookup failed: {failed}")
if subject is None:
return _rejected("the assertion names no user the gateway knows")
if isinstance(subject, SubjectTokenRefusal):
return subject
return AssertionIdentity(
user_id=subject.user_id,
team_id=subject.team_id,
replay_key=replay_key,
replay_ttl_seconds=replay_ttl_seconds,
scope=scope.scope,
allowed_routes=scope.allowed_routes,
)
def _replay_marker(
verified: Mapping[str, object], redemption: AssertionRedemption
) -> tuple[str, int] | SubjectTokenRefusal:
"""The single-use key and its lifetime for a verified assertion, once the assertion is bound to
the presenting app, carries no proof-of-possession the gateway cannot check, and is still alive."""
if verified.get("client_id") != redemption.client_id:
return _rejected("the assertion's client_id is not the requesting app that presented it")
if "cnf" in verified:
return _rejected("the assertion is proof-bound and this gateway does not check proof of possession")
issuer: Final = verified.get("iss")
jti: Final = verified.get("jti")
exp: Final = verified.get("exp")
if not isinstance(issuer, str) or not isinstance(jti, str) or not jti or not isinstance(exp, (int, float)):
return _rejected("the assertion's iss, jti, or exp claim has the wrong shape")
replay_ttl_seconds: Final = math.ceil(exp + redemption.leeway_seconds - redemption.now)
if replay_ttl_seconds <= 0:
return _rejected("the assertion has expired")
return _REPLAY_KEY_PREFIX + hashlib.sha256(f"{issuer}\n{jti}".encode()).hexdigest(), replay_ttl_seconds
def _rejected(reason: str) -> SubjectTokenRefusal:
verbose_proxy_logger.warning("jwt-bearer grant refused an assertion: %s", log_line(reason))
return SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION)
def _client_authenticated(app: IdJagRequestingApp, form_secret: str | None, authorization_header: str | None) -> bool:
"""A public app (no secret configured) is identified by ``client_id`` alone. A confidential
app must present its secret, by HTTP Basic (RFC 6749 section 2.3.1, the credentials
form-urlencoded inside the base64) or the ``client_secret`` form field, compared in constant
time."""
if app.client_secret is None:
return True
basic: Final = _basic_credentials(authorization_header)
if basic is not None:
return basic[0] == app.client_id and _same_secret(basic[1], app.client_secret)
return form_secret is not None and _same_secret(form_secret, app.client_secret)
def _same_secret(presented: str, expected: str) -> bool:
return hmac.compare_digest(presented.encode(), expected.encode())
def basic_client_id(authorization_header: str | None) -> str | None:
"""The ``client_id`` an HTTP Basic ``Authorization`` header carries, for a confidential app
that leaves the form field out as RFC 6749 section 2.3.1 allows."""
credentials: Final = _basic_credentials(authorization_header)
return None if credentials is None else credentials[0]
def _basic_credentials(authorization_header: str | None) -> tuple[str, str] | None:
if authorization_header is None:
return None
scheme, _, encoded = authorization_header.strip().partition(" ")
if scheme.lower() != "basic" or not encoded:
return None
try:
decoded: Final = base64.b64decode(encoded.strip(), validate=True).decode()
except (binascii.Error, UnicodeDecodeError):
return None
client_id, separator, secret = decoded.partition(":")
if not separator:
return None
return unquote_plus(client_id), unquote_plus(secret)

41
litellm/proxy/_experimental/mcp_server/idp_token_exchange.py
View file

@ -20,13 +20,15 @@ from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
EXCHANGE_ROUTE: Final = "/token"
REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth"
SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = (
"the gateway could not verify subject_token because its identity provider or database is unavailable; retry"
_CHECK_UNAVAILABLE: Final = (
"the gateway could not verify {parameter} because its identity provider or database is unavailable; retry"
)
SUBJECT_TOKEN_CHECK_FAULTED: Final = (
"the gateway could not verify subject_token because its database reported a fault that is not a transient "
_CHECK_FAULTED: Final = (
"the gateway could not verify {parameter} because its database reported a fault that is not a transient "
"outage; retrying will not help until the gateway deployment is repaired"
)
SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = _CHECK_UNAVAILABLE.format(parameter="subject_token")
SUBJECT_TOKEN_CHECK_FAULTED: Final = _CHECK_FAULTED.format(parameter="subject_token")
GatewayOutage = Literal["retryable", "faulted"]
@ -182,20 +184,37 @@ async def identity_from_subject_token(
def _refusal_for(denied: Exception, reason: object) -> SubjectTokenRefusal:
outage: Final = gateway_outage_refusal(denied, reason, what="a subject_token", parameter="subject_token")
if outage is not None:
return outage
verbose_proxy_logger.warning("token exchange refused a subject_token: %s", log_line(reason))
return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN)
def gateway_outage_refusal(denied: Exception, reason: object, what: str, parameter: str) -> SubjectTokenRefusal | None:
"""The ``temporarily_unavailable`` refusal when ``denied`` is the gateway failing to check
``what`` (its database or the IdP's JWKS down), naming the grant ``parameter`` it could not verify,
else ``None`` because the token itself was bad."""
outage: Final = _gateway_could_not_verify(denied)
if outage is None:
verbose_proxy_logger.warning("token exchange refused a subject_token: %s", reason)
return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN)
verbose_proxy_logger.error("token exchange could not verify a subject_token, %s: %s", outage, reason)
return SubjectTokenRefusal(error="temporarily_unavailable", description=_check_unavailable_description(outage))
return None
verbose_proxy_logger.error("token endpoint could not verify %s, %s: %s", what, outage, log_line(reason))
return SubjectTokenRefusal(
error="temporarily_unavailable", description=_check_unavailable_description(outage, parameter)
)
def _check_unavailable_description(outage: GatewayOutage) -> str:
def log_line(reason: object) -> str:
"""``reason`` as one log line, so text a caller controls cannot forge log entries."""
return str(reason).replace("\r", " ").replace("\n", " ")
def _check_unavailable_description(outage: GatewayOutage, parameter: str) -> str:
match outage:
case "retryable":
return SUBJECT_TOKEN_CHECK_UNAVAILABLE
return _CHECK_UNAVAILABLE.format(parameter=parameter)
case "faulted":
return SUBJECT_TOKEN_CHECK_FAULTED
return _CHECK_FAULTED.format(parameter=parameter)
case _:
assert_never(outage)

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

@ -455,6 +455,7 @@ class MCPServerConfig(TypedDict, total=False):
upstream_token_header: ReadOnly[str]
id_jag_resource_token_endpoint: str
id_jag_resource: str
id_jag_subject_source: ReadOnly[str]
id_jag_resource_client_id: ReadOnly[str]
id_jag_resource_client_secret: ReadOnly[str]
client_private_key: str
@ -2789,6 +2790,7 @@ class MCPServerManager:
# ID-JAG fields
id_jag_resource_token_endpoint=server_config.get("id_jag_resource_token_endpoint", None),
id_jag_resource=server_config.get("id_jag_resource", None),
id_jag_subject_source=server_config.get("id_jag_subject_source", None),
id_jag_resource_client_id=server_config.get("id_jag_resource_client_id", None),
id_jag_resource_client_secret=server_config.get("id_jag_resource_client_secret", None),
client_private_key=server_config.get("client_private_key", None),
@ -3380,6 +3382,7 @@ class MCPServerManager:
credentials_dict.get("id_jag_resource_token_endpoint") if credentials_dict else None
),
id_jag_resource=(credentials_dict.get("id_jag_resource") if credentials_dict else None),
id_jag_subject_source=(credentials_dict.get("id_jag_subject_source") if credentials_dict else None),
id_jag_resource_client_id=(credentials_dict.get("id_jag_resource_client_id") if credentials_dict else None),
id_jag_resource_client_secret=self._decrypt_credential_field(
credentials_dict.get("id_jag_resource_client_secret") if credentials_dict else None,

10
litellm/proxy/_experimental/mcp_server/oauth_utils.py
View file

@ -19,6 +19,9 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import (
build_token_endpoint_client_auth,
normalize_token_endpoint_auth_method,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
GOOGLE_REFRESH_TOKEN_PARAMS,
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.middleware.per_request_root_path_middleware import get_request_root_path
@ -59,7 +62,7 @@ _DEFAULT_NATIVE_REDIRECT_URIS: Final[list[str]] = [
]
_GOOGLE_AUTHORIZATION_HOSTS: Final = frozenset({"accounts.google.com"})
_GOOGLE_OFFLINE_ACCESS_PARAMS: Final = MappingProxyType({"access_type": "offline", "prompt": "consent"})
_GOOGLE_OFFLINE_ACCESS_PARAMS: Final = GOOGLE_REFRESH_TOKEN_PARAMS
_warned_invalid_proxy_base_url: str | None = None
@ -128,6 +131,11 @@ def build_upstream_authorize_url(authorization_url: str, params: Mapping[str, st
return urlunparse(parsed._replace(query=urlencode(merged)))
def proxy_base_url_from_env() -> str | None:
"""The gateway's fixed public URL from ``PROXY_BASE_URL``, ``None`` when unset or malformed."""
return _resolve_proxy_base_url_env()
def _resolve_proxy_base_url_env() -> str | None:
global _warned_invalid_proxy_base_url
configured: Final = os.environ.get("PROXY_BASE_URL", "").strip()

2
litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py
View file

@ -39,6 +39,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
HeaderCarrier,
IdJagConfig,
IdJagResourceClient,
IdJagSubjectSource,
NoneConfig,
PassthroughConfig,
PrivateKeyJwtAuth,
@ -71,6 +72,7 @@ __all__ = [
"HeaderCarrier",
"IdJagConfig",
"IdJagResourceClient",
"IdJagSubjectSource",
"NoOpAuth",
"NoneConfig",
"Ok",

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

@ -34,6 +34,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
CredError,
IdJagConfig,
IdJagResourceClient,
IdJagSubjectSource,
NoneConfig,
PassthroughConfig,
PrivateKeyJwtAuth,
@ -236,7 +237,8 @@ def _id_jag_spec(server: MCPServer, resource: str) -> ServerSpec | None:
if not org_token_endpoint or not resource_token_endpoint or not client_id or client_auth is None:
return None
resource_client: Final = _id_jag_resource_client(server, client_id, client_auth)
if resource_client is None:
subject_source: Final = _id_jag_subject_source(server)
if resource_client is None or subject_source is None:
return None
return ServerSpec(
server_id=server.server_id,
@ -249,6 +251,7 @@ def _id_jag_spec(server: MCPServer, resource: str) -> ServerSpec | None:
client_auth=client_auth,
resource_client=resource_client,
subject_token_type=_id_jag_subject_token_type(server),
subject_source=subject_source,
audience=server.audience,
resource=server.id_jag_resource,
scopes=tuple(server.scopes or ()),
@ -282,6 +285,18 @@ def _id_jag_client_auth(server: MCPServer) -> ClientAuth | None:
return None
def _id_jag_subject_source(server: MCPServer) -> IdJagSubjectSource | None:
"""``None`` for a value that names no source, so the half-configured server defers like a missing
endpoint does instead of silently presenting the wrong token."""
configured: Final = server.id_jag_subject_source
if not configured:
return IdJagSubjectSource.login_assertion
try:
return IdJagSubjectSource(configured)
except ValueError:
return None
def _id_jag_subject_token_type(server: MCPServer) -> str:
"""ID-JAG asserts the user's id_token, so the token-exchange access_token default maps to id_token;
an explicitly configured value (e.g. a SAML2 assertion type) is honored verbatim."""

264
litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py Normal file
View file

@ -0,0 +1,264 @@
"""Exchange the identity assertion captured at SSO login for a refresh token, at login.
An ``oauth2_id_jag`` server configured with ``id_jag_subject_source: refresh_token`` does not present the
login assertion on ID-JAG leg 1. It presents a refresh token that the IdP issued for that assertion through
one RFC 8693 exchange (``requested_token_type`` refresh_token), which is how a SAML requesting app works at
an IdP whose ID-JAG exchange does not take a ``saml2`` subject directly. That exchange has to happen at
login: a SAML assertion is valid for minutes, the first MCP call that needs it can come hours later, and
the refresh token outlives the assertion by as long as the IdP decides. An OIDC login signed in through the very
IdP client the ID-JAG server uses already holds such a refresh token: it is bound to that client, so it is
presented as is and no further exchange is posted, which Okta refuses for an id_token subject anyway.
The servers consulted are the two configuration authorities, the pod-local config declaration and the DB
rows, like the retention gate, not the in-memory registry, which is a per-process snapshot that can lag a
server added on another pod. The exchange authenticates as the ID-JAG server's own IdP client, since the
refresh token is bound to the client that later spends it. One row holds one refresh token per user, so every server in this mode has
to share one IdP client; a second client is logged and skipped rather than silently handed a token it
cannot redeem. A refused or unreachable IdP never fails the login, the row keeps the plain assertion and
the ID-JAG call reports the missing refresh token as the precondition it is.
"""
from __future__ import annotations
import asyncio
import os
from collections.abc import Iterable, Mapping, Sequence
from datetime import datetime, timezone
from typing import Final, Protocol
from litellm._logging import verbose_proxy_logger
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import (
SSOClientConfig,
sso_client_config,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
SSOIdentityAssertion,
with_login_refresh_token,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ExchangedToken
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
ClientAuth,
CredError,
IdJagConfig,
IdJagSubjectSource,
ServerSpec,
)
from litellm.proxy._types import LiteLLM_MCPServerTable
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
_TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"
REFRESH_TOKEN_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:refresh_token"
_LOGIN_EXCHANGE_DEADLINE_SECONDS: Final = 15.0
_REFRESH_TOKEN_SCOPE: Final = "openid offline_access"
_RUNTIME_APPROVAL_STATES: Final = (None, "active", "approved") # the states get_runtime_mcp_server_rows serves
class LoginExchangeCatalog(Protocol):
"""The slice of the MCP server manager the login-time exchange reads."""
@property
def config_mcp_servers(self) -> Mapping[str, MCPServer]: ...
async def build_mcp_server_from_table(
self,
mcp_server: LiteLLM_MCPServerTable,
*,
credentials_are_encrypted: bool = True,
env_vars_are_encrypted: bool | None = None,
register_oauth_discovery: bool = True,
) -> MCPServer: ...
class TokenExchangeEndpoint(Protocol):
async def fetch(
self, endpoint: str, client_id: str, grant_params: Mapping[str, str], client_auth: ClientAuth
) -> Result[ExchangedToken, CredError]: ...
def _is_runtime_id_jag_row(row: LiteLLM_MCPServerTable) -> bool:
"""An ID-JAG row the runtime registry would serve. A pending or rejected submission never gets a
login assertion posted to the token endpoint it names."""
return row.auth_type == MCPAuth.oauth2_id_jag and row.approval_status in _RUNTIME_APPROVAL_STATES
async def login_exchange_servers(
catalog: LoginExchangeCatalog, db_rows: Sequence[LiteLLM_MCPServerTable]
) -> tuple[MCPServer, ...]:
"""Every server the exchange has to consider: the config-declared ones plus the approved ID-JAG rows
in the DB, built the same way the catalog builds them so their credentials are decrypted. A DB row
hides a config server with the same ``server_id``, exactly as the live catalog serves them."""
id_jag_rows: Final = tuple(row for row in db_rows if _is_runtime_id_jag_row(row))
db_servers: Final = tuple([await _built_or_skipped(catalog, row) for row in id_jag_rows])
db_by_id: Final = {server.server_id: server for server in db_servers if server is not None}
return tuple({**catalog.config_mcp_servers, **db_by_id}.values())
async def _built_or_skipped(catalog: LoginExchangeCatalog, row: LiteLLM_MCPServerTable) -> MCPServer | None:
"""One row that cannot be built (undecryptable credentials, a half-configured submission) is skipped
with a warning. It must not cost every other refresh-token server its login exchange."""
try:
return await catalog.build_mcp_server_from_table(row, register_oauth_discovery=False)
except Exception as exc: # noqa: BLE001 # a per-row build failure is logged and isolated, not propagated
verbose_proxy_logger.warning(
"ID-JAG: MCP server row server_id=%s could not be built for the login-time refresh token exchange "
"and is skipped: %s",
row.server_id,
exc,
)
return None
def refresh_token_exchange_targets(specs: Iterable[ServerSpec | None]) -> tuple[IdJagConfig, ...]:
"""The ID-JAG configs that present a login refresh token, one per IdP client, in registration order."""
configs: Final = tuple(
spec.config
for spec in specs
if spec is not None
and isinstance(spec.config, IdJagConfig)
and spec.config.subject_source is IdJagSubjectSource.refresh_token
)
clients: Final = tuple((config.org_token_endpoint, config.client_id) for config in configs)
return tuple(config for index, config in enumerate(configs) if clients.index(clients[index]) == index)
async def exchange_login_assertion_for_refresh_token(
assertion: SSOIdentityAssertion,
targets: tuple[IdJagConfig, ...],
token_endpoint: TokenExchangeEndpoint,
now: datetime,
deadline_seconds: float = _LOGIN_EXCHANGE_DEADLINE_SECONDS,
sso_client: SSOClientConfig | None = None,
) -> SSOIdentityAssertion:
"""The assertion as captured when no server wants a refresh token, otherwise the assertion plus the
refresh token one target's IdP client can redeem. An OIDC login's own refresh token is that token
when the SSO client that received it is the target's client at the same token endpoint; otherwise it
is bound to a client the ID-JAG client is not, and the assertion is exchanged for one at the first
target's client."""
if not targets or assertion.exchanged_refresh_token is not None:
return assertion
reusable: Final = _target_redeeming_the_login_refresh_token(assertion, targets, sso_client)
if reusable is not None and assertion.refresh_token is not None:
verbose_proxy_logger.info(
"ID-JAG: the refresh token issued at OIDC login belongs to client_id=%s at %s, the IdP client of a server "
"with id_jag_subject_source: refresh_token, so it is kept as that subject and no further exchange is posted",
reusable.client_id,
reusable.org_token_endpoint,
)
return with_login_refresh_token(
assertion,
assertion.refresh_token.get_secret_value(),
None,
now,
endpoint=reusable.org_token_endpoint,
client_id=reusable.client_id,
)
config: Final = targets[0]
if len(targets) > 1:
verbose_proxy_logger.warning(
"ID-JAG: %d IdP clients are configured with id_jag_subject_source refresh_token; the login refresh "
"token is minted for client_id=%s at %s and servers using another client cannot redeem it.",
len(targets),
config.client_id,
config.org_token_endpoint,
)
params: Final = {
"grant_type": _TOKEN_EXCHANGE_GRANT_TYPE,
"requested_token_type": REFRESH_TOKEN_TOKEN_TYPE,
"subject_token": assertion.id_token.get_secret_value(),
"subject_token_type": assertion.subject_token_type or config.subject_token_type,
"scope": _REFRESH_TOKEN_SCOPE,
}
try:
outcome: Final = await asyncio.wait_for(
token_endpoint.fetch(config.org_token_endpoint, config.client_id, params, config.client_auth),
timeout=deadline_seconds,
)
except (TimeoutError, asyncio.TimeoutError):
verbose_proxy_logger.warning(
"ID-JAG: the IdP at %s did not answer the login-time refresh token exchange within %.0fs, so the login "
"finishes with the captured assertion only. Servers with id_jag_subject_source: refresh_token will ask "
"this user to sign in again.",
config.org_token_endpoint,
deadline_seconds,
)
return assertion
match outcome:
case Error(err):
verbose_proxy_logger.warning(
"ID-JAG: the IdP at %s refused to exchange the login assertion for a refresh token (%s); the "
"login goes on with the assertion alone and servers presenting a refresh token will ask the "
"user to sign in again.",
config.org_token_endpoint,
err.summary,
)
return assertion
case Ok(token):
verbose_proxy_logger.info(
"ID-JAG: exchanged the login assertion for a refresh token at %s (expires_in=%s)",
config.org_token_endpoint,
token.expires_in,
)
return with_login_refresh_token(
assertion,
token.access_token,
token.expires_in,
now,
endpoint=config.org_token_endpoint,
client_id=config.client_id,
)
def _target_redeeming_the_login_refresh_token(
assertion: SSOIdentityAssertion, targets: tuple[IdJagConfig, ...], sso_client: SSOClientConfig | None
) -> IdJagConfig | None:
"""The target whose IdP client is the SSO client the login's refresh token was issued to. A refresh
token is redeemable only by the client that received it at the endpoint that issued it, so both have
to match exactly before the login token is presented anywhere."""
if assertion.refresh_token is None or sso_client is None:
return None
return next(
(
config
for config in targets
if config.org_token_endpoint == sso_client.token_endpoint and config.client_id == sso_client.client_id
),
None,
)
async def exchange_login_assertion_if_configured(
assertion: SSOIdentityAssertion | None,
) -> SSOIdentityAssertion | None:
"""The SSO-callback hook: consults the configured and DB-declared MCP servers and performs the
exchange when one of them presents a login refresh token. Nothing here can fail the login."""
if assertion is None:
return None
from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers # noqa: PLC0415 # import cycle
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # avoids import cycle
global_mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 # import cycle
to_server_spec,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( # noqa: PLC0415 # cycle
TokenEndpointClient,
)
from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime global
try:
db_rows: Final = () if prisma_client is None else await get_all_mcp_servers(prisma_client)
servers: Final = await login_exchange_servers(global_mcp_server_manager, db_rows)
targets: Final = refresh_token_exchange_targets(to_server_spec(server) for server in servers)
return await exchange_login_assertion_for_refresh_token(
assertion,
targets,
TokenEndpointClient(),
datetime.now(timezone.utc),
sso_client=sso_client_config(os.environ),
)
except Exception as exc: # noqa: BLE001 # the login itself must not fail on an egress-side exchange
verbose_proxy_logger.warning("ID-JAG: the login-time refresh token exchange failed unexpectedly: %s", exc)
return assertion

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

@ -20,6 +20,7 @@ v2: no imports from v1.
from __future__ import annotations
import hashlib
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import partial
from typing import Final
@ -53,7 +54,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_r
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
AssertionStoreUnavailable,
SSOAssertionStore,
SSOIdentityAssertion,
assertion_expired,
exchanged_refresh_token_issued_for,
login_refresh_token_expired,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import (
ExchangedToken,
@ -75,6 +79,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
ClientSecretAuth,
CredError,
IdJagConfig,
IdJagSubjectSource,
NoneConfig,
PassthroughConfig,
PrivateKeyJwtAuth,
@ -88,6 +93,15 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
_TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"
_JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer"
_ID_JAG_REQUESTED_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:id-jag"
_REFRESH_TOKEN_SUBJECT_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:refresh_token"
@dataclass(frozen=True, slots=True)
class _IdJagSubjectToken:
"""What ID-JAG leg 1 asserts and the RFC 8693 type it is declared as."""
value: str
token_type: str
class _NullOAuthTokenStore:
@ -200,13 +214,15 @@ class UpstreamCredentialProvider:
async def _id_jag(
self, subject: Subject, server: ServerSpec, config: IdJagConfig
) -> Result[httpx2.Auth, CredError]:
match await self._id_jag_subject_token(subject):
match await self._id_jag_subject_token(subject, config):
case Error(err):
return Error(err)
case Ok(subject_token):
return await self._id_jag_exchange(subject, subject_token, server, config)
async def _id_jag_subject_token(self, subject: Subject) -> Result[str, CredError]:
async def _id_jag_subject_token(
self, subject: Subject, config: IdJagConfig
) -> Result[_IdJagSubjectToken, CredError]:
"""The identity token ID-JAG leg 1 asserts, from the request or from the SSO login it was captured at.
A caller that presents its own IdP identity token wins: that is the strongest available
@ -223,7 +239,7 @@ class UpstreamCredentialProvider:
a raised error so a DB blip cannot 500 the egress or the upstream-401 retry.
"""
if subject.inbound_token is not None:
return Ok(subject.inbound_token.get_secret_value())
return Ok(_IdJagSubjectToken(subject.inbound_token.get_secret_value(), config.subject_token_type))
if not subject.subject_id:
return Error(
CredError.of_precondition_required(
@ -232,7 +248,11 @@ class UpstreamCredentialProvider:
)
)
try:
assertion: Final = await self._sso_assertion_store.fetch(subject.subject_id)
assertion: Final = await (
self._sso_assertion_store.fetch_stored(subject.subject_id)
if config.subject_source is IdJagSubjectSource.refresh_token
else self._sso_assertion_store.fetch(subject.subject_id)
)
except AssertionStoreUnavailable as exc:
# The driver's message can name hosts, schemas or connection details, and this summary
# is returned to the caller verbatim as a 503 body. Operators get it from the log.
@ -253,6 +273,8 @@ class UpstreamCredentialProvider:
"Sign in through LiteLLM SSO so the gateway captures one."
)
)
if config.subject_source is IdJagSubjectSource.refresh_token:
return _login_refresh_token_subject(assertion, config, datetime.now(timezone.utc))
if assertion_expired(assertion, datetime.now(timezone.utc)):
return Error(
CredError.of_precondition_required(
@ -260,20 +282,24 @@ class UpstreamCredentialProvider:
"LiteLLM SSO again to capture a current one."
)
)
return Ok(assertion.id_token.get_secret_value())
return Ok(
_IdJagSubjectToken(
assertion.id_token.get_secret_value(), assertion.subject_token_type or config.subject_token_type
)
)
async def _id_jag_exchange(
self, subject: Subject, token: str, server: ServerSpec, config: IdJagConfig
self, subject: Subject, token: _IdJagSubjectToken, server: ServerSpec, config: IdJagConfig
) -> Result[httpx2.Auth, CredError]:
slot: Final = _id_jag_slot_key(subject, server)
fingerprint: Final = _id_jag_fingerprint(token, server.server_id, config)
fingerprint: Final = _id_jag_fingerprint(token.value, server.server_id, config)
async def _exchange() -> Result[ExchangedToken, CredError]:
leg1_params: Final = {
"grant_type": _TOKEN_EXCHANGE_GRANT_TYPE,
"requested_token_type": _ID_JAG_REQUESTED_TOKEN_TYPE,
"subject_token": token,
"subject_token_type": config.subject_token_type,
"subject_token": token.value,
"subject_token_type": token.token_type,
**({"audience": config.audience} if config.audience else {}),
**({"resource": config.resource} if config.resource else {}),
**({"scope": " ".join(config.scopes)} if config.scopes else {}),
@ -399,6 +425,42 @@ class UpstreamCredentialProvider:
return None
def _login_refresh_token_subject(
assertion: SSOIdentityAssertion, config: IdJagConfig, now: datetime
) -> Result[_IdJagSubjectToken, CredError]:
"""The subject for a server configured with ``id_jag_subject_source: refresh_token``: the refresh
token the ID-JAG client holds from the user's sign-in, exchanged from the login assertion or, when
the SSO client is this server's IdP client, the one the OIDC login itself returned. The login-time
hook records which client and endpoint it is bound to, and a token bound elsewhere is never sent
here. The assertion's own expiry does not apply, a SAML assertion is spent minutes after login
while the refresh token it bought lives as long as the IdP said."""
if assertion.exchanged_refresh_token is None:
return Error(
CredError.of_precondition_required(
"This server mints its ID-JAG from a refresh token exchanged at SSO login and none is stored "
"for this user. Sign in through LiteLLM SSO again; if the login log shows the IdP refused the "
"exchange, the IdP client has to allow a refresh token for this requesting app."
)
)
if not exchanged_refresh_token_issued_for(assertion, config.org_token_endpoint, config.client_id):
return Error(
CredError.of_precondition_required(
"The refresh token exchanged at SSO login was minted for a different IdP client than this server "
"uses, so it cannot be presented here. Point every id_jag_subject_source: refresh_token server at "
"the same IdP client, then sign in through LiteLLM SSO again."
)
)
if login_refresh_token_expired(assertion, now):
return Error(
CredError.of_precondition_required(
"The refresh token exchanged at SSO login for this user has expired. Sign in through LiteLLM SSO again."
)
)
return Ok(
_IdJagSubjectToken(assertion.exchanged_refresh_token.get_secret_value(), _REFRESH_TOKEN_SUBJECT_TOKEN_TYPE)
)
def _id_jag_slot_key(subject: Subject, server: ServerSpec) -> str:
"""Which cache slot this caller's bearer for this upstream lives in.

112
litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py
View file

@ -12,15 +12,18 @@ not the size of the user table) and single-flighted per user through the same ``
the ``authorization_code`` arm uses, because an IdP that rotates refresh tokens treats two concurrent
redemptions of one token as replay and can revoke the whole grant chain.
The refresh is redeemed against the generic-OIDC client the login itself used
(``GENERIC_TOKEN_ENDPOINT`` / ``GENERIC_CLIENT_ID`` / ``GENERIC_CLIENT_SECRET``, which the proxy
reconciles from the stored SSO row into the process environment at startup), authenticated the way
that login authenticated: the non-PKCE path always sends HTTP Basic, while the PKCE path sends the
credentials in the body when ``GENERIC_INCLUDE_CLIENT_ID`` is set, and an IdP application may accept
only one of the two. An assertion can only exist if that client minted it, so no other client could
redeem its refresh token, and no other method is known to be accepted. A deployment whose
``GENERIC_SCOPE`` omits ``offline_access`` captures no refresh token at all, which is why that miss
logs the scope by name rather than failing silently.
The refresh is redeemed against the SSO client the login itself used, resolved in the same order
the login callback dispatches (Google, then Microsoft, then generic OIDC, from the client ids the
proxy reconciles from the stored SSO row into the process environment at startup), authenticated the
way that login authenticated: the Google and Microsoft logins and the non-PKCE generic path send
HTTP Basic, while the PKCE generic path sends the credentials in the body when
``GENERIC_INCLUDE_CLIENT_ID`` is set, and an IdP application may accept only one of the two. An
assertion can only exist if that client minted it, so no other client could redeem its refresh
token, and no other method is known to be accepted. A login whose authorization request never asked
for a refresh token captures none at all: the Google and Microsoft logins ask (``access_type=offline``,
``offline_access``) only while an ``oauth2_id_jag`` server exists, and generic OIDC asks only when
``GENERIC_SCOPE`` carries ``offline_access``, which is why that miss names the cause rather than
failing silently.
Failures are values internally (``Result[_, RefreshFailure]``). At the store boundary they collapse
onto the protocol's existing two-outcome contract: a refusal returns the expired assertion unchanged
@ -77,7 +80,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_s
SSOIdentityAssertion,
assertion_expired,
assertion_from_sso_login,
carry_exchanged_refresh_token,
fetch_sso_identity_assertion,
is_saml_assertion,
persist_sso_identity_assertion,
)
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -92,6 +97,10 @@ _SINGLE_FLIGHT_KEY: Final = "sso_identity_assertion"
# Renew this far ahead of ``exp`` so a token that would die between resolution and the second leg of
# the exchange is replaced first. Matches the sibling per-user token store's skew.
_DEFAULT_EXPIRY_SKEW_SECONDS: Final = 60.0
# Where fastapi-sso's Google login redeemed the authorization code, read from Google's discovery
# document; Microsoft's mirrors ``CustomMicrosoftSSO``, override first, tenant template second.
_GOOGLE_TOKEN_ENDPOINT: Final = "https://oauth2.googleapis.com/token"
_MICROSOFT_TOKEN_ENDPOINT_TEMPLATE: Final = "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
class AssertionRead(Protocol):
@ -122,7 +131,7 @@ class FormPost(Protocol):
@dataclass(frozen=True, slots=True)
class SSOClientConfig:
"""The generic-OIDC client credentials a refresh_token grant has to authenticate as, and how."""
"""The SSO client credentials a refresh_token grant has to authenticate as, and how."""
token_endpoint: str
client_id: str
@ -131,32 +140,58 @@ class SSOClientConfig:
def sso_client_config(env: Mapping[str, str]) -> SSOClientConfig | None:
"""The configured generic-OIDC client, or ``None`` when the deployment has none.
"""The configured SSO client the login used, or ``None`` when the deployment has none.
Read from the process environment because that is where the login path reads it
(``_setup_generic_sso_env_vars``) and where the proxy materializes the stored ``sso_config`` row
at startup, so this resolves to the same client that minted the assertion. ``None`` is an
ordinary state, not an error: a deployment signing in through a provider that captures no
Read from the process environment because that is where the login callback reads it and where
the proxy materializes the stored ``sso_config`` row at startup, and resolved in the callback's
own order (``GOOGLE_CLIENT_ID``, then ``MICROSOFT_CLIENT_ID``, then ``GENERIC_CLIENT_ID``, each
selecting its branch when set at all), so this is the same client that minted the assertion: a
deployment carrying both a Google and a generic client signed the user in through Google. ``None``
is an ordinary state, not an error: a deployment signing in through a provider that captures no
assertion has nothing here to renew, and a client with no secret is not a confidential client
that could redeem one.
``auth_method`` is derived from the same ``GENERIC_INCLUDE_CLIENT_ID`` the login reads, because
the two login paths do not agree: the non-PKCE path always authenticates with HTTP Basic, while
the PKCE path puts the credentials in the body when that flag is set. Both capture assertions, so
a constant here would authenticate the renewal differently from the sign-in that produced the
``auth_method`` follows how each login authenticated at the token endpoint. The Google and
Microsoft logins (fastapi-sso) and the non-PKCE generic path send HTTP Basic, while the PKCE
generic path puts the credentials in the body when ``GENERIC_INCLUDE_CLIENT_ID`` is set. A
constant here would authenticate the renewal differently from the sign-in that produced the
refresh token and 401 against an IdP application registered for only one of the two.
"""
token_endpoint: Final = env.get("GENERIC_TOKEN_ENDPOINT")
client_id: Final = env.get("GENERIC_CLIENT_ID")
client_secret: Final = env.get("GENERIC_CLIENT_SECRET")
if env.get("GOOGLE_CLIENT_ID") is not None:
return _confidential_client(
_GOOGLE_TOKEN_ENDPOINT, env.get("GOOGLE_CLIENT_ID"), env.get("GOOGLE_CLIENT_SECRET"), "client_secret_basic"
)
if env.get("MICROSOFT_CLIENT_ID") is not None:
tenant: Final = env.get("MICROSOFT_TENANT")
return _confidential_client(
env.get("MICROSOFT_TOKEN_ENDPOINT")
or (_MICROSOFT_TOKEN_ENDPOINT_TEMPLATE.format(tenant=tenant) if tenant else None),
env.get("MICROSOFT_CLIENT_ID"),
env.get("MICROSOFT_CLIENT_SECRET"),
"client_secret_basic",
)
includes_client_id: Final = env.get("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true"
return _confidential_client(
env.get("GENERIC_TOKEN_ENDPOINT"),
env.get("GENERIC_CLIENT_ID"),
env.get("GENERIC_CLIENT_SECRET"),
"client_secret_post" if includes_client_id else "client_secret_basic",
)
def _confidential_client(
token_endpoint: str | None,
client_id: str | None,
client_secret: str | None,
auth_method: MCPTokenEndpointAuthMethod,
) -> SSOClientConfig | None:
if not token_endpoint or not client_id or not client_secret:
return None
includes_client_id: Final = env.get("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true"
return SSOClientConfig(
token_endpoint=token_endpoint,
client_id=client_id,
client_secret=SecretStr(client_secret),
auth_method="client_secret_post" if includes_client_id else "client_secret_basic",
auth_method=auth_method,
)
@ -260,11 +295,19 @@ class SSOAssertionRefresher:
async def refresh(
self, user_id: str, assertion: SSOIdentityAssertion
) -> Result[SSOIdentityAssertion, RefreshFailure]:
if is_saml_assertion(assertion):
verbose_proxy_logger.warning(
"ID-JAG: the stored SAML assertion for user_id=%s has reached its NotOnOrAfter. A SAML assertion "
"cannot be renewed, so the user has to sign in through SAML SSO again.",
user_id,
)
return Error(RefreshFailure.of_rejected("a SAML assertion cannot be renewed; sign in again"))
if assertion.refresh_token is None:
verbose_proxy_logger.warning(
"ID-JAG: the stored IdP identity assertion for user_id=%s has expired and no refresh token was "
"captured with it, so it cannot be renewed without another interactive sign-in. Add "
"'offline_access' to GENERIC_SCOPE so the SSO login captures one.",
"captured with it, so it cannot be renewed without another interactive sign-in. Google and Microsoft "
"logins ask for one once an oauth2_id_jag MCP server is registered, so a user who signed in before "
"that has to sign in again. Generic OIDC asks only with 'offline_access' in GENERIC_SCOPE.",
user_id,
)
return Error(RefreshFailure.of_rejected("no refresh token was captured at sign-in"))
@ -272,11 +315,12 @@ class SSOAssertionRefresher:
if config is None:
verbose_proxy_logger.warning(
"ID-JAG: the stored IdP identity assertion for user_id=%s has expired and cannot be renewed "
"because the generic SSO client is not configured (GENERIC_TOKEN_ENDPOINT, GENERIC_CLIENT_ID, "
"GENERIC_CLIENT_SECRET).",
"because the active SSO provider has no confidential client configured (GOOGLE_CLIENT_ID and "
"GOOGLE_CLIENT_SECRET, MICROSOFT_CLIENT_ID with MICROSOFT_CLIENT_SECRET and MICROSOFT_TENANT, or "
"GENERIC_TOKEN_ENDPOINT with GENERIC_CLIENT_ID and GENERIC_CLIENT_SECRET).",
user_id,
)
return Error(RefreshFailure.of_rejected("the generic SSO client is not configured"))
return Error(RefreshFailure.of_rejected("the SSO client is not configured"))
carried_refresh_token: Final = assertion.refresh_token.get_secret_value()
# Whichever method the SSO login used for this client, since that is the one the IdP
@ -323,10 +367,11 @@ class SSOAssertionRefresher:
user_id,
)
return Error(RefreshFailure.of_rejected("the IdP's refresh response carried no usable id_token"))
failure: Final = await self._store_renewal(user_id, previous, renewed)
carried: Final = carry_exchanged_refresh_token(previous, renewed)
failure: Final = await self._store_renewal(user_id, previous, carried)
if failure is not None:
return Error(failure)
return Ok(renewed)
return Ok(carried)
async def _store_renewal(
self, user_id: str, previous: SSOIdentityAssertion, renewed: SSOIdentityAssertion
@ -401,9 +446,12 @@ class RefreshingSSOAssertionStore:
self._skew = timedelta(seconds=expiry_skew_seconds)
self._clock = clock
async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
return await self._inner.fetch(user_id)
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
assertion: Final = await self._inner.fetch(user_id)
if not self._expiring(assertion):
if assertion is None or is_saml_assertion(assertion) or not self._expiring(assertion):
return assertion
await self._coordinator().run(
user_id,

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

@ -1,9 +1,9 @@
"""Store for the enterprise IdP identity assertion captured at SSO login (EMA).
The ``oauth2_id_jag`` egress arm needs the user's IdP ``id_token`` as its RFC 8693
``subject_token``. A front-door client holds an identity-only ``llm_session_`` bearer, not an
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
The ``oauth2_id_jag`` egress arm needs the user's IdP identity assertion as its RFC 8693
``subject_token``: the OIDC ``id_token``, or the signed SAML assertion when the login was SAML.
A front-door client holds an identity-only ``llm_session_`` bearer, not an 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.
@ -17,9 +17,11 @@ TTL ``MCP_SSO_ASSERTION_CACHE_TTL_SECONDS``; invalidation also guards against st
from __future__ import annotations
import base64
import json
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Protocol
import jwt
@ -35,7 +37,15 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
"""What a Google or Microsoft authorization request has to carry for the IdP to issue a refresh token
beside the id_token. Asked for only while an ``oauth2_id_jag`` server exists to renew an assertion for, so
a deployment without one sees no extra consent."""
MICROSOFT_REFRESH_TOKEN_SCOPE: Final = "offline_access"
GOOGLE_REFRESH_TOKEN_PARAMS: Final = MappingProxyType({"access_type": "offline", "prompt": "consent"})
_ASSERTION_DECRYPT_LOG_KEY: Final = "sso_identity_assertion"
SAML2_SUBJECT_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:saml2"
_STR_ADAPTER: Final[TypeAdapter[str]] = TypeAdapter(str)
_MAYBE_STR_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)
@ -69,8 +79,10 @@ def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable:
class SSOIdentityAssertion(LiteLLMBaseModel):
"""The IdP material an EMA exchange needs: ``id_token`` is the RFC 8693 subject token,
``expires_at`` bounds its usefulness, and the refresh token renews it without re-login."""
"""The IdP material an EMA exchange needs: ``id_token`` is the RFC 8693 subject token (the OIDC
id_token, or the base64 signed ``<saml:Assertion>`` when ``subject_token_type`` names SAML),
``expires_at`` bounds its usefulness, and the refresh token renews it without re-login.
``subject_token_type`` is ``None`` when the exchange should send the operator-configured type."""
model_config = ConfigDict(frozen=True)
@ -78,6 +90,11 @@ class SSOIdentityAssertion(LiteLLMBaseModel):
refresh_token: SecretStr | None = None
issuer: str | None = None
expires_at: datetime | None = None
subject_token_type: str | None = None
exchanged_refresh_token: SecretStr | None = None
exchanged_refresh_token_expires_at: datetime | None = None
exchanged_refresh_token_endpoint: str | None = None
exchanged_refresh_token_client_id: str | None = None
class SSOAssertionCache:
@ -130,6 +147,11 @@ class _StoredAssertionPayload(LiteLLMBaseModel):
refresh_token: str | None = None
issuer: str | None = None
expires_at: datetime | None = None
subject_token_type: str | None = None
exchanged_refresh_token: str | None = None
exchanged_refresh_token_expires_at: datetime | None = None
exchanged_refresh_token_endpoint: str | None = None
exchanged_refresh_token_client_id: str | None = None
def assertion_from_sso_login(id_token: object, refresh_token: object) -> SSOIdentityAssertion | None:
@ -159,6 +181,94 @@ def assertion_from_sso_login(id_token: object, refresh_token: object) -> SSOIden
)
def assertion_from_saml_login(
assertion_xml: bytes, not_on_or_after: int | None, issuer: str | None
) -> SSOIdentityAssertion:
"""The carrier for a SAML login: the signed ``<saml:Assertion>`` element the IdP issued, base64
encoded the way RFC 8693 carries a ``saml2`` subject token, with its own validity window as
the expiry and no refresh token, since SAML has none and the user has to sign in again."""
return SSOIdentityAssertion(
id_token=SecretStr(base64.urlsafe_b64encode(assertion_xml).rstrip(b"=").decode("ascii")),
issuer=issuer,
expires_at=datetime.fromtimestamp(not_on_or_after, tz=timezone.utc) if not_on_or_after is not None else None,
subject_token_type=SAML2_SUBJECT_TOKEN_TYPE,
)
def is_saml_assertion(assertion: SSOIdentityAssertion) -> bool:
return assertion.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE
def with_login_refresh_token(
assertion: SSOIdentityAssertion,
refresh_token: str,
expires_in: int | None,
now: datetime,
*,
endpoint: str | None = None,
client_id: str | None = None,
) -> SSOIdentityAssertion:
"""The same login, now also carrying the refresh token the ID-JAG client exchanged the assertion
for, with the lifetime the IdP gave that token. It is kept apart from ``refresh_token``, which an
OIDC login holds for the SSO client and which only that client can redeem. The assertion and its
own expiry stay as they were, so a server that presents the login assertion directly is unaffected."""
return assertion.model_copy(
update={
"exchanged_refresh_token": SecretStr(refresh_token),
"exchanged_refresh_token_expires_at": now + timedelta(seconds=expires_in)
if expires_in is not None
else None,
"exchanged_refresh_token_endpoint": endpoint,
"exchanged_refresh_token_client_id": client_id,
}
)
def carry_exchanged_refresh_token(
previous: SSOIdentityAssertion, renewed: SSOIdentityAssertion
) -> SSOIdentityAssertion:
"""The renewed OIDC assertion still carrying the refresh token the ID-JAG client exchanged at login.
Renewal only replaces the id_token and the SSO refresh token, the exchanged one is unrelated to
them and stays valid for as long as the IdP said. When the login's own refresh token was kept as
that subject (the SSO client is the ID-JAG client), a rotation the renewal just received replaces
it too, since the IdP has retired the token the row held."""
login_token_reused: Final = (
previous.refresh_token is not None
and previous.exchanged_refresh_token is not None
and previous.exchanged_refresh_token.get_secret_value() == previous.refresh_token.get_secret_value()
)
return renewed.model_copy(
update={
"exchanged_refresh_token": renewed.refresh_token
if login_token_reused
else previous.exchanged_refresh_token,
"exchanged_refresh_token_expires_at": previous.exchanged_refresh_token_expires_at,
"exchanged_refresh_token_endpoint": previous.exchanged_refresh_token_endpoint,
"exchanged_refresh_token_client_id": previous.exchanged_refresh_token_client_id,
}
)
def exchanged_refresh_token_issued_for(assertion: SSOIdentityAssertion, endpoint: str, client_id: str) -> bool:
"""Whether the exchanged refresh token was minted at this token endpoint for this client. A refresh
token is bound to the client that received it, so one minted for another server's IdP client must
never be posted to this server's endpoint."""
return (
assertion.exchanged_refresh_token_endpoint == endpoint
and assertion.exchanged_refresh_token_client_id == client_id
)
def login_refresh_token_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool:
"""Whether the refresh token exchanged at login has passed the lifetime the IdP gave it. One
with no recorded lifetime is left for the IdP to judge, like an assertion with no ``exp``."""
expires_at: Final = assertion.exchanged_refresh_token_expires_at
if expires_at is None:
return False
normalized: Final = expires_at if expires_at.tzinfo is not None else expires_at.replace(tzinfo=timezone.utc)
return normalized <= now
def assertion_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool:
"""Whether the assertion's ``exp`` has passed at ``now``. An assertion carrying no expiry is
treated as usable and left for the IdP to reject, since the store records what the id_token
@ -212,6 +322,27 @@ async def persist_sso_identity_assertion(
**({"refresh_token": assertion.refresh_token.get_secret_value()} if assertion.refresh_token else {}),
**({"issuer": assertion.issuer} if assertion.issuer else {}),
**({"expires_at": assertion.expires_at.isoformat()} if assertion.expires_at else {}),
**({"subject_token_type": assertion.subject_token_type} if assertion.subject_token_type else {}),
**(
{"exchanged_refresh_token": assertion.exchanged_refresh_token.get_secret_value()}
if assertion.exchanged_refresh_token
else {}
),
**(
{"exchanged_refresh_token_expires_at": assertion.exchanged_refresh_token_expires_at.isoformat()}
if assertion.exchanged_refresh_token_expires_at
else {}
),
**(
{"exchanged_refresh_token_endpoint": assertion.exchanged_refresh_token_endpoint}
if assertion.exchanged_refresh_token_endpoint
else {}
),
**(
{"exchanged_refresh_token_client_id": assertion.exchanged_refresh_token_client_id}
if assertion.exchanged_refresh_token_client_id
else {}
),
}
encoded: Final = _STR_ADAPTER.validate_python(encrypt_value_helper(json.dumps(payload)))
await _assertion_table(prisma_client).upsert(
@ -250,6 +381,11 @@ async def _read_assertion_from_db(user_id: str) -> SSOIdentityAssertion | None:
refresh_token=SecretStr(payload.refresh_token) if payload.refresh_token else None,
issuer=payload.issuer,
expires_at=payload.expires_at,
subject_token_type=payload.subject_token_type,
exchanged_refresh_token=SecretStr(payload.exchanged_refresh_token) if payload.exchanged_refresh_token else None,
exchanged_refresh_token_expires_at=payload.exchanged_refresh_token_expires_at,
exchanged_refresh_token_endpoint=payload.exchanged_refresh_token_endpoint,
exchanged_refresh_token_client_id=payload.exchanged_refresh_token_client_id,
)
@ -290,6 +426,12 @@ class SSOAssertionStore(Protocol):
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: ...
async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
"""The row as stored, never renewed on the way out. For a server that presents the refresh
token exchanged at login, the SSO id_token's own lifetime is irrelevant and an SSO renewal
failing must not stand between that server and a still valid refresh token."""
...
class DbSSOAssertionStore:
"""The live store: the row the SSO callback wrote, read back by ``user_id``.
@ -308,6 +450,9 @@ class DbSSOAssertionStore:
except Exception as exc: # noqa: BLE001 # any driver/storage failure is an outage, not an absence
raise AssertionStoreUnavailable(str(exc)) from exc
async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
return await self.fetch(user_id)
async def fetch_uncached(self, user_id: str) -> SSOIdentityAssertion | None:
try:
return await _read_assertion_from_db(user_id)

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

@ -349,6 +349,14 @@ class IdJagResourceClient(LiteLLMBaseModel):
client_auth: ClientAuth
class IdJagSubjectSource(str, Enum):
"""What ID-JAG leg 1 presents as the RFC 8693 subject_token: the identity assertion captured at SSO
login, or the refresh token that assertion was exchanged for at the IdP when the user logged in."""
login_assertion = "login_assertion"
refresh_token = "refresh_token"
class IdJagConfig(HeaderCarrier):
"""draft-ietf-oauth-identity-assertion-authz-grant (Okta "AI agent token exchange").
@ -368,6 +376,7 @@ class IdJagConfig(HeaderCarrier):
client_auth: ClientAuth
resource_client: IdJagResourceClient
subject_token_type: str = "urn:ietf:params:oauth:token-type:id_token"
subject_source: IdJagSubjectSource = IdJagSubjectSource.login_assertion
audience: str | None = None
resource: str | None = None
scopes: tuple[str, ...] = ()

18
litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py
View file

@ -25,7 +25,16 @@ from litellm.proxy.management_endpoints.ui_sso import (
async def lookup_consent_teams(user_id: str) -> tuple[ConsentTeam, ...] | ReloadUserFailure:
user: Final = await load_active_user_by_id(user_id)
return await _consent_teams(await load_active_user_by_id(user_id))
async def lookup_consent_teams_from_database(user_id: str) -> tuple[ConsentTeam, ...] | ReloadUserFailure:
"""The same teams read from the database row, for a grant that is consumed once it is settled: the
membership the minter will re-check must be the membership the team was picked from."""
return await _consent_teams(await load_active_user_by_id(user_id, source="database"))
async def _consent_teams(user: LiteLLM_UserTable | ReloadUserFailure) -> tuple[ConsentTeam, ...] | ReloadUserFailure:
if isinstance(user, str):
return user
details: Final = await _team_details(user.teams)
@ -39,7 +48,7 @@ async def lookup_consent_teams(user_id: str) -> tuple[ConsentTeam, ...] | Reload
async def mint_proxy_credential(
user_id: str, team_id: str | None
user_id: str, team_id: str | None, *, allowed_routes: Sequence[str] | None = None
) -> MintedProxyCredential | ProxyCredentialMintFailure:
"""Mint the ``lite login`` credential for a consented grant. Membership is checked
live against the database row, so a team the user left between consent and redemption (or between refreshes)
@ -53,7 +62,9 @@ async def mint_proxy_credential(
user row handed to the minter carries no team list, exactly like ``lite login``'s, so
the minter's own first-team fallback stays inert. The credential carries the role the
proxy already enforces for the user on every request, so a row with no role (JWT auth's
upsert writes none) mints as an internal user instead of being refused."""
upsert writes none) mints as an internal user instead of being refused. ``allowed_routes``
narrows the credential to those route groups, the way a virtual key's ``allowed_routes`` does,
for a grant whose scope asked for less than the full credential."""
user: Final = await load_active_user_by_id(user_id, source="database")
if isinstance(user, str):
return user
@ -75,6 +86,7 @@ async def mint_proxy_credential(
team_alias=selected.team_alias,
team_models=selected.team_models,
team_model_aliases=selected.team_model_aliases,
allowed_routes=allowed_routes,
)
return MintedProxyCredential(
key=key,

19
litellm/proxy/_lazy_features.py
View file

@ -155,18 +155,13 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
module_path="litellm.proxy.management_endpoints.mcp_management_endpoints",
path_prefixes=("/v1/mcp/",),
),
LazyFeature(
# Also serves /.well-known/oauth-* (OAuth metadata discovery).
# No /mcp/oauth prefix here: the mounted /mcp sub-app would
# shadow it, and there are no actual routes there anyway.
name="mcp_byok_oauth",
module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints",
path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"),
),
LazyFeature(
# Serves OAuth dance endpoints (/authorize, /token, /callback,
# /register) plus several /.well-known/ discovery URLs at the proxy
# root — needed for MCP-over-OAuth flows even before /mcp is hit.
# Registered before mcp_byok_oauth, which also claims the root
# /.well-known/oauth-authorization-server: the root document is the
# gateway's own authorization server metadata (token endpoint /token).
name="mcp_discoverable",
module_path="litellm.proxy._experimental.mcp_server.discoverable_endpoints",
path_prefixes=(
@ -185,6 +180,14 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
# Catches the /{mcp_server_name}/authorize|token|register variants.
path_suffixes=("/authorize", "/token", "/register"),
),
LazyFeature(
# Also serves /.well-known/oauth-* (OAuth metadata discovery).
# No /mcp/oauth prefix here: the mounted /mcp sub-app would
# shadow it, and there are no actual routes there anyway.
name="mcp_byok_oauth",
module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints",
path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"),
),
LazyFeature(
name="mcp_rest",
module_path="litellm.proxy._experimental.mcp_server.rest_endpoints",

130
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -34541,6 +34541,17 @@
],
"title": "Id Jag Resource Token Endpoint"
},
"id_jag_subject_source": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Id Jag Subject Source"
},
"redirect_uris": {
"anyOf": [
{
@ -35771,7 +35782,7 @@
"/.well-known/oauth-authorization-server": {
"get": {
"description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.",
"operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get",
"operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get_2",
"parameters": [
{
"in": "query",
@ -35819,7 +35830,7 @@
"/.well-known/oauth-authorization-server/mcp": {
"get": {
"description": "OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414\npath-inserted form for a client that treats {base}/mcp as its authorization base URL.\n\nThe single-segment /mcp is reserved for the aggregate so the discovery chain stays\nconsistent: the aggregate protected-resource document advertises {base}/mcp as its\nauthorization server, so the document served here must have issuer {base}/mcp. A server\nliterally named ``mcp`` therefore does not take this route; it keeps its standard\ntwo-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp. Letting the\nper-server row win here instead would serve an issuer of {base} against a resource that\nadvertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door.",
"operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get",
"operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get_2",
"responses": {
"200": {
"content": {
@ -35839,7 +35850,7 @@
"/.well-known/oauth-authorization-server/mcp/{mcp_server_name}": {
"get": {
"description": "OAuth authorization server discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-authorization-server/mcp/{server_name}",
"operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get",
"operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get_2",
"parameters": [
{
"in": "path",
@ -35880,7 +35891,7 @@
"/.well-known/oauth-authorization-server/{mcp_server_name}": {
"get": {
"description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.",
"operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get",
"operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get_2",
"parameters": [
{
"in": "path",
@ -35928,7 +35939,7 @@
"/.well-known/oauth-authorization-server/{mcp_server_name}/mcp": {
"get": {
"description": "OAuth authorization server discovery for legacy /{server_name}/mcp pattern.",
"operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get",
"operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get_2",
"parameters": [
{
"in": "path",
@ -35968,7 +35979,7 @@
},
"/.well-known/oauth-protected-resource": {
"get": {
"operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get",
"operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get_2",
"responses": {
"200": {
"content": {
@ -36004,7 +36015,7 @@
"/.well-known/oauth-protected-resource/mcp": {
"get": {
"description": "OAuth protected resource discovery for the aggregate /mcp endpoint.\n\nThe single-segment ``/mcp`` path does not collide with any per-server PRM pattern\n(those are two-segment: ``/mcp/{server}`` or ``/{server}/mcp``), so this unambiguously\ndescribes the aggregate resource.",
"operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get",
"operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get_2",
"responses": {
"200": {
"content": {
@ -36024,7 +36035,7 @@
"/.well-known/oauth-protected-resource/mcp/{mcp_server_name}": {
"get": {
"description": "OAuth protected resource discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-protected-resource/mcp/{server_name}\n\nThis endpoint is compliant with MCP specification and works with standard\nMCP clients like mcp-inspector and VSCode Copilot.",
"operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get",
"operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get_2",
"parameters": [
{
"in": "path",
@ -36065,7 +36076,7 @@
"/.well-known/oauth-protected-resource/{mcp_server_name}/mcp": {
"get": {
"description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.",
"operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get",
"operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get_2",
"parameters": [
{
"in": "path",
@ -36194,9 +36205,27 @@
},
"Body_token_endpoint__mcp_server_name__token_post": {
"properties": {
"assertion": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Assertion"
},
"client_id": {
"title": "Client Id",
"type": "string"
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Client Id"
},
"client_secret": {
"anyOf": [
@ -36293,17 +36322,34 @@
}
},
"required": [
"grant_type",
"client_id"
"grant_type"
],
"title": "Body_token_endpoint__mcp_server_name__token_post",
"type": "object"
},
"Body_token_endpoint_token_post": {
"properties": {
"assertion": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Assertion"
},
"client_id": {
"title": "Client Id",
"type": "string"
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Client Id"
},
"client_secret": {
"anyOf": [
@ -36400,8 +36446,7 @@
}
},
"required": [
"grant_type",
"client_id"
"grant_type"
],
"title": "Body_token_endpoint_token_post",
"type": "object"
@ -37324,6 +37369,17 @@
],
"title": "Id Jag Resource Token Endpoint"
},
"id_jag_subject_source": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Id Jag Subject Source"
},
"redirect_uris": {
"anyOf": [
{
@ -38151,7 +38207,7 @@
"/.well-known/oauth-authorization-server": {
"get": {
"description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.",
"operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get_2",
"operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get",
"parameters": [
{
"in": "query",
@ -38199,7 +38255,7 @@
"/.well-known/oauth-authorization-server/mcp": {
"get": {
"description": "OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414\npath-inserted form for a client that treats {base}/mcp as its authorization base URL.\n\nThe single-segment /mcp is reserved for the aggregate so the discovery chain stays\nconsistent: the aggregate protected-resource document advertises {base}/mcp as its\nauthorization server, so the document served here must have issuer {base}/mcp. A server\nliterally named ``mcp`` therefore does not take this route; it keeps its standard\ntwo-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp. Letting the\nper-server row win here instead would serve an issuer of {base} against a resource that\nadvertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door.",
"operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get_2",
"operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get",
"responses": {
"200": {
"content": {
@ -38219,7 +38275,7 @@
"/.well-known/oauth-authorization-server/mcp/{mcp_server_name}": {
"get": {
"description": "OAuth authorization server discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-authorization-server/mcp/{server_name}",
"operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get_2",
"operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get",
"parameters": [
{
"in": "path",
@ -38260,7 +38316,7 @@
"/.well-known/oauth-authorization-server/{mcp_server_name}": {
"get": {
"description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.",
"operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get_2",
"operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get",
"parameters": [
{
"in": "path",
@ -38308,7 +38364,7 @@
"/.well-known/oauth-authorization-server/{mcp_server_name}/mcp": {
"get": {
"description": "OAuth authorization server discovery for legacy /{server_name}/mcp pattern.",
"operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get_2",
"operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get",
"parameters": [
{
"in": "path",
@ -38348,7 +38404,7 @@
},
"/.well-known/oauth-protected-resource": {
"get": {
"operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get_2",
"operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get",
"responses": {
"200": {
"content": {
@ -38384,7 +38440,7 @@
"/.well-known/oauth-protected-resource/mcp": {
"get": {
"description": "OAuth protected resource discovery for the aggregate /mcp endpoint.\n\nThe single-segment ``/mcp`` path does not collide with any per-server PRM pattern\n(those are two-segment: ``/mcp/{server}`` or ``/{server}/mcp``), so this unambiguously\ndescribes the aggregate resource.",
"operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get_2",
"operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get",
"responses": {
"200": {
"content": {
@ -38404,7 +38460,7 @@
"/.well-known/oauth-protected-resource/mcp/{mcp_server_name}": {
"get": {
"description": "OAuth protected resource discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-protected-resource/mcp/{server_name}\n\nThis endpoint is compliant with MCP specification and works with standard\nMCP clients like mcp-inspector and VSCode Copilot.",
"operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get_2",
"operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get",
"parameters": [
{
"in": "path",
@ -38445,7 +38501,7 @@
"/.well-known/oauth-protected-resource/{mcp_server_name}/mcp": {
"get": {
"description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.",
"operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get_2",
"operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get",
"parameters": [
{
"in": "path",
@ -40729,6 +40785,17 @@
],
"title": "Id Jag Resource Token Endpoint"
},
"id_jag_subject_source": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Id Jag Subject Source"
},
"redirect_uris": {
"anyOf": [
{
@ -44809,6 +44876,17 @@
],
"title": "Id Jag Resource Token Endpoint"
},
"id_jag_subject_source": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Id Jag Subject Source"
},
"redirect_uris": {
"anyOf": [
{

25
litellm/proxy/_types.py
View file

@ -551,6 +551,7 @@ class LiteLLMRoutes(enum.Enum):
+ list(agent_inference_routes)
+ model_info_routes
)
llm_api_routes_without_mcp = sorted(set(llm_api_routes) - set(mcp_inference_routes))
info_routes = [
"/key/info",
"/key/health",
@ -4762,6 +4763,26 @@ class JWTIssuerConfig(LiteLLMBaseModel):
DEFAULT_JWKS_STALE_TTL: Final = 3600
class IdJagRequestingApp(LiteLLMBaseModel):
"""A requesting application the operator lets redeem an ID-JAG (Identity Assertion
Authorization Grant) at the MCP gateway's ``POST /token``. ``client_id`` is the value
the identity provider puts in the assertion's ``client_id`` claim and the app sends as
the ``client_id`` form field. A ``client_secret`` makes the app a confidential client
that must authenticate with it (HTTP Basic or the ``client_secret`` form field); without
one the app is a public client."""
model_config = ConfigDict(frozen=True)
client_id: str = Field(
min_length=1, description="The requesting app's client_id as the identity provider asserts it."
)
client_secret: str | None = Field(
default=None,
min_length=1,
description="Shared secret the app must present at the gateway's token endpoint. Omit for a public client.",
)
class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
"""
A class to define the roles and permissions for a LiteLLM Proxy w/ JWT Auth.
@ -4942,6 +4963,10 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
default=None,
description="Optional issuer-bound JWT validation rules. When a token's `iss` matches a configured issuer, validation uses that issuer's JWKS, audience, and claim mappings. Tokens with an unlisted `iss` fall back to the global JWT_AUDIENCE/JWT_ISSUER validation path — this is additive routing, not an allow-list.",
)
id_jag_requesting_apps: tuple[IdJagRequestingApp, ...] | None = Field(
default=None,
description="Requesting apps allowed to redeem an ID-JAG at the MCP gateway's POST /token with the RFC 7523 jwt-bearer grant. Unset or empty turns that grant off.",
)
#########################################################
def __init__(self, **kwargs: Any) -> None:

4
litellm/proxy/auth/auth_checks.py
View file

@ -3960,6 +3960,7 @@ class ExperimentalUIJWTToken:
team_models: Sequence[str] | None = None,
team_model_aliases: Mapping[str, str] | None = None,
max_budget: float | None = None,
allowed_routes: Sequence[str] | None = None,
) -> str:
"""
Generate a JWT token for CLI authentication with configurable expiration.
@ -3973,6 +3974,8 @@ class ExperimentalUIJWTToken:
team_alias: Team alias for the selected team, if available
team_models: Model allowlist granted by the selected team
team_model_aliases: Team model aliases for the selected team
allowed_routes: Route groups the session may call, enforced like a virtual key's
``allowed_routes`` on every route; None leaves every route the role allows
Returns:
Encrypted JWT token string
@ -4017,6 +4020,7 @@ class ExperimentalUIJWTToken:
max_parallel_requests=None,
user_role=LitellmUserRoles(user_info.user_role),
is_session_token=True,
allowed_routes=list(allowed_routes) if allowed_routes is not None else None,
)
return encrypt_bearer_token(valid_token.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX)

96
litellm/proxy/auth/handle_jwt.py
View file

@ -35,6 +35,7 @@ from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
from litellm.proxy._types import (
DEFAULT_JWKS_STALE_TTL,
RBAC_ROLES,
IdJagRequestingApp,
JWKKeyValue,
JWTAuthBuilderResult,
JWTIssuerConfig,
@ -99,6 +100,10 @@ class NoMatchingJWTPublicKeyError(Exception):
"""Raised when a JWKS endpoint returns no key matching the requested ``kid``."""
class IdJagIsNotABearerTokenError(Exception):
"""Raised when bearer auth is handed an ID-JAG, which only the token endpoint may redeem."""
class JWKSUnreachableError(Exception):
"""Raised when an IdP's JWKS / OIDC discovery endpoint is unreachable and no cached copy is left to fall back on."""
@ -114,9 +119,23 @@ _CACHED_VALUE: Final = TypeAdapter(object)
_CachedValueT = TypeVar("_CachedValueT", bound=JWKKeyValue | str)
ID_JAG_TYP: Final = "oauth-id-jag+jwt"
@dataclass(frozen=True, slots=True)
class AssertionRejection:
"""Why a JWT assertion failed verification. The reason is for the proxy log only, since
PyJWT's wording can quote the token's claims."""
reason: str
class _JWTAuthSettings(Protocol):
"""The JWT auth settings block this handler reads back through ``getattr``, when one is configured."""
@property
def id_jag_requesting_apps(self) -> Sequence[IdJagRequestingApp] | None: ...
@property
def issuers(self) -> Sequence[JWTIssuerConfig] | None: ...
@ -282,7 +301,7 @@ class JWTHandler:
return is_jwt_token(token=token)
@staticmethod
def get_unverified_claims(token: str) -> dict | None:
def get_unverified_claims(token: str) -> dict[str, object] | None:
"""
Decode JWT claims without signature verification.
Used for routing decisions before selecting validation path.
@ -949,7 +968,7 @@ class JWTHandler:
keys: Final = response_json["keys"] if "keys" in response_json else response_json
return cast(JWKKeyValue, keys) # cast-ok: JWTKeyItem declares only `kid`, validating would drop key material
async def _get_public_key_from_jwks_url(self, jwks_url: str, kid: str | None) -> dict:
async def _get_public_key_from_jwks_url(self, jwks_url: str, kid: str | None) -> JWTKeyItem:
resolved_jwks_url: Final = await self._resolve_jwks_url(jwks_url)
keys: Final = await self._cached_with_stale_fallback(
cache_key=f"litellm_jwt_auth_keys_{resolved_jwks_url}",
@ -960,11 +979,11 @@ class JWTHandler:
public_key: Final = self.parse_keys(keys=keys, kid=kid)
if public_key is not None:
return cast(dict, public_key)
return public_key
raise NoMatchingJWTPublicKeyError(f"No matching public key found. keys={resolved_jwks_url}, kid={kid}")
async def get_public_key(self, kid: str | None) -> dict:
async def get_public_key(self, kid: str | None) -> JWTKeyItem:
keys_url: Final = os.getenv("JWT_PUBLIC_KEY_URL")
if keys_url is None:
@ -1151,7 +1170,7 @@ class JWTHandler:
# its jwks_uri, matching JWTIssuerConfig.jwks_url's documented fallback.
return f"{issuer_config.issuer.rstrip('/')}/.well-known/openid-configuration"
def _get_claim_value_for_issuer_mapping(self, token: dict, claim_field: str) -> object:
def _get_claim_value_for_issuer_mapping(self, token: Mapping[str, object], claim_field: str) -> object:
"""Resolve a mapped claim from ``token``.
Returns ``None`` when the field is absent or empty so that mapped claims
@ -1168,8 +1187,10 @@ class JWTHandler:
return None
return claim_value
def _apply_issuer_claim_mappings(self, token: dict, issuer_config: JWTIssuerConfig) -> dict:
normalized: Final[dict] = {k: v for k, v in token.items() if k not in self.LITELLM_INTERNAL_CLAIMS}
def _apply_issuer_claim_mappings(
self, token: Mapping[str, object], issuer_config: JWTIssuerConfig
) -> dict[str, object]:
normalized: Final[dict[str, object]] = {k: v for k, v in token.items() if k not in self.LITELLM_INTERNAL_CLAIMS}
normalized[self.LITELLM_JWT_ISSUER_CLAIM] = issuer_config.issuer
claim_mappings: Final = [
(issuer_config.user_id_jwt_field, self.LITELLM_USER_ID_CLAIM),
@ -1192,8 +1213,8 @@ class JWTHandler:
return normalized
def _get_jwk_from_public_key(self, public_key: dict) -> dict:
jwk: Final = {}
def _get_jwk_from_public_key(self, public_key: Mapping[str, object]) -> dict[str, object]:
jwk: Final[dict[str, object]] = {}
for key in ["kty", "kid", "n", "e", "x", "y", "crv"]:
if key in public_key:
jwk[key] = public_key[key]
@ -1223,12 +1244,12 @@ class JWTHandler:
def _decode_jwt_with_public_key(
self,
token: str,
public_key: dict | str,
public_key: JWTKeyItem | str,
audience: str | list[str] | None,
issuer: str | None = None,
options: dict | None = None,
disable_audience_validation: bool = False,
) -> dict:
) -> dict[str, object]:
decode_options: Final = (
options
if options is not None
@ -1298,11 +1319,64 @@ class JWTHandler:
issuer_config=issuer_config,
)
async def decode_trusted_assertion(
self, token: str, audience: Sequence[str], required_claims: Sequence[str]
) -> Mapping[str, object] | AssertionRejection:
"""Verify a JWT assertion handed to the gateway's token endpoint with the issuer plumbing
JWT auth already trusts, a matching ``issuers`` entry or ``JWT_ISSUER`` with
``JWT_PUBLIC_KEY_URL``, but against the caller's audience instead of the API audience.
An assertion whose ``iss`` nothing trusts is refused rather than checked unscoped. A JWKS
outage propagates as the 503 ``ProxyException`` so the caller does not call the token bad."""
try:
header: Final = jwt.get_unverified_header(token)
except jwt.PyJWTError as malformed:
return AssertionRejection(reason=f"malformed assertion: {malformed}")
claims: Final = self.get_unverified_claims(token=token)
issuer_config: Final = self._get_configured_issuer(token=token)
trusted_issuer: Final = issuer_config.issuer if issuer_config is not None else os.getenv("JWT_ISSUER")
if not trusted_issuer or claims is None or claims.get("iss") != trusted_issuer:
return AssertionRejection(reason="iss is not an issuer this gateway trusts")
kid: Final = header.get("kid")
try:
public_key: Final = (
await self._get_public_key_from_jwks_url(self._get_jwks_url_for_issuer(issuer_config), kid)
if issuer_config is not None
else await self.get_public_key(kid=kid)
)
except JWKSUnreachableError as e:
raise jwks_unavailable_exception(e) from e
except Exception as no_key: # noqa: BLE001 # the key lookups raise plain Exceptions for a missing key
return AssertionRejection(reason=f"no signing key for the assertion: {no_key}")
try:
payload: Final = self._decode_jwt_with_public_key(
token=token,
public_key=public_key,
audience=list(audience),
issuer=trusted_issuer,
options={"require": list(required_claims)},
)
except jwt.PyJWTError as invalid:
return AssertionRejection(reason=f"assertion failed verification: {invalid}")
verified: Final = (
payload
if issuer_config is None
else self._apply_issuer_claim_mappings(token=payload, issuer_config=issuer_config)
)
validate: Final = self.litellm_jwtauth.custom_validate
if validate is not None and not validate(dict(verified)):
return AssertionRejection(reason="custom_validate rejected the assertion's claims")
return verified
async def auth_jwt(self, token: str) -> dict:
header: Final = jwt.get_unverified_header(token)
verbose_proxy_logger.debug("header: %s", header)
if header.get("typ") == ID_JAG_TYP:
raise IdJagIsNotABearerTokenError(
"an ID-JAG is an authorization grant, not a bearer token; redeem it at the token endpoint"
)
kid: Final = header.get("kid", None)
issuer_config: Final = self._get_configured_issuer(token=token)

11
litellm/proxy/auth/route_checks.py
View file

@ -69,7 +69,8 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES: Final = frozenset(
# paths directly because the request route carries the resolved key id.
_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES: Final = ("/regenerate", "/reset_spend")
_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", "llm_api_routes"))
_LLM_API_ROUTE_GROUPS: Final = frozenset(("llm_api_routes", "llm_api_routes_without_mcp"))
_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", *_LLM_API_ROUTE_GROUPS))
# Method-scoped because the generic credential CRUD handlers share these paths:
# the route check cannot admit PATCH/DELETE on a name like "user_connections"
@ -154,7 +155,7 @@ class RouteChecks:
################################################
# For llm_api_routes, also check registered pass-through endpoints
################################################
if allowed_route == "llm_api_routes":
if allowed_route in _LLM_API_ROUTE_GROUPS:
if route == "/auto_router/session" and RouteChecks.get_request_method(request) == "GET":
return True
@ -185,7 +186,9 @@ class RouteChecks:
# headers, env, credentials). POST/PUT/DELETE on
# these paths are admin-only management writes and
# are intentionally not covered.
if RouteChecks._is_get_mcp_server_discovery_route(route=route, request=request):
if allowed_route == "llm_api_routes" and RouteChecks._is_get_mcp_server_discovery_route(
route=route, request=request
):
return True
# Agent registry CRUD moved from llm_api_routes into
@ -195,7 +198,7 @@ class RouteChecks:
# could reach these paths, so keep them reachable here;
# the handlers in agent_endpoints/endpoints.py still
# enforce proxy-admin on writes and scope reads by role.
if RouteChecks.check_route_access(
if allowed_route == "llm_api_routes" and RouteChecks.check_route_access(
route=route,
allowed_routes=LiteLLMRoutes.agent_management_routes.value,
):

29
litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py
View file

@ -1,25 +1,26 @@
"""Whether the SSO provider the login callback dispatches to can capture an IdP identity assertion.
An ``oauth2_id_jag`` MCP server spends the ``id_token`` captured at SSO login as its RFC 8693
subject token. Only the generic OIDC login path reaches a token response the gateway retains one
from, so a deployment whose SSO runs through Google, Microsoft or SAML never stores an assertion
and every store-sourced ID-JAG exchange fails for every user, however many times they sign in.
Neither side can see that alone: the MCP registration knows nothing about SSO and the login knows
nothing about MCP. This module is the one shared answer both warn from.
subject token. The Google, Microsoft and generic OIDC login paths all reach a token response the
gateway retains one from, and the SAML login retains the signed assertion itself as a ``saml2``
subject token. A deployment with no SSO never signs anyone in, so every store-sourced ID-JAG
exchange fails for every user. Neither side can see that alone: the MCP registration knows nothing
about SSO and the login knows nothing about MCP. This module is the one shared answer both warn from.
"""
from __future__ import annotations
import os
from enum import Enum
from typing import Final
from typing_extensions import assert_never
from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler
_GENERIC_OIDC_REMEDY = (
"Point SSO at the generic OIDC provider (GENERIC_CLIENT_ID), the one login path whose token "
"response the gateway retains an id_token from"
_OIDC_REMEDY: Final = (
"Point SSO at an OIDC provider (GOOGLE_CLIENT_ID, MICROSOFT_CLIENT_ID or GENERIC_CLIENT_ID) or at SAML, "
"the login paths the gateway retains an identity assertion from"
)
@ -55,18 +56,14 @@ def id_jag_assertion_capture_gap() -> str | None:
or ``None`` when that provider does capture an assertion."""
provider = active_sso_provider()
match provider:
case ActiveSSOProvider.generic:
case (
ActiveSSOProvider.google | ActiveSSOProvider.microsoft | ActiveSSOProvider.generic | ActiveSSOProvider.saml
):
return None
case ActiveSSOProvider.none:
return (
"no SSO provider is configured, so no IdP identity assertion is ever captured and "
f"ID-JAG credential resolution fails for every user. {_GENERIC_OIDC_REMEDY}"
)
case ActiveSSOProvider.google | ActiveSSOProvider.microsoft | ActiveSSOProvider.saml:
return (
f"the active SSO provider ({provider.value}) has no identity-assertion capture path, so no "
"IdP id_token is ever stored and ID-JAG credential resolution fails for every user no matter "
f"how often they sign in. {_GENERIC_OIDC_REMEDY}"
f"ID-JAG credential resolution fails for every user. {_OIDC_REMEDY}"
)
case _:
assert_never(provider)

58
litellm/proxy/management_endpoints/sso/saml_sso.py
View file

@ -27,6 +27,7 @@ import hashlib
import os
import secrets
import time
from collections.abc import Mapping
from typing import Final, cast
from urllib.parse import parse_qsl
@ -37,6 +38,10 @@ from pydantic import ValidationError
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
SSOIdentityAssertion,
assertion_from_saml_login,
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.management_endpoints.sso_helper_utils import SSO_SESSIONS_TARGET
from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role
@ -212,7 +217,7 @@ class SAMLAuthHandler:
return OneLogin_Saml2_IdPMetadataParser.merge_settings(sp_settings, idp_settings)
@staticmethod
def _prepare_request_data(request: Request, post_data: dict[str, str] | None = None) -> dict[str, object]:
def _prepare_request_data(request: Request, post_data: Mapping[str, str] | None = None) -> dict[str, object]:
base: Final = SAMLAuthHandler._base_url(request)
scheme, _, host_part = base.partition("://")
host: Final = host_part.split("/", 1)[0]
@ -228,7 +233,7 @@ class SAMLAuthHandler:
async def _build_auth(
request: Request,
cache: DualCache,
post_data: dict[str, str] | None = None,
post_data: Mapping[str, str] | None = None,
) -> "OneLogin_Saml2_Auth":
if not SAML_AVAILABLE:
raise _saml_unavailable_error()
@ -310,7 +315,9 @@ class SAMLAuthHandler:
return dict(parse_qsl(body.decode("utf-8", "replace")))
@staticmethod
async def handle_acs(request: Request, cache: DualCache, post_data: dict[str, str]) -> CustomOpenID:
async def handle_acs(
request: Request, cache: DualCache, post_data: Mapping[str, str]
) -> tuple[CustomOpenID, SSOIdentityAssertion | None]:
auth: Final = await SAMLAuthHandler._build_auth(request, cache, post_data=post_data)
browser_request_id: Final = request.cookies.get(_SAML_AUTHN_STATE_COOKIE)
try:
@ -330,7 +337,50 @@ class SAMLAuthHandler:
)
await SAMLAuthHandler._enforce_response_binding(auth, cache, browser_request_id)
return SAMLAuthHandler._result_from_auth(auth)
return SAMLAuthHandler._result_from_auth(auth), SAMLAuthHandler._retained_assertion(auth)
@staticmethod
def _retained_assertion(auth: "OneLogin_Saml2_Auth") -> SSOIdentityAssertion | None:
"""The one signed ``<saml:Assertion>`` the validated response carried, serialized on its own
so its enveloped signature still verifies at the IdP's token endpoint. Retention is optional
egress material, so a failure here is logged and the validated login goes on without it."""
try:
return SAMLAuthHandler._parse_retained_assertion(auth)
except Exception as e: # noqa: BLE001 - lxml and python3-saml expose no common base; capture must not block login
verbose_proxy_logger.warning(
"SAML login validated but its assertion could not be retained for oauth2_id_jag MCP servers: %s", e
)
return None
@staticmethod
def _parse_retained_assertion(auth: "OneLogin_Saml2_Auth") -> SSOIdentityAssertion | None:
xml: Final = cast(bytes | None, auth.get_last_response_xml()) # cast-ok: untyped python3-saml
if not xml:
verbose_proxy_logger.warning(
"SAML login validated but the response XML was not retained, so this user has no identity "
"assertion for oauth2_id_jag MCP servers."
)
return None
root: Final = OneLogin_Saml2_XML.to_etree(xml)
assertions: Final = OneLogin_Saml2_XML.query(root, "//saml:Assertion")
if len(assertions) != 1:
encrypted: Final = OneLogin_Saml2_XML.query(root, "//saml:EncryptedAssertion")
verbose_proxy_logger.warning(
"SAML login validated but the response carried %d plaintext saml:Assertion elements and %d "
"saml:EncryptedAssertion elements, so no identity assertion was retained for oauth2_id_jag MCP "
"servers. Those servers need exactly one signed, unencrypted assertion from the IdP.",
len(assertions),
len(encrypted),
)
return None
issuers: Final = OneLogin_Saml2_XML.query(assertions[0], "./saml:Issuer")
issuer: Final = cast(str | None, issuers[0].text) if issuers else None # cast-ok: untyped python3-saml
not_on_or_after: Final = auth.get_last_assertion_not_on_or_after()
return assertion_from_saml_login(
assertion_xml=cast(bytes, OneLogin_Saml2_XML.to_string(assertions[0], encoding="utf-8")), # cast-ok: saml
not_on_or_after=not_on_or_after if isinstance(not_on_or_after, int) else None,
issuer=issuer,
)
@staticmethod
def _replay_guard_ttl(auth: "OneLogin_Saml2_Auth") -> int:

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

@ -37,10 +37,12 @@ from urllib.parse import parse_qs, urlencode, urlparse
if TYPE_CHECKING:
import httpx
from litellm.caching.redis_cache import RedisCache
import jwt
from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status
from fastapi.responses import RedirectResponse
from pydantic import ConfigDict, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, field_validator
import litellm
from litellm._internal_context import with_service_target
@ -69,7 +71,12 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.login_assertion_exchange import (
exchange_login_assertion_if_configured,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
GOOGLE_REFRESH_TOKEN_PARAMS,
MICROSOFT_REFRESH_TOKEN_SCOPE,
SSOIdentityAssertion,
assertion_from_sso_login,
ema_assertion_retention_enabled,
@ -1791,6 +1798,19 @@ async def warn_if_id_jag_assertion_uncaptured(
)
async def login_asks_for_refresh_token(retention_enabled: RetentionCheck | None = None) -> bool:
"""A refresh token is only issued when the authorization request asks for one, and asking costs the
user an offline-access consent, so the Google and Microsoft logins ask only while an ``oauth2_id_jag``
server exists whose assertion the refresher would renew. A store outage answers no: a login must not
fail over a renewal that is only an optimisation on top of signing in again."""
try:
check: Final = retention_enabled if retention_enabled is not None else ema_assertion_retention_enabled
return await check()
except Exception as exc: # noqa: BLE001 # the login must not fail over the renewal gate
verbose_proxy_logger.debug("Could not check for oauth2_id_jag MCP servers before SSO login: %s", exc)
return False
async def warn_if_id_jag_capture_gap(*, retention_enabled: RetentionCheck | None = None) -> None:
gap: Final = id_jag_assertion_capture_gap()
if gap is None:
@ -2130,6 +2150,50 @@ async def check_and_update_if_proxy_admin_id(user_role: str, user_id: str, prism
return user_role
async def _sso_provider_callback(
request: Request,
redirect_url: str,
*,
google_client_id: str | None,
microsoft_client_id: str | None,
generic_client_id: str | None,
jwt_handler: JWTHandler,
sso_jwt_handler: JWTHandler | None,
return_raw_sso_response: bool = False,
) -> tuple[
OpenID | dict[str, object] | None, dict[str, object] | None, dict[str, object] | None, SSOIdentityAssertion | None
]:
"""Dispatch the callback to the configured provider in the order the login redirect used (Google,
Microsoft, generic) and hand back what every branch has to agree on: the parsed identity, the raw
userinfo and access-token claims only the generic branch exposes, and the IdP identity assertion
an ``oauth2_id_jag`` MCP server spends. ``None`` identity means no provider is configured."""
if google_client_id is not None:
google_result, google_assertion = await GoogleSSOHandler.get_google_callback_response(
request=request,
google_client_id=google_client_id,
redirect_url=redirect_url,
return_raw_sso_response=return_raw_sso_response,
)
return google_result, None, None, google_assertion
if microsoft_client_id is not None:
microsoft_result, microsoft_assertion = await MicrosoftSSOHandler.get_microsoft_callback_response(
request=request,
microsoft_client_id=microsoft_client_id,
redirect_url=redirect_url,
return_raw_sso_response=return_raw_sso_response,
)
return microsoft_result, None, None, microsoft_assertion
if generic_client_id is not None:
return await get_generic_sso_response(
request=request,
jwt_handler=jwt_handler,
generic_client_id=generic_client_id,
redirect_url=redirect_url,
sso_jwt_handler=sso_jwt_handler,
)
return None, None, None, None
@router.get("/sso/callback", tags=["experimental"], include_in_schema=False)
async def auth_callback(request: Request, state: str | None = None):
"""Verify login"""
@ -2178,9 +2242,6 @@ async def auth_callback(request: Request, state: str | None = None):
microsoft_client_id: Final = os.getenv("MICROSOFT_CLIENT_ID", None)
google_client_id: Final = os.getenv("GOOGLE_CLIENT_ID", None)
generic_client_id: Final = os.getenv("GENERIC_CLIENT_ID", None)
received_response: dict | None = None
access_token_payload: dict | None = None
sso_assertion: SSOIdentityAssertion | None = None
# get url from request
if master_key is None:
raise ProxyException(
@ -2192,33 +2253,15 @@ async def auth_callback(request: Request, state: str | None = None):
redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(request=request, sso_callback_route="sso/callback")
verbose_proxy_logger.info("Redirecting to %s", redirect_url)
result = None
if google_client_id is not None:
result = await GoogleSSOHandler.get_google_callback_response(
request=request,
google_client_id=google_client_id,
redirect_url=redirect_url,
)
elif microsoft_client_id is not None:
result = await MicrosoftSSOHandler.get_microsoft_callback_response(
request=request,
microsoft_client_id=microsoft_client_id,
redirect_url=redirect_url,
)
elif generic_client_id is not None:
(
result,
received_response,
access_token_payload,
sso_assertion,
) = await get_generic_sso_response(
request=request,
jwt_handler=jwt_handler,
generic_client_id=generic_client_id,
redirect_url=redirect_url,
sso_jwt_handler=sso_jwt_handler,
)
result, received_response, access_token_payload, sso_assertion = await _sso_provider_callback(
request,
redirect_url,
google_client_id=google_client_id,
microsoft_client_id=microsoft_client_id,
generic_client_id=generic_client_id,
jwt_handler=jwt_handler,
sso_jwt_handler=sso_jwt_handler,
)
if result is None:
raise HTTPException(
@ -2308,7 +2351,7 @@ async def saml_callback(request: Request):
if "SAMLResponse" not in post_data:
raise HTTPException(status_code=400, detail="Missing SAMLResponse in callback request.")
result: Final = await SAMLAuthHandler.handle_acs(request=request, cache=saml_sso_cache, post_data=post_data)
user, retained = await SAMLAuthHandler.handle_acs(request=request, cache=saml_sso_cache, post_data=post_data)
await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client)
@ -2321,7 +2364,7 @@ async def saml_callback(request: Request):
)
return await SSOAuthenticationHandler.get_redirect_response_from_openid(
result=result,
result=user,
request=request,
received_response=None,
generic_client_id=None,
@ -2329,6 +2372,7 @@ async def saml_callback(request: Request):
access_token_payload=None,
jwt_handler=jwt_handler,
return_to=cp_return_to,
sso_assertion=retained,
)
@ -2445,7 +2489,8 @@ async def _complete_cli_sso_callback_session(
raise HTTPException(status_code=500, detail="Failed to retrieve user information from SSO")
resolved_user_id: Final = _require_sso_user_id(user_info.user_id)
await retain_sso_identity_assertion_for_ema(user_id=resolved_user_id, assertion=sso_assertion)
retained_assertion: Final = await exchange_login_assertion_if_configured(sso_assertion)
await retain_sso_identity_assertion_for_ema(user_id=resolved_user_id, assertion=retained_assertion)
await warn_if_id_jag_assertion_uncaptured(sso_assertion)
teams: list[str] = []
@ -2966,8 +3011,8 @@ def _is_same_origin_return_path(return_to: str) -> bool:
async def sso_return_to_redirect(
return_to: str | None,
jwt_token: str,
redis_usage_cache,
user_api_key_cache,
redis_usage_cache: "RedisCache | None",
user_api_key_cache: UserApiKeyCache,
request: Request,
) -> RedirectResponse | None:
"""Resolve the post-SSO redirect for a ``return_to``, or None to fall through to the dashboard.
@ -3114,6 +3159,7 @@ class SSOAuthenticationHandler:
generic_client_id: str | None = None,
state: str | None = None,
request: Request | None = None,
retention_enabled: RetentionCheck | None = None,
) -> RedirectResponse | None:
"""
Step 1. Call Get Login Redirect for the SSO provider. Send the redirect response to `redirect_url`
@ -3125,10 +3171,13 @@ class SSOAuthenticationHandler:
generic_client_id (Optional[str], optional): The Generic Client ID. Defaults to None.
request: Optional FastAPI request, used to drive the ``Secure``
attribute on the ``litellm_oauth_state`` CSRF cookie.
retention_enabled: Whether an ``oauth2_id_jag`` server exists, deciding if the Google and
Microsoft logins ask the IdP for a refresh token. Defaults to the store's own check.
Returns:
RedirectResponse: The redirect response from the SSO provider.
"""
asks_for_refresh_token: Final = await login_asks_for_refresh_token(retention_enabled)
# Google SSO Auth
if google_client_id is not None:
google_client_secret: Final = os.getenv("GOOGLE_CLIENT_SECRET", None)
@ -3150,7 +3199,10 @@ class SSOAuthenticationHandler:
google_client_id,
)
with google_sso:
return await google_sso.get_login_redirect(state=state)
return await google_sso.get_login_redirect(
state=state,
params=dict(GOOGLE_REFRESH_TOKEN_PARAMS) if asks_for_refresh_token else None,
)
# Microsoft SSO Auth
elif microsoft_client_id is not None:
microsoft_client_secret: Final = os.getenv("MICROSOFT_CLIENT_SECRET", None)
@ -3168,6 +3220,7 @@ class SSOAuthenticationHandler:
tenant=microsoft_tenant,
redirect_uri=redirect_url,
allow_insecure_http=True,
scope=[*CustomMicrosoftSSO.scope, MICROSOFT_REFRESH_TOKEN_SCOPE] if asks_for_refresh_token else None,
)
with microsoft_sso:
return await microsoft_sso.get_login_redirect(state=state)
@ -3867,7 +3920,8 @@ class SSOAuthenticationHandler:
)
if isinstance(user_id, str) and user_id:
await retain_sso_identity_assertion_for_ema(user_id=user_id, assertion=sso_assertion)
retained_assertion: Final = await exchange_login_assertion_if_configured(sso_assertion)
await retain_sso_identity_assertion_for_ema(user_id=user_id, assertion=retained_assertion)
await warn_if_id_jag_assertion_uncaptured(sso_assertion)
disabled_non_admin_personal_key_creation: Final = get_disabled_non_admin_personal_key_creation()
@ -4451,6 +4505,21 @@ class SSOAuthenticationHandler:
return userinfo
class _SSOTokenResponse(BaseModel):
"""The id_token and refresh token fastapi-sso keeps on the SSO instance after the code exchange,
read through a validated view because the package ships no type information."""
model_config = ConfigDict(frozen=True, from_attributes=True)
id_token: str | None = None
refresh_token: str | None = None
@field_validator("id_token", "refresh_token", mode="before")
@classmethod
def _strings_only(cls, value: object) -> str | None:
return value if isinstance(value, str) else None
class MicrosoftSSOHandler:
"""
Handles Microsoft SSO callback response and returns a CustomOpenID object
@ -4485,9 +4554,10 @@ class MicrosoftSSOHandler:
microsoft_client_id: str,
redirect_url: str,
return_raw_sso_response: bool = False,
) -> CustomOpenID | OpenID | dict:
) -> tuple[CustomOpenID | OpenID | dict[str, object], SSOIdentityAssertion | None]:
"""
Get the Microsoft SSO callback response
Get the Microsoft SSO callback response, with the IdP identity assertion the token
response carried (the id_token and refresh token an ``oauth2_id_jag`` server spends)
Args:
return_raw_sso_response: If True, return the raw SSO response
@ -4522,6 +4592,8 @@ class MicrosoftSSOHandler:
)
or {}
)
microsoft_tokens: Final = _SSOTokenResponse.model_validate(microsoft_sso)
sso_assertion: Final = assertion_from_sso_login(microsoft_tokens.id_token, microsoft_tokens.refresh_token)
if not return_raw_sso_response and "error" in original_msft_result and original_msft_result["error"]:
verbose_proxy_logger.warning("Microsoft SSO login failed: Graph /me returned an error instead of a profile")
raise HTTPException(
@ -4546,7 +4618,7 @@ class MicrosoftSSOHandler:
if return_raw_sso_response:
original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = user_team_ids
original_msft_result["app_roles"] = app_roles
return original_msft_result or {}
return original_msft_result or {}, sso_assertion
from litellm.proxy.management_endpoints.sso.agent_subject_enrollment import microsoft_interactive_subject
@ -4569,7 +4641,7 @@ class MicrosoftSSOHandler:
team_ids=user_team_ids,
user_role=user_role,
)
return result
return result, sso_assertion
@staticmethod
def openid_from_response(
@ -4845,9 +4917,10 @@ class GoogleSSOHandler:
google_client_id: str,
redirect_url: str,
return_raw_sso_response: bool = False,
) -> OpenID | dict:
) -> tuple[OpenID | dict[str, object], SSOIdentityAssertion | None]:
"""
Get the Google SSO callback response
Get the Google SSO callback response, with the IdP identity assertion the token
response carried (the id_token and refresh token an ``oauth2_id_jag`` server spends)
Args:
return_raw_sso_response: If True, return the raw SSO response
@ -4868,17 +4941,13 @@ class GoogleSSOHandler:
)
# if user is trying to get the raw sso response for debugging, return the raw sso response
if return_raw_sso_response:
return (
await google_sso.verify_and_process(
request=request,
convert_response=False,
)
or {}
)
result: Final = await google_sso.verify_and_process(request)
return result or {}
result: Final = (
await google_sso.verify_and_process(request=request, convert_response=False)
if return_raw_sso_response
else await google_sso.verify_and_process(request)
)
google_tokens: Final = _SSOTokenResponse.model_validate(google_sso)
return result or {}, assertion_from_sso_login(google_tokens.id_token, google_tokens.refresh_token)
def _raise_if_sso_debug_disabled() -> None:
@ -4980,37 +5049,16 @@ async def debug_sso_callback(request: Request):
else:
redirect_url += "/sso/debug/callback"
result = None
received_response: dict | None = None
access_token_payload: dict | None = None
if google_client_id is not None:
result = await GoogleSSOHandler.get_google_callback_response(
request=request,
google_client_id=google_client_id,
redirect_url=redirect_url,
return_raw_sso_response=True,
)
elif microsoft_client_id is not None:
result = await MicrosoftSSOHandler.get_microsoft_callback_response(
request=request,
microsoft_client_id=microsoft_client_id,
redirect_url=redirect_url,
return_raw_sso_response=True,
)
elif generic_client_id is not None:
(
result,
received_response,
access_token_payload,
_sso_assertion,
) = await get_generic_sso_response(
request=request,
jwt_handler=jwt_handler,
generic_client_id=generic_client_id,
redirect_url=redirect_url,
sso_jwt_handler=sso_jwt_handler,
)
result, received_response, access_token_payload, _sso_assertion = await _sso_provider_callback(
request,
redirect_url,
google_client_id=google_client_id,
microsoft_client_id=microsoft_client_id,
generic_client_id=generic_client_id,
jwt_handler=jwt_handler,
sso_jwt_handler=sso_jwt_handler,
return_raw_sso_response=True,
)
# If result is None, return a basic error message
if result is None:

9
litellm/types/mcp.py
View file

@ -248,6 +248,15 @@ class MCPCredentials(TypedDict, total=False):
Optional RFC 8707 resource indicator sent on ID-JAG leg 1
"""
id_jag_subject_source: ReadOnly[str | None]
"""
Which token ID-JAG leg 1 presents as its RFC 8693 subject_token. ``login_assertion`` (the default) is the
identity assertion captured at SSO login, an OIDC id_token or a SAML assertion. ``refresh_token`` exchanges
that assertion at the IdP token endpoint once, at login, for a refresh token and presents the refresh token
on every leg 1 instead, which is how a SAML requesting app works at IdPs whose ID-JAG exchange does not take
a saml2 subject directly.
"""
id_jag_resource_client_id: ReadOnly[str | None]
"""OAuth client identifier used at the ID-JAG resource authorization server."""

1
litellm/types/mcp_server/mcp_server_manager.py
View file

@ -182,6 +182,7 @@ class MCPServer(LiteLLMBaseModel):
# posts the ID-JAG assertion to id_jag_resource_token_endpoint.
id_jag_resource_token_endpoint: str | None = None
id_jag_resource: str | None = None
id_jag_subject_source: str | None = None
id_jag_resource_client_id: str | None = Field(
default=None,
description="OAuth client identifier used at the ID-JAG resource authorization server.",

117
tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -4,6 +4,7 @@ import os
from collections.abc import Sequence
from collections.abc import Set as AbstractSet
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
@ -12,6 +13,7 @@ from fastapi import HTTPException
from fastapi.testclient import TestClient
from starlette.datastructures import Headers
from litellm.caching.caching import DualCache
from litellm.constants import MCP_ALL_TOOLS_WILDCARD
from litellm.models.access_group import LiteLLM_AccessGroupTable
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
@ -23,6 +25,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
from litellm.proxy._types import (
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamTable,
LiteLLM_UserTable,
ProxyException,
SpecialHeaders,
SpecialMCPServerNames,
@ -11549,3 +11552,117 @@ class TestAccessGroupToolOwnershipAndFaults:
patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups, missing={"typo-group"})),
):
assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"]
class _EmptyTables:
"""Storage edge for a session credential's catalog refresh: the user row the real
``reload_admitted_user`` reads, and every other table (keys, MCP servers, config) empty, so the
real ``_reload_admitted_key`` fails closed, since a session credential has no key row."""
def __init__(self, live_user: LiteLLM_UserTable) -> None:
self.litellm_usertable = SimpleNamespace(find_unique=AsyncMock(return_value=live_user))
def __getattr__(self, table: str) -> SimpleNamespace:
empty: Final = SimpleNamespace(
find_unique=AsyncMock(return_value=None),
find_first=AsyncMock(return_value=None),
find_many=AsyncMock(return_value=[]),
count=AsyncMock(return_value=0),
)
setattr(self, table, empty)
return empty
def _session_credential_storage(live_user: LiteLLM_UserTable) -> SimpleNamespace:
tables: Final = _EmptyTables(live_user)
return SimpleNamespace(db=tables, writer_db=tables)
@pytest.mark.asyncio
async def test_catalog_refresh_reloads_session_credential_through_its_user_and_keeps_token_grants(monkeypatch):
"""A ``/token``-minted or ``lite login`` credential has no key row: the catalog refresh must read the live
user policy, never the key table, and keep the team, routes and models sealed in the token."""
from litellm.proxy import proxy_server
database = _session_credential_storage(
LiteLLM_UserTable(user_id="session-user", organization_id="live-org", user_role="internal_user", teams=[])
)
monkeypatch.setattr(proxy_server, "prisma_client", database)
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
caller = UserAPIKeyAuth(
api_key="session-token-hash",
user_id="session-user",
team_id="token-team",
org_id="stale-org",
user_role="proxy_admin",
models=["token-model"],
allowed_routes=["mcp_inference_routes"],
is_session_token=True,
)
caller.via_virtual_key = True
refreshed = await MCPRequestHandler.refresh_catalog_authority(caller)
assert database.writer_db.litellm_usertable.find_unique.await_count >= 1
database.writer_db.litellm_verificationtoken.find_unique.assert_not_awaited()
assert (refreshed.object_permission, refreshed.object_permission_id) == (None, None)
assert (refreshed.user_role, refreshed.org_id) == ("internal_user", "live-org")
assert (refreshed.team_id, refreshed.models, refreshed.allowed_routes) == (
"token-team",
["token-model"],
["mcp_inference_routes"],
)
assert refreshed.is_session_token and refreshed.requires_fresh_policy
assert caller.org_id == "stale-org" and not caller.requires_fresh_policy
@pytest.mark.asyncio
@pytest.mark.parametrize("scope", ["litellm:mcp litellm:llm", "litellm:mcp", None])
async def test_token_minted_credential_lists_tools_for_every_scope_shape(monkeypatch, scope):
"""``tools/list`` with the credential ``/token`` mints (a session token the shared auth marks
``via_virtual_key`` at the key-cache boundary) returns the tool list instead of the JSON-RPC
credential error, for the combined, MCP-only and unscoped ID-JAG shapes."""
from mcp.types import ListToolsRequest, ListToolsResult, Tool
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import operations
from litellm.proxy._experimental.mcp_server.id_jag_grant import SubjectTokenRefusal, granted_scope
from litellm.proxy._experimental.mcp_server.operations import GatewayOperations, prepare_context
from litellm.proxy._types import MCPTransport
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.types.mcp_server.mcp_server_manager import MCPServer
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-session-credential-tools-list")
granted = granted_scope(scope, mcp_resource=False)
assert not isinstance(granted, SubjectTokenRefusal)
user = LiteLLM_UserTable(user_id="okta-user-1", user_email="okta@example.com", user_role="internal_user")
minted = ExperimentalUIJWTToken.get_cli_jwt_auth_token(
user, team_id="team-eng", team_alias="eng", allowed_routes=granted.allowed_routes
)
credential = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(minted)
assert credential is not None and credential.is_session_token
credential.api_key = "hash-of-" + minted[:16]
credential.via_virtual_key = True
server = MCPServer(
server_id="eng-server",
name="eng_server",
server_name="eng_server",
alias="eng_server",
url="https://upstream.example.com/mcp",
transport=MCPTransport.http,
)
database = _session_credential_storage(user)
monkeypatch.setattr(proxy_server, "prisma_client", database)
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
tool = Tool(name="eng_server-list_repos", description="list repos", inputSchema={"type": "object"})
with (
patch.object(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])),
patch.object(
operations.global_mcp_server_manager,
"get_tools_page",
AsyncMock(return_value=ListToolsResult(tools=[tool])),
),
):
result = await GatewayOperations().execute(ListToolsRequest(), prepare_context(user_api_key_auth=credential))
assert [listed.name for listed in result.tools] == ["eng_server-list_repos"]
assert database.writer_db.litellm_usertable.find_unique.await_count >= 1
database.writer_db.litellm_verificationtoken.find_unique.assert_not_awaited()
assert credential.team_id == "team-eng" and credential.allowed_routes == list(granted.allowed_routes or ())

23
tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py
View file

@ -31,6 +31,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
ClientSecretAuth,
CredError,
IdJagConfig,
IdJagSubjectSource,
NoneConfig,
PassthroughConfig,
PrivateKeyJwtAuth,
@ -908,3 +909,25 @@ def test_incomplete_obo_keeps_exchange_ownership(client_secret: str | None, is_b
assert isinstance(spec.config, TokenExchangeConfig)
assert spec.config.client_id == "client"
assert spec.config.client_secret is None
assert spec.config.client_id == "client"
assert spec.config.client_secret is None
@pytest.mark.parametrize(
("configured", "expected"),
[
(None, IdJagSubjectSource.login_assertion),
("login_assertion", IdJagSubjectSource.login_assertion),
("refresh_token", IdJagSubjectSource.refresh_token),
],
)
def test_id_jag_subject_source_maps_to_config(configured: str | None, expected: IdJagSubjectSource) -> None:
spec = to_server_spec(_id_jag_server(id_jag_subject_source=configured))
assert spec is not None
assert isinstance(spec.config, IdJagConfig)
assert spec.config.subject_source is expected
def test_id_jag_unknown_subject_source_defers_like_a_half_configured_server() -> None:
assert to_server_spec(_id_jag_server(id_jag_subject_source="assertion")) is None

412
tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py Normal file
View file

@ -0,0 +1,412 @@
"""The login-time exchange of a captured identity assertion for an IdP refresh token."""
import asyncio
import logging
from datetime import datetime, timedelta, timezone
import pytest
from pydantic import SecretStr
from litellm.proxy._experimental.mcp_server.outbound_credentials.login_assertion_exchange import (
exchange_login_assertion_for_refresh_token,
login_exchange_servers,
refresh_token_exchange_targets,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import SSOClientConfig
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
SAML2_SUBJECT_TOKEN_TYPE,
SSOIdentityAssertion,
assertion_from_saml_login,
login_refresh_token_expired,
with_login_refresh_token,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ExchangedToken
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
ClientAuth,
ClientSecretAuth,
CredError,
IdJagConfig,
IdJagResourceClient,
IdJagSubjectSource,
NoneConfig,
ServerSpec,
)
from litellm.proxy._types import LiteLLM_MCPServerTable
from litellm.types.mcp import MCPAuth, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
_NOW = datetime(2026, 10, 16, 12, 0, tzinfo=timezone.utc)
def _config(
*, source: IdJagSubjectSource = IdJagSubjectSource.refresh_token, client_id: str = "agent", endpoint: str = ""
) -> IdJagConfig:
return IdJagConfig(
org_token_endpoint=endpoint or "https://idp.example.com/token",
resource_token_endpoint="https://mcp-as.example.com/token",
client_id=client_id,
client_auth=ClientSecretAuth(client_secret=SecretStr("s")),
resource_client=IdJagResourceClient(
client_id=client_id, client_auth=ClientSecretAuth(client_secret=SecretStr("s"))
),
subject_source=source,
)
def _spec(config: IdJagConfig | NoneConfig, server_id: str = "srv") -> ServerSpec:
return ServerSpec(server_id=server_id, resource="https://mcp.example.com", config=config)
def _saml() -> SSOIdentityAssertion:
return assertion_from_saml_login(b"<saml:Assertion/>", int(_NOW.timestamp()) + 300, "https://idp.example.com")
class _Endpoint:
def __init__(self, result: Result[ExchangedToken, CredError]) -> None:
self._result = result
self.calls: list[tuple[str, str, dict[str, str], ClientAuth]] = []
async def fetch(
self, endpoint: str, client_id: str, grant_params: dict[str, str], client_auth: ClientAuth
) -> Result[ExchangedToken, CredError]:
self.calls.append((endpoint, client_id, dict(grant_params), client_auth))
return self._result
def test_targets_keep_only_refresh_token_sourced_id_jag_servers_one_per_idp_client():
direct = _config(source=IdJagSubjectSource.login_assertion)
first = _config(client_id="agent")
twin = _config(client_id="agent")
other = _config(client_id="other-agent")
targets = refresh_token_exchange_targets(
[_spec(direct, "a"), None, _spec(NoneConfig(), "b"), _spec(first, "c"), _spec(other, "d"), _spec(twin, "e")]
)
assert targets == (first, other)
@pytest.mark.asyncio
async def test_exchange_presents_the_captured_assertion_and_keeps_the_refresh_token():
endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600)))
config = _config()
saml = _saml()
exchanged = await exchange_login_assertion_for_refresh_token(saml, (config,), endpoint, _NOW)
((url, client_id, params, auth),) = endpoint.calls
assert (url, client_id, auth) == (config.org_token_endpoint, config.client_id, config.client_auth)
assert params == {
"grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
"requested_token_type": "urn:ietf:params:oauth:token-type:refresh_token",
"subject_token": saml.id_token.get_secret_value(),
"subject_token_type": SAML2_SUBJECT_TOKEN_TYPE,
"scope": "openid offline_access",
}
assert exchanged.exchanged_refresh_token is not None
assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1"
assert exchanged.exchanged_refresh_token_endpoint == "https://idp.example.com/token"
assert exchanged.exchanged_refresh_token_client_id == "agent"
assert exchanged.exchanged_refresh_token_expires_at == _NOW + timedelta(seconds=3600)
assert exchanged.id_token == saml.id_token
assert exchanged.expires_at == saml.expires_at
assert exchanged.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE
assert not login_refresh_token_expired(exchanged, _NOW + timedelta(seconds=3599))
assert login_refresh_token_expired(exchanged, _NOW + timedelta(seconds=3600))
@pytest.mark.asyncio
async def test_exchange_without_a_lifetime_leaves_expiry_to_the_idp():
endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=None)))
exchanged = await exchange_login_assertion_for_refresh_token(_saml(), (_config(),), endpoint, _NOW)
assert exchanged.exchanged_refresh_token_expires_at is None
assert not login_refresh_token_expired(exchanged, _NOW + timedelta(days=365))
@pytest.mark.asyncio
async def test_a_refused_exchange_keeps_the_plain_assertion(caplog: pytest.LogCaptureFixture):
endpoint = _Endpoint(Error(CredError.of_upstream_unavailable("token exchange failed with status 400")))
saml = _saml()
with caplog.at_level("WARNING"):
exchanged = await exchange_login_assertion_for_refresh_token(saml, (_config(),), endpoint, _NOW)
assert exchanged == saml
assert exchanged.exchanged_refresh_token is None
assert any("refused to exchange the login assertion" in message for message in caplog.messages)
@pytest.mark.asyncio
async def test_no_exchange_without_a_refresh_token_sourced_server():
endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600)))
saml = _saml()
assert await exchange_login_assertion_for_refresh_token(saml, (), endpoint, _NOW) == saml
assert endpoint.calls == []
@pytest.mark.asyncio
async def test_an_oidc_login_is_exchanged_at_the_id_jag_client_and_keeps_its_sso_refresh_token():
"""The SSO refresh token belongs to the SSO client and cannot be redeemed by the ID-JAG client, so
an OIDC login still exchanges its id_token and both tokens live side by side."""
endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600)))
oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt"))
exchanged = await exchange_login_assertion_for_refresh_token(oidc, (_config(),), endpoint, _NOW)
assert endpoint.calls[0][2]["subject_token"] == "id-token"
assert endpoint.calls[0][2]["subject_token_type"] == "urn:ietf:params:oauth:token-type:id_token"
assert exchanged.refresh_token is not None and exchanged.refresh_token.get_secret_value() == "sso-rt"
assert exchanged.exchanged_refresh_token is not None
assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1"
def _sso_client(client_id: str = "agent", endpoint: str = "https://idp.example.com/token") -> SSOClientConfig:
return SSOClientConfig(
token_endpoint=endpoint, client_id=client_id, client_secret=SecretStr("s"), auth_method="client_secret_basic"
)
@pytest.mark.asyncio
async def test_an_oidc_login_through_the_id_jag_client_keeps_its_own_refresh_token_without_an_exchange():
"""Okta refuses requested_token_type=refresh_token for an id_token subject. When the SSO app is the
ID-JAG client itself, the refresh token the login returned is already that client's, so it is bound
to the target and nothing is posted to the IdP."""
endpoint = _Endpoint(
Error(
CredError.of_upstream_unavailable(
"token exchange failed with status 400: 'requested_token_type' is invalid or not supported."
)
)
)
oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt"))
kept = await exchange_login_assertion_for_refresh_token(
oidc, (_config(),), endpoint, _NOW, sso_client=_sso_client()
)
assert endpoint.calls == []
assert kept.exchanged_refresh_token is not None and kept.exchanged_refresh_token.get_secret_value() == "sso-rt"
assert kept.exchanged_refresh_token_expires_at is None
assert kept.exchanged_refresh_token_endpoint == "https://idp.example.com/token"
assert kept.exchanged_refresh_token_client_id == "agent"
assert kept.id_token == oidc.id_token and kept.refresh_token == oidc.refresh_token
@pytest.mark.asyncio
@pytest.mark.parametrize(
"sso_client",
[
pytest.param(_sso_client(client_id="sso-app"), id="other-client"),
pytest.param(_sso_client(endpoint="https://other-idp.example.com/token"), id="other-endpoint"),
pytest.param(None, id="unknown-client"),
],
)
async def test_a_login_refresh_token_from_another_client_or_issuer_is_never_presented(
sso_client: SSOClientConfig | None,
):
endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600)))
oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt"))
exchanged = await exchange_login_assertion_for_refresh_token(
oidc, (_config(),), endpoint, _NOW, sso_client=sso_client
)
assert [(call[0], call[1], call[2]["subject_token"]) for call in endpoint.calls] == [
("https://idp.example.com/token", "agent", "id-token")
]
assert exchanged.exchanged_refresh_token is not None
assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1"
@pytest.mark.asyncio
async def test_the_login_refresh_token_binds_to_the_target_whose_client_issued_it_not_the_first():
endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600)))
oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt"))
targets = (_config(client_id="other-agent"), _config(client_id="agent"))
kept = await exchange_login_assertion_for_refresh_token(oidc, targets, endpoint, _NOW, sso_client=_sso_client())
assert endpoint.calls == []
assert kept.exchanged_refresh_token_client_id == "agent"
assert kept.exchanged_refresh_token is not None and kept.exchanged_refresh_token.get_secret_value() == "sso-rt"
@pytest.mark.asyncio
async def test_a_saml_login_through_the_id_jag_client_still_exchanges_its_assertion():
endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600)))
exchanged = await exchange_login_assertion_for_refresh_token(
_saml(), (_config(),), endpoint, _NOW, sso_client=_sso_client()
)
assert endpoint.calls[0][2]["subject_token_type"] == "urn:ietf:params:oauth:token-type:saml2"
assert exchanged.exchanged_refresh_token is not None
assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1"
@pytest.mark.asyncio
async def test_a_login_already_exchanged_is_not_exchanged_again():
endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-2", expires_in=3600)))
done = with_login_refresh_token(_saml(), "rt-1", 3600, _NOW)
assert await exchange_login_assertion_for_refresh_token(done, (_config(),), endpoint, _NOW) == done
assert endpoint.calls == []
@pytest.mark.asyncio
async def test_several_idp_clients_mint_for_the_first_and_warn(caplog: pytest.LogCaptureFixture):
endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600)))
first, second = _config(client_id="agent"), _config(client_id="other-agent")
with caplog.at_level("WARNING"):
exchanged = await exchange_login_assertion_for_refresh_token(_saml(), (first, second), endpoint, _NOW)
assert [call[1] for call in endpoint.calls] == ["agent"]
assert exchanged.exchanged_refresh_token is not None
assert any("2 IdP clients" in message and "client_id=agent" in message for message in caplog.messages)
class _FakeCatalog:
def __init__(self, config_servers: dict[str, MCPServer]) -> None:
self.config_mcp_servers = config_servers
self.built: list[tuple[str, bool]] = []
async def build_mcp_server_from_table(
self,
mcp_server: LiteLLM_MCPServerTable,
*,
credentials_are_encrypted: bool = True,
env_vars_are_encrypted: bool | None = None,
register_oauth_discovery: bool = True,
) -> MCPServer:
self.built.append((mcp_server.server_id, register_oauth_discovery))
return MCPServer(
server_id=mcp_server.server_id, name=mcp_server.server_id, url=mcp_server.url, transport=MCPTransport.http
)
def _row(server_id: str, auth_type: MCPAuth, approval_status: str | None = None) -> LiteLLM_MCPServerTable:
return LiteLLM_MCPServerTable(
server_id=server_id,
alias=server_id,
url=f"https://{server_id}.example.com/mcp",
transport=MCPTransport.http,
auth_type=auth_type,
approval_status=approval_status,
created_at=_NOW,
updated_at=_NOW,
)
@pytest.mark.asyncio
async def test_login_exchange_servers_reads_config_and_db_id_jag_rows_not_the_registry():
"""A server added on another pod is in the DB before this pod's registry reloads, so the login
exchange has to judge against the DB rows, like the retention gate, and only the ID-JAG ones."""
config_server = MCPServer(
server_id="from-config", name="from-config", url="https://cfg.example.com/mcp", transport=MCPTransport.http
)
catalog = _FakeCatalog({"from-config": config_server})
servers = await login_exchange_servers(
catalog,
[
_row("db-id-jag", MCPAuth.oauth2_id_jag),
_row("db-id-jag-approved", MCPAuth.oauth2_id_jag, "approved"),
_row("db-oauth", MCPAuth.oauth2),
],
)
assert [server.server_id for server in servers] == ["from-config", "db-id-jag", "db-id-jag-approved"]
assert catalog.built == [("db-id-jag", False), ("db-id-jag-approved", False)]
@pytest.mark.asyncio
async def test_a_db_row_hides_the_config_server_with_the_same_id_like_the_live_catalog():
"""The catalog serves ``config_mcp_servers | registry``, so a DB row wins its server_id. The login
exchange must mint for that live server, not for the hidden config entry's IdP client."""
hidden = MCPServer(
server_id="shared", name="shared", url="https://cfg.example.com/mcp", transport=MCPTransport.http
)
catalog = _FakeCatalog({"shared": hidden})
servers = await login_exchange_servers(catalog, [_row("shared", MCPAuth.oauth2_id_jag)])
assert [(server.server_id, server.url) for server in servers] == [("shared", "https://shared.example.com/mcp")]
@pytest.mark.asyncio
@pytest.mark.parametrize("approval_status", ["pending_review", "rejected", "draft"])
async def test_login_exchange_servers_skip_unapproved_submissions(approval_status: str):
"""A team key can submit an ID-JAG server naming any token endpoint. Until an admin approves it,
no login assertion is posted there."""
catalog = _FakeCatalog({})
servers = await login_exchange_servers(catalog, [_row("submitted", MCPAuth.oauth2_id_jag, approval_status)])
assert servers == ()
assert catalog.built == []
class _StalledEndpoint(_Endpoint):
"""An IdP that accepts the connection and never answers."""
async def fetch(
self, endpoint: str, client_id: str, grant_params: dict[str, str], client_auth: ClientAuth
) -> Result[ExchangedToken, CredError]:
await asyncio.sleep(0.5)
return await super().fetch(endpoint, client_id, grant_params, client_auth)
@pytest.mark.asyncio
async def test_a_stalled_idp_does_not_hold_up_the_login(caplog: pytest.LogCaptureFixture):
"""The exchange is optional. A user who only wants the dashboard must not wait on an MCP IdP that
stopped answering, so the exchange has its own short deadline and the login keeps the plain assertion."""
saml = _saml()
with caplog.at_level(logging.WARNING, logger="litellm.proxy"):
exchanged = await exchange_login_assertion_for_refresh_token(
saml,
(_config(),),
_StalledEndpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))),
_NOW,
deadline_seconds=0.05,
)
assert exchanged == saml
assert any("did not answer the login-time refresh token exchange" in r.getMessage() for r in caplog.records)
class _OneBadRowCatalog(_FakeCatalog):
async def build_mcp_server_from_table(
self,
mcp_server: LiteLLM_MCPServerTable,
*,
credentials_are_encrypted: bool = True,
env_vars_are_encrypted: bool | None = None,
register_oauth_discovery: bool = True,
) -> MCPServer:
if mcp_server.server_id == "broken":
raise ValueError("credentials cannot be decrypted")
return await super().build_mcp_server_from_table(
mcp_server,
credentials_are_encrypted=credentials_are_encrypted,
env_vars_are_encrypted=env_vars_are_encrypted,
register_oauth_discovery=register_oauth_discovery,
)
@pytest.mark.asyncio
async def test_one_unbuildable_db_row_does_not_cost_the_other_servers_their_exchange(caplog: pytest.LogCaptureFixture):
catalog = _OneBadRowCatalog({})
with caplog.at_level(logging.WARNING, logger="litellm.proxy"):
servers = await login_exchange_servers(
catalog, [_row("broken", MCPAuth.oauth2_id_jag), _row("healthy", MCPAuth.oauth2_id_jag)]
)
assert [server.server_id for server in servers] == ["healthy"]
assert any("server_id=broken" in r.getMessage() for r in caplog.records)

354
tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py
View file

@ -32,6 +32,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import (
Error,
IdJagConfig,
IdJagResourceClient,
IdJagSubjectSource,
NoneConfig,
NoOpAuth,
Ok,
@ -55,8 +56,11 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_r
SSOClientConfig,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
SAML2_SUBJECT_TOKEN_TYPE,
AssertionStoreUnavailable,
SSOIdentityAssertion,
assertion_from_saml_login,
with_login_refresh_token,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import (
ExchangedToken,
@ -111,6 +115,10 @@ class _FakeAssertionStore:
self._assertions = dict(assertions or {})
self.lookups: list[str] = []
async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
return await self.fetch(user_id)
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
self.lookups.append(user_id)
return self._assertions.get(user_id)
@ -539,12 +547,8 @@ async def test_id_jag_runs_both_legs_and_returns_the_leg2_bearer():
assert leg1_client_id == "litellm"
assert isinstance(endpoint.auths[0], ClientSecretAuth)
assert endpoint.auths[0].client_secret.get_secret_value() == "s"
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["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"
leg2_endpoint, leg2_client_id, leg2_params = endpoint.calls[1]
@ -563,9 +567,7 @@ async def test_id_jag_without_inbound_token_or_stored_assertion_is_precondition_
endpoint = _FakeTokenEndpoint([])
store = _FakeAssertionStore()
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store)
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"
@ -581,9 +583,7 @@ async def test_id_jag_exchanges_the_stored_sso_assertion_when_the_caller_present
store = _FakeAssertionStore({"alice": _assertion("alice-id-token")})
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store)
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, Ok)
assert _emitted(result.ok)["Authorization"] == "Bearer final-access"
@ -593,6 +593,54 @@ async def test_id_jag_exchanges_the_stored_sso_assertion_when_the_caller_present
assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag"
@pytest.mark.asyncio
async def test_id_jag_declares_a_stored_saml_assertion_as_a_saml2_subject_token():
"""A SAML login retains the signed assertion, not an id_token, so leg 1 has to tell the IdP what
it is being handed even though the server is configured with the id_token default."""
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
saml = assertion_from_saml_login(b"<saml:Assertion/>", int(time.time()) + 300, "https://idp.example.com")
store = _FakeAssertionStore({"alice": saml})
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store)
result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()))
assert isinstance(result, Ok)
_, _, leg1_params = endpoint.calls[0]
assert leg1_params["subject_token"] == saml.id_token.get_secret_value()
assert leg1_params["subject_token_type"] == SAML2_SUBJECT_TOKEN_TYPE
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"
@pytest.mark.asyncio
async def test_id_jag_sends_the_configured_type_for_a_stored_oidc_assertion():
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
store = _FakeAssertionStore({"alice": _assertion("alice-id-token")})
config = _id_jag_config().model_copy(update={"subject_token_type": "urn:ietf:params:oauth:token-type:jwt"})
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store)
result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(config))
assert isinstance(result, Ok)
_, _, leg1_params = endpoint.calls[0]
assert leg1_params["subject_token_type"] == "urn:ietf:params:oauth:token-type:jwt"
@pytest.mark.asyncio
async def test_id_jag_refuses_a_saml_assertion_past_its_not_on_or_after_without_calling_the_idp():
endpoint = _FakeTokenEndpoint([])
saml = assertion_from_saml_login(b"<saml:Assertion/>", int(time.time()) - 1, "https://idp.example.com")
provider = UpstreamCredentialProvider(
token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": saml})
)
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"
assert endpoint.calls == []
@pytest.mark.asyncio
async def test_id_jag_prefers_the_callers_own_token_over_the_stored_assertion():
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
@ -613,9 +661,7 @@ async def test_id_jag_refuses_an_expired_stored_assertion_without_calling_the_id
store = _FakeAssertionStore({"alice": _assertion("stale-id-token", expires_in=-timedelta(seconds=1))})
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store)
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"
@ -646,6 +692,9 @@ async def test_id_jag_renews_an_expired_stored_assertion_instead_of_challenging(
rows[user_id] = assertion
class _Inner:
async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
return await self.fetch(user_id)
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
return await _read(user_id)
@ -672,9 +721,7 @@ async def test_id_jag_renews_an_expired_stored_assertion_instead_of_challenging(
),
)
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, Ok)
_, _, leg1_params = endpoint.calls[0]
@ -686,7 +733,7 @@ def test_the_resolver_defaults_to_the_renewing_assertion_store():
the plain database reader would strand every agent an id_token lifetime after its user's login."""
provider = UpstreamCredentialProvider()
assert isinstance(provider._sso_assertion_store, RefreshingSSOAssertionStore) # noqa: SLF001 # the wiring is the assertion
assert isinstance(provider._sso_assertion_store, RefreshingSSOAssertionStore) # the wiring is the assertion
@pytest.mark.asyncio
@ -695,9 +742,7 @@ async def test_id_jag_accepts_a_stored_assertion_that_declares_no_expiry():
store = _FakeAssertionStore({"alice": _assertion("undated-id-token", expires_in=None)})
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store)
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, Ok)
_, _, leg1_params = endpoint.calls[0]
@ -730,14 +775,10 @@ async def test_id_jag_keeps_store_sourced_bearers_partitioned_per_user():
Ok(ExchangedToken(access_token="bob-bearer", expires_in=3600)),
]
)
store = _FakeAssertionStore(
{"alice": _assertion("alice-id-token"), "bob": _assertion("bob-id-token")}
)
store = _FakeAssertionStore({"alice": _assertion("alice-id-token"), "bob": _assertion("bob-id-token")})
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store)
alice = await provider.resolve_credentials(
Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())
)
alice = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()))
bob = await provider.resolve_credentials(Subject(tenant_id="", subject_id="bob"), _spec(_id_jag_config()))
assert isinstance(alice, Ok) and isinstance(bob, Ok)
@ -754,6 +795,10 @@ class _OutageAssertionStore:
def __init__(self) -> None:
self.lookups: list[str] = []
async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
return await self.fetch(user_id)
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
self.lookups.append(user_id)
raise AssertionStoreUnavailable(_DRIVER_DETAIL)
@ -766,9 +811,7 @@ async def test_id_jag_maps_an_assertion_store_outage_to_upstream_unavailable():
endpoint = _FakeTokenEndpoint([])
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=_OutageAssertionStore())
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 == "upstream_unavailable"
@ -784,9 +827,7 @@ async def test_id_jag_store_outage_does_not_leak_driver_detail_to_the_caller(cap
)
with caplog.at_level(logging.WARNING):
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 _DRIVER_DETAIL not in result.error.summary
@ -803,9 +844,7 @@ async def test_id_jag_invalidation_survives_an_assertion_store_outage():
token_endpoint=_FakeTokenEndpoint([]), sso_assertion_store=_OutageAssertionStore()
)
await provider.invalidate_credentials(
Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())
)
await provider.invalidate_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()))
class _FlakyAssertionStore:
@ -815,6 +854,10 @@ class _FlakyAssertionStore:
self._assertion = assertion
self.down = False
async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
return await self.fetch(user_id)
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
if self.down:
raise AssertionStoreUnavailable("connection refused")
@ -855,6 +898,10 @@ class _SwitchableAssertionStore:
def __init__(self, id_token: str) -> None:
self.id_token = id_token
async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
return await self.fetch(user_id)
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
return _assertion(self.id_token)
@ -899,6 +946,10 @@ class _SequentialAssertionStore:
self.issued: list[str] = []
self._n = 0
async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
return await self.fetch(user_id)
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
await asyncio.sleep(0)
if self.pinned is not None:
@ -996,14 +1047,10 @@ async def test_id_jag_reads_a_naive_stored_expiry_as_utc():
comparison would be a TypeError on the egress path, turning a 412 into a 500."""
endpoint = _FakeTokenEndpoint([])
naive_past = datetime.now(timezone.utc).replace(tzinfo=None) - timedelta(hours=1)
store = _FakeAssertionStore(
{"alice": SSOIdentityAssertion(id_token=SecretStr("stale"), expires_at=naive_past)}
)
store = _FakeAssertionStore({"alice": SSOIdentityAssertion(id_token=SecretStr("stale"), expires_at=naive_past)})
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store)
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"
@ -1032,13 +1079,9 @@ async def test_invalidate_evicts_a_store_sourced_id_jag_bearer():
@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-id-token"), _spec(_id_jag_config()))
assert isinstance(result, Error)
assert result.error.tag == "upstream_unavailable"
@ -1055,9 +1098,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-id-token"), _spec(_id_jag_config()))
assert isinstance(result, Error)
assert result.error.tag == "upstream_unavailable"
@ -1065,7 +1106,7 @@ async def test_id_jag_propagates_a_leg2_error():
assert len(endpoint.calls) == 2
def _two_leg_ok(bearer: str) -> list:
def _two_leg_ok(bearer: str) -> list[Result[ExchangedToken, CredError]]:
return [
Ok(ExchangedToken(access_token="the-id-jag", expires_in=300)),
Ok(ExchangedToken(access_token=bearer, expires_in=3600)),
@ -1256,8 +1297,18 @@ async def test_passthrough_ignores_the_carrier_and_keeps_the_callers_slot():
(NoneConfig(), _SUBJECT, "no-auth", None),
(PassthroughConfig(), _SUBJECT, "no-auth", None),
(PassthroughConfig(), _with_inbound("Bearer caller-token"), "oauth2-passthrough", "Bearer caller-token"),
(ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-key"))), _SUBJECT, "static-token", "Bearer static-key"),
(AuthorizationCodeConfig(), Subject(tenant_id="", subject_id="alice"), "stored-user-token", "Bearer stored-alice"),
(
ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-key"))),
_SUBJECT,
"static-token",
"Bearer static-key",
),
(
AuthorizationCodeConfig(),
Subject(tenant_id="", subject_id="alice"),
"stored-user-token",
"Bearer stored-alice",
),
],
)
async def test_resolved_source_matches_the_credential_sent_upstream(
@ -1279,7 +1330,9 @@ async def test_resolved_source_matches_the_credential_sent_upstream(
async def test_resolved_source_preserves_missing_user_token_error() -> None:
from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source
result = await resolve_credentials_with_source(UpstreamCredentialProvider(), _SUBJECT, _spec(AuthorizationCodeConfig()))
result = await resolve_credentials_with_source(
UpstreamCredentialProvider(), _SUBJECT, _spec(AuthorizationCodeConfig())
)
assert isinstance(result, Error)
assert result.error.tag == "unauthorized"
@ -1289,7 +1342,9 @@ async def test_minted_token_sources_match_egress_and_do_not_fetch_twice() -> Non
from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source
source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at")))
m2m = await resolve_credentials_with_source(UpstreamCredentialProvider(client_credentials_source=source), _SUBJECT, _spec(_M2M))
m2m = await resolve_credentials_with_source(
UpstreamCredentialProvider(client_credentials_source=source), _SUBJECT, _spec(_M2M)
)
assert isinstance(m2m, Ok)
headers, _ = await _emitted_async(m2m.ok.auth)
assert headers["Authorization"] == "Bearer m2m-at"
@ -1297,7 +1352,9 @@ async def test_minted_token_sources_match_egress_and_do_not_fetch_twice() -> Non
assert source.gets == ["s"]
exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at")))
exchanged = await resolve_credentials_with_source(UpstreamCredentialProvider(token_exchanger=exchanger), _with_inbound("subject"), _spec(_OBO))
exchanged = await resolve_credentials_with_source(
UpstreamCredentialProvider(token_exchanger=exchanger), _with_inbound("subject"), _spec(_OBO)
)
assert isinstance(exchanged, Ok)
assert _emitted(exchanged.ok.auth)["Authorization"] == "Bearer exchanged-at"
assert exchanged.ok.source.value == "token-exchange"
@ -1321,9 +1378,14 @@ async def test_id_jag_source_describes_final_token_after_both_exchanges() -> Non
async def test_client_credentials_subscription_identity_preserves_resolved_grant_scopes():
async def resolve(token):
source = _FakeM2MSource(Ok(token))
result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials(_SUBJECT, _spec(_M2M))
result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials(
_SUBJECT, _spec(_M2M)
)
assert isinstance(result, Ok)
from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ClientCredentialsBearerAuth
from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import (
ClientCredentialsBearerAuth,
)
assert isinstance(result.ok, ClientCredentialsBearerAuth)
return result.ok.subscription_identity
@ -1333,3 +1395,175 @@ async def test_client_credentials_subscription_identity_preserves_resolved_grant
assert original is not None
assert original == renewed
assert original != reduced
def _login_with_refresh_token(expires_in: int | None) -> SSOIdentityAssertion:
saml = assertion_from_saml_login(b"<saml:Assertion/>", int(time.time()) + 300, "https://idp.example.com")
return with_login_refresh_token(
saml,
"rt-from-login",
expires_in,
datetime.now(timezone.utc),
endpoint="https://idp.example.com/token",
client_id="litellm",
)
def _refresh_token_source_config() -> IdJagConfig:
return _id_jag_config().model_copy(update={"subject_source": IdJagSubjectSource.refresh_token})
@pytest.mark.asyncio
async def test_id_jag_presents_the_login_refresh_token_when_the_server_asks_for_it():
"""``id_jag_subject_source: refresh_token`` presents the refresh token the login assertion was
exchanged for, typed as one, never the SAML assertion itself."""
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
store = _FakeAssertionStore({"alice": _login_with_refresh_token(3600)})
provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store)
result = await provider.resolve_credentials(
Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config())
)
assert isinstance(result, Ok)
_, _, leg1_params = endpoint.calls[0]
assert leg1_params["subject_token"] == "rt-from-login"
assert leg1_params["subject_token_type"] == "urn:ietf:params:oauth:token-type:refresh_token"
assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag"
@pytest.mark.asyncio
async def test_id_jag_refresh_token_source_never_presents_the_sso_refresh_token():
"""An OIDC login holds a refresh token for the SSO client. Only the token the ID-JAG client
exchanged at login is presented, the SSO one cannot be redeemed there."""
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt"))
row = with_login_refresh_token(
oidc,
"rt-from-login",
3600,
datetime.now(timezone.utc),
endpoint="https://idp.example.com/token",
client_id="litellm",
)
provider = UpstreamCredentialProvider(
token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": row})
)
result = await provider.resolve_credentials(
Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config())
)
assert isinstance(result, Ok)
assert endpoint.calls[0][2]["subject_token"] == "rt-from-login"
@pytest.mark.asyncio
async def test_id_jag_refresh_token_source_refuses_a_token_minted_for_another_idp_client():
"""A refresh token is bound to the client that received it. One minted for another server's IdP
client is never posted to this server's endpoint, the user is told to sign in again instead."""
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
row = with_login_refresh_token(
_login_with_refresh_token(3600),
"rt-for-other",
3600,
datetime.now(timezone.utc),
endpoint="https://other-idp.example.com/token",
client_id="other-agent",
)
provider = UpstreamCredentialProvider(
token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": row})
)
result = await provider.resolve_credentials(
Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config())
)
assert isinstance(result, Error)
assert result.error.tag == "precondition_required"
assert "different IdP client" in str(result.error)
assert endpoint.calls == []
class _RenewalFailsStore(_FakeAssertionStore):
"""An SSO renewal that is timing out: ``fetch`` raises the way RefreshingSSOAssertionStore does,
while the row itself is still readable."""
async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
return self._assertions.get(user_id)
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
raise AssertionStoreUnavailable("IdP token endpoint timed out")
@pytest.mark.asyncio
async def test_id_jag_refresh_token_source_does_not_wait_on_an_sso_renewal():
"""The exchanged refresh token has its own lifetime. A failing renewal of the unrelated SSO
id_token must not turn into a 503 for a server that never presents that id_token."""
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
provider = UpstreamCredentialProvider(
token_endpoint=endpoint, sso_assertion_store=_RenewalFailsStore({"alice": _login_with_refresh_token(3600)})
)
result = await provider.resolve_credentials(
Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config())
)
assert isinstance(result, Ok)
assert endpoint.calls[0][2]["subject_token"] == "rt-from-login"
@pytest.mark.asyncio
async def test_id_jag_refresh_token_source_outlives_the_saml_assertion():
"""The SAML assertion is spent minutes after login; the refresh token it bought is what counts."""
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
expired_saml = assertion_from_saml_login(b"<saml:Assertion/>", int(time.time()) - 60, "https://idp.example.com")
row = with_login_refresh_token(
expired_saml,
"rt-from-login",
3600,
datetime.now(timezone.utc),
endpoint="https://idp.example.com/token",
client_id="litellm",
)
provider = UpstreamCredentialProvider(
token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": row})
)
result = await provider.resolve_credentials(
Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config())
)
assert isinstance(result, Ok)
assert endpoint.calls[0][2]["subject_token"] == "rt-from-login"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("row", "expected"),
[
(
assertion_from_saml_login(b"<saml:Assertion/>", int(time.time()) + 300, "https://idp.example.com"),
"none is stored",
),
(_login_with_refresh_token(-1), "has expired"),
],
ids=["no refresh token", "expired refresh token"],
)
async def test_id_jag_refresh_token_source_is_a_precondition_without_a_live_refresh_token(
row: SSOIdentityAssertion, expected: str
):
endpoint = _FakeTokenEndpoint([])
provider = UpstreamCredentialProvider(
token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": row})
)
result = await provider.resolve_credentials(
Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config())
)
assert isinstance(result, Error)
assert result.error.tag == "precondition_required"
assert expected in result.error.summary
assert "Sign in through LiteLLM SSO" in result.error.summary
assert endpoint.calls == []

150
tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py
View file

@ -35,8 +35,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_r
sso_client_config,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
with_login_refresh_token,
AssertionStoreUnavailable,
SSOIdentityAssertion,
assertion_from_saml_login,
)
SIGNING_KEY = "test-idp-signing-key-32-bytes-long-xxxx"
@ -318,6 +320,53 @@ async def test_a_missing_refresh_token_names_the_scope_the_operator_has_to_set(c
assert "offline_access" in caplog.text
@pytest.mark.asyncio
async def test_an_expired_saml_assertion_is_never_redeemed_and_names_the_re_login(caplog: pytest.LogCaptureFixture):
"""SAML has no refresh grant. Even a row that somehow carries a refresh token must not be posted
to the generic OIDC client, and the operator hint about GENERIC_SCOPE does not apply."""
saml = assertion_from_saml_login(b"<saml:Assertion/>", int(time.time()) - 1, "https://idp.example.com")
with_refresh_token = saml.model_copy(update={"refresh_token": SecretStr("rt_saml")})
rows = _FakeRows({"alice": with_refresh_token})
transport = _FakeTransport(_renewal(_id_token()))
served = await _store(rows, transport).fetch("alice")
refresher = SSOAssertionRefresher(transport, client_config=lambda: _CLIENT, read=rows.fetch, write=rows.write)
with caplog.at_level(logging.WARNING):
outcome = await refresher.refresh("alice", with_refresh_token)
assert served == with_refresh_token
assert transport.calls == []
assert rows.writes == []
match outcome:
case Error(failure):
assert failure.kind == "rejected"
assert "sign in" in failure.detail
case Ok(_):
pytest.fail("a SAML assertion must not be renewable")
assert "sign in" in caplog.text
assert "GENERIC_SCOPE" not in caplog.text
class _ForbiddenCoordinator:
async def run(self, user_id: str, key: str, *, refresh: object, reread: object) -> None:
raise AssertionError(f"a SAML assertion must never enter the refresh path (user_id={user_id}, key={key})")
@pytest.mark.asyncio
async def test_a_saml_assertion_inside_the_expiry_skew_is_served_as_is_without_electing_a_renewal():
"""SAML cannot be renewed, so a still-valid assertion near its NotOnOrAfter is handed out untouched
instead of running the single-flight election that would 503 the losers."""
saml = assertion_from_saml_login(b"<saml:Assertion/>", int(time.time()) + 30, "https://idp.example.com")
rows = _FakeRows({"alice": saml})
transport = _FakeTransport(_renewal(_id_token()))
served = await _store(rows, transport, coordinator_factory=_ForbiddenCoordinator).fetch("alice")
assert served == saml
assert transport.calls == []
assert rows.writes == []
@pytest.mark.asyncio
async def test_an_unconfigured_sso_client_never_calls_the_idp(caplog):
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
@ -694,6 +743,9 @@ async def test_the_in_process_coordinator_is_retried_until_redis_appears():
{"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, "GENERIC_CLIENT_SECRET": "s"},
{"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, "GENERIC_CLIENT_ID": "litellm"},
{"GENERIC_TOKEN_ENDPOINT": "", "GENERIC_CLIENT_ID": "litellm", "GENERIC_CLIENT_SECRET": "s"},
{"GOOGLE_CLIENT_ID": "google-cid"},
{"MICROSOFT_CLIENT_ID": "ms-cid", "MICROSOFT_CLIENT_SECRET": "s"},
{"MICROSOFT_CLIENT_ID": "ms-cid", "MICROSOFT_TENANT": "tenant"},
],
)
def test_a_partial_sso_client_is_no_client(env):
@ -717,6 +769,64 @@ def test_the_configured_sso_client_is_the_one_the_login_used():
assert config.client_secret.get_secret_value() == "s3cret"
def test_a_google_login_is_renewed_at_googles_token_endpoint_as_the_google_client():
"""fastapi-sso's Google login redeems the code at Google's token endpoint with HTTP Basic, so
that is where and how the refresh token it handed back is redeemed."""
config = sso_client_config({"GOOGLE_CLIENT_ID": "google-cid", "GOOGLE_CLIENT_SECRET": "google-secret"})
assert config is not None
assert config.token_endpoint == "https://oauth2.googleapis.com/token"
assert config.client_id == "google-cid"
assert config.client_secret.get_secret_value() == "google-secret"
assert config.auth_method == "client_secret_basic"
def test_a_microsoft_login_is_renewed_at_the_tenant_token_endpoint_as_the_microsoft_client():
config = sso_client_config(
{"MICROSOFT_CLIENT_ID": "ms-cid", "MICROSOFT_CLIENT_SECRET": "ms-secret", "MICROSOFT_TENANT": "tenant-123"}
)
assert config is not None
assert config.token_endpoint == "https://login.microsoftonline.com/tenant-123/oauth2/v2.0/token"
assert config.client_id == "ms-cid"
assert config.client_secret.get_secret_value() == "ms-secret"
assert config.auth_method == "client_secret_basic"
def test_a_custom_microsoft_token_endpoint_wins_over_the_tenant_template():
"""``MICROSOFT_TOKEN_ENDPOINT`` is what the Microsoft login itself honours for sovereign
clouds, so the renewal has to go to the same place."""
config = sso_client_config(
{
"MICROSOFT_CLIENT_ID": "ms-cid",
"MICROSOFT_CLIENT_SECRET": "ms-secret",
"MICROSOFT_TENANT": "tenant-123",
"MICROSOFT_TOKEN_ENDPOINT": "https://login.microsoftonline.us/tenant-123/oauth2/v2.0/token",
}
)
assert config is not None
assert config.token_endpoint == "https://login.microsoftonline.us/tenant-123/oauth2/v2.0/token"
def test_the_client_follows_the_login_callbacks_provider_order():
"""A deployment carrying a Google and a generic client signs users in through Google, so the
stored refresh token belongs to the Google client and only that client can redeem it."""
config = sso_client_config(
{
"GOOGLE_CLIENT_ID": "google-cid",
"GOOGLE_CLIENT_SECRET": "google-secret",
"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT,
"GENERIC_CLIENT_ID": "litellm",
"GENERIC_CLIENT_SECRET": "s3cret",
}
)
assert config is not None
assert config.client_id == "google-cid"
assert config.token_endpoint == "https://oauth2.googleapis.com/token"
def test_the_live_store_renews_over_the_database_reader():
"""The composition root has to produce a renewing store, or none of this runs in production."""
assert isinstance(default_sso_assertion_store(), RefreshingSSOAssertionStore)
@ -792,3 +902,43 @@ async def test_a_successful_grant_is_handed_back_as_the_parsed_body():
assert isinstance(outcome, Ok)
assert outcome.ok["id_token"] == "idt"
@pytest.mark.asyncio
async def test_a_renewal_keeps_the_refresh_token_exchanged_at_login_in_the_stored_row():
"""Renewal replaces the id_token and the SSO refresh token. The refresh token the ID-JAG client
exchanged at login is unrelated to either and must still be in the row the next MCP call reads."""
stale, fresh = _id_token(exp_offset=-1), _id_token()
row = with_login_refresh_token(
_stored(stale, expires_in=-1),
"rt-from-login",
3600,
datetime.now(timezone.utc),
endpoint="https://idp.example.com/token",
client_id="agent",
)
rows = _FakeRows({"alice": row})
served = await _store(rows, _FakeTransport(_renewal(fresh, "rt_2"))).fetch("alice")
stored = rows.rows["alice"]
assert stored.id_token.get_secret_value() == fresh
assert stored.refresh_token is not None and stored.refresh_token.get_secret_value() == "rt_2"
assert stored.exchanged_refresh_token is not None
assert stored.exchanged_refresh_token.get_secret_value() == "rt-from-login"
assert stored.exchanged_refresh_token_client_id == "agent"
assert served is not None and served.exchanged_refresh_token is not None
@pytest.mark.asyncio
async def test_fetch_stored_hands_back_the_row_without_renewing_it():
"""The read a refresh-token server uses: the row as stored, no IdP round trip even when the SSO
id_token is about to expire."""
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(_renewal(_id_token()))
served = await _store(rows, transport).fetch_stored("alice")
assert served is not None
assert served.id_token.get_secret_value() == rows.rows["alice"].id_token.get_secret_value()
assert transport.calls == []

173
tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py
View file

@ -8,9 +8,11 @@ rotation re-encrypts stored rows like the sibling per-user credential tables.
"""
import asyncio
import base64
import json
import os
import time
from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock, MagicMock, patch
import jwt as pyjwt
@ -18,15 +20,20 @@ import pytest
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
_ASSERTION_CACHE,
SAML2_SUBJECT_TOKEN_TYPE,
AssertionStoreUnavailable,
DbSSOAssertionStore,
SSOAssertionCache,
assertion_from_saml_login,
assertion_from_sso_login,
carry_exchanged_refresh_token,
ema_assertion_retention_enabled,
fetch_sso_identity_assertion,
login_refresh_token_expired,
persist_sso_identity_assertion,
retain_sso_identity_assertion_for_ema,
rotate_sso_identity_assertions_master_key,
with_login_refresh_token,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper
from litellm.types.mcp import MCPAuth
@ -111,6 +118,41 @@ def test_assertion_from_sso_login_happy_path():
assert assertion.expires_at.timestamp() == pytest.approx(time.time() + 3600, abs=5)
SAML_ASSERTION_XML = (
b'<saml:Assertion xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion" ID="_a1" Version="2.0">'
b"<saml:Issuer>https://idp.example.com/metadata</saml:Issuer></saml:Assertion>"
)
def _saml_subject_bytes(encoded: str) -> bytes:
assert set(encoded) <= set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"), encoded[:40]
return base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4))
def test_assertion_from_saml_login_carries_the_base64url_assertion_as_a_saml2_subject_token():
not_on_or_after = int(time.time()) + 300
assertion = assertion_from_saml_login(SAML_ASSERTION_XML, not_on_or_after, "https://idp.example.com/metadata")
assert _saml_subject_bytes(assertion.id_token.get_secret_value()) == SAML_ASSERTION_XML
assert assertion.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE
assert assertion.refresh_token is None
assert assertion.issuer == "https://idp.example.com/metadata"
assert assertion.expires_at is not None
assert assertion.expires_at.tzinfo is not None
assert assertion.expires_at.timestamp() == not_on_or_after
def test_assertion_from_saml_login_without_not_on_or_after_has_no_expiry():
assertion = assertion_from_saml_login(SAML_ASSERTION_XML, None, None)
assert assertion.expires_at is None
assert assertion.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE
def test_oidc_login_assertion_leaves_the_subject_token_type_to_the_operator():
assertion = assertion_from_sso_login(_make_id_token(), None)
assert assertion is not None
assert assertion.subject_token_type is None
def test_assertion_repr_never_leaks_token_material():
token = _make_id_token()
assertion = assertion_from_sso_login(token, "rt_secret_value")
@ -226,6 +268,37 @@ async def test_persist_and_fetch_round_trip_encrypted_at_rest():
assert json.loads(decrypted)["id_token"] == token
@pytest.mark.asyncio
async def test_saml_assertion_round_trips_with_its_subject_token_type():
stored: dict[str, str] = {}
prisma = _make_prisma(stored)
assertion = assertion_from_saml_login(SAML_ASSERTION_XML, int(time.time()) + 300, "https://idp.example.com")
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
await persist_sso_identity_assertion("user-saml", assertion)
fetched = await fetch_sso_identity_assertion("user-saml")
assert fetched is not None
assert fetched.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE
assert _saml_subject_bytes(fetched.id_token.get_secret_value()) == SAML_ASSERTION_XML
assert fetched.expires_at == assertion.expires_at
assert fetched.refresh_token is None
assert SAML_ASSERTION_XML.decode() not in stored["user-saml"]
@pytest.mark.asyncio
async def test_oidc_row_written_before_subject_token_types_reads_back_without_one():
stored: dict[str, str] = {}
prisma = _make_prisma(stored)
token = _make_id_token()
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(token, None))
fetched = await fetch_sso_identity_assertion("user-a")
assert fetched is not None
assert fetched.subject_token_type is None
decrypted = decrypt_value_helper(stored["user-a"], "test", exception_type="debug")
assert decrypted is not None
assert "subject_token_type" not in json.loads(decrypted)
@pytest.mark.asyncio
async def test_persist_overwrites_previous_login():
stored = {}
@ -248,7 +321,9 @@ async def test_fetch_serves_second_read_from_cache_without_db_read():
cache = SSOAssertionCache()
token = _make_id_token()
assertion = assertion_from_sso_login(token, "rt_1")
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
with patch(
"litellm.proxy.proxy_server.prisma_client", prisma
): # test-quality-ok: fake DB seam has no injection boundary
await persist_sso_identity_assertion("user-a", assertion, cache=cache)
first = await fetch_sso_identity_assertion("user-a", cache=cache)
second = await fetch_sso_identity_assertion("user-a", cache=cache)
@ -266,7 +341,9 @@ async def test_persist_busts_cache_so_relogin_is_visible_immediately():
cache = SSOAssertionCache()
first_token = _make_id_token(exp_offset=100)
second_token = _make_id_token(exp_offset=7200)
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
with patch(
"litellm.proxy.proxy_server.prisma_client", prisma
): # test-quality-ok: fake DB seam has no injection boundary
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache)
first = await fetch_sso_identity_assertion("user-a", cache=cache)
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(second_token, "rt_new"), cache=cache)
@ -295,7 +372,9 @@ async def test_fetch_racing_a_relogin_does_not_cache_the_previous_assertion():
return row
prisma.db.litellm_ssoidentityassertion.find_unique = _paused_find_unique
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
with patch(
"litellm.proxy.proxy_server.prisma_client", prisma
): # test-quality-ok: fake DB seam has no injection boundary
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache)
racing_fetch = asyncio.create_task(fetch_sso_identity_assertion("user-a", cache=cache))
await db_read_started.wait()
@ -324,7 +403,9 @@ async def test_fetch_missing_row_returns_none():
async def test_fetch_does_not_cache_a_missing_row():
prisma = _make_prisma({})
cache = SSOAssertionCache()
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
with patch(
"litellm.proxy.proxy_server.prisma_client", prisma
): # test-quality-ok: fake DB seam has no injection boundary
first = await fetch_sso_identity_assertion("nobody", cache=cache)
second = await fetch_sso_identity_assertion("nobody", cache=cache)
assert first is None
@ -353,7 +434,9 @@ async def test_cached_assertion_expires_after_ttl():
cache = SSOAssertionCache(ttl_seconds=1)
first_token = _make_id_token(exp_offset=100)
second_token = _make_id_token(exp_offset=7200)
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
with patch(
"litellm.proxy.proxy_server.prisma_client", prisma
): # test-quality-ok: fake DB seam has no injection boundary
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache)
first = await fetch_sso_identity_assertion("user-a", cache=cache)
await persist_sso_identity_assertion(
@ -477,7 +560,9 @@ async def test_db_store_uses_injected_cache():
prisma = _make_prisma(stored)
cache = SSOAssertionCache()
token = _make_id_token()
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
with patch(
"litellm.proxy.proxy_server.prisma_client", prisma
): # test-quality-ok: fake DB seam has no injection boundary
await persist_sso_identity_assertion("alice", assertion_from_sso_login(token, None), cache=cache)
store = DbSSOAssertionStore(cache=cache)
first = await store.fetch("alice")
@ -496,3 +581,79 @@ async def test_db_store_returns_none_for_a_user_with_no_stored_assertion():
with patch.dict(os.environ, {"LITELLM_SALT_KEY": SALT_KEY}):
with patch("litellm.proxy.proxy_server.prisma_client", _make_prisma({})):
assert await DbSSOAssertionStore().fetch("nobody") is None
@pytest.mark.asyncio
async def test_login_refresh_token_round_trips_beside_the_saml_assertion():
"""A SAML login exchanged for a refresh token keeps the assertion, its type and its expiry, and adds
the refresh token with the lifetime the IdP gave it, all encrypted at rest."""
stored: dict[str, str] = {}
prisma = _make_prisma(stored)
now = datetime.now(timezone.utc).replace(microsecond=0)
saml = assertion_from_saml_login(SAML_ASSERTION_XML, int(time.time()) + 300, "https://idp.example.com")
assertion = with_login_refresh_token(
saml, "rt-from-login", 3600, now, endpoint="https://idp.example.com/token", client_id="agent"
)
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
await persist_sso_identity_assertion("user-rt", assertion)
fetched = await fetch_sso_identity_assertion("user-rt")
assert fetched is not None
assert fetched.exchanged_refresh_token is not None
assert fetched.exchanged_refresh_token.get_secret_value() == "rt-from-login"
assert fetched.exchanged_refresh_token_expires_at == now + timedelta(seconds=3600)
assert fetched.exchanged_refresh_token_endpoint == "https://idp.example.com/token"
assert fetched.exchanged_refresh_token_client_id == "agent"
assert fetched.refresh_token is None
assert fetched.id_token == saml.id_token
assert fetched.expires_at == saml.expires_at
assert fetched.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE
assert "rt-from-login" not in stored["user-rt"]
assert not login_refresh_token_expired(fetched, now)
assert login_refresh_token_expired(fetched, now + timedelta(seconds=3600))
def test_renewal_follows_a_rotated_login_refresh_token_that_is_the_id_jag_subject():
"""When the OIDC login's own refresh token is the ID-JAG subject (the SSO client is the ID-JAG
client), a renewal that rotated it leaves the old token revoked at the IdP, so the subject moves
to the token the renewal returned."""
now = datetime.now(timezone.utc).replace(microsecond=0)
login = assertion_from_sso_login(_make_id_token(), "sso-rt")
assert login is not None
previous = with_login_refresh_token(
login, "sso-rt", None, now, endpoint="https://idp.example.com/token", client_id="agent"
)
renewed = assertion_from_sso_login(_make_id_token(7200), "sso-rt-2")
assert renewed is not None
carried = carry_exchanged_refresh_token(previous, renewed)
assert carried.exchanged_refresh_token is not None
assert carried.exchanged_refresh_token.get_secret_value() == "sso-rt-2"
assert carried.exchanged_refresh_token_endpoint == "https://idp.example.com/token"
assert carried.exchanged_refresh_token_client_id == "agent"
def test_renewal_carries_the_exchanged_refresh_token_forward():
"""An OIDC renewal rebuilds the row from the new id_token. The refresh token the ID-JAG client
exchanged at login is unrelated to that renewal and must survive it."""
now = datetime.now(timezone.utc).replace(microsecond=0)
previous = with_login_refresh_token(
assertion_from_sso_login(_make_id_token(), "sso-rt"),
"rt-from-login",
3600,
now,
endpoint="https://idp.example.com/token",
client_id="agent",
)
renewed = assertion_from_sso_login(_make_id_token(7200), "sso-rt-2")
assert previous is not None and renewed is not None
carried = carry_exchanged_refresh_token(previous, renewed)
assert carried.id_token == renewed.id_token
assert carried.refresh_token is not None and carried.refresh_token.get_secret_value() == "sso-rt-2"
assert carried.exchanged_refresh_token is not None
assert carried.exchanged_refresh_token.get_secret_value() == "rt-from-login"
assert carried.exchanged_refresh_token_expires_at == now + timedelta(seconds=3600)
assert carried.exchanged_refresh_token_endpoint == "https://idp.example.com/token"
assert carried.exchanged_refresh_token_client_id == "agent"

218
tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -102,9 +102,8 @@ def isolate_global_mcp_registry(monkeypatch):
restore, entries leak into other test modules sharing the same worker and break
assertions over the full registry contents.
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
from litellm.proxy._experimental.mcp_server.catalog import CatalogSnapshots
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
monkeypatch.setattr(global_mcp_server_manager, "catalog", CatalogSnapshots(global_mcp_server_manager))
snapshot = dict(global_mcp_server_manager.registry)
@ -8133,9 +8132,9 @@ async def test_resolve_active_litellm_key_resolves_key_without_user_id(proxy_glo
blocked and expiry, and the key hash (not the user) is what the mint seals. The per-user token
store still gets no user for such a key, since there is none to key a stored credential by."""
from litellm.proxy._experimental.mcp_server.bridge_token_flow import (
extract_user_id_from_request,
_ResolvedKey,
_resolve_active_litellm_key,
_ResolvedKey,
extract_user_id_from_request,
)
from litellm.proxy._types import UserAPIKeyAuth, hash_token
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@ -8417,8 +8416,8 @@ async def test_load_active_user_by_id_serves_a_cached_row_without_a_database_rea
cached row answers without a database read, and only a caller that asks for the database row pays for
one."""
from litellm.proxy._experimental.mcp_server.bridge_token_flow import (
reload_active_user_by_id,
load_active_user_by_id,
reload_active_user_by_id,
)
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@ -11428,11 +11427,10 @@ async def test_upstream_resource_sent_on_dcr_bridge_relay_authorize():
"""The DCR-bridge relay arm builds its own upstream authorize params, so it needs the resource
too. Without it the relayed authorize would omit the resource while the gateway's token leg
still sent one, which is itself an invalid_target."""
from litellm.types.mcp import MCPAuth
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_dcr_bridge_relays_client_registration,
)
from litellm.types.mcp import MCPAuth
server = _resource_server(
client_id=None,
@ -11693,7 +11691,7 @@ def _native_client_app(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", _NATIVE_CLIENT_MASTER_KEY, raising=False)
minted = []
async def fake_mint(user_id, team_id):
async def fake_mint(user_id, team_id, allowed_routes=None):
minted.append((user_id, team_id))
return MintedProxyCredential(key=f"sk-cli-{len(minted)}", expires_in=3600, user_id=user_id, team_id=team_id)
@ -12136,7 +12134,9 @@ def test_introspect_route_answers_for_authenticated_caller(monkeypatch):
async def test_identity_bound_authorization_carries_nonce_and_caller_through_callback(monkeypatch):
from http.cookies import SimpleCookie
from urllib.parse import parse_qs, urlparse
from fastapi import Request
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_oauth_state_cookie_name,
authorize_with_server,
@ -12211,6 +12211,7 @@ async def test_identity_bound_authorization_carries_nonce_and_caller_through_cal
@pytest.mark.asyncio
async def test_enforced_login_warms_verified_token_readable_without_database_lookup(monkeypatch):
from types import SimpleNamespace
from litellm.caching.caching import DualCache
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import db, mcp_server_manager
@ -12279,9 +12280,10 @@ async def test_enforced_login_warms_verified_token_readable_without_database_loo
async def test_dcr_refusal_is_actionable_without_upstream_body(
upstream_status: int, auth_type: MCPAuth, dcr_bridge: bool, flow: str, monkeypatch: pytest.MonkeyPatch
) -> None:
import httpx
from typing import Final
import httpx
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
mint_ephemeral_dcr_client,
register_client_with_server,
@ -12522,8 +12524,8 @@ async def test_oauth_exchange_stores_token_for_validated_jwt_user(
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy import proxy_server
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy import proxy_server
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -12618,8 +12620,8 @@ async def test_oauth_jwt_identity_rejects_untrusted_or_inactive_owner(
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import mcp_server_manager
from litellm.proxy._experimental.mcp_server.bridge_token_flow import (
extract_user_id_from_request,
authorize_oauth_credential_request,
extract_user_id_from_request,
)
allowed_servers: Final = AsyncMock(return_value=["server-a"])
@ -12688,7 +12690,7 @@ async def test_oauth_jwt_uses_configured_virtual_key_owner(
) -> None:
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy._experimental.mcp_server.bridge_token_flow import extract_user_id_from_request
from litellm.proxy._types import UserAPIKeyAuth, UnregisteredJWTClientBehavior, hash_token
from litellm.proxy._types import UnregisteredJWTClientBehavior, UserAPIKeyAuth, hash_token
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
handler, signing_key = jwt_oauth_identity
@ -13069,8 +13071,8 @@ async def test_signed_oauth_callback_honors_credential_write_policy(
cimd: bool,
) -> None:
import httpx
import litellm
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy import proxy_server
@ -13186,17 +13188,16 @@ async def test_identity_bound_authorize_preserves_presented_jwt_permissions(
allowed: bool,
credential: str,
) -> None:
import jwt
from datetime import datetime, timedelta, timezone
from urllib.parse import parse_qs, urlparse
import jwt
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy import proxy_server
from litellm.proxy._types import JWTIssuerConfig, UserAPIKeyAuth, hash_token
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.proxy._experimental.mcp_server import discoverable_endpoints, mcp_server_manager
from litellm.proxy._types import MCPTransport
from litellm.proxy._types import JWTIssuerConfig, MCPTransport, UserAPIKeyAuth, hash_token
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
handler, signing_key = jwt_oauth_identity
@ -13330,9 +13331,10 @@ async def test_identity_bound_authorize_unrelated_bearer_uses_browser_session(
credential: str,
cookie_state: str,
) -> None:
import jwt
from urllib.parse import parse_qs, urlparse
import jwt
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import discoverable_endpoints, mcp_server_manager
@ -13664,6 +13666,7 @@ def test_persisted_client_is_not_applied_to_another_upstream(issuer, url):
async def test_registration_does_not_write_client_after_upstream_edit(monkeypatch):
monkeypatch.setenv("LITELLM_SALT_KEY", "test-mcp-registration-salt")
from prisma import models
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _persist_dcr_client_registration
@ -13750,6 +13753,7 @@ async def test_callback_does_not_forward_error_from_another_issuer(monkeypatch):
@pytest.mark.parametrize("operation", ["authorize", "token"])
async def test_bound_client_cannot_be_sent_to_a_different_issuer(monkeypatch, operation):
from fastapi import Request
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -13796,6 +13800,7 @@ async def test_bound_client_cannot_be_sent_to_a_different_issuer(monkeypatch, op
@pytest.mark.asyncio
async def test_failed_registration_preserves_cached_credentials(monkeypatch):
from fastapi import Request
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints
server = _dcr_redirect_test_server("old-client").model_copy(
@ -13819,8 +13824,13 @@ async def test_failed_registration_preserves_cached_credentials(monkeypatch):
request.headers = {}
with pytest.raises(HTTPException) as error:
await endpoints.register_client_with_server(
request=request, mcp_server=server, client_name="app", grant_types=["authorization_code"],
response_types=["code"], token_endpoint_auth_method="none", replace_stale_client=True,
request=request,
mcp_server=server,
client_name="app",
grant_types=["authorization_code"],
response_types=["code"],
token_endpoint_auth_method="none",
replace_stale_client=True,
)
assert error.value.status_code == 503
assert "could not be saved" in error.value.detail
@ -13835,6 +13845,7 @@ async def test_failed_registration_preserves_cached_credentials(monkeypatch):
@pytest.mark.asyncio
async def test_registration_without_database_keeps_client_in_temporary_server(monkeypatch):
from fastapi import Request
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints
@ -13850,8 +13861,13 @@ async def test_registration_without_database_keeps_client_in_temporary_server(mo
request.base_url = "https://gateway.example/"
request.headers = {}
response = await endpoints.register_client_with_server(
request=request, mcp_server=server, client_name="app", grant_types=["authorization_code"],
response_types=["code"], token_endpoint_auth_method="none", replace_stale_client=True,
request=request,
mcp_server=server,
client_name="app",
grant_types=["authorization_code"],
response_types=["code"],
token_endpoint_auth_method="none",
replace_stale_client=True,
)
assert response.status_code == 200
assert json.loads(response.body)["client_id"] == "temporary-client"
@ -13864,9 +13880,10 @@ async def test_registration_without_database_keeps_client_in_temporary_server(mo
@pytest.mark.asyncio
@pytest.mark.parametrize("config_store", [False, True])
async def test_registration_does_not_overwrite_credentials_after_failed_identity_read(monkeypatch, config_store):
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints, db
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
from litellm.proxy import utils
from litellm.proxy._experimental.mcp_server import db
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
server = _dcr_redirect_test_server(None)
monkeypatch.setattr(utils, "get_prisma_client_or_throw", lambda _: MagicMock())
@ -13895,9 +13912,10 @@ async def test_registration_does_not_overwrite_credentials_after_failed_identity
@pytest.mark.asyncio
@pytest.mark.parametrize("winner_available", [False, True])
async def test_registration_losing_conditional_write_reuses_only_a_matching_winner(monkeypatch, winner_available):
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints, db
from litellm.proxy._types import LiteLLM_MCPServerTable
from litellm.proxy import utils
from litellm.proxy._experimental.mcp_server import db
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints
from litellm.proxy._types import LiteLLM_MCPServerTable
server = _dcr_redirect_test_server(None).model_copy(update={"url": "https://upstream.example/mcp"})
row = LiteLLM_MCPServerTable.model_validate(server.model_dump(exclude_none=True))
@ -14503,7 +14521,8 @@ async def test_managed_authorize_retains_the_callers_previous_grant(
from urllib.parse import parse_qs, urlsplit
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints
from litellm.proxy._experimental.mcp_server import db
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
_, signing_key = jwt_oauth_identity
@ -14548,7 +14567,8 @@ async def test_managed_code_exchange_preserves_requested_scope_only_when_omitted
from urllib.parse import parse_qs, urlsplit
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints
from litellm.proxy._experimental.mcp_server import db
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
_, signing_key = jwt_oauth_identity
@ -14692,7 +14712,8 @@ async def test_managed_refresh_retains_only_the_actual_grant(
from types import SimpleNamespace
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints
from litellm.proxy._experimental.mcp_server import db
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
_, signing_key = jwt_oauth_identity
@ -14774,7 +14795,8 @@ async def test_managed_resource_metadata_excludes_refresh_permission(monkeypatch
],
)
async def test_connect_requires_the_requested_vendor_grant(monkeypatch, granted, required, expected):
from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints
from litellm.proxy._experimental.mcp_server import db
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints
credential = {**_stored_grant(), "scopes": granted}
read = AsyncMock(return_value=credential)
@ -14836,3 +14858,137 @@ def test_aggregate_protected_resource_advertises_header_server_scopes(
)
with patch.object(mcp_server_manager, "global_mcp_server_manager", manager):
assert _build_aggregate_protected_resource_response(request)["scopes_supported"] == expected
@pytest.mark.asyncio
async def test_token_endpoint_routes_the_jwt_bearer_grant_to_the_aggregate_gateway():
"""A requesting app presents its identity provider's client_id, never a gateway DCR one, so
the aggregate dispatch keys on the grant type for RFC 7523 and hands the assertion to the
ID-JAG redeemer. A per-server token endpoint never serves that grant."""
from litellm.proxy._experimental.mcp_server import discoverable_endpoints
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import JWT_BEARER_GRANT_TYPE
from litellm.proxy._experimental.mcp_server.id_jag_grant import IdJagPrerequisites
request = _mock_callback_request("https://litellm.example.com/")
with patch("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit9670"):
response = await discoverable_endpoints.token_endpoint(
request=request,
grant_type=JWT_BEARER_GRANT_TYPE,
client_id="0oa-requesting-app",
client_secret="app-secret",
assertion="eyJ.assertion.sig",
resource=None,
)
assert response.status_code == 400
refused = json.loads(bytes(response.body))
assert refused["error"] == "unsupported_grant_type"
unmet = IdJagPrerequisites(
jwt_auth_enabled=False,
has_database=False,
licensed=False,
maps_jwts_to_virtual_keys=False,
shared_replay_store=False,
base_url=None,
requesting_apps=(),
).refusal()
assert unmet is not None and refused["error_description"] == unmet.description
with pytest.raises(HTTPException):
await discoverable_endpoints.token_endpoint(
request=request,
grant_type=JWT_BEARER_GRANT_TYPE,
client_id="0oa-requesting-app",
assertion="eyJ.assertion.sig",
mcp_server_name="no-such-server",
)
@pytest.mark.asyncio
async def test_token_endpoint_takes_the_client_id_from_http_basic_when_the_form_leaves_it_out():
"""RFC 6749 section 2.3.1: a confidential app may authenticate with HTTP Basic alone, so the
aggregate dispatch reads client_id from that header. With neither the form field nor the
header the request is a 400 invalid_request, not a 422."""
from base64 import b64encode
from litellm.proxy._experimental.mcp_server import discoverable_endpoints
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import JWT_BEARER_GRANT_TYPE
request = _mock_callback_request("https://litellm.example.com/")
request.headers = {"authorization": "Basic " + b64encode(b"0oa-requesting-app:app-secret").decode()}
with patch("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit9670"):
response = await discoverable_endpoints.token_endpoint(
request=request,
grant_type=JWT_BEARER_GRANT_TYPE,
client_id=None,
assertion="eyJ.assertion.sig",
resource=None,
)
assert response.status_code == 400
assert json.loads(bytes(response.body))["error"] == "unsupported_grant_type"
request.headers = {}
with pytest.raises(HTTPException) as refused:
await discoverable_endpoints.token_endpoint(
request=request, grant_type=JWT_BEARER_GRANT_TYPE, client_id=None, assertion="eyJ.assertion.sig"
)
assert refused.value.status_code == 400
assert refused.value.detail["error"] == "invalid_request"
@pytest.mark.parametrize("available", [True, False])
def test_root_metadata_is_the_gateways_own_and_lists_the_jwt_bearer_grant_when_it_is_served(available):
"""With no OAuth2 MCP server resolving at the root, the bare-origin document is the gateway's
own authorization server metadata: issuer is the origin an ID-JAG names as ``aud``, the token
endpoint is ``{base}/token``, and the jwt-bearer grant (with HTTP Basic client auth) is listed
exactly when the gateway serves it, under the same gates as the aggregate ``/mcp`` document."""
from fastapi import Request
from litellm.proxy._experimental.mcp_server import discoverable_endpoints
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import JWT_BEARER_GRANT_TYPE
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
global_mcp_server_manager.registry.clear()
request = MagicMock(spec=Request)
request.base_url = "https://llm.example.com/"
request.headers = {}
with (
patch.object(discoverable_endpoints, "jwt_bearer_available", return_value=available), # test-quality-ok: deployment gate seam
patch.object(discoverable_endpoints, "token_exchange_available", return_value=True), # test-quality-ok: deployment gate seam
):
document = discoverable_endpoints._build_oauth_authorization_server_response(request=request, mcp_server_name=None)
assert document["issuer"] == "https://llm.example.com"
assert document["token_endpoint"] == "https://llm.example.com/token"
assert (JWT_BEARER_GRANT_TYPE in document["grant_types_supported"]) is available
assert ("client_secret_basic" in document["token_endpoint_auth_methods_supported"]) is available
assert ("none" in document["token_endpoint_auth_methods_supported"]) is available
assert "urn:ietf:params:oauth:grant-type:token-exchange" in document["grant_types_supported"]
def test_root_metadata_owned_by_a_resolved_oauth2_server_never_lists_the_gateways_grant():
"""When exactly one OAuth2 MCP server resolves at the root, the root document describes that
upstream's token endpoint, which cannot redeem an ID-JAG, so the jwt-bearer grant is not
advertised there even though the gateway serves it; a requesting app uses the aggregate
``/.well-known/oauth-authorization-server/mcp`` document in that deployment."""
from fastapi import Request
from litellm.proxy._experimental.mcp_server import discoverable_endpoints
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import JWT_BEARER_GRANT_TYPE
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
server = _create_oauth2_server()
global_mcp_server_manager.registry.clear()
global_mcp_server_manager.registry[server.server_id] = server
request = MagicMock(spec=Request)
request.base_url = "https://llm.example.com/"
request.headers = {}
try:
with (
patch.object(discoverable_endpoints, "jwt_bearer_available", return_value=True), # test-quality-ok: deployment gate seam
patch.object(discoverable_endpoints, "token_exchange_available", return_value=True), # test-quality-ok: deployment gate seam
patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=server), # test-quality-ok: resolver seam
):
document = discoverable_endpoints._build_oauth_authorization_server_response(request=request, mcp_server_name=None)
finally:
global_mcp_server_manager.registry.clear()
assert document["token_endpoint"] == "https://llm.example.com/test_oauth/token"
assert JWT_BEARER_GRANT_TYPE not in document["grant_types_supported"]
assert document["grant_types_supported"] == ["authorization_code", "refresh_token"]
assert document["token_endpoint_auth_methods_supported"] == ["client_secret_post"]

204
tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
View file

@ -20,10 +20,12 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
CONNECT_FLOW_COOKIE_PREFIX,
GATEWAY_AUTH_CODE_PREFIX,
GATEWAY_AUTH_CODE_TTL_SECONDS,
JWT_BEARER_GRANT_TYPE,
MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS,
MAX_CLIENT_ID_LENGTH,
SUBJECT_TOKEN_TYPES,
TOKEN_EXCHANGE_GRANT_TYPE,
AssertionIdentity,
ConsentTeam,
MintedProxyCredential,
SubjectIdentity,
@ -44,6 +46,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
open_gateway_dcr_client,
register_aggregate_client,
revoke_refresh_token,
supported_grant_types,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
SessionBearerAdmitted,
@ -1771,10 +1774,12 @@ CONSENT_TEAMS = (ConsentTeam(team_id="team-a", team_alias="Team A"), ConsentTeam
class _Minter:
def __init__(self, result=None):
self.calls = []
self.allowed_routes = []
self.result = result
async def __call__(self, user_id, team_id):
async def __call__(self, user_id, team_id, allowed_routes=None):
self.calls.append((user_id, team_id))
self.allowed_routes.append(allowed_routes)
if self.result is not None:
return self.result
return MintedProxyCredential(key=f"sk-cli-{user_id}", expires_in=3600, user_id=user_id, team_id=team_id)
@ -3296,3 +3301,200 @@ async def test_gateway_scope_still_narrows_when_the_resource_does_not_resolve()
_, cookies = _flow_cookie_from(response)
sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")
assert sealed_flow["resource_server_id"] == "alpha-id"
ID_JAG = "eyJhbGciOiJSUzI1NiIsInR5cCI6Im9hdXRoLWlkLWphZytqd3QifQ.eyJzdWIiOiJ1MSJ9.idp-signature"
REQUESTING_APP = "0oa-requesting-app"
ASSERTED = AssertionIdentity(
user_id="u1",
team_id="team-b",
replay_key="mcp_gateway_dcr:id_jag:jti-1",
replay_ttl_seconds=300,
scope="litellm:mcp",
allowed_routes=("mcp_inference_routes",),
)
class _Redeemer:
def __init__(self, result=None):
self.calls = []
self.result = result
async def __call__(self, assertion, client_id, client_secret, request):
self.calls.append((assertion, client_id, client_secret, request.url.path))
return self.result if self.result is not None else ASSERTED
async def _redeem_assertion_native(minter, redeemer, cache=None, client_id=REQUESTING_APP, **overrides):
arguments = {
"grant_type": JWT_BEARER_GRANT_TYPE,
"assertion": ID_JAG,
"client_secret": "app-secret",
"redeem_assertion": redeemer,
}
return await _redeem_native(None, client_id, minter, cache=cache, **{**arguments, **overrides})
@pytest.mark.asyncio
async def test_jwt_bearer_grant_mints_the_proxy_credential_for_the_asserted_user_once_per_assertion():
"""RFC 7523: an allowlisted requesting app redeems the ID-JAG its identity provider minted for
the same credential the consent flow mints, attributed to the user the assertion names and
the team asked for, with the assertion's scope stated and no refresh token (RFC 7521 section
4.1, the app renews with a fresh assertion). The assertion's jti is single-use: presenting it again is a replay and mints
nothing, while a fresh assertion for the same user is served."""
minter, redeemer, cache = _Minter(), _Redeemer(), DualCache()
response = await _redeem_assertion_native(minter, redeemer, cache=cache)
assert response.status_code == 200
assert response.headers["cache-control"] == "no-store"
body = json.loads(response.body)
assert redeemer.calls == [(ID_JAG, REQUESTING_APP, "app-secret", "/token")]
assert minter.calls == [("u1", "team-b")]
assert body["access_token"] == "sk-cli-u1"
assert body["token_type"] == "Bearer"
assert body["scope"] == "litellm:mcp"
assert "issued_token_type" not in body
assert (body["user_id"], body["team_id"]) == ("u1", "team-b")
assert "refresh_token" not in body
replayed = await _redeem_assertion_native(minter, redeemer, cache=cache)
assert replayed.status_code == 400
assert json.loads(replayed.body)["error"] == "invalid_grant"
assert minter.calls == [("u1", "team-b")]
fresh = _Redeemer(ASSERTED.model_copy(update={"replay_key": "mcp_gateway_dcr:id_jag:jti-2", "scope": None}))
again = await _redeem_assertion_native(minter, fresh, cache=cache)
assert again.status_code == 200
assert "scope" not in json.loads(again.body)
assert minter.calls == [("u1", "team-b"), ("u1", "team-b")]
@pytest.mark.asyncio
async def test_jwt_bearer_grant_claims_the_jti_before_it_mints():
"""A mint the gateway could not complete still consumed the assertion: the client gets the
503 and comes back with a fresh assertion, never with the same one."""
minter, cache = _Minter(result="unavailable"), DualCache()
first = await _redeem_assertion_native(minter, _Redeemer(), cache=cache)
assert first.status_code == 503
retried = await _redeem_assertion_native(_Minter(), _Redeemer(), cache=cache)
assert retried.status_code == 400
assert json.loads(retried.body)["error"] == "invalid_grant"
@pytest.mark.parametrize(
"error, status",
[
("invalid_client", 401),
("invalid_grant", 400),
("invalid_request", 400),
("unsupported_grant_type", 400),
("temporarily_unavailable", 503),
],
)
@pytest.mark.asyncio
async def test_jwt_bearer_grant_refusals_keep_their_oauth_error_and_status(error, status):
minter, cache = _Minter(), DualCache()
refusal = SubjectTokenRefusal(error=error, description="why")
response = await _redeem_assertion_native(minter, _Redeemer(refusal), cache=cache)
assert response.status_code == status
assert json.loads(response.body) == {"error": error, "error_description": "why"}
assert minter.calls == []
served = await _redeem_assertion_native(minter, _Redeemer(), cache=cache)
assert served.status_code == 200
@pytest.mark.asyncio
async def test_jwt_bearer_grant_requires_an_assertion_and_the_proxy_api_target():
minter, redeemer = _Minter(), _Redeemer()
missing = await _redeem_assertion_native(minter, redeemer, assertion=None)
assert missing.status_code == 400
assert json.loads(missing.body)["error"] == "invalid_request"
foreign = await _redeem_assertion_native(minter, redeemer, resource="https://elsewhere.example/api")
assert foreign.status_code == 400
assert json.loads(foreign.body)["error"] == "invalid_target"
assert redeemer.calls == []
assert minter.calls == []
@pytest.mark.asyncio
async def test_jwt_bearer_grant_is_refused_where_no_redeemer_is_wired():
minter = _Minter()
response = await _redeem_native(
None, REQUESTING_APP, minter, grant_type=JWT_BEARER_GRANT_TYPE, assertion=ID_JAG, client_secret=None
)
assert response.status_code == 400
assert json.loads(response.body)["error"] == "unsupported_grant_type"
assert minter.calls == []
@pytest.mark.asyncio
async def test_an_unknown_grant_type_names_the_jwt_bearer_grant_among_the_served_ones():
response = await _redeem_native(None, REQUESTING_APP, _Minter(), grant_type="urn:example:nope")
assert response.status_code == 400
assert JWT_BEARER_GRANT_TYPE in json.loads(response.body)["error_description"]
@pytest.mark.parametrize(
"token_exchange, jwt_bearer, expected",
[
(False, False, ("authorization_code", "refresh_token")),
(True, False, ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE)),
(False, True, ("authorization_code", "refresh_token", JWT_BEARER_GRANT_TYPE)),
(True, True, ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE, JWT_BEARER_GRANT_TYPE)),
],
)
def test_the_jwt_bearer_grant_is_advertised_exactly_when_it_is_served(token_exchange, jwt_bearer, expected):
assert supported_grant_types(token_exchange, jwt_bearer) == expected
contract = native_client_auth_contract(_request("/.well-known/litellm-cli-auth"), token_exchange, jwt_bearer)
assert contract["grant_types_supported"] == expected
@pytest.mark.asyncio
async def test_jwt_bearer_grant_narrows_the_credential_to_the_routes_the_scope_bought():
"""The route groups the assertion's scope resolved to reach the minter, so the credential is
enforced on every proxy route the way a virtual key's allowed_routes is, and the response states
exactly the scope that was enforced. No scope on the assertion mints the full credential."""
minter = _Minter()
both = ASSERTED.model_copy(
update={"scope": "litellm:mcp litellm:llm", "allowed_routes": ("mcp_inference_routes", "openai_routes")}
)
scoped = await _redeem_assertion_native(minter, _Redeemer(both), cache=DualCache())
assert scoped.status_code == 200
assert json.loads(scoped.body)["scope"] == "litellm:mcp litellm:llm"
full = await _redeem_assertion_native(
minter, _Redeemer(ASSERTED.model_copy(update={"scope": None, "allowed_routes": None})), cache=DualCache()
)
assert full.status_code == 200
assert "scope" not in json.loads(full.body)
assert minter.allowed_routes == [("mcp_inference_routes", "openai_routes"), None]
@pytest.mark.asyncio
async def test_jwt_bearer_grant_accepts_the_mcp_resource_indicator_as_its_target():
"""RFC 8707 clients send the ID-JAG's audience as ``resource``, and that audience may be the MCP
resource ``{base}/mcp``, so it is a valid target next to the proxy origin. Anything else stays
``invalid_target``."""
minter = _Minter()
served = await _redeem_assertion_native(
minter, _Redeemer(), cache=DualCache(), resource=PROXY_API_RESOURCE.rstrip("/") + "/mcp"
)
assert served.status_code == 200
foreign = await _redeem_assertion_native(
minter, _Redeemer(), cache=DualCache(), resource=PROXY_API_RESOURCE.rstrip("/") + "/other/mcp"
)
assert json.loads(foreign.body)["error"] == "invalid_target"
assert minter.calls == [("u1", "team-b")]
@pytest.mark.parametrize("jwt_bearer_available", [True, False])
def test_aggregate_metadata_advertises_basic_client_auth_only_with_the_grant_that_reads_it(jwt_bearer_available):
"""Only the jwt-bearer redemption reads an HTTP Basic client secret, the other grants take
``client_secret`` from the form, so ``client_secret_basic`` is advertised only when that grant
is on. A client following the document otherwise sends Basic alone and authenticates with no
secret."""
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_build_aggregate_authorization_server_response,
)
document = _build_aggregate_authorization_server_response(_request(), True, jwt_bearer_available)
methods = document["token_endpoint_auth_methods_supported"]
assert isinstance(methods, list)
assert ("client_secret_basic" in methods) is jwt_bearer_available
assert "none" in methods and "client_secret_post" in methods

738
tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py Normal file
View file

@ -0,0 +1,738 @@
import hashlib
import logging
import math
from base64 import b64encode
from collections.abc import Mapping, Sequence
from typing import Final
import jwt
import pytest
from fastapi import HTTPException
from litellm.caching.caching import DualCache
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
AssertionIdentity,
ConsentTeam,
ReloadUserFailure,
SubjectIdentity,
SubjectTokenRefusal,
)
from litellm.proxy._experimental.mcp_server.id_jag_grant import (
ID_JAG_REQUIRED_CLAIMS,
REJECTED_ASSERTION,
TEAM_HEADER,
AssertionRedemption,
GrantedScope,
IdJagPrerequisites,
LookupTeams,
_claim_reader,
basic_client_id,
gateway_audiences,
granted_scope,
identity_from_id_jag,
jwt_bearer_available,
names_mcp_resource,
proxy_user_lookup,
select_team,
subject_resolver,
)
from litellm.proxy._types import IdJagRequestingApp, LiteLLM_JWTAuth, LiteLLMRoutes
from litellm.proxy.auth.handle_jwt import (
ID_JAG_TYP,
AssertionRejection,
JWKSUnreachableError,
JWTHandler,
jwks_unavailable_exception,
)
ISSUER = "https://idp.example.com"
GATEWAY = "https://gateway.example.com"
AUDIENCES = (f"{GATEWAY}/mcp", GATEWAY)
NOW = 1_800_000_000.0
PUBLIC_APP = IdJagRequestingApp(client_id="0oa-public")
CONFIDENTIAL_APP = IdJagRequestingApp(client_id="0oa-confidential", client_secret="s3cret")
SPACED_APP = IdJagRequestingApp(client_id="0oa spaced", client_secret="app secret")
EVERY_GATE_HOLDS = {
"jwt_auth_enabled": True,
"has_database": True,
"licensed": True,
"maps_jwts_to_virtual_keys": False,
"shared_replay_store": True,
"base_url": "https://gateway.example.com",
"requesting_apps": (PUBLIC_APP, CONFIDENTIAL_APP, SPACED_APP),
}
JWKS_DOWN = jwks_unavailable_exception(JWKSUnreachableError("ConnectError fetching the IdP's JWKS after 3 attempts"))
def _claims(**overrides: object) -> dict[str, object]:
return {
"iss": ISSUER,
"sub": "okta-user-1",
"aud": AUDIENCES[1],
"exp": NOW + 300,
"iat": NOW - 5,
"jti": "jti-1",
"client_id": PUBLIC_APP.client_id,
"scope": "litellm:mcp",
**overrides,
}
def _assertion(typ: str = ID_JAG_TYP, **claims: object) -> str:
return jwt.encode(
_claims(**claims),
"unused-by-the-fake-verifier-but-long-enough-for-hs256",
algorithm="HS256",
headers={"typ": typ},
)
class _Verifier:
"""Stands in for the injected ``DecodeAssertion`` port; records each call as a new tuple."""
def __init__(
self, result: Mapping[str, object] | AssertionRejection | None = None, raises: BaseException | None = None
) -> None:
self.calls: tuple[tuple[str, tuple[str, ...], tuple[str, ...]], ...] = ()
self.result = result
self.raises = raises
async def __call__(
self, assertion: str, audience: Sequence[str], required_claims: Sequence[str]
) -> Mapping[str, object] | AssertionRejection:
self.calls = (*self.calls, (assertion, tuple(audience), tuple(required_claims))) # rebind-ok: immutable record
if self.raises is not None:
raise self.raises
return self.result if self.result is not None else jwt.decode(assertion, options={"verify_signature": False})
class _Resolver:
"""Stands in for the injected ``ResolveSubject`` port; records each call as a new tuple."""
def __init__(
self,
user_id: str | None = "u1",
raises: BaseException | None = None,
result: SubjectIdentity | SubjectTokenRefusal | None = None,
) -> None:
self.calls: tuple[tuple[Mapping[str, object], str | None], ...] = ()
self.user_id = user_id
self.raises = raises
self.result = result
async def __call__(
self, claims: Mapping[str, object], team_header: str | None
) -> SubjectIdentity | SubjectTokenRefusal | None:
self.calls = (*self.calls, (dict(claims), team_header)) # rebind-ok: immutable record
if self.raises is not None:
raise self.raises
if self.result is not None:
return self.result
return None if self.user_id is None else SubjectIdentity(user_id=self.user_id, team_id=team_header)
def _redemption(
assertion: str,
client_id: str = PUBLIC_APP.client_id,
client_secret: str | None = None,
authorization: str | None = None,
team_id: str | None = "team-b",
leeway: float = 0.0,
) -> AssertionRedemption:
return AssertionRedemption(
assertion=assertion,
client_id=client_id,
client_secret=client_secret,
authorization_header=authorization,
team_id=team_id,
audiences=AUDIENCES,
now=NOW,
leeway_seconds=leeway,
)
@pytest.mark.asyncio
async def test_the_replay_marker_outlives_the_leeway_the_verifier_grants_exp():
"""PyJWT accepts an assertion until ``exp`` plus the handler's leeway, so the marker must
live that long too, including for an assertion already past ``exp`` but inside the leeway."""
within_leeway = await _identity(_redemption(_assertion(), leeway=30), _Verifier(), _Resolver())
assert isinstance(within_leeway, AssertionIdentity)
assert within_leeway.replay_ttl_seconds == 330
just_expired = await _identity(
_redemption(_assertion(), leeway=30), _Verifier(result=_claims(exp=NOW - 10)), _Resolver()
)
assert isinstance(just_expired, AssertionIdentity)
assert just_expired.replay_ttl_seconds == 20
async def _identity(
redemption: AssertionRedemption,
verifier: _Verifier | None = None,
resolver: _Resolver | None = None,
**unmet: object,
) -> AssertionIdentity | SubjectTokenRefusal:
return await identity_from_id_jag(
redemption,
prerequisites=IdJagPrerequisites(**{**EVERY_GATE_HOLDS, **unmet}),
decode=verifier if verifier is not None else _Verifier(),
resolve_subject=resolver if resolver is not None else _Resolver(),
)
@pytest.mark.asyncio
async def test_a_valid_id_jag_names_the_user_its_single_use_key_and_its_scope():
"""The verifier is asked for the gateway's own issuer values as audience and for every claim
the ID-JAG profile requires, the subject lookup sees the verified claims, and the identity
carries the team asked for, a replay key derived from iss and jti, a marker TTL that outlives
the assertion, and the assertion's scope."""
assertion, verifier, resolver = _assertion(), _Verifier(), _Resolver()
identity = await _identity(_redemption(assertion), verifier, resolver)
assert isinstance(identity, AssertionIdentity)
assert verifier.calls == ((assertion, AUDIENCES, ID_JAG_REQUIRED_CLAIMS),)
assert {"iss", "sub", "aud", "exp", "iat", "jti", "client_id"} <= set(ID_JAG_REQUIRED_CLAIMS)
assert resolver.calls == ((_claims(), "team-b"),)
assert (identity.user_id, identity.team_id, identity.scope) == ("u1", "team-b", "litellm:mcp")
assert identity.allowed_routes == ("mcp_inference_routes",)
assert identity.replay_key == "mcp_gateway_dcr:id_jag:" + hashlib.sha256(f"{ISSUER}\njti-1".encode()).hexdigest()
assert identity.replay_ttl_seconds == math.ceil(300)
other_jti = await _identity(_redemption(_assertion(jti="jti-2")), _Verifier(), _Resolver())
assert isinstance(other_jti, AssertionIdentity)
assert other_jti.replay_key != identity.replay_key
teamless = await _identity(_redemption(_assertion(scope=""), team_id=None))
assert isinstance(teamless, AssertionIdentity)
assert (teamless.team_id, teamless.scope, teamless.allowed_routes) == (None, None, None)
@pytest.mark.asyncio
async def test_a_requesting_app_the_operator_did_not_allowlist_learns_nothing_about_the_assertion():
verifier, resolver = _Verifier(), _Resolver()
refused = await _identity(_redemption(_assertion(), client_id="0oa-stranger"), verifier, resolver)
assert refused == SubjectTokenRefusal(
error="invalid_client", description="client_id is not an allowlisted requesting app"
)
assert verifier.calls == ()
assert resolver.calls == ()
def _basic(client_id: str, secret: str) -> str:
return "Basic " + b64encode(f"{client_id}:{secret}".encode()).decode()
@pytest.mark.parametrize(
"client_secret, authorization, accepted",
[
(None, None, False),
("wrong", None, False),
("s3cr\u00e9t", None, False),
("s3cret", None, True),
(None, _basic(CONFIDENTIAL_APP.client_id, "s3cr\u00e9t"), False),
(None, _basic(CONFIDENTIAL_APP.client_id, "s3cret"), True),
(None, _basic(CONFIDENTIAL_APP.client_id, "wrong"), False),
(None, _basic(PUBLIC_APP.client_id, "s3cret"), False),
("s3cret", _basic(CONFIDENTIAL_APP.client_id, "wrong"), False),
(None, "Basic not-base64!", False),
(None, "Bearer s3cret", False),
],
)
@pytest.mark.asyncio
async def test_a_confidential_app_must_authenticate_with_its_secret(client_secret, authorization, accepted):
verifier = _Verifier()
assertion = _assertion(client_id=CONFIDENTIAL_APP.client_id)
redemption = _redemption(assertion, CONFIDENTIAL_APP.client_id, client_secret, authorization)
identity = await _identity(redemption, verifier)
if accepted:
assert isinstance(identity, AssertionIdentity)
return
assert identity == SubjectTokenRefusal(error="invalid_client", description="client authentication failed")
assert verifier.calls == ()
@pytest.mark.parametrize(
"client_secret, authorization, accepted",
[
("app secret", None, True),
(None, "Basic " + b64encode(b"0oa+spaced:app+secret").decode(), True),
(None, "Basic " + b64encode(b"0oa%20spaced:app%20secret").decode(), True),
(None, "Basic " + b64encode(b"0oa spaced:app secret").decode(), True),
(None, "Basic " + b64encode(b"0oa+spaced:app+wrong").decode(), False),
(None, "Basic " + b64encode(b"0oa-spaced:app+secret").decode(), False),
],
)
@pytest.mark.asyncio
async def test_basic_credentials_are_form_decoded_so_a_space_in_an_id_or_secret_still_matches(
client_secret, authorization, accepted
):
"""RFC 6749 section 2.3.1 form-encodes the id and secret before base64, so ``+`` is a space. The
endpoint hands over the ``client_id`` it read from the header, so the header must name the same app."""
assertion = _assertion(client_id=SPACED_APP.client_id)
identity = await _identity(_redemption(assertion, SPACED_APP.client_id, client_secret, authorization), _Verifier())
if accepted:
assert isinstance(identity, AssertionIdentity)
return
assert identity == SubjectTokenRefusal(error="invalid_client", description="client authentication failed")
@pytest.mark.asyncio
async def test_a_public_app_is_identified_by_client_id_alone():
identity = await _identity(_redemption(_assertion(), client_secret="anything", authorization="Bearer x"))
assert isinstance(identity, AssertionIdentity)
@pytest.mark.parametrize(
"assertion, description",
[
("not-a-jwt", "assertion is not a JWT"),
(_assertion(typ="JWT"), f"assertion typ must be {ID_JAG_TYP}"),
(_assertion(typ="at+jwt"), f"assertion typ must be {ID_JAG_TYP}"),
],
)
@pytest.mark.asyncio
async def test_a_bearer_or_id_token_is_not_an_id_jag(assertion, description):
verifier = _Verifier()
refused = await _identity(_redemption(assertion), verifier)
assert refused == SubjectTokenRefusal(error="invalid_grant", description=description)
assert verifier.calls == ()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"verifier, reason",
[
(_Verifier(result=AssertionRejection(reason="Signature verification failed")), "Signature verification"),
(_Verifier(result=AssertionRejection(reason="Audience doesn't match")), "Audience"),
(_Verifier(result=AssertionRejection(reason='Token is missing the "jti" claim')), "jti"),
(_Verifier(result=AssertionRejection(reason="Signature has expired")), "expired"),
(_Verifier(result=AssertionRejection(reason="iss is not an issuer this gateway trusts")), "iss"),
(_Verifier(raises=Exception("Validation fails: Invalid audience")), "Invalid audience"),
(_Verifier(result=AssertionRejection(reason="bad kid\nWARNING forged line")), "bad kid WARNING forged line"),
],
)
async def test_an_assertion_the_verifier_rejects_is_refused_with_the_reason_kept_in_the_log(verifier, reason, caplog):
resolver = _Resolver()
with caplog.at_level(logging.WARNING, logger="litellm.proxy"):
refused = await _identity(_redemption(_assertion()), verifier, resolver)
assert refused == SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION)
assert reason not in refused.description
assert reason in caplog.text
assert resolver.calls == ()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"claims, reason",
[
({"client_id": "0oa-confidential"}, "client_id"),
({"exp": NOW}, "expired"),
({"exp": NOW - 1}, "expired"),
({"jti": ""}, "jti"),
({"jti": 7}, "jti"),
({"iss": None}, "iss"),
({"cnf": {"jkt": "thumbprint"}}, "proof"),
],
)
async def test_verified_claims_that_do_not_fit_the_profile_are_refused(claims, reason, caplog):
resolver = _Resolver()
with caplog.at_level(logging.WARNING, logger="litellm.proxy"):
refused = await _identity(_redemption(_assertion()), _Verifier(result=_claims(**claims)), resolver)
assert refused == SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION)
assert reason in caplog.text
assert resolver.calls == ()
@pytest.mark.asyncio
async def test_an_assertion_that_names_no_user_the_gateway_knows_cannot_be_redeemed(caplog):
with caplog.at_level(logging.WARNING, logger="litellm.proxy"):
refused = await _identity(_redemption(_assertion()), resolver=_Resolver(user_id=None))
assert refused == SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION)
assert "names no user" in caplog.text
missing = await _identity(
_redemption(_assertion()), resolver=_Resolver(raises=ValueError("User doesn't exist in db. 'user_id'=u1"))
)
assert missing == SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"verifier, resolver",
[
(_Verifier(raises=JWKS_DOWN), _Resolver()),
(_Verifier(), _Resolver(raises=HTTPException(status_code=503, detail="database unavailable"))),
],
)
async def test_a_gateway_outage_is_reported_as_retryable_not_as_a_bad_assertion(verifier, resolver):
refused = await _identity(_redemption(_assertion()), verifier, resolver)
assert isinstance(refused, SubjectTokenRefusal)
assert refused.error == "temporarily_unavailable"
assert "retry" in refused.description
assert "assertion" in refused.description
assert "subject_token" not in refused.description
@pytest.mark.parametrize(
"unmet, description",
[
({"jwt_auth_enabled": False}, "JWT auth is not enabled"),
({"has_database": False}, "no database"),
({"licensed": False}, "enterprise only"),
({"maps_jwts_to_virtual_keys": True}, "maps IdP tokens to virtual keys"),
({"shared_replay_store": False}, "Redis"),
({"base_url": None}, "PROXY_BASE_URL"),
({"requesting_apps": ()}, "no requesting apps allowlisted"),
],
)
@pytest.mark.asyncio
async def test_the_deployment_gates_refuse_before_any_verification(unmet, description):
verifier = _Verifier()
refused = await _identity(_redemption(_assertion()), verifier, **unmet)
assert isinstance(refused, SubjectTokenRefusal)
assert refused.error == "unsupported_grant_type"
assert description in refused.description
assert verifier.calls == ()
assert IdJagPrerequisites(**{**EVERY_GATE_HOLDS, **unmet}).available is False
assert IdJagPrerequisites(**EVERY_GATE_HOLDS).available is True
def _running_jwt_handler(litellm_jwtauth: LiteLLM_JWTAuth | None) -> JWTHandler:
handler: Final = JWTHandler()
if litellm_jwtauth is not None:
handler.update_environment(prisma_client=None, user_api_key_cache=DualCache(), litellm_jwtauth=litellm_jwtauth)
return handler
@pytest.mark.parametrize(
"general_settings, prisma_client, premium_user, litellm_jwtauth, expected",
[
({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), True),
(
{"enable_jwt_auth": True, "redis": None},
object(),
True,
LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]),
False,
),
(
{"enable_jwt_auth": True, "base_url": ""},
object(),
True,
LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]),
False,
),
({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(id_jag_requesting_apps=[]), False),
(
{"enable_jwt_auth": True},
object(),
True,
LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP], virtual_key_claim_field="vkey"),
False,
),
({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(), False),
({"enable_jwt_auth": True}, object(), True, None, False),
({}, object(), True, LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), False),
({"enable_jwt_auth": True}, None, True, LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), False),
({"enable_jwt_auth": True}, object(), False, LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), False),
],
)
def test_availability_is_read_from_the_running_proxy(
monkeypatch, general_settings, prisma_client, premium_user, litellm_jwtauth, expected
):
monkeypatch.setenv("PROXY_BASE_URL", general_settings.get("base_url", "https://gateway.example.com"))
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", premium_user)
monkeypatch.setattr("litellm.proxy.proxy_server.redis_usage_cache", general_settings.get("redis", object()))
monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", _running_jwt_handler(litellm_jwtauth))
assert jwt_bearer_available() is expected
def test_the_assertion_audience_is_the_fixed_gateway_url_not_the_request_host():
"""An ID-JAG must name this gateway's MCP resource or its authorization server under the
operator-fixed PROXY_BASE_URL. Without it there is no audience, so there is no grant."""
assert gateway_audiences("https://gateway.example.com") == (
"https://gateway.example.com/mcp",
"https://gateway.example.com",
)
assert gateway_audiences(None) == ()
@pytest.mark.asyncio
async def test_a_team_the_subject_lookup_cannot_settle_is_refused_before_the_assertion_is_consumed():
"""The resolver's refusal (a user on several teams who named none, or a header naming a team
they are not on) comes back as is, so the client fixes its request and presents the same
assertion again instead of finding its jti burned."""
refusal = SubjectTokenRefusal(error="invalid_request", description=f"send the {TEAM_HEADER} header")
refused = await _identity(_redemption(_assertion()), resolver=_Resolver(result=refusal))
assert refused is refusal
teamed = await _identity(
_redemption(_assertion()), resolver=_Resolver(result=SubjectIdentity(user_id="u1", team_id="t9"))
)
assert isinstance(teamed, AssertionIdentity)
assert teamed.team_id == "t9"
ENG = ConsentTeam(team_id="team-eng", team_alias="engineering")
OPS = ConsentTeam(team_id="team-ops", team_alias=None)
@pytest.mark.parametrize(
"team_header, claim_teams, teams, expected",
[
(None, (), (), None),
(None, (), (ENG,), "team-eng"),
("team-eng", (), (ENG, OPS), "team-eng"),
("engineering", (), (ENG, OPS), "team-eng"),
("team-ops", ("team-eng",), (ENG, OPS), "team-ops"),
(None, ("team-ops",), (ENG, OPS), "team-ops"),
(None, ("team-sales", "team-ops"), (ENG, OPS), "team-ops"),
(None, ("0oa-xaa-public-app",), (ENG,), "team-eng"),
(None, ("0oa-xaa-public-app",), (), None),
],
)
def test_select_team_prefers_the_header_then_the_assertion_then_the_only_team(
team_header, claim_teams, teams, expected
):
assert select_team(team_header, claim_teams, teams) == expected
class _Reader:
def __init__(self, user_id=None, email=None, team_ids=(), allowed_domain=None):
self.user_id, self.email, self.team_ids, self.allowed_domain = user_id, email, team_ids, allowed_domain
def is_allowed_domain(self, user_email):
return self.allowed_domain is None or user_email.split("@")[-1] == self.allowed_domain
def get_user_id(self, token, default_value):
return self.user_id or default_value
def get_user_email(self, token, default_value):
return self.email or default_value
def get_all_jwt_team_ids(self, token):
return list(self.team_ids)
async def _known_user(user_id: str, user_email: str | None) -> str | None:
return user_id if user_id in {"okta-user-1", "u-mapped"} else None
async def _two_teams(user_id: str) -> tuple[ConsentTeam, ...]:
return (ENG, OPS)
async def _teams_unreadable(user_id: str) -> ReloadUserFailure:
return "unavailable"
def _teams_failing(failure: ReloadUserFailure) -> LookupTeams:
async def lookup_teams(user_id: str) -> tuple[ConsentTeam, ...] | ReloadUserFailure:
return failure
return lookup_teams
@pytest.mark.asyncio
async def test_the_subject_resolver_binds_the_claims_to_a_known_user_and_their_asserted_team():
"""The user comes from the JWT auth user mapping, else ``sub``, and must exist. The team the
assertion asserts settles a multi-team user, the header still wins, and unreadable teams are a
retryable refusal rather than a credential for the wrong team."""
asserted = subject_resolver(_Reader(team_ids=("team-ops",)), _known_user, _two_teams)
assert await asserted(_claims(), None) == SubjectIdentity(user_id="okta-user-1", team_id="team-ops")
assert await asserted(_claims(), "team-eng") == SubjectIdentity(user_id="okta-user-1", team_id="team-eng")
mapped = subject_resolver(_Reader(user_id="u-mapped", email="m@example.com"), _known_user, _two_teams)
assert await mapped(_claims(), "engineering") == SubjectIdentity(user_id="u-mapped", team_id="team-eng")
unknown = subject_resolver(_Reader(), _known_user, _two_teams)
assert await unknown(_claims(sub="nobody"), None) is None
unreadable = await subject_resolver(_Reader(), _known_user, _teams_unreadable)(_claims(), None)
assert isinstance(unreadable, SubjectTokenRefusal)
assert unreadable.error == "temporarily_unavailable"
@pytest.mark.parametrize("failure", ["no_active_key", "unresolvable"])
@pytest.mark.asyncio
async def test_a_blocked_or_unmapped_user_is_a_bad_grant_not_an_outage_to_retry(failure):
"""A deactivated user, or one the resolver cannot map, can retry forever and never succeed, so
the refusal is invalid_grant. Only a database the gateway could not reach is retryable."""
refused = await subject_resolver(_Reader(), _known_user, _teams_failing(failure))(_claims(), None)
assert isinstance(refused, SubjectTokenRefusal)
assert refused.error == "invalid_grant"
assert "retry" not in refused.description
outage = await subject_resolver(_Reader(), _known_user, _teams_failing("faulted"))(_claims(), None)
assert isinstance(outage, SubjectTokenRefusal)
assert outage.error == "temporarily_unavailable"
def test_the_assertion_team_claims_ignore_team_id_default():
"""``team_id_default`` is how bearer JWT auth picks a team for a request, not a claim the
assertion made, so a user outside that team must still get their only team or no team."""
handler = JWTHandler()
handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_default="default-team")
claim_teams = tuple(_claim_reader(handler).get_all_jwt_team_ids(token={"iss": ISSUER, "sub": "u1"}))
assert claim_teams == ()
assert select_team(None, claim_teams, (ENG,)) == "team-eng"
assert select_team(None, claim_teams, ()) is None
@pytest.mark.parametrize(
"team_header, claim_teams, teams, error, mention",
[
("team-sales", (), (ENG, OPS), "invalid_request", TEAM_HEADER),
("engineering", (), (), "invalid_request", TEAM_HEADER),
(None, (), (ENG, OPS), "invalid_request", TEAM_HEADER),
(None, ("team-eng", "team-ops"), (ENG, OPS), "invalid_request", TEAM_HEADER),
(None, ("0oa-xaa-public-app",), (ENG, OPS), "invalid_request", TEAM_HEADER),
],
)
def test_select_team_refuses_a_team_the_user_is_not_on_or_an_unnamed_choice(
team_header, claim_teams, teams, error, mention
):
refused = select_team(team_header, claim_teams, teams)
assert isinstance(refused, SubjectTokenRefusal)
assert refused.error == error
assert mention in refused.description
@pytest.mark.parametrize(
"authorization, expected",
[
(None, None),
("Bearer abc", None),
("Basic not-base64!", None),
(_basic("0oa-confidential", "s3cret"), "0oa-confidential"),
("basic " + b64encode(b"0oa%20x:s3cret").decode(), "0oa x"),
("Basic " + b64encode(b"0oa+x:s3cret").decode(), "0oa x"),
("Basic " + b64encode(b"no-separator").decode(), None),
],
)
def test_basic_client_id_reads_the_app_from_the_authorization_header(authorization, expected):
assert basic_client_id(authorization) == expected
MCP_ROUTE_GROUPS = ("mcp_inference_routes",)
LLM_ROUTE_GROUPS = ("llm_api_routes_without_mcp",)
@pytest.mark.parametrize(
"claim, scope, allowed_routes",
[
(None, None, None),
("", None, None),
("litellm:mcp", "litellm:mcp", MCP_ROUTE_GROUPS),
("litellm:llm", "litellm:llm", LLM_ROUTE_GROUPS),
("litellm:mcp litellm:llm", "litellm:mcp litellm:llm", MCP_ROUTE_GROUPS + LLM_ROUTE_GROUPS),
("litellm:llm litellm:mcp litellm:llm", "litellm:mcp litellm:llm", MCP_ROUTE_GROUPS + LLM_ROUTE_GROUPS),
],
)
def test_the_scope_claim_buys_exactly_the_route_groups_it_names(claim, scope, allowed_routes):
"""``litellm:mcp`` is the MCP gateway's data-plane routes, ``litellm:llm`` the LLM API routes
without their MCP members, both together are both, and no claim is the full credential. The
scope echoed back is canonical: deduplicated, in vocabulary order, so it is exactly what is
enforced. Every group named must be one the route check resolves."""
granted = granted_scope(claim, mcp_resource=False)
assert (granted.scope, granted.allowed_routes) == (scope, allowed_routes)
if allowed_routes is not None:
assert set(allowed_routes) <= set(LiteLLMRoutes._member_names_)
mcp_members = set(LiteLLMRoutes.mcp_inference_routes.value)
llm_members = set().union(*(LiteLLMRoutes[group].value for group in LLM_ROUTE_GROUPS))
assert "/chat/completions" in llm_members and "/mcp" not in llm_members
assert "/mcp" in mcp_members and "/chat/completions" not in mcp_members
@pytest.mark.parametrize("claim", ["openid", "litellm:mcp openid", "litellm:admin", 7, ["litellm:mcp"]])
@pytest.mark.asyncio
async def test_a_scope_the_gateway_does_not_know_is_refused_as_invalid_scope_before_any_lookup(claim):
refusal = granted_scope(claim, mcp_resource=False)
assert isinstance(refusal, SubjectTokenRefusal)
assert refusal.error == "invalid_scope"
assert "litellm:mcp" in refusal.description and "litellm:llm" in refusal.description
resolver = _Resolver()
refused = await _identity(_redemption(_assertion()), _Verifier(result=_claims(scope=claim)), resolver)
assert refused == refusal
assert resolver.calls == ()
@pytest.mark.asyncio
async def test_a_linked_sso_account_resolves_by_its_provider_subject_without_an_email_claim():
"""An SSO login stores the provider subject as ``sso_user_id`` on a user whose ``user_id`` is
something else, so the lookup must offer the subject as both, with no email to fall back on."""
from litellm.models.user import LiteLLM_UserTable
async def get_user(*, user_id, user_email, sso_user_id):
if sso_user_id == "00u-okta-subject" and user_email is None:
return LiteLLM_UserTable(user_id="local-user-7", user_email=None)
return None
lookup = proxy_user_lookup(get_user)
assert await lookup("00u-okta-subject", None) == "local-user-7"
assert await lookup("nobody", None) is None
@pytest.mark.parametrize("email", ["eve@evil.example", None])
@pytest.mark.asyncio
async def test_the_allowed_email_domain_gates_the_assertions_subject_like_bearer_jwt_auth(email):
"""With ``user_allowed_email_domain`` set, bearer JWT auth refuses an identity whose email is
outside the domain or missing. An assertion for an existing user gets the same refusal before
any lookup, as invalid_grant, and a matching email still resolves."""
refused = await subject_resolver(_Reader(email=email, allowed_domain="corp.example"), _known_user, _two_teams)(
_claims(), "team-eng"
)
assert isinstance(refused, SubjectTokenRefusal)
assert refused.error == "invalid_grant"
allowed = subject_resolver(
_Reader(email="ann@corp.example", allowed_domain="corp.example"), _known_user, _two_teams
)
assert await allowed(_claims(), "team-eng") == SubjectIdentity(user_id="okta-user-1", team_id="team-eng")
@pytest.mark.parametrize("claim", [None, "", "litellm:mcp", "litellm:mcp litellm:mcp"])
def test_an_assertion_for_the_mcp_resource_is_ceilinged_to_the_mcp_routes(claim):
"""``aud = {base}/mcp`` means the identity provider named the MCP resource, so no scope or
``litellm:mcp`` both mint an MCP-only credential and echo ``litellm:mcp``."""
granted = granted_scope(claim, mcp_resource=True)
assert granted == GrantedScope(scope="litellm:mcp", allowed_routes=("mcp_inference_routes",))
@pytest.mark.parametrize("claim", ["litellm:llm", "litellm:mcp litellm:llm"])
def test_llm_routes_cannot_be_granted_under_the_mcp_resource_audience(claim):
refusal = granted_scope(claim, mcp_resource=True)
assert isinstance(refusal, SubjectTokenRefusal)
assert refusal.error == "invalid_scope"
assert "/mcp audience" in refusal.description and "bare origin" in refusal.description
@pytest.mark.parametrize(
"aud, expected",
[
(AUDIENCES[0], True),
(AUDIENCES[1], False),
([AUDIENCES[0]], True),
([AUDIENCES[0], AUDIENCES[1]], False),
("https://other.example", False),
(None, False),
],
)
def test_the_mcp_resource_audience_is_the_mcp_path_without_the_bare_origin(aud, expected):
assert names_mcp_resource(aud, GATEWAY) is expected
assert names_mcp_resource(aud, None) is False
@pytest.mark.asyncio
async def test_redeeming_for_the_mcp_resource_narrows_the_identity_and_refuses_llm_scope():
"""End to end through identity_from_id_jag: the same verified claims, with ``aud`` the MCP
resource, come out ceilinged, and asking for LLM routes under that audience is a scope error."""
narrowed = await _identity(_redemption(_assertion(aud=AUDIENCES[0])), _Verifier(), _Resolver())
assert isinstance(narrowed, AssertionIdentity)
assert (narrowed.scope, narrowed.allowed_routes) == ("litellm:mcp", ("mcp_inference_routes",))
refused = await _identity(_redemption(_assertion(aud=AUDIENCES[0], scope="litellm:llm")), _Verifier(), _Resolver())
assert isinstance(refused, SubjectTokenRefusal)
assert refused.error == "invalid_scope"
unscoped = jwt.encode(
{key: value for key, value in _claims(aud=AUDIENCES[1]).items() if key != "scope"},
"unused-by-the-fake-verifier-but-long-enough-for-hs256",
algorithm="HS256",
headers={"typ": ID_JAG_TYP},
)
full = await _identity(_redemption(unscoped), _Verifier(), _Resolver())
assert isinstance(full, AssertionIdentity)
assert (full.scope, full.allowed_routes) == (None, None)
scoped = await _identity(_redemption(_assertion(aud=AUDIENCES[1], scope="litellm:llm")), _Verifier(), _Resolver())
assert isinstance(scoped, AssertionIdentity)
assert scoped.scope == "litellm:llm"

28
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -1641,9 +1641,10 @@ class TestMCPServerManager:
assert server.oauth2_flow is None
@pytest.mark.asyncio
async def test_load_servers_from_config_warns_for_id_jag_with_google_sso(self, monkeypatch, caplog):
async def test_load_servers_from_config_is_silent_for_id_jag_with_saml_sso(self, monkeypatch, caplog):
"""A SAML login retains the signed assertion as the saml2 subject token, so SAML is not a gap."""
self._clear_sso_env(monkeypatch)
monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid")
monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata")
manager = MCPServerManager()
with (
patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)),
@ -1652,10 +1653,25 @@ class TestMCPServerManager:
):
await manager.load_servers_from_config(self._id_jag_config())
warnings = [message for message in caplog.messages if "oauth2_id_jag" in message]
assert len(warnings) == 1
assert "idjag_server" in warnings[0]
assert "GENERIC_CLIENT_ID" in warnings[0]
assert [message for message in caplog.messages if "oauth2_id_jag" in message] == []
@pytest.mark.asyncio
@pytest.mark.parametrize("client_id_env", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"])
async def test_load_servers_from_config_is_silent_for_id_jag_under_oidc_logins(
self, monkeypatch, caplog, client_id_env
):
"""Every OIDC login captures the IdP id_token an id_jag server spends, so none of them is a gap."""
self._clear_sso_env(monkeypatch)
monkeypatch.setenv(client_id_env, "cid")
manager = MCPServerManager()
with (
patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)),
patch.object(manager, "hydrate_config_servers_dcr_clients", new=AsyncMock()),
caplog.at_level(logging.WARNING, logger="LiteLLM"),
):
await manager.load_servers_from_config(self._id_jag_config())
assert not any("oauth2_id_jag" in message for message in caplog.messages)
@pytest.mark.asyncio
async def test_load_servers_from_config_does_not_warn_for_id_jag_without_sso(self, monkeypatch, caplog):

89
tests/unit/proxy/_experimental/mcp_server/test_proxy_api_credentials.py
View file

@ -1,15 +1,21 @@
"""Tests for minting the ``lite login`` credential from a consented native-client grant."""
from unittest.mock import ANY, AsyncMock, MagicMock
from unittest.mock import ANY, AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from litellm.constants import CLI_JWT_EXPIRATION_HOURS
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ConsentTeam, MintedProxyCredential
from litellm.proxy._experimental.mcp_server.proxy_api_credentials import lookup_consent_teams, mint_proxy_credential
from litellm.proxy._experimental.mcp_server.proxy_api_credentials import (
lookup_consent_teams,
lookup_consent_teams_from_database,
mint_proxy_credential,
)
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail
@ -235,3 +241,82 @@ async def test_lookup_consent_teams_passes_user_lookup_failures_through(failure,
async def test_lookup_consent_teams_reports_unavailable_when_details_cannot_load(load_user, fetch_teams):
fetch_teams.return_value = None
assert await lookup_consent_teams("u1") == "unavailable"
def _route_status(minted: MintedProxyCredential, route: str, method: str | None = None) -> int:
request = None if method is None else MagicMock(method=method)
try:
RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=_decoded(minted), request=request)
except HTTPException as refused:
return refused.status_code
return 200
@pytest.mark.parametrize(
"allowed_routes, chat_completions, mcp",
[
(None, True, True),
(("mcp_inference_routes",), False, True),
(("llm_api_routes_without_mcp",), True, False),
(("mcp_inference_routes", "llm_api_routes_without_mcp"), True, True),
],
ids=["full credential", "mcp only", "llm only", "both"],
)
@pytest.mark.asyncio
async def test_mint_narrows_the_credential_to_the_allowed_routes_the_route_check_enforces(
load_user, fetch_teams, allowed_routes, chat_completions, mcp
):
"""The allowed_routes a grant's scope resolved to ride inside the credential, so the same route
check a virtual key goes through on every request refuses /chat/completions for an MCP-only
credential and /mcp for an LLM-only one, and a credential minted without them is unchanged."""
minted = await mint_proxy_credential("u1", "team-a", allowed_routes=allowed_routes)
assert isinstance(minted, MintedProxyCredential)
assert _decoded(minted).allowed_routes == list(allowed_routes or ())
chat_status, mcp_status = (200 if chat_completions else 403), (200 if mcp else 403)
assert _route_status(minted, "/chat/completions") == chat_status
assert _route_status(minted, "/v1/chat/completions") == chat_status
assert _route_status(minted, "/mcp") == mcp_status
assert _route_status(minted, "/mcp/tools/call") == mcp_status
@pytest.mark.asyncio
async def test_lookup_consent_teams_from_database_reads_the_row_the_minter_will_recheck(load_user, fetch_teams):
"""A grant consumed once a team is picked must pick it from the same membership the minter re-reads,
or a stale cache would burn the assertion on a request the caller could have fixed."""
teams = await lookup_consent_teams_from_database("u1")
load_user.assert_awaited_once_with("u1", source="database")
assert teams == (ConsentTeam(team_id="team-a", team_alias="Team A"), ConsentTeam(team_id="team-b"))
@pytest.mark.parametrize(
"allowed_routes, registered_pass_through, auto_router_session, mcp_server_listing, agent_registry",
[
(("llm_api_routes",), 200, 200, 200, 200),
(("llm_api_routes_without_mcp",), 200, 200, 403, 403),
(("mcp_inference_routes",), 403, 403, 403, 403),
],
ids=["llm_api_routes key", "llm only", "mcp only"],
)
@pytest.mark.asyncio
async def test_an_llm_only_credential_keeps_the_llm_api_carve_outs_except_mcp_discovery(
load_user,
fetch_teams,
allowed_routes,
registered_pass_through,
auto_router_session,
mcp_server_listing,
agent_registry,
):
"""Operator-registered pass-through endpoints and the auto-router session read are LLM API
surface that only the ``llm_api_routes`` group used to reach, so an LLM-only credential keeps
them, while the MCP server listing that group also carves out stays MCP surface."""
minted = await mint_proxy_credential("u1", "team-a", allowed_routes=allowed_routes)
assert isinstance(minted, MintedProxyCredential)
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route",
side_effect=lambda route: route.startswith("/my-llm-endpoint"),
):
assert _route_status(minted, "/my-llm-endpoint/v1/chat", "POST") == registered_pass_through
assert _route_status(minted, "/auto_router/session", "GET") == auto_router_session
assert _route_status(minted, "/v1/mcp/server", "GET") == mcp_server_listing
assert _route_status(minted, "/v1/agents", "GET") == agent_registry

198
tests/unit/proxy/auth/test_handle_jwt.py
View file

@ -13,6 +13,7 @@ from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import (
DEFAULT_JWKS_STALE_TTL,
JWTLiteLLMRoleMap,
JWTIssuerConfig,
LiteLLMRoutes,
LiteLLM_JWTAuth,
LiteLLM_ModelTable,
@ -29,9 +30,11 @@ from litellm.proxy._types import (
from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
from litellm.proxy.auth.auth_checks import TeamNotFoundError
from litellm.proxy.auth.handle_jwt import (
ID_JAG_TYP,
JWKS_FETCH_ATTEMPTS,
STALE_CACHE_KEY_PREFIX,
STALE_WRITTEN_AT_CACHE_KEY_PREFIX,
AssertionRejection,
HeaderTeam,
JWKSUnreachableError,
JWTAuthManager,
@ -8303,3 +8306,198 @@ def test_team_passthrough_route_denial_names_the_matched_deny_entry(
assert exc_info.value.status_code == 403
assert expected_detail in str(exc_info.value.detail)
ASSERTION_ISSUER = "https://idp.example.com/oauth2/default"
ASSERTION_JWKS_URL = f"{ASSERTION_ISSUER}/v1/keys"
GATEWAY_AUDIENCES = ("https://gateway.example.com/mcp", "https://gateway.example.com")
ID_JAG_CLAIMS = ("iss", "sub", "aud", "exp", "iat", "jti", "client_id")
def _assertion_signer():
import json
import jwt as pyjwt
from cryptography.hazmat.primitives.asymmetric import rsa
private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
jwk: Final = {**json.loads(pyjwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())), "kid": "k1"}
def sign(claims: Mapping[str, object], typ: str = ID_JAG_TYP) -> str:
return pyjwt.encode(dict(claims), private_key, algorithm="RS256", headers={"kid": "k1", "typ": typ})
return jwk, sign
def _id_jag_claims(**overrides):
now: Final = int(time.time())
return {
"iss": ASSERTION_ISSUER,
"sub": "okta-user-1",
"aud": GATEWAY_AUDIENCES[0],
"exp": now + 300,
"iat": now - 5,
"jti": "jti-1",
"client_id": "0oa-requesting-app",
**overrides,
}
def _assertion_handler(jwk, litellm_jwtauth, jwks_url=ASSERTION_JWKS_URL):
"""A handler whose JWKS for the issuer is already cached, so no HTTP is needed."""
handler: Final = JWTHandler()
cache: Final = DualCache()
cache.set_cache(f"litellm_jwt_auth_keys_{jwks_url}", [jwk])
handler.update_environment(prisma_client=None, user_api_key_cache=cache, litellm_jwtauth=litellm_jwtauth)
return handler
@pytest.mark.asyncio
async def test_decode_trusted_assertion_verifies_against_the_configured_issuer_with_the_callers_audience():
"""The assertion is checked with the issuer's JWKS and claim mappings, but for the audience the
caller names (the gateway's own issuer values), not the API audience the issuer entry carries,
since an ID-JAG is addressed to the authorization server, never to the API."""
jwk, sign = _assertion_signer()
handler = _assertion_handler(
jwk,
LiteLLM_JWTAuth(
issuers=[
JWTIssuerConfig(
issuer=ASSERTION_ISSUER,
jwks_url=ASSERTION_JWKS_URL,
audience="api://litellm",
user_id_jwt_field="sub",
)
]
),
)
claims = _id_jag_claims()
verified = await handler.decode_trusted_assertion(sign(claims), GATEWAY_AUDIENCES, ID_JAG_CLAIMS)
assert isinstance(verified, dict)
assert {key: verified[key] for key in claims} == claims
assert handler.get_user_id(token=verified, default_value=None) == "okta-user-1"
second = await handler.decode_trusted_assertion(
sign(_id_jag_claims(aud=GATEWAY_AUDIENCES[1])), GATEWAY_AUDIENCES, ID_JAG_CLAIMS
)
assert isinstance(second, dict)
@pytest.mark.asyncio
async def test_decode_trusted_assertion_falls_back_to_the_env_scoped_issuer(monkeypatch):
jwk, sign = _assertion_signer()
monkeypatch.setenv("JWT_ISSUER", ASSERTION_ISSUER)
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", ASSERTION_JWKS_URL)
handler = _assertion_handler(jwk, LiteLLM_JWTAuth())
verified = await handler.decode_trusted_assertion(sign(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS)
assert isinstance(verified, dict)
assert verified["jti"] == "jti-1"
monkeypatch.setenv("JWT_ISSUER", "https://someone-else.example.com")
foreign = await handler.decode_trusted_assertion(sign(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS)
assert isinstance(foreign, AssertionRejection)
assert "trusts" in foreign.reason
@pytest.mark.asyncio
@pytest.mark.parametrize(
"fault, reason",
[
({"aud": "api://litellm"}, "udience"),
({"jti": None}, "jti"),
({"client_id": None}, "client_id"),
({"exp": -10}, "xpired"),
({"iat": 600}, "iat"),
({"iss": "https://untrusted.example.com"}, "trusts"),
],
)
async def test_decode_trusted_assertion_rejects_a_claim_outside_the_profile(fault: dict[str, object], reason: str):
"""The claims are built inside the test, relative to the clock the verifier reads, so a slow
collection cannot turn an intended claim fault into an unrelated expiry (``None`` drops a claim,
an int on ``exp`` or ``iat`` is an offset from now)."""
now: Final = int(time.time())
offsets: Final = {k: now + v for k, v in fault.items() if k in ("exp", "iat") and isinstance(v, int)}
claims: Final = {k: v for k, v in _id_jag_claims(**offsets).items() if fault.get(k, k) is not None}
claims.update({k: v for k, v in fault.items() if v is not None and k not in offsets})
jwk, sign = _assertion_signer()
handler = _assertion_handler(
jwk,
LiteLLM_JWTAuth(issuers=[JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience="x")]),
)
rejected = await handler.decode_trusted_assertion(sign(claims), GATEWAY_AUDIENCES, ID_JAG_CLAIMS)
assert isinstance(rejected, AssertionRejection)
assert reason in rejected.reason
@pytest.mark.asyncio
async def test_decode_trusted_assertion_rejects_a_forged_signature_and_garbage(monkeypatch):
jwk, _ = _assertion_signer()
_, forge = _assertion_signer()
handler = _assertion_handler(
jwk,
LiteLLM_JWTAuth(issuers=[JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience="x")]),
)
forged = await handler.decode_trusted_assertion(forge(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS)
assert isinstance(forged, AssertionRejection)
assert "verification" in forged.reason
garbage = await handler.decode_trusted_assertion("not.a.jwt", GATEWAY_AUDIENCES, ID_JAG_CLAIMS)
assert isinstance(garbage, AssertionRejection)
monkeypatch.delenv("JWT_ISSUER", raising=False)
unconfigured = await _assertion_handler(jwk, LiteLLM_JWTAuth()).decode_trusted_assertion(
forge(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS
)
assert isinstance(unconfigured, AssertionRejection)
assert "trusts" in unconfigured.reason
@pytest.mark.asyncio
async def test_decode_trusted_assertion_reports_an_unreachable_jwks_as_an_outage_not_a_bad_assertion():
jwk, sign = _assertion_signer()
handler = _assertion_handler(
jwk,
LiteLLM_JWTAuth(issuers=[JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience="x")]),
jwks_url="https://a-different-cached-url.example.com/keys",
)
with patch.object(handler, "_fetch_jwks_keys", new_callable=AsyncMock, side_effect=JWKSUnreachableError("down")):
with pytest.raises(ProxyException) as outage:
await handler.decode_trusted_assertion(sign(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS)
assert outage.value.code == "503"
@pytest.mark.asyncio
async def test_auth_jwt_never_accepts_an_id_jag_as_a_bearer_token():
"""An assertion typed oauth-id-jag+jwt is a grant for the token endpoint, so bearer auth
refuses it before any key lookup, even where it would otherwise verify."""
jwk, sign = _assertion_signer()
handler = _assertion_handler(
jwk,
LiteLLM_JWTAuth(
issuers=[
JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience=GATEWAY_AUDIENCES[0])
]
),
)
accepted = await handler.auth_jwt(sign(_id_jag_claims(), typ="JWT"))
assert accepted["jti"] == "jti-1"
with pytest.raises(Exception, match="token endpoint"):
await handler.auth_jwt(sign(_id_jag_claims()))
@pytest.mark.asyncio
@pytest.mark.parametrize("verdict", [True, False])
async def test_decode_trusted_assertion_runs_the_configured_custom_validate_on_the_verified_claims(verdict):
"""The operator's ``custom_validate`` gates every JWT the proxy admits, so an assertion redeemed
at the token endpoint goes through it too, after signature and claim verification."""
jwk, sign = _assertion_signer()
seen: Final = []
handler = _assertion_handler(
jwk,
LiteLLM_JWTAuth(
issuers=[JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience="api://litellm")]
),
)
handler.litellm_jwtauth.custom_validate = lambda claims: seen.append(dict(claims)) or verdict
claims = _id_jag_claims()
verified = await handler.decode_trusted_assertion(sign(claims), GATEWAY_AUDIENCES, ID_JAG_CLAIMS)
assert [{key: call[key] for key in claims} for call in seen] == [claims]
if verdict:
assert isinstance(verified, dict)
else:
assert isinstance(verified, AssertionRejection)
assert "custom_validate" in verified.reason

60
tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py
View file

@ -49,8 +49,8 @@ class TestActiveSSOProviderMirrorsTheCallback:
assert active_sso_provider() is ActiveSSOProvider.none
def test_google_outranks_a_co_configured_generic_client(self, monkeypatch):
"""The callback tests GOOGLE_CLIENT_ID first, so the generic arm never runs here and
no assertion is captured; reporting generic would clear a gap that is still open."""
"""The callback tests GOOGLE_CLIENT_ID first, so the generic arm never runs here and the
assertion the refresher later renews was minted by the Google client, not the generic one."""
monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid")
monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid")
assert active_sso_provider() is ActiveSSOProvider.google
@ -67,51 +67,41 @@ class TestActiveSSOProviderMirrorsTheCallback:
class TestIdJagAssertionCaptureGap:
def test_generic_oidc_has_no_gap(self, monkeypatch):
monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid")
@pytest.mark.parametrize("env_var", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"])
def test_every_oidc_login_captures_so_has_no_gap(self, monkeypatch, env_var):
"""Google, Microsoft and generic OIDC all read the id_token off the token response, so a
deployment on any of them must not be told to switch providers."""
monkeypatch.setenv(env_var, "configured")
assert id_jag_assertion_capture_gap() is None
@pytest.mark.parametrize(
"env_var, provider_label",
[
("GOOGLE_CLIENT_ID", "google"),
("MICROSOFT_CLIENT_ID", "microsoft"),
("SAML_IDP_METADATA_URL", "saml"),
],
)
def test_non_capturing_provider_is_named_with_the_remedy(self, monkeypatch, env_var, provider_label):
monkeypatch.setenv(env_var, "configured")
gap = id_jag_assertion_capture_gap()
assert gap is not None
assert provider_label in gap
assert "GENERIC_CLIENT_ID" in gap
def test_saml_has_no_gap(self, monkeypatch: pytest.MonkeyPatch):
"""The SAML callback retains the signed assertion as the saml2 subject token, so a SAML
deployment is a capturing one and must not be told to switch to an OIDC provider."""
monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata")
assert id_jag_assertion_capture_gap() is None
def test_no_sso_configured_reports_a_gap(self):
gap = id_jag_assertion_capture_gap()
assert gap is not None
assert "no SSO provider is configured" in gap
def test_google_beside_generic_still_reports_a_gap(self, monkeypatch):
"""The precedence trap in operator terms: adding a generic client id without removing
GOOGLE_CLIENT_ID does not fix the deployment, so the gap must not clear."""
monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid")
monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid")
gap = id_jag_assertion_capture_gap()
assert gap is not None
assert "google" in gap
def test_saml_beside_an_oidc_client_has_no_gap(self, monkeypatch):
"""The callback dispatches to the OIDC client before SAML, so that is the login that runs
and it captures; reporting SAML's gap would tell the operator to fix a path nobody takes."""
monkeypatch.setenv("MICROSOFT_CLIENT_ID", "ms-cid")
monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata")
assert id_jag_assertion_capture_gap() is None
class TestIdJagAssertionCaptureGapAtStartup:
def test_no_provider_at_startup_is_not_yet_a_gap(self):
assert id_jag_assertion_capture_gap_at_startup() is None
def test_google_provider_at_startup_reports_the_capture_gap(self, monkeypatch):
monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid")
startup_gap = id_jag_assertion_capture_gap_at_startup()
callback_gap = id_jag_assertion_capture_gap()
assert startup_gap is not None
assert startup_gap == callback_gap
def test_generic_provider_at_startup_has_no_gap(self, monkeypatch):
monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid")
def test_saml_at_startup_has_no_gap(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata")
assert id_jag_assertion_capture_gap_at_startup() is None
@pytest.mark.parametrize("env_var", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"])
def test_an_oidc_provider_at_startup_has_no_gap(self, monkeypatch, env_var):
monkeypatch.setenv(env_var, "configured")
assert id_jag_assertion_capture_gap_at_startup() is None

19
tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py
View file

@ -4458,9 +4458,9 @@ class TestAddMCPServerAtomicity:
class TestIdJagRegistrationWarnsAboutTheSSOGap:
"""An `oauth2_id_jag` server only ever works when the login path captures an IdP identity
assertion, and only the generic OIDC arm does. Registering one under Google or Microsoft
succeeds and then fails for every user on every call, so the mismatch has to be said at
registration time, while the admin is still looking at the configuration."""
assertion, and only the OIDC logins (Google, Microsoft, generic) do. Registering one under SAML
or with no SSO succeeds and then fails for every user on every call, so the mismatch has to be
said at registration time, while the admin is still looking at the configuration."""
@staticmethod
def _clear_sso_env(monkeypatch):
@ -4530,9 +4530,6 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap:
@pytest.mark.parametrize(
"provider_env, expected_fragment",
[
({"GOOGLE_CLIENT_ID": "cid"}, "google"),
({"MICROSOFT_CLIENT_ID": "cid"}, "microsoft"),
({"SAML_IDP_METADATA_URL": "https://idp.example.com/metadata"}, "saml"),
({}, "no SSO provider is configured"),
],
)
@ -4546,8 +4543,9 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap:
assert "ema-1" in str(warnings[0])
@pytest.mark.asyncio
async def test_create_is_silent_under_generic_oidc(self, monkeypatch, caplog):
await self._run_create(monkeypatch, {"GENERIC_CLIENT_ID": "cid"}, MCPAuth.oauth2_id_jag, caplog)
@pytest.mark.parametrize("client_id_env", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"])
async def test_create_is_silent_under_oidc_logins(self, monkeypatch, caplog, client_id_env):
await self._run_create(monkeypatch, {client_id_env: "cid"}, MCPAuth.oauth2_id_jag, caplog)
assert self._id_jag_warnings(caplog) == []
@pytest.mark.asyncio
@ -4565,7 +4563,8 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap:
)
self._clear_sso_env(monkeypatch)
monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid")
for var in ("MICROSOFT_CLIENT_ID", "GOOGLE_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"):
monkeypatch.delenv(var, raising=False)
mock_manager = MagicMock()
mock_manager.update_server = AsyncMock()
@ -4602,7 +4601,7 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap:
warnings = self._id_jag_warnings(caplog)
assert len(warnings) == 1
assert "google" in str(warnings[0])
assert "no SSO provider is configured" in str(warnings[0])
class TestHealthCheckServers:

221
tests/unit/proxy/management_endpoints/test_saml_sso.py
View file

@ -9,15 +9,17 @@ makes a test fail.
import base64
import datetime
import logging
import re
import time
from http.cookies import SimpleCookie
import pytest
from fastapi import HTTPException, Request
pytest.importorskip(
"onelogin", reason="python3-saml (saml extra) is required for SAML SSO tests"
)
pytest.importorskip("onelogin", reason="python3-saml (saml extra) is required for SAML SSO tests")
from typing import cast
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
@ -26,11 +28,10 @@ from cryptography.x509.oid import NameOID
from onelogin.saml2.utils import OneLogin_Saml2_Utils
from starlette.datastructures import URL
from typing import cast
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import SSOIdentityAssertion
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.management_endpoints.sso.saml_sso import (
_SAML_AUTHN_REQUEST_CACHE_PREFIX,
@ -40,15 +41,17 @@ from litellm.proxy.management_endpoints.sso.saml_sso import (
_SAML_REPLAY_GUARD_MAX_TTL_SECONDS,
SAMLAuthHandler,
)
from litellm.proxy.management_endpoints.types import CustomOpenID
def _shared_cache(store=None):
def _shared_cache(store: InMemoryCache | None = None) -> DualCache:
"""A DualCache whose replay guard is backed by a shared, atomic store.
An InMemoryCache instance stands in for Redis; passing the same instance to
two DualCaches simulates two workers sharing one atomic backend."""
return DualCache(redis_cache=cast(RedisCache, store or InMemoryCache()))
IDP_ENTITY = "https://idp.example.com/metadata"
SP_ENTITY = "https://proxy.example.com/sso/saml/metadata"
ACS = "https://proxy.example.com/sso/saml/callback"
@ -56,7 +59,7 @@ SSO_URL = "https://idp.example.com/sso"
PROXY_BASE_URL = "https://proxy.example.com"
def _make_idp_keypair():
def _make_idp_keypair() -> tuple[str, str]:
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "idp.example.com")])
cert = (
@ -79,9 +82,7 @@ def _make_idp_keypair():
def _idp_metadata_xml(cert_pem):
cert_body = "".join(
line for line in cert_pem.splitlines() if "CERTIFICATE" not in line
)
cert_body = "".join(line for line in cert_pem.splitlines() if "CERTIFICATE" not in line)
return (
'<?xml version="1.0"?>'
f'<EntityDescriptor xmlns="urn:oasis:names:tc:SAML:2.0:metadata" entityID="{IDP_ENTITY}">'
@ -101,17 +102,17 @@ def _saml_time(delta_seconds):
def _build_signed_response(
key_pem,
cert_pem,
key_pem: str,
cert_pem: str,
*,
in_response_to=None,
response_level_in_response_to=True,
email="alice@example.com",
attributes=None,
not_before_delta=-60,
not_on_or_after_delta=300,
sign=True,
):
in_response_to: str | None = None,
response_level_in_response_to: bool = True,
email: str = "alice@example.com",
attributes: dict[str, list[str]] | None = None,
not_before_delta: int = -60,
not_on_or_after_delta: int = 300,
sign: bool = True,
) -> str:
if attributes is None:
attributes = {
"email": [email],
@ -157,9 +158,7 @@ def _build_signed_response(
if sign:
signed = OneLogin_Saml2_Utils.add_sign(assertion, key_pem, cert_pem)
assertion = (signed.decode() if isinstance(signed, bytes) else signed).replace(
'<?xml version="1.0"?>', ""
)
assertion = (signed.decode() if isinstance(signed, bytes) else signed).replace('<?xml version="1.0"?>', "")
return (
'<?xml version="1.0"?>'
@ -174,11 +173,11 @@ def _build_signed_response(
)
def _b64(xml):
def _b64(xml: str) -> str:
return base64.b64encode(xml.encode()).decode()
def _fake_request(cookies=None):
def _fake_request(cookies: dict[str, str] | None = None):
return type(
"Req",
(),
@ -190,14 +189,18 @@ def _fake_request(cookies=None):
)()
async def _acs_login(
b64: str, cache: DualCache, cookies: dict[str, str] | None = None
) -> tuple[CustomOpenID, SSOIdentityAssertion | None]:
return await SAMLAuthHandler.handle_acs(_fake_request(cookies), cache, {"SAMLResponse": b64})
async def _acs(b64, cache, cookies=None):
return await SAMLAuthHandler.handle_acs(
_fake_request(cookies), cache, {"SAMLResponse": b64}
)
return (await _acs_login(b64, cache, cookies))[0]
@pytest.fixture
def saml_env(monkeypatch):
def saml_env(monkeypatch: pytest.MonkeyPatch) -> tuple[str, str]:
key_pem, cert_pem = _make_idp_keypair()
monkeypatch.setenv("SAML_IDP_METADATA_XML", _idp_metadata_xml(cert_pem))
monkeypatch.setenv("SAML_SP_ENTITY_ID", SP_ENTITY)
@ -214,7 +217,7 @@ def saml_env(monkeypatch):
@pytest.fixture
def saml_env_idp_initiated(saml_env, monkeypatch):
def saml_env_idp_initiated(saml_env: tuple[str, str], monkeypatch: pytest.MonkeyPatch) -> tuple[str, str]:
monkeypatch.setenv("SAML_ALLOW_UNSOLICITED", "true")
return saml_env
@ -234,6 +237,89 @@ async def test_valid_idp_initiated_login_maps_assertion_to_user(saml_env_idp_ini
assert result.provider == "saml"
@pytest.mark.asyncio
async def test_valid_login_retains_the_signed_assertion_as_a_saml2_subject_token(
saml_env_idp_initiated: tuple[str, str],
):
"""The retained material is the inner <saml:Assertion>, still carrying a verifiable signature, with
the assertion's own NotOnOrAfter as its expiry and no refresh token, since SAML has none."""
key_pem, cert_pem = saml_env_idp_initiated
resp = _build_signed_response(key_pem, cert_pem, not_on_or_after_delta=300)
user, assertion = await _acs_login(_b64(resp), _shared_cache())
assert assertion is not None
assert assertion.subject_token_type == "urn:ietf:params:oauth:token-type:saml2"
assert assertion.refresh_token is None
assert assertion.issuer == IDP_ENTITY
retained = _retained_bytes(assertion.id_token.get_secret_value())
assert retained.startswith(b"<saml:Assertion")
assert b"<samlp:Response" not in retained
assert _signature_valid(retained, cert_pem)
assert assertion.expires_at == datetime.datetime.strptime(
_conditions_not_on_or_after(resp), "%Y-%m-%dT%H:%M:%SZ"
).replace(tzinfo=datetime.timezone.utc)
def _conditions_not_on_or_after(response_xml: str) -> str:
match = re.search(r'<saml:Conditions[^>]*NotOnOrAfter="([^"]+)"', response_xml)
assert match is not None
return match.group(1)
def _signature_valid(xml: bytes, cert_pem: str) -> bool:
return bool(OneLogin_Saml2_Utils.validate_sign(xml, cert_pem, raise_exceptions=True, xpath="//ds:Signature"))
def _retained_bytes(encoded: str) -> bytes:
assert set(encoded) <= set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"), encoded[:40]
return base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4))
@pytest.mark.asyncio
async def test_retained_assertion_keeps_non_ascii_subject_bytes_verbatim(saml_env_idp_initiated: tuple[str, str]):
key_pem, cert_pem = saml_env_idp_initiated
resp = _build_signed_response(key_pem, cert_pem, email="zo\u00eb@example.com")
user, assertion = await _acs_login(_b64(resp), _shared_cache())
assert assertion is not None
retained = _retained_bytes(assertion.id_token.get_secret_value())
assert "zo\u00eb@example.com".encode() in retained
assert b"&#" not in retained
assert _signature_valid(retained, cert_pem)
class _UnparseableResponseAuth:
def get_last_response_xml(self) -> bytes:
return b"<samlp:Response><saml:Assertion"
def get_last_assertion_not_on_or_after(self) -> int | None:
return None
def test_assertion_retention_failure_does_not_fail_the_validated_login(caplog: pytest.LogCaptureFixture):
"""Retention is optional egress material, so a response the retention step cannot parse logs and yields
no assertion instead of turning an IdP-validated login into a 500."""
with caplog.at_level(logging.WARNING):
retained = SAMLAuthHandler._retained_assertion(_UnparseableResponseAuth())
assert retained is None
assert "could not be retained for oauth2_id_jag" in caplog.text
@pytest.mark.asyncio
async def test_retained_assertion_is_the_one_the_user_was_mapped_from(saml_env_idp_initiated: tuple[str, str]):
key_pem, cert_pem = saml_env_idp_initiated
resp = _build_signed_response(key_pem, cert_pem, email="bob@example.com")
user, assertion = await _acs_login(_b64(resp), _shared_cache())
assert user.email == "bob@example.com"
assert assertion is not None
assert b"bob@example.com" in _retained_bytes(assertion.id_token.get_secret_value())
@pytest.mark.asyncio
async def test_tampered_assertion_is_rejected(saml_env):
key_pem, cert_pem = saml_env
@ -269,9 +355,7 @@ async def test_signature_from_untrusted_key_is_rejected(saml_env):
@pytest.mark.asyncio
async def test_expired_assertion_is_rejected(saml_env):
key_pem, cert_pem = saml_env
resp = _build_signed_response(
key_pem, cert_pem, not_before_delta=-7200, not_on_or_after_delta=-3600
)
resp = _build_signed_response(key_pem, cert_pem, not_before_delta=-7200, not_on_or_after_delta=-3600)
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), DualCache())
@ -293,9 +377,7 @@ async def test_sp_initiated_known_request_succeeds_once_then_replay_rejected(sam
key_pem, cert_pem = saml_env
cache = DualCache()
request_id = "_authn_req_known"
cache.set_cache(
key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600
)
cache.set_cache(key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600)
resp = _build_signed_response(key_pem, cert_pem, in_response_to=request_id)
cookies = {_SAML_AUTHN_STATE_COOKIE: request_id}
@ -312,9 +394,7 @@ async def test_sp_initiated_response_not_bound_to_browser_is_rejected(saml_env):
key_pem, cert_pem = saml_env
cache = DualCache()
request_id = "_authn_req_known"
cache.set_cache(
key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600
)
cache.set_cache(key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600)
resp = _build_signed_response(key_pem, cert_pem, in_response_to=request_id)
with pytest.raises(HTTPException) as exc:
@ -322,9 +402,7 @@ async def test_sp_initiated_response_not_bound_to_browser_is_rejected(saml_env):
assert exc.value.status_code == 401
with pytest.raises(HTTPException) as exc:
await _acs(
_b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: "_attacker_request"}
)
await _acs(_b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: "_attacker_request"})
assert exc.value.status_code == 401
@ -338,9 +416,7 @@ async def test_subjectconfirmation_only_in_response_to_without_cookie_is_rejecte
key_pem, cert_pem = saml_env_idp_initiated
cache = DualCache()
request_id = "_authn_req_known"
cache.set_cache(
key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600
)
cache.set_cache(key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600)
resp = _build_signed_response(
key_pem,
cert_pem,
@ -358,9 +434,7 @@ async def test_subjectconfirmation_only_in_response_to_with_cookie_succeeds(saml
key_pem, cert_pem = saml_env
cache = DualCache()
request_id = "_authn_req_known"
cache.set_cache(
key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600
)
cache.set_cache(key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600)
resp = _build_signed_response(
key_pem,
cert_pem,
@ -368,9 +442,7 @@ async def test_subjectconfirmation_only_in_response_to_with_cookie_succeeds(saml
response_level_in_response_to=False,
)
result = await _acs(
_b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: request_id}
)
result = await _acs(_b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: request_id})
assert result.email == "alice@example.com"
@ -415,9 +487,7 @@ async def test_assertion_without_id_is_rejected(saml_env_idp_initiated):
return None
with pytest.raises(HTTPException) as exc:
await SAMLAuthHandler._enforce_response_binding(
_AuthNoAssertionId(), _shared_cache(), None
)
await SAMLAuthHandler._enforce_response_binding(_AuthNoAssertionId(), _shared_cache(), None)
assert exc.value.status_code == 401
assert "ID" in exc.value.detail
@ -450,9 +520,7 @@ async def test_invalid_email_in_assertion_is_rejected_cleanly(saml_env_idp_initi
@pytest.mark.asyncio
async def test_email_less_assertion_rejected_when_domain_restriction_configured(
saml_env_idp_initiated, monkeypatch
):
async def test_email_less_assertion_rejected_when_domain_restriction_configured(saml_env_idp_initiated, monkeypatch):
key_pem, cert_pem = saml_env_idp_initiated
monkeypatch.setenv("ALLOWED_EMAIL_DOMAINS", "example.com")
resp = _build_signed_response(
@ -558,11 +626,7 @@ async def test_build_login_redirect_targets_idp_and_caches_request_id(saml_env):
location = redirect.headers["location"]
assert location.startswith(SSO_URL)
assert "SAMLRequest=" in location
cached = [
k
for k in cache.in_memory_cache.cache_dict
if k.startswith(_SAML_AUTHN_REQUEST_CACHE_PREFIX)
]
cached = [k for k in cache.in_memory_cache.cache_dict if k.startswith(_SAML_AUTHN_REQUEST_CACHE_PREFIX)]
assert len(cached) == 1
request_id = cached[0].split(":", 1)[1]
@ -761,9 +825,7 @@ class TestSAMLAuthnCookieSecureFlag:
{"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]},
)
cache = DualCache()
request = _fake_request_with_scheme(
"http", headers={"X-Forwarded-Proto": "https"}, client_host="10.0.0.5"
)
request = _fake_request_with_scheme("http", headers={"X-Forwarded-Proto": "https"}, client_host="10.0.0.5")
redirect = await SAMLAuthHandler.build_login_redirect(request, cache)
cookie = redirect.headers["set-cookie"]
assert "Secure" in cookie
@ -773,9 +835,32 @@ class TestSAMLAuthnCookieSecureFlag:
monkeypatch.delenv("PROXY_BASE_URL", raising=False)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
cache = DualCache()
request = _fake_request_with_scheme(
"http", headers={"X-Forwarded-Proto": "https"}, client_host="203.0.113.5"
)
request = _fake_request_with_scheme("http", headers={"X-Forwarded-Proto": "https"}, client_host="203.0.113.5")
redirect = await SAMLAuthHandler.build_login_redirect(request, cache)
cookie = redirect.headers["set-cookie"]
assert "Secure" not in cookie
def test_an_encrypted_assertion_is_not_retained_and_the_login_log_says_why(caplog: pytest.LogCaptureFixture):
"""An IdP that encrypts its assertion still signs the user in, but the ID-JAG arm has nothing to
present. The login log has to name that, otherwise every later 412 looks like a bug re-login cannot fix."""
class _Encrypted:
def get_last_response_xml(self) -> bytes:
return (
b'<samlp:Response xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol" '
b'xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion">'
b'<saml:EncryptedAssertion><xenc:EncryptedData xmlns:xenc="http://www.w3.org/2001/04/xmlenc#"/>'
b"</saml:EncryptedAssertion></samlp:Response>"
)
def get_last_assertion_not_on_or_after(self) -> int | None:
return None
with caplog.at_level(logging.WARNING, logger="litellm.proxy"):
retained = SAMLAuthHandler._retained_assertion(_Encrypted()) # pyright: ignore[reportArgumentType] # stub auth # pyright: ignore[reportPrivateUsage] # the retention seam under test
assert retained is None
assert any(
"1 saml:EncryptedAssertion" in r.getMessage() and "oauth2_id_jag" in r.getMessage() for r in caplog.records
)

321
tests/unit/proxy/management_endpoints/test_ui_sso.py
View file

@ -10,6 +10,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from pydantic import SecretStr
import respx
from fastapi import HTTPException, Request
@ -245,17 +246,24 @@ def test_get_microsoft_callback_response(stubbed_graph_api):
"surname": "User",
}
id_token = _ema_id_token(sub="msft123")
async def fake_verify(self, *args, **kwargs):
self._id_token = id_token
self._refresh_token = "msft-refresh-token"
return mock_response
with patch.dict(
os.environ,
{"MICROSOFT_CLIENT_SECRET": "mock_secret", "MICROSOFT_TENANT": "mock_tenant"},
):
mock_verify = AsyncMock(return_value=mock_response)
with patch(
with patch( # test-quality-ok: stubs the upstream Microsoft IdP boundary, where the token response is read
"fastapi_sso.sso.microsoft.MicrosoftSSO.verify_and_process",
new=mock_verify,
autospec=True,
side_effect=fake_verify,
):
# Act
result = asyncio.run(
result, sso_assertion = asyncio.run(
MicrosoftSSOHandler.get_microsoft_callback_response(
request=mock_request,
microsoft_client_id="mock_client_id",
@ -271,6 +279,11 @@ def test_get_microsoft_callback_response(stubbed_graph_api):
assert result.id == "msft123"
assert result.first_name == "Microsoft"
assert result.last_name == "User"
assert sso_assertion is not None
assert sso_assertion.id_token.get_secret_value() == id_token
assert sso_assertion.refresh_token is not None
assert sso_assertion.refresh_token.get_secret_value() == "msft-refresh-token"
assert sso_assertion.issuer == "https://idp.example.com"
GRAPH_ERROR_BODIES: Final = (
@ -319,7 +332,7 @@ def test_get_microsoft_callback_response_signs_in_a_profile_with_an_empty_error_
"fastapi_sso.sso.microsoft.MicrosoftSSO.verify_and_process",
new=AsyncMock(return_value={**profile, "error": empty_error}),
):
result = asyncio.run(
result, _ = asyncio.run(
MicrosoftSSOHandler.get_microsoft_callback_response(
request=mock_request,
microsoft_client_id="mock_client_id",
@ -344,7 +357,7 @@ def test_get_microsoft_callback_response_raw_sso_response_keeps_the_graph_error(
"fastapi_sso.sso.microsoft.MicrosoftSSO.verify_and_process",
new=AsyncMock(return_value=dict(graph_error)),
):
result = asyncio.run(
result, _ = asyncio.run(
MicrosoftSSOHandler.get_microsoft_callback_response(
request=mock_request,
microsoft_client_id="mock_client_id",
@ -378,7 +391,7 @@ def test_get_microsoft_callback_response_raw_sso_response(stubbed_graph_api):
new=mock_verify,
):
# Act
result = asyncio.run(
result, _ = asyncio.run(
MicrosoftSSOHandler.get_microsoft_callback_response(
request=mock_request,
microsoft_client_id="mock_client_id",
@ -407,13 +420,19 @@ def test_get_google_callback_response():
"family_name": "User",
}
id_token = _ema_id_token(sub="google123")
async def fake_verify(self, *args, **kwargs):
self._id_token = id_token
self._refresh_token = None
return mock_response
with patch.dict(os.environ, {"GOOGLE_CLIENT_SECRET": "mock_secret"}):
mock_verify = AsyncMock(return_value=mock_response)
with patch(
"fastapi_sso.sso.google.GoogleSSO.verify_and_process", new=mock_verify
with patch( # test-quality-ok: stubs the upstream Google IdP boundary, where the token response is read
"fastapi_sso.sso.google.GoogleSSO.verify_and_process", autospec=True, side_effect=fake_verify
):
# Act
result = asyncio.run(
result, sso_assertion = asyncio.run(
GoogleSSOHandler.get_google_callback_response(
request=mock_request,
google_client_id="mock_client_id",
@ -428,6 +447,9 @@ def test_get_google_callback_response():
assert result.get("sub") == "google123"
assert result.get("given_name") == "Google"
assert result.get("family_name") == "User"
assert sso_assertion is not None
assert sso_assertion.id_token.get_secret_value() == id_token
assert sso_assertion.refresh_token is None
@pytest.mark.asyncio
@ -4096,7 +4118,13 @@ class TestCLIKeyRegenerationFlow:
cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-session-key-456"
# Mock the CLI callback and required proxy server components
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
assertion_from_sso_login,
)
mock_result = {"user_id": "test-user", "email": "test@example.com"}
google_assertion = assertion_from_sso_login(_ema_id_token(sub="test-user"), None)
assert google_assertion is not None
with (
patch(
@ -4110,7 +4138,7 @@ class TestCLIKeyRegenerationFlow:
patch.dict(os.environ, {"GOOGLE_CLIENT_ID": "test-google-id"}, clear=True),
patch(
"litellm.proxy.management_endpoints.ui_sso.GoogleSSOHandler.get_google_callback_response",
return_value=mock_result,
return_value=(mock_result, google_assertion),
),
):
mock_cli_callback.return_value = MagicMock()
@ -4124,7 +4152,7 @@ class TestCLIKeyRegenerationFlow:
prefill_user_code=None,
result=mock_result,
received_response=None,
sso_assertion=None,
sso_assertion=google_assertion,
)
@pytest.mark.asyncio
@ -4152,7 +4180,7 @@ class TestCLIKeyRegenerationFlow:
patch.dict(os.environ, {"GOOGLE_CLIENT_ID": "test-google-id"}, clear=True),
patch(
"litellm.proxy.management_endpoints.ui_sso.GoogleSSOHandler.get_google_callback_response",
return_value=mock_result,
return_value=(mock_result, None),
),
):
mock_cli_callback.return_value = MagicMock()
@ -8914,7 +8942,7 @@ async def test_debug_sso_callback_handles_missing_raw_response():
)
async def fake_microsoft_callback(**kwargs):
return parsed_openid
return parsed_openid, None
with (
patch.dict(
@ -8964,7 +8992,7 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False)
return parsed, {"sub": "user_123"}, {"scope": "openid"}, None
async def fake_google(**kwargs):
return parsed
return parsed, None
stack = [
patch.dict(os.environ, {**provider_env, "ENABLE_SSO_DEBUG": "true"}, clear=False),
@ -9003,14 +9031,14 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False)
@pytest.mark.asyncio
async def test_debug_page_logs_the_capture_gap_but_never_renders_it(caplog):
async def test_debug_page_logs_no_capture_gap_for_a_google_login(caplog):
"""Google reads the id_token off its token response like generic OIDC does, so the debug page
under GOOGLE_CLIENT_ID with an oauth2_id_jag server registered has no gap to log and nothing
to render about it."""
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
body = await _render_debug_page({"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=True)
warnings = _id_jag_gap_warnings(caplog)
assert len(warnings) == 1
assert "google" in warnings[0]
assert "GENERIC_CLIENT_ID" in warnings[0]
assert _id_jag_gap_warnings(caplog) == []
assert "id_jag" not in body
assert "GENERIC_CLIENT_ID" not in body
@ -9055,7 +9083,8 @@ async def test_debug_page_survives_a_store_outage(monkeypatch, caplog):
not the page."""
from litellm.proxy.management_endpoints.ui_sso import warn_if_id_jag_capture_gap
monkeypatch.setenv("GOOGLE_CLIENT_ID", _GOOGLE_DEBUG_CLIENT_ID)
for var in ("MICROSOFT_CLIENT_ID", "GOOGLE_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"):
monkeypatch.delenv(var, raising=False)
retention_check = AsyncMock(side_effect=Exception("db down"))
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
assert await warn_if_id_jag_capture_gap(retention_enabled=retention_check) is None
@ -9181,7 +9210,7 @@ async def test_saml_callback_enforces_free_sso_user_limit_after_validation():
async def _fake_handle_acs(**kwargs):
call_order.append("validate")
return CustomOpenID(
return (CustomOpenID(
id="dana@litellm.ai",
email="dana@litellm.ai",
first_name=None,
@ -9191,7 +9220,7 @@ async def test_saml_callback_enforces_free_sso_user_limit_after_validation():
provider="saml",
team_ids=[],
user_role=None,
)
), None)
async def _fake_count_billable_users():
call_order.append("count")
@ -9713,6 +9742,233 @@ async def test_redirect_from_openid_persists_assertion_under_canonical_user_id()
assert response is not None
@pytest.mark.asyncio
async def test_redirect_from_openid_persists_the_login_exchanged_assertion():
"""A server with ``id_jag_subject_source: refresh_token`` needs the refresh token the login
assertion was exchanged for, so the funnel persists what the login-time exchange returns, not
the raw capture."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
assertion_from_sso_login,
)
assertion = assertion_from_sso_login(_ema_id_token(), "rt_1")
assert assertion is not None
mock_request = MagicMock(spec=Request)
mock_request.scope = {}
mock_request.base_url = "http://localhost:4000/"
mock_request.cookies = {}
retain_mock = AsyncMock()
exchanged = assertion.model_copy(update={"exchanged_refresh_token": SecretStr("rt-from-login")})
exchange_mock = AsyncMock(return_value=exchanged)
with (
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()),
patch("litellm.proxy.proxy_server.master_key", "sk-master"),
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.premium_user", False),
patch("litellm.proxy.proxy_server.user_custom_sso", None),
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
patch("litellm.proxy.proxy_server.redis_usage_cache", None),
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
patch(
"litellm.proxy.proxy_server.generate_key_helper_fn",
AsyncMock(
return_value={"token": "sk-ui-key", "user_id": "canonical-user-id"}
),
),
patch(
"litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db",
AsyncMock(return_value=None),
),
patch(
"litellm.proxy.management_endpoints.ui_sso.check_and_update_if_proxy_admin_id",
AsyncMock(return_value="internal_user"),
),
patch(
"litellm.proxy.management_endpoints.ui_sso.retain_sso_identity_assertion_for_ema",
retain_mock,
),
patch(
"litellm.proxy.management_endpoints.ui_sso.exchange_login_assertion_if_configured",
exchange_mock,
),
):
response = await SSOAuthenticationHandler.get_redirect_response_from_openid(
result=CustomOpenID(
id="raw-idp-subject",
email="u@example.com",
first_name="U",
last_name="Ser",
display_name="U Ser",
provider="generic",
team_ids=[],
user_role=None,
),
request=mock_request,
received_response=None,
generic_client_id="cid",
ui_access_mode=None,
access_token_payload=None,
jwt_handler=None,
sso_assertion=assertion,
)
exchange_mock.assert_awaited_once_with(assertion)
retain_mock.assert_awaited_once()
assert retain_mock.await_args.kwargs["assertion"] is exchanged
@pytest.mark.asyncio
@pytest.mark.parametrize(
"provider_env, sso_class, payload, needs_graph",
[
(
{"MICROSOFT_CLIENT_ID": "ms-cid", "MICROSOFT_CLIENT_SECRET": "ms-secret", "MICROSOFT_TENANT": "tenant"},
"fastapi_sso.sso.microsoft.MicrosoftSSO",
{"mail": "u@example.com", "displayName": "U Ser", "id": "idp-subject"},
True,
),
(
{"GOOGLE_CLIENT_ID": "google-cid", "GOOGLE_CLIENT_SECRET": "google-secret"},
"fastapi_sso.sso.google.GoogleSSO",
{"email": "u@example.com", "name": "U Ser", "sub": "idp-subject"},
False,
),
],
)
async def test_ui_login_through_microsoft_or_google_hands_its_assertion_to_the_redirect(
request, provider_env, sso_class, payload, needs_graph
):
"""The UI funnel used to capture an assertion only on the generic branch, so an ``oauth2_id_jag``
server had nothing to spend for a deployment signing in through Entra or Google. The real
Microsoft and Google handlers run here, only the IdP token exchange is stubbed, so the id_token
and refresh token the IdP returned must reach the redirect that persists them."""
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler, auth_callback
if needs_graph:
request.getfixturevalue("stubbed_graph_api")
id_token = _ema_id_token(sub="idp-subject")
async def fake_verify(self, *args, **kwargs):
self._id_token = id_token
self._refresh_token = "idp-refresh"
return payload
mock_request = MagicMock(spec=Request)
mock_request.query_params = {"code": "auth-code"}
mock_request.cookies = {}
mock_request.scope = {}
mock_request.base_url = "http://localhost:4000/"
redirect_mock = AsyncMock(return_value=MagicMock())
with (
patch.dict(os.environ, provider_env, clear=True),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.master_key", "sk-master"),
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.jwt_handler", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
patch( # test-quality-ok: stubs the upstream IdP token exchange, where the token response is read
f"{sso_class}.verify_and_process", autospec=True, side_effect=fake_verify
),
patch.object( # test-quality-ok: the DB and key-mint boundary after the login, asserted on its inputs
SSOAuthenticationHandler, "get_redirect_response_from_openid", redirect_mock
),
):
await auth_callback(request=mock_request, state=None)
redirect_mock.assert_awaited_once()
assert redirect_mock.await_args.kwargs["result"] is not None
sso_assertion = redirect_mock.await_args.kwargs["sso_assertion"]
assert sso_assertion is not None
assert sso_assertion.id_token.get_secret_value() == id_token
assert sso_assertion.refresh_token is not None
assert sso_assertion.refresh_token.get_secret_value() == "idp-refresh"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"retention_enabled, expected_params",
[
(AsyncMock(return_value=True), {"access_type": "offline", "prompt": "consent"}),
(AsyncMock(return_value=False), None),
(AsyncMock(side_effect=RuntimeError("store down")), None),
],
)
async def test_google_login_asks_for_a_refresh_token_only_while_an_id_jag_server_exists(
retention_enabled, expected_params
):
"""Google issues a refresh token only when the authorization request carries ``access_type=offline``,
and the refresher cannot renew a Google assertion without one. The login asks only while an
``oauth2_id_jag`` server exists, and a store outage must not fail the login or add the consent."""
from fastapi.responses import RedirectResponse
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
seen: dict[str, object] = {}
async def fake_redirect(self, *, redirect_uri=None, params=None, state=None):
seen["params"] = params
return RedirectResponse(url="https://accounts.google.com/o/oauth2/v2/auth")
with (
patch.dict(os.environ, {"GOOGLE_CLIENT_SECRET": "google-secret"}),
patch( # test-quality-ok: stubs the upstream Google authorization boundary, where the request leaves
"fastapi_sso.sso.google.GoogleSSO.get_login_redirect", autospec=True, side_effect=fake_redirect
),
):
response = await SSOAuthenticationHandler.get_sso_login_redirect(
redirect_url="http://localhost:4000/sso/callback",
google_client_id="google-cid",
retention_enabled=retention_enabled,
)
assert response is not None
assert seen["params"] == expected_params
@pytest.mark.asyncio
@pytest.mark.parametrize(
"retention_enabled, asks_offline_access",
[
(AsyncMock(return_value=True), True),
(AsyncMock(return_value=False), False),
(AsyncMock(side_effect=RuntimeError("store down")), False),
],
)
async def test_microsoft_login_asks_for_offline_access_only_while_an_id_jag_server_exists(
retention_enabled, asks_offline_access
):
"""Entra issues a refresh token only when ``offline_access`` is in the requested scope, and the
refresher cannot renew an Entra assertion without one. The default sign-in scopes stay in place so
the login keeps resolving the same profile."""
from fastapi.responses import RedirectResponse
from fastapi_sso.sso.microsoft import MicrosoftSSO
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
seen: dict[str, object] = {}
async def fake_redirect(self, *, redirect_uri=None, params=None, state=None):
seen["scope"] = list(self._scope)
return RedirectResponse(url="https://login.microsoftonline.com/tenant/oauth2/v2.0/authorize")
with (
patch.dict(os.environ, {"MICROSOFT_CLIENT_SECRET": "ms-secret", "MICROSOFT_TENANT": "tenant"}),
patch( # test-quality-ok: stubs the upstream Entra authorization boundary, where the request leaves
"fastapi_sso.sso.microsoft.MicrosoftSSO.get_login_redirect", autospec=True, side_effect=fake_redirect
),
):
response = await SSOAuthenticationHandler.get_sso_login_redirect(
redirect_url="http://localhost:4000/sso/callback",
microsoft_client_id="ms-cid",
retention_enabled=retention_enabled,
)
assert response is not None
assert all(default in seen["scope"] for default in MicrosoftSSO.scope), seen["scope"]
assert ("offline_access" in seen["scope"]) is asks_offline_access, seen["scope"]
@pytest.mark.asyncio
async def test_cli_completion_persists_assertion_under_db_user_id():
"""The CLI funnel persists the captured assertion under the DB-resolved user_id."""
@ -9788,8 +10044,6 @@ def _id_jag_gap_warnings(caplog) -> list[str]:
@pytest.mark.parametrize(
"provider_env, expected_fragment",
[
({"GOOGLE_CLIENT_ID": "cid"}, "google"),
({"MICROSOFT_CLIENT_ID": "cid", "MICROSOFT_TENANT": "t"}, "microsoft"),
({}, "no SSO provider is configured"),
],
)
@ -9816,16 +10070,17 @@ async def test_uncaptured_assertion_warns_when_an_id_jag_server_is_registered(
@pytest.mark.asyncio
async def test_generic_provider_that_returned_no_id_token_still_warns(monkeypatch, caplog):
"""Generic OIDC has a capture path, so there is no configuration gap to report; the login
still handed the id_jag arm nothing, and that must not pass silently."""
@pytest.mark.parametrize("provider_env_var", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"])
async def test_a_capturing_provider_that_returned_no_id_token_still_warns(monkeypatch, caplog, provider_env_var):
"""Google, Microsoft and generic OIDC all have a capture path, so there is no configuration gap
to report; the login still handed the id_jag arm nothing, and that must not pass silently."""
from litellm.proxy.management_endpoints.ui_sso import (
warn_if_id_jag_assertion_uncaptured,
)
for name in ("GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "SAML_IDP_METADATA_URL"):
for name in ("GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"):
monkeypatch.delenv(name, raising=False)
monkeypatch.setenv("GENERIC_CLIENT_ID", "cid")
monkeypatch.setenv(provider_env_var, "cid")
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
await warn_if_id_jag_assertion_uncaptured(None, retention_enabled=AsyncMock(return_value=True))
@ -9953,7 +10208,7 @@ async def test_browser_funnel_reports_an_uncaptured_assertion(monkeypatch, caplo
warnings = _id_jag_gap_warnings(caplog)
assert len(warnings) == 1
assert "google" in str(warnings[0])
assert "no usable id_token" in str(warnings[0])
@pytest.mark.asyncio
@ -10017,7 +10272,7 @@ async def test_cli_funnel_reports_an_uncaptured_assertion(monkeypatch, caplog):
warnings = _id_jag_gap_warnings(caplog)
assert len(warnings) == 1
assert "microsoft" in str(warnings[0])
assert "no usable id_token" in str(warnings[0])
def _cli_callback_kwargs(flow):

22
tests/unit/proxy/test__lazy_features.py
View file

@ -401,3 +401,25 @@ def test_without_the_flag_the_warmup_route_registers_a_feature_and_returns_its_p
assert body.paths["/alpha/list"]["get"].tags == ("alpha",)
assert loaded_lazy_modules(app) == {features[0].module_path}
assert "/alpha/list" in _paths(app) and "/beta/list" not in _paths(app)
def test_root_authorization_server_metadata_is_the_gateways_not_the_byok_connector_document(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Two features claim the root ``/.well-known/oauth-authorization-server``. The one a client
gets must be the gateway's own authorization server metadata, whose ``/token`` serves the
gateway's grants, not the BYOK connector form flow whose token endpoint is
``/v1/mcp/oauth/token``; the connector keeps its path-appended document."""
monkeypatch.setenv(FLAG, "false")
app: Final = FastAPI()
attach_lazy_features(app)
with TestClient(app, base_url="https://gateway.example.com") as client:
root: Final = client.get("/.well-known/oauth-authorization-server")
connector: Final = client.get("/.well-known/oauth-authorization-server/v1/mcp/oauth")
assert root.status_code == 200
assert root.json()["issuer"] == "https://gateway.example.com"
assert root.json()["token_endpoint"] == "https://gateway.example.com/token"
assert connector.status_code == 200
assert connector.json()["token_endpoint"] == "https://gateway.example.com/v1/mcp/oauth/token"

10
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -27480,8 +27480,10 @@ export interface components {
};
/** Body_token_endpoint__mcp_server_name__token_post */
Body_token_endpoint__mcp_server_name__token_post: {
/** Assertion */
assertion?: string | null;
/** Client Id */
client_id: string;
client_id?: string | null;
/** Client Secret */
client_secret?: string | null;
/** Code */
@ -27507,8 +27509,10 @@ export interface components {
};
/** Body_token_endpoint_token_post */
Body_token_endpoint_token_post: {
/** Assertion */
assertion?: string | null;
/** Client Id */
client_id: string;
client_id?: string | null;
/** Client Secret */
client_secret?: string | null;
/** Code */
@ -37288,6 +37292,8 @@ export interface components {
id_jag_resource_client_secret?: string | null;
/** Id Jag Resource Token Endpoint */
id_jag_resource_token_endpoint?: string | null;
/** Id Jag Subject Source */
id_jag_subject_source?: string | null;
/** Redirect Uris */
redirect_uris?: string[] | null;
/** Scopes */