From 110f654f342897ea438a6c71e91f0078bb4d76fa Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 5 Sep 2026 12:43:02 -0700 Subject: [PATCH] feat(mcp): renew the stored SSO identity assertion behind ID-JAG (#35401) * 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 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../per_user_oauth_store.py | 21 +- .../outbound_credentials/resolver.py | 23 +- .../runtime_refresh_coordinator.py | 41 + .../sso_assertion_refresher.py | 469 +++++++++++ .../sso_assertion_store.py | 39 +- .../outbound_credentials/test_resolver.py | 74 ++ .../test_sso_assertion_refresher.py | 794 ++++++++++++++++++ type-discipline-budget.json | 6 +- 8 files changed, 1427 insertions(+), 40 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py index 9273ddda9cf..5e28396dcfb 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 404baa14350..8328aae01ab 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -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. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py new file mode 100644 index 00000000000..e799838b5b7 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py @@ -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) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py new file mode 100644 index 00000000000..7600fd7ab8a --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py @@ -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, + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index 6552008ca54..f7b92df5ba3 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -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}, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 5e2f2cf97d7..6b3098d9e60 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -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")) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py new file mode 100644 index 00000000000..d80913f8d33 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py @@ -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="maintenance", 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" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 481e3591ce9..225134b4e2b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -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