scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-17 04:34:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 172

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

Browse source 
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam 2026-10-16 12:16:56 -07:00 • committed by GitHub
parent 3501e36fa9
commit f02449ad38
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
25 changed files with 1720 additions and 200 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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,

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."""

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

@ -0,0 +1,216 @@
"""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.
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
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_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,
) -> SSOIdentityAssertion:
"""The assertion as captured when no server wants a refresh token, otherwise the assertion plus the
refresh token the first target's IdP client exchanged it for. An OIDC login's own SSO refresh token
does not count, it is bound to the SSO client and the ID-JAG client cannot redeem it."""
if not targets or assertion.exchanged_refresh_token is not None:
return assertion
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,
)
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)
)
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

79
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,41 @@ 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 exchanged the login assertion for when the user signed in. Never the SSO
refresh token an OIDC login holds, that one is bound to the SSO client. The assertion's own expiry
does not apply here, 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.

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

@ -80,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
@ -293,6 +295,13 @@ 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 "
@ -358,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
@ -436,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,

142
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,10 @@ 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
@ -44,6 +45,7 @@ GOOGLE_REFRESH_TOKEN_PARAMS: Final = MappingProxyType({"access_type": "offline",
_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)
@ -77,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)
@ -86,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:
@ -138,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:
@ -167,6 +181,85 @@ 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."""
return renewed.model_copy(
update={
"exchanged_refresh_token": 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
@ -220,6 +313,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(
@ -258,6 +372,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,
)
@ -298,6 +417,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``.
@ -316,6 +441,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, ...] = ()

44
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": [
{
@ -37358,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": [
{
@ -40763,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": [
{
@ -44843,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": [
{

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

@ -2,10 +2,10 @@
An ``oauth2_id_jag`` MCP server spends the ``id_token`` captured at SSO login as its RFC 8693
subject token. The Google, Microsoft and generic OIDC login paths all reach a token response the
gateway retains one from; SAML has no id_token to capture, and a deployment with no SSO never signs
anyone in, so under either 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.
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
@ -19,8 +19,8 @@ from typing_extensions import assert_never
from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler
_OIDC_REMEDY: Final = (
"Point SSO at an OIDC provider (GOOGLE_CLIENT_ID, MICROSOFT_CLIENT_ID or GENERIC_CLIENT_ID), the login "
"paths whose token response the gateway retains an id_token from"
"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"
)
@ -56,19 +56,15 @@ 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.google | ActiveSSOProvider.microsoft | 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. {_OIDC_REMEDY}"
)
case 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. {_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:

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

@ -71,6 +71,9 @@ 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,
@ -2343,7 +2346,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)
@ -2356,7 +2359,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,
@ -2364,6 +2367,7 @@ async def saml_callback(request: Request):
access_token_payload=None,
jwt_handler=jwt_handler,
return_to=cp_return_to,
sso_assertion=retained,
)
@ -2480,7 +2484,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] = []
@ -3863,7 +3868,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()

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.",

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

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

@ -0,0 +1,327 @@
"""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_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"
@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 == []

89
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)})
@ -853,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 == []

152
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,58 @@ 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_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"

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

@ -1641,7 +1641,8 @@ class TestMCPServerManager:
assert server.oauth2_flow is None
@pytest.mark.asyncio
async def test_load_servers_from_config_warns_for_id_jag_with_saml_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("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata")
manager = MCPServerManager()
@ -1652,11 +1653,7 @@ 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 "saml" in warnings[0]
assert "MICROSOFT_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"])

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

@ -74,14 +74,11 @@ class TestIdJagAssertionCaptureGap:
monkeypatch.setenv(env_var, "configured")
assert id_jag_assertion_capture_gap() is None
def test_saml_is_named_with_the_oidc_remedy(self, monkeypatch):
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")
gap = id_jag_assertion_capture_gap()
assert gap is not None
assert "saml" in gap
assert "GOOGLE_CLIENT_ID" in gap
assert "MICROSOFT_CLIENT_ID" in gap
assert "GENERIC_CLIENT_ID" in gap
assert id_jag_assertion_capture_gap() is None
def test_no_sso_configured_reports_a_gap(self):
gap = id_jag_assertion_capture_gap()
@ -100,12 +97,9 @@ 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_saml_at_startup_reports_the_capture_gap(self, monkeypatch):
def test_saml_at_startup_has_no_gap(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata")
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
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):

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

@ -4530,7 +4530,6 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap:
@pytest.mark.parametrize(
"provider_env, expected_fragment",
[
({"SAML_IDP_METADATA_URL": "https://idp.example.com/metadata"}, "saml"),
({}, "no SSO provider is configured"),
],
)
@ -4564,7 +4563,8 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap:
)
self._clear_sso_env(monkeypatch)
monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata")
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()
@ -4601,7 +4601,7 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap:
warnings = self._id_jag_warnings(caplog)
assert len(warnings) == 1
assert "saml" 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
)

84
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
@ -9004,7 +9005,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("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata")
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
@ -9130,7 +9132,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,
@ -9140,7 +9142,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")
@ -9662,6 +9664,81 @@ 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",
@ -9889,7 +9966,6 @@ def _id_jag_gap_warnings(caplog) -> list[str]:
@pytest.mark.parametrize(
"provider_env, expected_fragment",
[
({"SAML_IDP_METADATA_URL": "https://idp.example.com/metadata"}, "saml"),
({}, "no SSO provider is configured"),
],
)

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

@ -37292,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 */