scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

feat(mcp): renew the stored SSO identity assertion behind ID-JAG (#35401)

Browse source 
* feat(mcp): renew the stored SSO identity assertion behind ID-JAG

The oauth2_id_jag arm asserts the id_token captured at the user's last
interactive SSO login, and nothing ever renewed it, so an agent holding a
brokered LiteLLM key could act for that user only until that token's exp.
The assertion already carried the IdP refresh token beside it; this
redeems it.

RefreshingSSOAssertionStore wraps the database reader and satisfies the
same protocol, so the egress arm is unchanged. Renewal is lazy and
single-flighted per user through the same RefreshCoordinator the
authorization_code arm uses, since an IdP that rotates refresh tokens
treats two concurrent redemptions as replay. A refusal leaves the expired
assertion in place so the reader still challenges the user; an
unreachable IdP surfaces as a store outage instead.

* fix(mcp): let a cross-replica loser settle the SSO assertion renewal itself

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

* fix(mcp): satisfy type discipline lint budget

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

* chore(ci): rerun checks after docs main added the missing router setting row

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

* fix(mcp): answer a cross-replica loser retryable instead of re-electing it

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

* fix(mcp): bypass stale assertion cache during renewal

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

* chore: ratchet type-discipline budget after merge

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

---------

Co-authored-by: Yassin Kortam <yassin.kortam@gmail.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam 2026-09-05 12:43:02 -07:00 • committed by GitHub
parent b3f28a77d8
commit 110f654f34
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 1427 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py
View file

@ -31,11 +31,8 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto
TokenCacheBackend,
TokenStoreUnavailable,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_distributed_lock import (
RedisDistributedLock,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_coordinator import (
RedisRefreshCoordinator,
from litellm.proxy._experimental.mcp_server.outbound_credentials.runtime_refresh_coordinator import (
runtime_refresh_coordinator,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import (
OAuthTokenCacheCodec,
@ -131,23 +128,17 @@ def _runtime_backend_and_coordinator() -> tuple[TokenCacheBackend | None, Refres
)
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
redis_cache: Final = user_api_key_cache.redis_cache
if redis_cache is None:
coordinator: Final = runtime_refresh_coordinator()
if coordinator is None:
return None, None, False
codec: Final = OAuthTokenCacheCodec(
encrypt_value_helper,
lambda blob: decrypt_value_helper(blob, "mcp_per_user_token", exception_type="debug"),
)
# user_api_key_cache satisfies the AsyncCache slice (DualCache types ttl via **kwargs) and the
# Redis client from init_async_client() is partially typed - both are untyped-boundary casts.
# user_api_key_cache satisfies the AsyncCache slice (DualCache types ttl via **kwargs) - an
# untyped-boundary cast.
cache: Final[AsyncCache] = user_api_key_cache # pyright: ignore
redis_client: Final = redis_cache.init_async_client() # pyright: ignore
lock: Final = RedisDistributedLock(
redis_client, # pyright: ignore
namespace_key=redis_cache.check_and_fix_namespace,
)
backend: Final = DualCacheTokenCacheBackend(cache, codec)
coordinator: Final = RedisRefreshCoordinator(lock)
return backend, coordinator, True

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

@ -46,11 +46,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.result import (
Ok,
Result,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import (
default_sso_assertion_store,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
AssertionStoreUnavailable,
DbSSOAssertionStore,
SSOAssertionStore,
SSOIdentityAssertion,
assertion_expired,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import (
ExchangedToken,
@ -129,7 +131,7 @@ class UpstreamCredentialProvider:
self._token_endpoint: TokenEndpointClient = token_endpoint or TokenEndpointClient()
self._exchanged_tokens: ExchangedTokenCache = exchanged_tokens or ExchangedTokenCache()
self._client_credentials_source = client_credentials_source or ClientCredentialsTokenSource()
self._sso_assertion_store: SSOAssertionStore = sso_assertion_store or DbSSOAssertionStore()
self._sso_assertion_store: SSOAssertionStore = sso_assertion_store or default_sso_assertion_store()
async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]:
match server.config:
@ -246,7 +248,7 @@ class UpstreamCredentialProvider:
"Sign in through LiteLLM SSO so the gateway captures one."
)
)
if _assertion_expired(assertion, datetime.now(timezone.utc)):
if assertion_expired(assertion, datetime.now(timezone.utc)):
return Error(
CredError.of_precondition_required(
"The stored IdP identity assertion for this user has expired. Sign in through "
@ -405,19 +407,6 @@ def _id_jag_slot_key(subject: Subject, server: ServerSpec) -> str:
return hashlib.sha256(material.encode()).hexdigest()
def _assertion_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool:
"""Whether the stored assertion's ``exp`` has passed. An assertion carrying no expiry is
treated as usable and left for the IdP to reject, since the store records what the id_token
claimed rather than imposing a lifetime of its own. A naive ``expires_at`` is read as UTC so a
stored value that lost its offset compares instead of raising.
"""
expires_at: Final = assertion.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 _id_jag_fingerprint(subject_token: str, server_id: str, config: IdJagConfig) -> str:
"""What the cached leg-2 bearer was minted from: the subject token, the server, and the config.

41
litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py Normal file
View file

@ -0,0 +1,41 @@
"""The runtime ``RefreshCoordinator``: cross-replica single-flight when Redis is wired.
Builds ``RedisRefreshCoordinator`` over the proxy's shared Redis so one refresh runs per key
across the fleet, or returns ``None`` when Redis is absent so the caller keeps the foundation's
in-process default (correct for a single replica). The proxy globals it reads are not ready at
import time, so this is called per composition rather than held as module state.
Shared by every credential arm that renews a stored grant: a rotating refresh token must be
redeemed once across all workers, so each arm electing its own winner with its own lock shape
would be a bug waiting to differ.
"""
from __future__ import annotations
from typing import Final
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
RefreshCoordinator,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_distributed_lock import (
RedisDistributedLock,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_coordinator import (
RedisRefreshCoordinator,
)
def runtime_refresh_coordinator() -> RefreshCoordinator | None:
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 # runtime global
redis_cache: Final = user_api_key_cache.redis_cache
if redis_cache is None:
return None
# The Redis client from init_async_client() is only partially typed; the lock validates every
# reply it depends on, so the untyped boundary is contained here.
redis_client: Final = redis_cache.init_async_client() # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # litellm redis wrapper is untyped
lock: Final = RedisDistributedLock(
redis_client, # pyright: ignore[reportArgumentType,reportUnknownArgumentType] # litellm redis wrapper is untyped
namespace_key=redis_cache.check_and_fix_namespace,
)
return RedisRefreshCoordinator(lock)

469
litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py Normal file
View file

@ -0,0 +1,469 @@
"""Renew the stored SSO identity assertion so an ID-JAG agent outlives one id_token.
The ``oauth2_id_jag`` arm asserts the id_token captured at the user's last interactive sign-in, so
without renewal an agent holding a brokered LiteLLM key can act for that user only until that token's
``exp``, typically an hour, and the sole recovery is another interactive login. The assertion already
carries the IdP refresh token beside it; this module is what redeems it.
``RefreshingSSOAssertionStore`` wraps any ``SSOAssertionStore`` and satisfies the same protocol, so
the egress arm is unchanged: it still reads one assertion and still judges expiry itself. Renewal is
lazy (only a read that finds a near-expiry assertion triggers one, so IdP traffic tracks actual use,
not the size of the user table) and single-flighted per user through the same ``RefreshCoordinator``
the ``authorization_code`` arm uses, because an IdP that rotates refresh tokens treats two concurrent
redemptions of one token as replay and can revoke the whole grant chain.
The refresh is redeemed against the generic-OIDC client the login itself used
(``GENERIC_TOKEN_ENDPOINT`` / ``GENERIC_CLIENT_ID`` / ``GENERIC_CLIENT_SECRET``, which the proxy
reconciles from the stored SSO row into the process environment at startup), authenticated the way
that login authenticated: the non-PKCE path always sends HTTP Basic, while the PKCE path sends the
credentials in the body when ``GENERIC_INCLUDE_CLIENT_ID`` is set, and an IdP application may accept
only one of the two. An assertion can only exist if that client minted it, so no other client could
redeem its refresh token, and no other method is known to be accepted. A deployment whose
``GENERIC_SCOPE`` omits ``offline_access`` captures no refresh token at all, which is why that miss
logs the scope by name rather than failing silently.
Failures are values internally (``Result[_, RefreshFailure]``). At the store boundary they collapse
onto the protocol's existing two-outcome contract: a refusal returns the expired assertion unchanged
so the reader's own guard challenges the user to sign in again, while a transient IdP failure raises
``AssertionStoreUnavailable`` so the reader answers 503 instead of blaming the user for an outage.
One ambiguity remains under Redis-coordinated renewal across replicas. A cross-replica loser that
finds the row still expiring after the holder finished cannot tell a refused refresh from a renewal
that could not be recorded. Redeeming itself could consume a refresh token the holder may already
have rotated, so it answers retryable 503 rather than guessing a sign-in challenge. The next
uncontended read settles the outcome itself: a refusal challenges, and a successful refresh persists.
If the holder rotated the token but its write failed, that rotation is lost and the next uncontended
read's refusal challenges, which is the only honest answer because the rotated token was never
recorded. On the refusal path, the loser pays for one retry before that challenge.
"""
from __future__ import annotations
import json
import os
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Final, Literal, Protocol
import httpx
from pydantic import SecretStr, TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import Timeout
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # litellm http handler is untyped
)
from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import (
build_token_endpoint_client_auth,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
InProcessRefreshCoordinator,
RefreshCoordinator,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import (
Error,
Ok,
Result,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.runtime_refresh_coordinator import (
runtime_refresh_coordinator,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
AssertionStoreUnavailable,
DbSSOAssertionStore,
SSOAssertionStore,
SSOIdentityAssertion,
assertion_expired,
assertion_from_sso_login,
fetch_sso_identity_assertion,
persist_sso_identity_assertion,
)
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.mcp import MCPTokenEndpointAuthMethod
_BODY_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(dict[str, object])
_REFRESH_GRANT_TYPE: Final = "refresh_token"
# The lock namespace for the one assertion row a user has; the sibling arm keys the same lock by
# server_id, and no server_id can collide with this literal.
_SINGLE_FLIGHT_KEY: Final = "sso_identity_assertion"
# Renew this far ahead of ``exp`` so a token that would die between resolution and the second leg of
# the exchange is replaced first. Matches the sibling per-user token store's skew.
_DEFAULT_EXPIRY_SKEW_SECONDS: Final = 60.0
class AssertionRead(Protocol):
"""Reads the user's stored assertion row."""
async def __call__(self, user_id: str) -> SSOIdentityAssertion | None: ...
class AssertionWrite(Protocol):
"""Replaces the user's stored assertion row."""
async def __call__(self, user_id: str, assertion: SSOIdentityAssertion) -> None: ...
class CoordinatorFactory(Protocol):
"""Builds the cross-replica coordinator, or ``None`` when there is no shared lock to build on."""
def __call__(self) -> RefreshCoordinator | None: ...
class FormPost(Protocol):
"""POSTs an OAuth form and hands back the raw response."""
async def __call__(
self, url: str, form: Mapping[str, str], headers: Mapping[str, str]
) -> httpx.Response | None: ...
@dataclass(frozen=True, slots=True)
class SSOClientConfig:
"""The generic-OIDC client credentials a refresh_token grant has to authenticate as, and how."""
token_endpoint: str
client_id: str
client_secret: SecretStr
auth_method: MCPTokenEndpointAuthMethod
def sso_client_config(env: Mapping[str, str]) -> SSOClientConfig | None:
"""The configured generic-OIDC client, or ``None`` when the deployment has none.
Read from the process environment because that is where the login path reads it
(``_setup_generic_sso_env_vars``) and where the proxy materializes the stored ``sso_config`` row
at startup, so this resolves to the same client that minted the assertion. ``None`` is an
ordinary state, not an error: a deployment signing in through a provider that captures no
assertion has nothing here to renew, and a client with no secret is not a confidential client
that could redeem one.
``auth_method`` is derived from the same ``GENERIC_INCLUDE_CLIENT_ID`` the login reads, because
the two login paths do not agree: the non-PKCE path always authenticates with HTTP Basic, while
the PKCE path puts the credentials in the body when that flag is set. Both capture assertions, so
a constant here would authenticate the renewal differently from the sign-in that produced the
refresh token and 401 against an IdP application registered for only one of the two.
"""
token_endpoint: Final = env.get("GENERIC_TOKEN_ENDPOINT")
client_id: Final = env.get("GENERIC_CLIENT_ID")
client_secret: Final = env.get("GENERIC_CLIENT_SECRET")
if not token_endpoint or not client_id or not client_secret:
return None
includes_client_id: Final = env.get("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true"
return SSOClientConfig(
token_endpoint=token_endpoint,
client_id=client_id,
client_secret=SecretStr(client_secret),
auth_method="client_secret_post" if includes_client_id else "client_secret_basic",
)
@dataclass(frozen=True, slots=True)
class RefreshFailure:
"""Why a renewal produced nothing, split by what the caller can do about it.
``rejected`` is settled: this refresh token will never work again, so the user has to sign in.
``unavailable`` is transient: the same attempt may succeed in a minute, so telling the user to
sign in again would be a lie about whose problem it is. Both arms carry the same payload, so
this is a ``Literal`` discriminant rather than a ``tagged_union``; consumers still ``match`` on
``kind`` with an ``assert_never`` tail.
"""
kind: Literal["rejected", "unavailable"]
detail: str
@staticmethod
def of_rejected(detail: str) -> RefreshFailure:
return RefreshFailure(kind="rejected", detail=detail)
@staticmethod
def of_unavailable(detail: str) -> RefreshFailure:
return RefreshFailure(kind="unavailable", detail=detail)
class TokenEndpointTransport(Protocol):
"""One form POST to the IdP token endpoint, with the refusal/outage split preserved.
That split is the whole reason this is not the resolver's ``TokenEndpointClient``: that
collaborator maps every non-2xx to ``upstream_unavailable``, which is right for an exchange leg
and wrong here, where a 400 ``invalid_grant`` means the stored refresh token is dead and the user
must act.
"""
async def post(
self, url: str, form: Mapping[str, str], headers: Mapping[str, str]
) -> Result[Mapping[str, object], RefreshFailure]: ...
async def post_form(url: str, form: Mapping[str, str], headers: Mapping[str, str]) -> httpx.Response | None:
# litellm's httpx handler is only partially typed; nothing but the response object crosses back,
# and the transport below validates its body, so the untyped boundary is contained here.
client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) # pyright: ignore[reportUnknownVariableType] # litellm http handler is untyped
return await client.post(url, data=form, headers=headers) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType,reportReturnType,reportArgumentType] # litellm http handler is untyped and its stub narrows data=/headers= to dict, which httpx itself does not require
class HttpxTokenEndpointTransport:
"""The live transport. 4xx is the IdP refusing this grant; anything else is an outage.
The POST itself is injected so that split, which decides whether the user is challenged or told
to wait, is testable without a live IdP.
"""
def __init__(self, post: FormPost = post_form) -> None:
self._post = post
async def post(
self, url: str, form: Mapping[str, str], headers: Mapping[str, str]
) -> Result[Mapping[str, object], RefreshFailure]:
try:
response: Final = await self._post(url, form, headers)
if response is None:
return Error(RefreshFailure.of_unavailable("the IdP token endpoint returned no response"))
response.raise_for_status()
body: Final = _BODY_ADAPTER.validate_python(response.json()) # pyright: ignore[reportAny] # untyped JSON; the adapter is the type gate
except httpx.HTTPStatusError as exc:
status: Final = exc.response.status_code
if 400 <= status < 500:
return Error(RefreshFailure.of_rejected(f"the IdP refused the refresh with status {status}"))
return Error(RefreshFailure.of_unavailable(f"the IdP token endpoint answered with status {status}"))
except (httpx.RequestError, Timeout) as exc:
return Error(RefreshFailure.of_unavailable(f"the IdP token endpoint is unreachable ({type(exc).__name__})"))
except json.JSONDecodeError:
return Error(RefreshFailure.of_unavailable("the IdP token endpoint returned a non-JSON response"))
except ValidationError:
return Error(RefreshFailure.of_unavailable("the IdP token endpoint returned a non-object response"))
return Ok(body)
class SSOAssertionRefresher:
"""Redeems the stored refresh token for a current id_token and writes the rotation back.
Collaborators are injected so the orchestration, the untyped response parsing and the
write-back race are all testable without an IdP or a database.
"""
def __init__(
self,
transport: TokenEndpointTransport,
*,
client_config: Callable[[], SSOClientConfig | None] = lambda: sso_client_config(os.environ),
read: AssertionRead = fetch_sso_identity_assertion,
write: AssertionWrite = persist_sso_identity_assertion,
) -> None:
self._transport = transport
self._client_config = client_config
self._read = read
self._write = write
async def refresh(
self, user_id: str, assertion: SSOIdentityAssertion
) -> Result[SSOIdentityAssertion, RefreshFailure]:
if assertion.refresh_token is None:
verbose_proxy_logger.warning(
"ID-JAG: the stored IdP identity assertion for user_id=%s has expired and no refresh token was "
"captured with it, so it cannot be renewed without another interactive sign-in. Add "
"'offline_access' to GENERIC_SCOPE so the SSO login captures one.",
user_id,
)
return Error(RefreshFailure.of_rejected("no refresh token was captured at sign-in"))
config: Final = self._client_config()
if config is None:
verbose_proxy_logger.warning(
"ID-JAG: the stored IdP identity assertion for user_id=%s has expired and cannot be renewed "
"because the generic SSO client is not configured (GENERIC_TOKEN_ENDPOINT, GENERIC_CLIENT_ID, "
"GENERIC_CLIENT_SECRET).",
user_id,
)
return Error(RefreshFailure.of_rejected("the generic SSO client is not configured"))
carried_refresh_token: Final = assertion.refresh_token.get_secret_value()
# Whichever method the SSO login used for this client, since that is the one the IdP
# application is known to accept: an assertion only exists to renew because a sign-in already
# authenticated this client that way.
client_auth: Final = build_token_endpoint_client_auth(
auth_method=config.auth_method,
client_id=config.client_id,
client_secret=config.client_secret.get_secret_value(),
)
form: Final = { # mutable-ok: the RFC 6749 form body is a wire format the HTTP client takes as a mapping
"grant_type": _REFRESH_GRANT_TYPE,
"refresh_token": carried_refresh_token,
**client_auth.body,
}
match await self._transport.post(config.token_endpoint, form, client_auth.headers):
case Error(failure):
return Error(failure)
case Ok(body):
return await self._renewed_from(user_id, assertion, body, carried_refresh_token)
async def _renewed_from(
self,
user_id: str,
previous: SSOIdentityAssertion,
body: Mapping[str, object],
carried_refresh_token: str,
) -> Result[SSOIdentityAssertion, RefreshFailure]:
"""The renewed assertion, built by the same validator the login path uses.
A rotated refresh token replaces the stored one; an omitted one carries forward, since an
IdP that does not rotate expects the original to keep working.
"""
rotated: Final = body.get("refresh_token")
renewed: Final = assertion_from_sso_login(
body.get("id_token"),
rotated if isinstance(rotated, str) and rotated else carried_refresh_token,
)
if renewed is None:
verbose_proxy_logger.warning(
"ID-JAG: the IdP accepted the refresh for user_id=%s but returned no usable id_token, so there "
"is nothing to assert upstream. The SSO client's grant needs the 'openid' scope for the token "
"endpoint to return one on a refresh.",
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)
if failure is not None:
return Error(failure)
return Ok(renewed)
async def _store_renewal(
self, user_id: str, previous: SSOIdentityAssertion, renewed: SSOIdentityAssertion
) -> RefreshFailure | None:
"""Write the renewal back, unless the row moved on while this renewal was in flight.
The row is one per user and last-write-wins, so an interactive sign-in landing mid-renewal
would otherwise be overwritten with a refresh token the IdP has already rotated away, costing
that user a sign-in later. Comparing against the id_token this renewal started from is what
detects that; skipping is safe because the newer row is the one the reader wants anyway.
A failed write is transient, not settled. The store, not this return value, is what every
caller reads, so a renewal that could not be recorded is a renewal nobody will see; saying so
keeps a database problem answering 503 rather than telling the user to sign in again over it.
"""
try:
current: Final = await self._read(user_id)
if current is not None and current.id_token.get_secret_value() != previous.id_token.get_secret_value():
verbose_proxy_logger.info(
"ID-JAG: a newer IdP identity assertion for user_id=%s was stored while this renewal was in "
"flight; keeping the stored one.",
user_id,
)
return None
await self._write(user_id, renewed)
except Exception as exc: # noqa: BLE001 # any storage failure is transient here, never the user's fault
verbose_proxy_logger.warning(
"ID-JAG: could not persist the renewed IdP identity assertion for user_id=%s, so the rotated "
"refresh token is lost and this user will have to sign in again once the renewed token expires: %s",
user_id,
exc,
)
return RefreshFailure.of_unavailable("the renewed IdP identity assertion could not be persisted")
return None
class RefreshingSSOAssertionStore:
"""An ``SSOAssertionStore`` that renews a near-expiry assertion before handing it back.
Reads the inner store; an assertion still comfortably inside its lifetime is returned untouched,
so the common path costs exactly what it did before. Otherwise one renewal runs per user through
the injected ``RefreshCoordinator`` and every caller then re-reads the inner store, which is the
authority: the winner's write is what they all observe, and a renewal the write-back guard
skipped yields the newer assertion that displaced it rather than a private copy.
A refusal leaves the expired assertion in place for the reader's own guard to reject, so the user
sees the same sign-in-again challenge as before this store existed. A transient IdP failure
raises ``AssertionStoreUnavailable``, the protocol's existing signal for "this is not the user's
fault"; concurrent in-process callers share that outcome, while a cross-replica loser answers 503
when its re-read still finds the row expiring. On the refusal path that costs the loser one retry,
which then challenges. If the holder rotated the token but its write failed, the rotation is lost
and the next uncontended read's refusal challenges, the only honest answer because that token was
never recorded.
"""
def __init__(
self,
inner: SSOAssertionStore,
refresher: SSOAssertionRefresher,
*,
fresh_read: AssertionRead,
coordinator_factory: CoordinatorFactory = runtime_refresh_coordinator,
expiry_skew_seconds: float = _DEFAULT_EXPIRY_SKEW_SECONDS,
clock: Callable[[], datetime] = lambda: datetime.now(timezone.utc),
) -> None:
self._inner = inner
self._refresher = refresher
self._fresh_read = fresh_read
self._coordinator_factory = coordinator_factory
self._in_process_coordinator = InProcessRefreshCoordinator()
self._distributed_coordinator: RefreshCoordinator | None = None
self._skew = timedelta(seconds=expiry_skew_seconds)
self._clock = clock
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
assertion: Final = await self._inner.fetch(user_id)
if not self._expiring(assertion):
return assertion
await self._coordinator().run(
user_id,
_SINGLE_FLIGHT_KEY,
refresh=lambda: self._renew(user_id),
reread=lambda: self._reread_renewed(user_id),
)
return await self._fresh_read(user_id)
def _expiring(self, assertion: SSOIdentityAssertion | None) -> bool:
return assertion is not None and assertion_expired(assertion, self._clock() + self._skew)
def _coordinator(self) -> RefreshCoordinator:
"""The cross-replica coordinator once Redis is reachable, else the in-process one.
Built on first use and kept, because the proxy's Redis client is not wired at import time;
retried while it is absent so a proxy that gains Redis later stops electing per-worker.
"""
if self._distributed_coordinator is None:
self._distributed_coordinator = self._coordinator_factory()
return self._distributed_coordinator or self._in_process_coordinator
async def _renew(self, user_id: str) -> None:
"""The elected renewal, judged from a fresh read so a rotation another replica just landed is
never redeemed again. Returns nothing: the inner store, not this return value, is what every
caller reads afterwards, so the winner and the losers cannot disagree."""
latest: Final = await self._fresh_read(user_id)
if latest is None or not self._expiring(latest):
return
match await self._refresher.refresh(user_id, latest):
case Ok(_):
return
case Error(failure):
match failure.kind:
case "rejected":
return
case "unavailable":
raise AssertionStoreUnavailable(failure.detail)
assert_never(failure.kind)
async def _reread_renewed(self, user_id: str) -> None:
"""A loser cannot distinguish refusal from an unrecorded renewal without risking token replay.
It answers retryable 503 instead of guessing a sign-in challenge; the retry runs uncontended
and settles the outcome itself.
"""
latest: Final = await self._fresh_read(user_id)
if self._expiring(latest):
raise AssertionStoreUnavailable(
f"the IdP identity assertion for user_id={user_id} was being renewed by another replica "
"and is not yet current; retry shortly"
)
def default_sso_assertion_store() -> SSOAssertionStore:
"""The live read seam for the ``id_jag`` arm: the stored assertion, renewed when it is stale."""
db_store: Final = DbSSOAssertionStore()
fresh_read: Final = db_store.fetch_uncached
return RefreshingSSOAssertionStore(
db_store,
SSOAssertionRefresher(HttpxTokenEndpointTransport(), read=fresh_read),
fresh_read=fresh_read,
)

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

@ -127,6 +127,23 @@ def assertion_from_sso_login(id_token: object, refresh_token: object) -> SSOIden
)
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
claimed rather than imposing a lifetime of its own. A naive ``expires_at`` is read as UTC so a
stored value that lost its offset compares instead of raising.
Lives beside the model rather than in either reader so the egress guard and the renewal
trigger judge the same field the same way; passing a ``now`` in the future is how a caller
asks "is this about to expire" without a second, driftable predicate.
"""
expires_at: Final = assertion.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
async def ema_assertion_retention_enabled() -> bool:
"""Whether any MCP server uses ``oauth2_id_jag``, evaluated per login so the gateway only
retains bearer material while an EMA upstream exists to spend it on. Judged against the two
@ -146,7 +163,9 @@ async def ema_assertion_retention_enabled() -> bool:
return True
if prisma_client is None:
return False
row = await prisma_client.db.litellm_mcpservertable.find_first(where={"auth_type": MCPAuth.oauth2_id_jag.value})
row: Final = await prisma_client.db.litellm_mcpservertable.find_first(
where={"auth_type": MCPAuth.oauth2_id_jag.value}
)
return row is not None
@ -158,7 +177,7 @@ async def persist_sso_identity_assertion(
if prisma_client is None:
return
payload: Final[dict[str, str]] = {
payload: Final = {
"id_token": assertion.id_token.get_secret_value(),
**({"refresh_token": assertion.refresh_token.get_secret_value()} if assertion.refresh_token else {}),
**({"issuer": assertion.issuer} if assertion.issuer else {}),
@ -220,11 +239,13 @@ async def fetch_sso_identity_assertion(
class AssertionStoreUnavailable(Exception):
"""Raised by ``fetch`` when the backing store is unreachable (e.g. the DB is down).
"""Raised by ``fetch`` when the assertion cannot be read for a transient reason: the DB is
down, or the IdP behind a renewing store could not be reached.
Distinct from returning ``None`` for "this user has no captured assertion": an outage must not
read as a definite absence, which would tell the user to sign in again over a transient failure,
and it must not escape as an unhandled error on the egress or retry path. Mirrors
and it must not escape as an unhandled error on the egress or retry path. The message names the
real component for the operator log; callers get the reader's generic 503. Mirrors
``TokenStoreUnavailable`` on the sibling per-user OAuth store.
"""
@ -257,6 +278,12 @@ 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_uncached(self, user_id: str) -> SSOIdentityAssertion | None:
try:
return await _read_assertion_from_db(user_id)
except Exception as exc: # noqa: BLE001 # any driver/storage failure is an outage, not an absence
raise AssertionStoreUnavailable(str(exc)) from exc
async def rotate_sso_identity_assertions_master_key(prisma_client: PrismaClient, new_master_key: str) -> None:
"""Re-encrypt every stored assertion under ``new_master_key`` during a salt-key rotation,
@ -280,7 +307,9 @@ async def rotate_sso_identity_assertions_master_key(prisma_client: PrismaClient,
row.user_id,
)
return False
re_encrypted = _STR_ADAPTER.validate_python(encrypt_value_helper(plaintext, new_encryption_key=new_master_key))
re_encrypted: Final = _STR_ADAPTER.validate_python(
encrypt_value_helper(plaintext, new_encryption_key=new_master_key)
)
await prisma_client.db.litellm_ssoidentityassertion.update(
where={"user_id": row.user_id},
data={"assertion_b64": re_encrypted},

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

@ -9,9 +9,11 @@ returning the stub.
import asyncio
import logging
import time
from datetime import datetime, timedelta, timezone
import httpx
import jwt as pyjwt
import pytest
from pydantic import SecretStr
@ -42,6 +44,11 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto
OAuthToken,
TokenStoreUnavailable,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import (
RefreshingSSOAssertionStore,
SSOAssertionRefresher,
SSOClientConfig,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
AssertionStoreUnavailable,
SSOIdentityAssertion,
@ -589,6 +596,73 @@ async def test_id_jag_refuses_an_expired_stored_assertion_without_calling_the_id
assert endpoint.calls == []
@pytest.mark.asyncio
async def test_id_jag_renews_an_expired_stored_assertion_instead_of_challenging():
"""The unattended-agent case end to end: the user last signed in more than an id_token lifetime
ago, so without renewal this is the 412 above. With the renewing store wired the arm resolves,
and leg 1 asserts the renewed token rather than the one that ran out."""
renewed_id_token = pyjwt.encode(
{"iss": "https://idp.example.com", "sub": "alice", "exp": int(time.time()) + 3600},
"test-idp-signing-key-32-bytes-long-xxxx",
algorithm="HS256",
)
expired = SSOIdentityAssertion(
id_token=SecretStr("stale-id-token"),
refresh_token=SecretStr("rt_1"),
expires_at=datetime.now(timezone.utc) - timedelta(seconds=1),
)
rows = {"alice": expired}
async def _read(user_id: str) -> SSOIdentityAssertion | None:
return rows.get(user_id)
async def _write(user_id: str, assertion: SSOIdentityAssertion) -> None:
rows[user_id] = assertion
class _Inner:
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
return await _read(user_id)
class _Transport:
async def post(self, url, form, headers):
return Ok({"access_token": "at", "id_token": renewed_id_token})
refresher = SSOAssertionRefresher(
_Transport(),
client_config=lambda: SSOClientConfig(
token_endpoint="https://idp.example.com/token",
client_id="litellm",
client_secret=SecretStr("s"),
auth_method="client_secret_basic",
),
read=_read,
write=_write,
)
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))
provider = UpstreamCredentialProvider(
token_endpoint=endpoint,
sso_assertion_store=RefreshingSSOAssertionStore(
_Inner(), refresher, fresh_read=_read, coordinator_factory=lambda: None
),
)
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"] == renewed_id_token
def test_the_resolver_defaults_to_the_renewing_assertion_store():
"""A resolver built without collaborators is what production gets, so the default has to renew;
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
@pytest.mark.asyncio
async def test_id_jag_accepts_a_stored_assertion_that_declares_no_expiry():
endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access"))

794
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py Normal file
View file

@ -0,0 +1,794 @@
"""Tests for renewing the stored SSO identity assertion behind the ID-JAG arm.
Pins the contract an unattended agent depends on: an assertion that has run out is renewed from the
refresh token captured beside it instead of stranding the agent until its user signs in again, the
IdP sees one redemption per user no matter how many tool calls arrive at once, a rotation is written
back without overwriting a sign-in that landed mid-renewal, and the two failure kinds stay
distinguishable - a dead refresh token still challenges the user, an unreachable IdP does not.
"""
import asyncio
import base64
import itertools
import logging
import time
from collections.abc import Awaitable, Callable, Mapping
from datetime import datetime, timedelta, timezone
import httpx
import jwt as pyjwt
import pytest
from pydantic import SecretStr
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import (
Error,
Ok,
Result,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import (
HttpxTokenEndpointTransport,
RefreshFailure,
RefreshingSSOAssertionStore,
SSOAssertionRefresher,
SSOClientConfig,
default_sso_assertion_store,
sso_client_config,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
AssertionStoreUnavailable,
SSOIdentityAssertion,
)
SIGNING_KEY = "test-idp-signing-key-32-bytes-long-xxxx"
ISSUER = "https://idp.example.com"
TOKEN_ENDPOINT = "https://idp.example.com/token"
_CLIENT = SSOClientConfig(
token_endpoint=TOKEN_ENDPOINT,
client_id="litellm",
client_secret=SecretStr("s3cret"),
auth_method="client_secret_basic",
)
_POST_CLIENT = SSOClientConfig(
token_endpoint=TOKEN_ENDPOINT,
client_id="litellm",
client_secret=SecretStr("s3cret"),
auth_method="client_secret_post",
)
_MINTED = itertools.count()
def _id_token(subject: str = "u1", exp_offset: int = 3600) -> str:
"""A distinct token per call. Two mints with the same claims in the same second would encode
identically, which would let a test that means "the renewed token replaced the old one" pass
while comparing a value to itself."""
return pyjwt.encode(
{"iss": ISSUER, "sub": subject, "exp": int(time.time()) + exp_offset, "jti": f"t{next(_MINTED)}"},
SIGNING_KEY,
algorithm="HS256",
)
def _stored(id_token: str, *, expires_in: int, refresh_token: str | None = "rt_1") -> SSOIdentityAssertion:
"""A row as the SSO callback wrote it: ``expires_in`` seconds from now, mirroring the id_token."""
return SSOIdentityAssertion(
id_token=SecretStr(id_token),
refresh_token=SecretStr(refresh_token) if refresh_token else None,
issuer=ISSUER,
expires_at=datetime.now(timezone.utc) + timedelta(seconds=expires_in),
)
class _FakeRows:
"""The one assertion row per user: the inner read seam and the refresher's read/write pair."""
def __init__(self, rows: dict[str, SSOIdentityAssertion] | None = None) -> None:
self.rows: dict[str, SSOIdentityAssertion] = dict(rows or {})
self.cached_rows: dict[str, SSOIdentityAssertion] = {}
self.reads: list[str] = []
self.writes: list[tuple[str, SSOIdentityAssertion]] = []
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
self.reads.append(user_id)
# A real suspension point, so concurrent callers interleave here instead of running to
# completion one at a time and never actually racing.
await asyncio.sleep(0)
return self.cached_rows.get(user_id, self.rows.get(user_id))
async def fetch_fresh(self, user_id: str) -> SSOIdentityAssertion | None:
self.reads.append(user_id)
await asyncio.sleep(0)
return self.rows.get(user_id)
async def write(self, user_id: str, assertion: SSOIdentityAssertion) -> None:
self.writes.append((user_id, assertion))
self.rows[user_id] = assertion
class _FakeTransport:
"""Answers every refresh with the same canned result, optionally holding until ``gate`` opens."""
def __init__(
self,
response: Result[Mapping[str, object], RefreshFailure],
*,
gate: asyncio.Event | None = None,
on_call: Callable[[], None] | None = None,
) -> None:
self._response = response
self._gate = gate
self._on_call = on_call
self.calls: list[tuple[str, dict[str, str]]] = []
self.headers: list[dict[str, str]] = []
async def post(
self, url: str, form: Mapping[str, str], headers: Mapping[str, str]
) -> Result[Mapping[str, object], RefreshFailure]:
self.calls.append((url, dict(form)))
self.headers.append(dict(headers))
if self._on_call is not None:
self._on_call()
if self._gate is not None:
await self._gate.wait()
return self._response
def _renewal(id_token: str, refresh_token: str | None = None) -> Result[Mapping[str, object], RefreshFailure]:
body: dict[str, object] = {"access_token": "at", "id_token": id_token, "token_type": "Bearer"}
return Ok({**body, "refresh_token": refresh_token} if refresh_token else body)
def _store(
rows: _FakeRows,
transport: _FakeTransport,
*,
client_config: Callable[[], SSOClientConfig | None] = lambda: _CLIENT,
coordinator_factory: Callable[[], object] = lambda: None,
) -> RefreshingSSOAssertionStore:
refresher = SSOAssertionRefresher(transport, client_config=client_config, read=rows.fetch, write=rows.write)
return RefreshingSSOAssertionStore(
rows,
refresher,
fresh_read=rows.fetch_fresh,
coordinator_factory=coordinator_factory, # pyright: ignore[reportArgumentType] # test doubles stand in for the runtime factory
)
async def _until(predicate: Callable[[], bool]) -> None:
for _ in range(2000):
if predicate():
return
await asyncio.sleep(0)
raise AssertionError("condition never became true")
@pytest.mark.asyncio
async def test_an_expiring_assertion_is_renewed_and_the_renewal_is_what_the_reader_gets():
"""The whole point: an agent calling after its user's id_token ran out keeps working."""
stale, fresh = _id_token(exp_offset=-1), _id_token()
rows = _FakeRows({"alice": _stored(stale, expires_in=-1)})
transport = _FakeTransport(_renewal(fresh))
served = await _store(rows, transport).fetch("alice")
assert served is not None
assert served.id_token.get_secret_value() == fresh
assert len(transport.calls) == 1
url, form = transport.calls[0]
assert url == TOKEN_ENDPOINT
assert form["grant_type"] == "refresh_token"
assert form["refresh_token"] == "rt_1"
@pytest.mark.asyncio
async def test_a_basic_auth_login_gets_a_basic_auth_refresh():
"""The non-PKCE login always sends HTTP Basic, so the renewal must too; credentials in the body
would 401 against an IdP application registered for Basic."""
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(_renewal(_id_token()))
await _store(rows, transport).fetch("alice")
expected = base64.b64encode(b"litellm:s3cret").decode()
assert transport.headers[0]["Authorization"] == f"Basic {expected}"
_url, form = transport.calls[0]
assert "client_secret" not in form
assert "client_id" not in form
@pytest.mark.asyncio
async def test_a_body_credential_login_gets_a_body_credential_refresh():
"""The mirror case. A PKCE deployment with GENERIC_INCLUDE_CLIENT_ID set signs in with the
credentials in the body, so Basic here would 401 against an application registered for post; the
renewal has to follow the login rather than a constant."""
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(_renewal(_id_token()))
await _store(rows, transport, client_config=lambda: _POST_CLIENT).fetch("alice")
assert "Authorization" not in transport.headers[0]
_url, form = transport.calls[0]
assert form["client_id"] == "litellm"
assert form["client_secret"] == "s3cret"
@pytest.mark.parametrize(
("include_client_id", "expected"),
[
(None, "client_secret_basic"),
("false", "client_secret_basic"),
("TRUE", "client_secret_post"),
("true", "client_secret_post"),
],
)
def test_the_auth_method_follows_the_flag_the_login_reads(include_client_id, expected):
"""``GENERIC_INCLUDE_CLIENT_ID`` is what the PKCE login branches on, parsed the same way it
parses it, so the renewal cannot pick a method the sign-in did not use."""
env = {
"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT,
"GENERIC_CLIENT_ID": "litellm",
"GENERIC_CLIENT_SECRET": "s3cret",
**({"GENERIC_INCLUDE_CLIENT_ID": include_client_id} if include_client_id is not None else {}),
}
config = sso_client_config(env)
assert config is not None
assert config.auth_method == expected
@pytest.mark.asyncio
async def test_an_assertion_well_inside_its_lifetime_never_reaches_the_idp():
"""The common path must cost exactly what it did before this store existed."""
current = _id_token()
rows = _FakeRows({"alice": _stored(current, expires_in=1800)})
transport = _FakeTransport(_renewal(_id_token()))
served = await _store(rows, transport).fetch("alice")
assert served is not None
assert served.id_token.get_secret_value() == current
assert transport.calls == []
assert rows.writes == []
@pytest.mark.asyncio
async def test_renewal_starts_inside_the_skew_rather_than_after_expiry():
"""A token that would die between resolution and the second exchange leg is replaced first."""
about_to_expire, fresh = _id_token(), _id_token()
assert about_to_expire != fresh
rows = _FakeRows({"alice": _stored(about_to_expire, expires_in=30)})
transport = _FakeTransport(_renewal(fresh))
served = await _store(rows, transport).fetch("alice")
assert served is not None
assert served.id_token.get_secret_value() == fresh
@pytest.mark.asyncio
async def test_a_user_with_no_stored_assertion_is_still_absent():
rows = _FakeRows()
transport = _FakeTransport(_renewal(_id_token()))
assert await _store(rows, transport).fetch("nobody") is None
assert transport.calls == []
@pytest.mark.asyncio
async def test_a_refused_refresh_leaves_the_expired_assertion_for_the_reader_to_reject():
"""A dead refresh token is the user's problem, and the reader's expiry guard is what tells them;
swapping in a renewed-looking value or hiding the row would break that challenge."""
stale = _id_token(exp_offset=-1)
rows = _FakeRows({"alice": _stored(stale, expires_in=-1)})
transport = _FakeTransport(Error(RefreshFailure.of_rejected("the IdP refused the refresh with status 400")))
served = await _store(rows, transport).fetch("alice")
assert served is not None
assert served.id_token.get_secret_value() == stale
assert rows.writes == []
@pytest.mark.asyncio
async def test_an_unreachable_idp_is_a_store_outage_not_a_sign_in_again_challenge():
"""503, not 412: the user has nothing to fix by signing in again while the IdP is down."""
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(Error(RefreshFailure.of_unavailable("the IdP token endpoint is unreachable")))
with pytest.raises(AssertionStoreUnavailable):
await _store(rows, transport).fetch("alice")
@pytest.mark.asyncio
async def test_a_missing_refresh_token_names_the_scope_the_operator_has_to_set(caplog):
"""Nothing to redeem is the default state of a deployment, so the log has to say what to change
or the feature stays silently inert."""
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1, refresh_token=None)})
transport = _FakeTransport(_renewal(_id_token()))
with caplog.at_level(logging.WARNING):
served = await _store(rows, transport).fetch("alice")
assert served is not None
assert transport.calls == []
assert "GENERIC_SCOPE" in caplog.text
assert "offline_access" in caplog.text
@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)})
transport = _FakeTransport(_renewal(_id_token()))
with caplog.at_level(logging.WARNING):
served = await _store(rows, transport, client_config=lambda: None).fetch("alice")
assert served is not None
assert transport.calls == []
assert "GENERIC_TOKEN_ENDPOINT" in caplog.text
@pytest.mark.asyncio
async def test_a_refresh_response_carrying_no_id_token_is_refused(caplog):
"""An access token is not an identity assertion, so there is nothing to assert upstream."""
stale = _id_token(exp_offset=-1)
rows = _FakeRows({"alice": _stored(stale, expires_in=-1)})
transport = _FakeTransport(Ok({"access_token": "at", "token_type": "Bearer"}))
with caplog.at_level(logging.WARNING):
served = await _store(rows, transport).fetch("alice")
assert served is not None
assert served.id_token.get_secret_value() == stale
assert rows.writes == []
assert "openid" in caplog.text
@pytest.mark.asyncio
async def test_a_rotated_refresh_token_replaces_the_stored_one():
"""An IdP that rotates invalidates the old token, so keeping it would cost a sign-in next time."""
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(_renewal(_id_token(), refresh_token="rt_2"))
await _store(rows, transport).fetch("alice")
stored = rows.rows["alice"]
assert stored.refresh_token is not None
assert stored.refresh_token.get_secret_value() == "rt_2"
@pytest.mark.asyncio
async def test_an_omitted_refresh_token_carries_the_previous_one_forward():
"""An IdP that does not rotate expects the original to keep working; dropping it would strand
the user after exactly one renewal."""
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(_renewal(_id_token()))
await _store(rows, transport).fetch("alice")
stored = rows.rows["alice"]
assert stored.refresh_token is not None
assert stored.refresh_token.get_secret_value() == "rt_1"
@pytest.mark.asyncio
async def test_the_renewed_expiry_moves_forward_so_the_next_read_does_not_refresh_again():
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(_renewal(_id_token(exp_offset=3600)))
store = _store(rows, transport)
await store.fetch("alice")
await store.fetch("alice")
assert len(transport.calls) == 1
async def _explode(user_id: str, assertion: SSOIdentityAssertion) -> None:
raise RuntimeError("write failed")
@pytest.mark.asyncio
async def test_a_renewal_that_cannot_be_recorded_is_reported_as_transient():
"""The store is what every caller reads, so a renewal nobody can see is not a success. Calling it
one would hand back a token the gateway failed to record."""
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
refresher = SSOAssertionRefresher(
_FakeTransport(_renewal(_id_token())), client_config=lambda: _CLIENT, read=rows.fetch, write=_explode
)
outcome = await refresher.refresh("alice", rows.rows["alice"])
assert isinstance(outcome, Error)
assert outcome.error.kind == "unavailable"
@pytest.mark.asyncio
async def test_a_failed_write_does_not_tell_the_user_to_sign_in_again():
"""A database that cannot take the write is not something signing in again fixes, so the reader
has to see an outage rather than the stale row's expiry."""
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(_renewal(_id_token()))
refresher = SSOAssertionRefresher(transport, client_config=lambda: _CLIENT, read=rows.fetch, write=_explode)
store = RefreshingSSOAssertionStore(
rows,
refresher,
fresh_read=rows.fetch_fresh,
coordinator_factory=lambda: None, # pyright: ignore[reportArgumentType] # test double stands in for the runtime factory
)
with pytest.raises(AssertionStoreUnavailable):
await store.fetch("alice")
assert len(transport.calls) == 1
@pytest.mark.asyncio
async def test_concurrent_reads_for_one_user_redeem_the_refresh_token_once():
"""A burst of tool calls must not replay one refresh token N times: an IdP that rotates reads
that as reuse and can revoke the whole grant chain."""
gate = asyncio.Event()
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
fresh = _id_token()
transport = _FakeTransport(_renewal(fresh), gate=gate)
store = _store(rows, transport)
callers = [asyncio.create_task(store.fetch("alice")) for _ in range(8)]
await _until(lambda: len(transport.calls) >= 1 and len(rows.reads) >= 8)
# Guards against a vacuous pass: every caller must have read the expired row and entered the
# renewal branch while the winner is still blocked, otherwise they never raced at all.
assert len(rows.reads) >= 8
assert not any(task.done() for task in callers)
gate.set()
served = await asyncio.gather(*callers)
assert len(transport.calls) == 1
assert {assertion.id_token.get_secret_value() for assertion in served if assertion is not None} == {fresh}
@pytest.mark.asyncio
async def test_concurrent_reads_for_different_users_each_get_their_own_refresh():
"""Single-flight is per user; collapsing across users would leave everyone but one stranded."""
gate = asyncio.Event()
rows = _FakeRows(
{
"alice": _stored(_id_token("alice", exp_offset=-1), expires_in=-1),
"bob": _stored(_id_token("bob", exp_offset=-1), expires_in=-1),
}
)
transport = _FakeTransport(_renewal(_id_token()), gate=gate)
store = _store(rows, transport)
callers = [asyncio.create_task(store.fetch(user)) for user in ("alice", "bob")]
await _until(lambda: len(transport.calls) >= 2)
gate.set()
await asyncio.gather(*callers)
assert len(transport.calls) == 2
assert {form["refresh_token"] for _url, form in transport.calls} == {"rt_1"}
@pytest.mark.asyncio
async def test_a_renewal_writes_back_when_the_row_did_not_move():
"""The refresh-then-sign-in ordering: nothing displaced the row, so the rotation must land."""
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
fresh = _id_token()
transport = _FakeTransport(_renewal(fresh, refresh_token="rt_2"))
served = await _store(rows, transport).fetch("alice")
assert [user_id for user_id, _assertion in rows.writes] == ["alice"]
assert rows.rows["alice"].id_token.get_secret_value() == fresh
assert served is not None
assert served.id_token.get_secret_value() == fresh
@pytest.mark.asyncio
async def test_a_sign_in_landing_mid_renewal_is_not_overwritten():
"""The sign-in-then-refresh ordering. The login wrote a newer assertion while the IdP call was in
flight; overwriting it would put back a refresh token the IdP has already rotated away, costing
that user a sign-in later."""
from_login = _stored(_id_token("alice"), expires_in=3600, refresh_token="rt_from_login")
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
def _login_lands() -> None:
rows.rows["alice"] = from_login
transport = _FakeTransport(_renewal(_id_token(), refresh_token="rt_2"), on_call=_login_lands)
served = await _store(rows, transport).fetch("alice")
assert rows.writes == []
stored = rows.rows["alice"]
assert stored.refresh_token is not None
assert stored.refresh_token.get_secret_value() == "rt_from_login"
assert served is not None
assert served.id_token.get_secret_value() == from_login.id_token.get_secret_value()
class _RecordingCoordinator:
"""Stands in for the cross-replica coordinator, running the winner's refresh inline."""
def __init__(self) -> None:
self.runs: list[tuple[str, str]] = []
async def run(
self,
user_id: str,
server_id: str,
refresh: Callable[[], Awaitable[None]],
reread: Callable[[], Awaitable[None]],
) -> None:
self.runs.append((user_id, server_id))
return await refresh()
class _ReplaceThenRefreshCoordinator:
"""Replaces the row before running the elected refresh."""
def __init__(self, replace: Callable[[], None]) -> None:
self._replace = replace
self.runs: list[tuple[str, str]] = []
async def run(
self,
user_id: str,
server_id: str,
refresh: Callable[[], Awaitable[None]],
reread: Callable[[], Awaitable[None]],
) -> None:
self.runs.append((user_id, server_id))
self._replace()
return await refresh()
class _HeldCoordinator:
"""Emulates a cross-replica holder finishing before the loser re-reads."""
def __init__(self, before_reread: Callable[[], None] | None = None) -> None:
self._before_reread = before_reread
self.runs: list[tuple[str, str]] = []
async def run(
self,
user_id: str,
server_id: str,
refresh: Callable[[], Awaitable[None]],
reread: Callable[[], Awaitable[None]],
) -> None:
self.runs.append((user_id, server_id))
if self._before_reread is not None:
self._before_reread()
return await reread()
@pytest.mark.asyncio
async def test_an_elected_renewal_redeems_the_row_it_re_reads_not_the_one_it_entered_with():
stale = _id_token(exp_offset=-1)
fresh = _stored(_id_token(), expires_in=3600, refresh_token="rt_2")
rows = _FakeRows({"alice": _stored(stale, expires_in=-1)})
transport = _FakeTransport(_renewal(_id_token()))
coordinator = _ReplaceThenRefreshCoordinator(lambda: rows.rows.__setitem__("alice", fresh))
served = await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice")
assert transport.calls == []
assert served is not None
assert served.id_token.get_secret_value() == fresh.id_token.get_secret_value()
@pytest.mark.asyncio
async def test_a_cross_replica_loser_whose_winner_renewed_reads_the_renewal_without_redeeming():
fresh = _stored(_id_token(), expires_in=3600, refresh_token="rt_2")
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(_renewal(_id_token()))
coordinator = _HeldCoordinator(before_reread=lambda: rows.rows.__setitem__("alice", fresh))
served = await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice")
assert served is fresh
assert transport.calls == []
assert len(coordinator.runs) == 1
@pytest.mark.asyncio
async def test_a_cross_replica_loser_rereads_past_a_stale_process_local_cache():
stale = _stored(_id_token(exp_offset=-1), expires_in=-1)
fresh = _stored(_id_token(), expires_in=3600, refresh_token="rt_2")
rows = _FakeRows({"alice": fresh})
rows.cached_rows["alice"] = stale
transport = _FakeTransport(_renewal(_id_token()))
coordinator = _HeldCoordinator()
served = await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice")
assert served is fresh
assert transport.calls == []
assert len(coordinator.runs) == 1
@pytest.mark.asyncio
async def test_a_cross_replica_loser_does_not_turn_a_write_failure_into_a_sign_in_challenge():
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(_renewal(_id_token()))
refresher = SSOAssertionRefresher(transport, client_config=lambda: _CLIENT, read=rows.fetch, write=_explode)
coordinator = _HeldCoordinator()
store = RefreshingSSOAssertionStore(
rows,
refresher,
fresh_read=rows.fetch_fresh,
coordinator_factory=lambda: coordinator, # pyright: ignore[reportArgumentType] # test double stands in for the runtime factory
)
with pytest.raises(AssertionStoreUnavailable):
await store.fetch("alice")
assert transport.calls == []
assert len(coordinator.runs) == 1
@pytest.mark.asyncio
async def test_a_cross_replica_loser_never_redeems_the_token_the_holder_may_have_rotated():
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(Error(RefreshFailure.of_rejected("dead")))
coordinator = _HeldCoordinator()
with pytest.raises(AssertionStoreUnavailable):
await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice")
assert transport.calls == []
assert len(coordinator.runs) == 1
@pytest.mark.asyncio
async def test_the_cross_replica_coordinator_is_used_and_built_once():
"""Redis elects one refresher across the fleet; rebuilding its client per renewal would open a
connection every time."""
coordinator = _RecordingCoordinator()
builds: list[int] = []
def _factory() -> object:
builds.append(1)
return coordinator
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(_renewal(_id_token(exp_offset=-1)))
store = _store(rows, transport, coordinator_factory=_factory)
await store.fetch("alice")
await store.fetch("alice")
assert len(builds) == 1
assert coordinator.runs == [("alice", "sso_identity_assertion"), ("alice", "sso_identity_assertion")]
@pytest.mark.asyncio
async def test_the_in_process_coordinator_is_retried_until_redis_appears():
"""A proxy that gains Redis after boot must stop electing a winner per worker."""
coordinator = _RecordingCoordinator()
available: list[bool] = [False]
def _factory() -> object | None:
return coordinator if available[0] else None
rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
transport = _FakeTransport(_renewal(_id_token(exp_offset=-1)))
store = _store(rows, transport, coordinator_factory=_factory)
await store.fetch("alice")
assert coordinator.runs == []
available[0] = True
await store.fetch("alice")
assert coordinator.runs == [("alice", "sso_identity_assertion")]
@pytest.mark.parametrize(
"env",
[
{},
{"GENERIC_CLIENT_ID": "litellm", "GENERIC_CLIENT_SECRET": "s"},
{"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, "GENERIC_CLIENT_SECRET": "s"},
{"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, "GENERIC_CLIENT_ID": "litellm"},
{"GENERIC_TOKEN_ENDPOINT": "", "GENERIC_CLIENT_ID": "litellm", "GENERIC_CLIENT_SECRET": "s"},
],
)
def test_a_partial_sso_client_is_no_client(env):
"""Redeeming against a half-configured client would post credentials nowhere useful; the arm
treats it as "cannot renew" and falls back to the sign-in challenge."""
assert sso_client_config(env) is None
def test_the_configured_sso_client_is_the_one_the_login_used():
config = sso_client_config(
{
"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT,
"GENERIC_CLIENT_ID": "litellm",
"GENERIC_CLIENT_SECRET": "s3cret",
}
)
assert config is not None
assert config.token_endpoint == TOKEN_ENDPOINT
assert config.client_id == "litellm"
assert config.client_secret.get_secret_value() == "s3cret"
def test_the_live_store_renews_over_the_database_reader():
"""The composition root has to produce a renewing store, or none of this runs in production."""
assert isinstance(default_sso_assertion_store(), RefreshingSSOAssertionStore)
def _responding(response: httpx.Response | None) -> HttpxTokenEndpointTransport:
async def _post(url: str, form: Mapping[str, str], headers: Mapping[str, str]) -> httpx.Response | None:
return response
return HttpxTokenEndpointTransport(_post)
def _json_response(status: int, payload: dict[str, object]) -> httpx.Response:
return httpx.Response(status, json=payload, request=httpx.Request("POST", TOKEN_ENDPOINT))
@pytest.mark.parametrize("status", [400, 401, 403])
@pytest.mark.asyncio
async def test_the_idp_declining_the_grant_is_a_refusal_the_user_must_act_on(status):
"""A 4xx means this refresh token is finished; calling that an outage would sit the user behind a
503 forever instead of telling them to sign in."""
outcome = await _responding(_json_response(status, {"error": "invalid_grant"})).post(TOKEN_ENDPOINT, {}, {})
assert isinstance(outcome, Error)
assert outcome.error.kind == "rejected"
@pytest.mark.parametrize("status", [500, 502, 503])
@pytest.mark.asyncio
async def test_a_failing_idp_is_an_outage_not_a_refusal(status):
"""The refresh token is probably fine; telling the user to sign in again would blame them for
someone else's outage, and would burn their session for nothing."""
outcome = await _responding(_json_response(status, {})).post(TOKEN_ENDPOINT, {}, {})
assert isinstance(outcome, Error)
assert outcome.error.kind == "unavailable"
@pytest.mark.asyncio
async def test_an_unreachable_endpoint_is_an_outage():
async def _post(url: str, form: Mapping[str, str], headers: Mapping[str, str]) -> httpx.Response | None:
raise httpx.ConnectError("connection refused")
outcome = await HttpxTokenEndpointTransport(_post).post(TOKEN_ENDPOINT, {}, {})
assert isinstance(outcome, Error)
assert outcome.error.kind == "unavailable"
@pytest.mark.asyncio
async def test_a_non_json_body_is_an_outage():
response = httpx.Response(200, text="<html>maintenance</html>", request=httpx.Request("POST", TOKEN_ENDPOINT))
outcome = await _responding(response).post(TOKEN_ENDPOINT, {}, {})
assert isinstance(outcome, Error)
assert outcome.error.kind == "unavailable"
@pytest.mark.asyncio
async def test_a_missing_response_is_an_outage():
outcome = await _responding(None).post(TOKEN_ENDPOINT, {}, {})
assert isinstance(outcome, Error)
assert outcome.error.kind == "unavailable"
@pytest.mark.asyncio
async def test_a_successful_grant_is_handed_back_as_the_parsed_body():
outcome = await _responding(_json_response(200, {"access_token": "at", "id_token": "idt"})).post(
TOKEN_ENDPOINT, {"grant_type": "refresh_token"}, {}
)
assert isinstance(outcome, Ok)
assert outcome.ok["id_token"] == "idt"

6
type-discipline-budget.json
View file

@ -1,6 +1,6 @@
{
"LIT001": {
"limit": 22181
"limit": 22180
},
"LIT002": {
"limit": 26745
@ -9,7 +9,7 @@
"limit": 261
},
"LIT004": {
"limit": 40
"limit": 38
},
"LIT005": {
"limit": 0
@ -27,7 +27,7 @@
"limit": 0
},
"LIT010": {
"limit": 16464
"limit": 16462
},
"LIT011": {
"limit": 5506