feat(mcp): source the ID-JAG subject from the user's stored SSO assertion

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -814,6 +814,20 @@ def _extract_upstream_auth_failure(
     return upstream_auth_challenge(exc)
 
 
+def _obo_retry_applies(server: MCPServer, subject_token: str | None) -> bool:
+    """Whether an upstream 401/403 should invalidate the minted credential and retry once.
+
+    ``oauth2_token_exchange`` can only mint from an inbound subject token, so with no token there is
+    nothing to re-mint and the plain single call is correct. ``oauth2_id_jag`` also sources its
+    subject from the identity assertion stored for the user at SSO login, so it qualifies whether or
+    not the caller presented a token of its own; gating it on the inbound token would leave a
+    store-sourced bearer un-invalidated and replayed until its TTL.
+    """
+    if server.auth_type == MCPAuth.oauth2_id_jag:
+        return True
+    return server.auth_type == MCPAuth.oauth2_token_exchange and bool(subject_token)
+
+
 def _warn_on_server_name_fields(
     *,
     server_id: str,
@@ -4786,7 +4800,7 @@ async def _call_regular_mcp_tool(
             arguments=arguments,
         )
 
-        if mcp_server.auth_type in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag) and subject_token:
+        if _obo_retry_applies(mcp_server, subject_token):
             # OBO / ID-JAG: the exchanged token may have been revoked/rotated upstream since it was
             # cached, so an upstream 401 gets one invalidate + re-mint + retry. Gated to these modes;
             # all others keep the plain single call below.

litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py

@@ -10,19 +10,24 @@
 `none`, `api_key` (shared-key source), and `passthrough` (forwards the caller's own inbound token)
 are live, as is `authorization_code`, which reads the user's token from the injected
 `OAuthTokenStore`, `token_exchange`, which swaps the caller's inbound token through the injected
-`TokenExchanger`, and `client_credentials`, which mints and caches the gateway's M2M token through
-the injected `ClientCredentialsTokenSource`. The remaining arms are `not_implemented` stubs that
-each land in a follow-up PR with their seam. Pure v2: no imports from v1.
+`TokenExchanger`, `client_credentials`, which mints and caches the gateway's M2M token through the
+injected `ClientCredentialsTokenSource`, and `id_jag`, which runs the two-leg identity-assertion
+grant against a subject token taken from the request or from the injected `SSOAssertionStore`. The
+remaining arms are `not_implemented` stubs that each land in a follow-up PR with their seam. Pure
+v2: no imports from v1.
 """
 
 from __future__ import annotations
 
 import hashlib
+from datetime import datetime, timezone
 from functools import partial
 
 import httpx
 from typing_extensions import assert_never
 
+from litellm._logging import verbose_proxy_logger
+
 from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import (
     ClientCredentialsBearerAuth,
     ClientCredentialsTokenSource,
@@ -41,6 +46,12 @@
     Ok,
     Result,
 )
+from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
+    AssertionStoreUnavailable,
+    DbSSOAssertionStore,
+    SSOAssertionStore,
+    SSOIdentityAssertion,
+)
 from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import (
     ExchangedToken,
     ExchangedTokenCache,
@@ -111,12 +122,14 @@ def __init__(
         token_endpoint: TokenEndpointClient | None = None,
         exchanged_tokens: ExchangedTokenCache | None = None,
         client_credentials_source: ClientCredentialsTokenSource | None = None,
+        sso_assertion_store: SSOAssertionStore | None = None,
     ) -> None:
         self._oauth_token_store: OAuthTokenStore = oauth_token_store or _NullOAuthTokenStore()
         self._token_exchanger: TokenExchanger = token_exchanger or _NullTokenExchanger()
         self._token_endpoint: TokenEndpointClient = token_endpoint or TokenEndpointClient()
         self._exchanged_tokens: ExchangedTokenCache = exchanged_tokens or ExchangedTokenCache()
         self._client_credentials_source = client_credentials_source or ClientCredentialsTokenSource()
+        self._sso_assertion_store: SSOAssertionStore = sso_assertion_store or DbSSOAssertionStore()
 
     async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]:
         match server.config:
@@ -171,15 +184,73 @@ def _api_key(self, config: ApiKeyConfig) -> Result[httpx.Auth, CredError]:
         assert_never(config.key_source)
 
     async def _id_jag(self, subject: Subject, server: ServerSpec, config: IdJagConfig) -> Result[httpx.Auth, CredError]:
-        if subject.inbound_token is None:
+        match await self._id_jag_subject_token(subject):
+            case Error(err):
+                return Error(err)
+            case Ok(subject_token):
+                return await self._id_jag_exchange(subject, subject_token, server, config)
+
+    async def _id_jag_subject_token(self, subject: Subject) -> Result[str, CredError]:
+        """The identity token ID-JAG leg 1 asserts, from the request or from the SSO login it was captured at.
+
+        A caller that presents its own IdP identity token wins: that is the strongest available
+        assertion of who is calling. Otherwise the subject is the assertion captured for this user
+        at LiteLLM SSO login, which is what lets an agent holding a brokered LiteLLM credential
+        reach an upstream as the user it was issued for. The user is always taken from the
+        authenticated principal, never from a caller-supplied field, so no caller can select whose
+        identity is asserted upstream.
+
+        Every miss is ``precondition_required`` (412) rather than a fall-through to a weaker
+        credential: ID-JAG exists to assert a specific user, so a missing subject has no safe
+        substitute. A store outage is the one exception: it is ``upstream_unavailable`` (503), not
+        412, because the user has nothing to fix by signing in again, and it is a value rather than
+        a raised error so a DB blip cannot 500 the egress or the upstream-401 retry.
+        """
+        if subject.inbound_token is not None:
+            return Ok(subject.inbound_token.get_secret_value())
+        if not subject.subject_id:
+            return Error(
+                CredError.of_precondition_required(
+                    "ID-JAG requires an identified caller; this request carries neither an "
+                    "identity token nor a resolved LiteLLM user."
+                )
+            )
+        try:
+            assertion = await self._sso_assertion_store.fetch(subject.subject_id)
+        except AssertionStoreUnavailable as exc:
+            # The driver's message can name hosts, schemas or connection details, and this summary
+            # is returned to the caller verbatim as a 503 body. Operators get it from the log.
+            verbose_proxy_logger.warning(
+                "ID-JAG: the IdP identity assertion store is unreachable for user_id=%s: %s",
+                subject.subject_id,
+                exc,
+            )
+            return Error(
+                CredError.of_upstream_unavailable(
+                    "The IdP identity assertion store is unreachable, so ID-JAG cannot resolve a subject."
+                )
+            )
+        if assertion is None:
+            return Error(
+                CredError.of_precondition_required(
+                    "ID-JAG requires an IdP identity assertion for this user and none is stored. "
+                    "Sign in through LiteLLM SSO so the gateway captures one."
+                )
+            )
+        if _assertion_expired(assertion, datetime.now(timezone.utc)):
             return Error(
                 CredError.of_precondition_required(
-                    "ID-JAG requires a caller identity token; it asserts the calling "
-                    "user's identity upstream and cannot use a static credential."
+                    "The stored IdP identity assertion for this user has expired. Sign in through "
+                    "LiteLLM SSO again to capture a current one."
                 )
             )
-        token = subject.inbound_token.get_secret_value()
-        cache_key = _id_jag_cache_key(token, server.server_id, config)
+        return Ok(assertion.id_token.get_secret_value())
+
+    async def _id_jag_exchange(
+        self, subject: Subject, token: str, server: ServerSpec, config: IdJagConfig
+    ) -> Result[httpx.Auth, CredError]:
+        slot = _id_jag_slot_key(subject, server)
+        fingerprint = _id_jag_fingerprint(token, server.server_id, config)
 
         async def _exchange() -> Result[ExchangedToken, CredError]:
             leg1_params = {
@@ -211,7 +282,7 @@ async def _exchange() -> Result[ExchangedToken, CredError]:
                         config.client_auth,
                     )
 
-        match await self._exchanged_tokens.get_or_compute(cache_key, _exchange):
+        match await self._exchanged_tokens.get_or_compute(slot, _exchange, fingerprint=fingerprint):
             case Ok(access_token):
                 return Ok(StaticHeaderAuth(f"Bearer {access_token}"))
             case Error(err):
@@ -273,17 +344,27 @@ async def invalidate_credentials(self, subject: Subject, server: ServerSpec) ->
         re-mintable cached credential here; `client_credentials` recovers inside its own auth flow
         (`ClientCredentialsBearerAuth` retries the 401'd request once with a fresh token), and
         other modes are a no-op.
+
+        `id_jag` evicts by a slot key derived from the principal, so it needs no lookup against the
+        assertion store on this path; the fingerprint stored beside the entry is what keeps a slot
+        shared between callers safe.
         """
-        if subject.inbound_token is None:
-            return
-        if isinstance(server.config, TokenExchangeConfig):
+        if isinstance(server.config, IdJagConfig):
+            self._invalidate_id_jag(subject, server)
+        elif isinstance(server.config, TokenExchangeConfig) and subject.inbound_token is not None:
             await self._token_exchanger.invalidate(
                 subject.inbound_token.get_secret_value(), server, server.config, tenant_id=subject.tenant_id
             )
-        if isinstance(server.config, IdJagConfig):
-            self._exchanged_tokens.invalidate(
-                _id_jag_cache_key(subject.inbound_token.get_secret_value(), server.server_id, server.config)
-            )
+
+    def _invalidate_id_jag(self, subject: Subject, server: ServerSpec) -> None:
+        """Evict the bearer this `(subject, server)` last resolved, without depending on the store.
+
+        The slot is addressed by the principal (plus the caller's own token when it presented one),
+        never by the credential material, so it stays computable when the assertion store is down.
+        The fingerprint stored with the entry is what keeps that safe: an entry minted for different
+        inputs reads as a miss rather than being served.
+        """
+        self._exchanged_tokens.invalidate(_id_jag_slot_key(subject, server))
 
     async def _authz_token(self, subject: Subject, server: ServerSpec) -> OAuthToken | None:
         """The user's authorization_code token, or None when absent or the store is unreachable.
@@ -297,8 +378,37 @@ async def _authz_token(self, subject: Subject, server: ServerSpec) -> OAuthToken
             return None
 
 
-def _id_jag_cache_key(subject_token: str, server_id: str, config: IdJagConfig) -> str:
-    """Bind the cached leg-2 bearer to the caller token, the server, AND the config that minted it.
+def _id_jag_slot_key(subject: Subject, server: ServerSpec) -> str:
+    """Which cache slot this caller's bearer for this upstream lives in.
+
+    Addressed by the principal, plus the caller's own token when it presented one so two callers
+    sharing an empty principal do not contend for one slot. Deliberately free of the stored
+    assertion, which is what lets invalidation compute this while the assertion store is down. The
+    entry's fingerprint, not this key, is what guarantees a cached bearer matches current inputs.
+    """
+    inbound = subject.inbound_token.get_secret_value() if subject.inbound_token is not None else ""
+    material = "\x00".join((subject.tenant_id, subject.subject_id, server.server_id, inbound))
+    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 = assertion.expires_at
+    if expires_at is None:
+        return False
+    normalized = 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.
+
+    Stored beside the bearer and compared on every read, so a rotated assertion or an edited server
+    config reads as a miss and re-mints instead of serving a bearer authorized under the old policy.
 
     Every exchange parameter derives from the config (endpoints, audience, resource, scopes, client
     auth), so a server update that changes any of them must change the key; otherwise the old bearer,

litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py

@@ -18,7 +18,7 @@
 
 import json
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol
 
 import jwt
 from pydantic import BaseModel, ConfigDict, SecretStr, TypeAdapter, ValidationError
@@ -160,6 +160,42 @@ async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | N
     )
 
 
+class AssertionStoreUnavailable(Exception):
+    """Raised by ``fetch`` when the backing store is unreachable (e.g. the DB is down).
+
+    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
+    ``TokenStoreUnavailable`` on the sibling per-user OAuth store.
+    """
+
+
+class SSOAssertionStore(Protocol):
+    """The read seam the ``id_jag`` egress arm depends on, so the arm takes a collaborator
+    rather than reaching for a module-level function and a proxy global at call time.
+
+    Returns the user's captured assertion, or ``None`` when they have never signed in. Raises
+    ``AssertionStoreUnavailable`` when the backing store is unreachable.
+    """
+
+    async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: ...
+
+
+class DbSSOAssertionStore:
+    """The live store: the row the SSO callback wrote, read back by ``user_id``.
+
+    A storage failure is re-raised as ``AssertionStoreUnavailable`` so the resolver can map it to a
+    typed fail-closed result; letting the raw driver error escape would surface a DB blip as a 500
+    from credential resolution and from the upstream-401 retry.
+    """
+
+    async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
+        try:
+            return await fetch_sso_identity_assertion(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,
     mirroring the sibling per-user credential tables; an unreadable row is skipped so one

litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py

@@ -23,7 +23,7 @@
 
 import httpx
 import jwt
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel, TypeAdapter, ValidationError
 from typing_extensions import assert_never
 
 from litellm._logging import verbose_proxy_logger
@@ -51,6 +51,9 @@
 )
 from litellm.types.llms.custom_http import httpxSpecialProvider
 
+# The cache stores (fingerprint, token); anything else in the slot is treated as absent.
+_CACHED_ENTRY_ADAPTER: TypeAdapter[tuple[str, str]] = TypeAdapter(tuple[str, str])
+
 CLIENT_ASSERTION_TYPE = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"
 CLIENT_ASSERTION_LIFETIME_SECONDS = 60
 
@@ -134,19 +137,28 @@ async def get_or_compute(
         self,
         cache_key: str,
         compute: Callable[[], Awaitable[Result[ExchangedToken, CredError]]],
+        *,
+        fingerprint: str = "",
     ) -> Result[str, CredError]:
-        cached = self._get(cache_key)
+        """The cached token for `cache_key`, minting one when absent.
+
+        `fingerprint` lets a caller address a slot by something stable (a principal) while still
+        guaranteeing the token it gets back was minted for the *current* inputs: a stored entry
+        whose fingerprint differs reads as a miss and is re-minted over. That keeps eviction
+        addressable without the key having to encode the credential material it protects.
+        """
+        cached = self._get(cache_key, fingerprint)
         if cached is not None:
             return Ok(cached)
         async with self._lock(cache_key):
-            cached = self._get(cache_key)
+            cached = self._get(cache_key, fingerprint)
             if cached is not None:
                 return Ok(cached)
             match await compute():
                 case Ok(token):
                     self._cache.set_cache(  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
                         cache_key,
-                        token.access_token,
+                        (fingerprint, token.access_token),
                         ttl=_cache_ttl_seconds(token.expires_in),
                     )
                     return Ok(token.access_token)
@@ -157,9 +169,18 @@ def invalidate(self, cache_key: str) -> None:
         """Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401)."""
         self._cache.delete_cache(cache_key)  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
 
-    def _get(self, cache_key: str) -> str | None:
-        value = self._cache.get_cache(cache_key)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]  # InMemoryCache is untyped; narrowed by isinstance below
-        return value if isinstance(value, str) else None
+    def _get(self, cache_key: str, fingerprint: str) -> str | None:
+        """The stored token, or None when absent or minted for different inputs.
+
+        The fingerprint comparison is what makes a shared slot safe: a mismatch never returns the
+        other party's token, it just reads as a miss.
+        """
+        value = self._cache.get_cache(cache_key)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]  # InMemoryCache is untyped; the adapter below is the type gate
+        try:
+            stored_fingerprint, token = _CACHED_ENTRY_ADAPTER.validate_python(value)
+        except ValidationError:
+            return None
+        return token if stored_fingerprint == fingerprint else None
 
     def _lock(self, cache_key: str) -> asyncio.Lock:
         lock = self._locks.get(cache_key)

litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py

@@ -391,7 +391,8 @@ class Subject(BaseModel):
 
     tenant_id: str
     subject_id: str
-    # Opaque, already-validated inbound identity. Only `token_exchange` / `passthrough` read it.
+    # Opaque, already-validated inbound identity. Read by `token_exchange`, `passthrough`, and
+    # `id_jag` (which falls back to the user's stored SSO assertion when it is absent).
     inbound_token: SecretStr | None = None