diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 82b820d8cd96..6cb38c53fe82 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/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. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 69984a56311a..b70db64ba943 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/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, 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 e0927cc4f648..d52c718c0b8a 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 @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py index 4bc5732ec0e4..3ed22732c905 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py +++ b/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) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 0f276cb8e5c6..f80954986e5e 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/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 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 a710da81962a..0d130767bd57 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 @@ -7,6 +7,10 @@ returning the stub. """ +import asyncio +import logging +from datetime import datetime, timedelta, timezone + import httpx import pytest from pydantic import SecretStr @@ -25,6 +29,7 @@ NoOpAuth, Ok, PassthroughConfig, + PrivateKeyJwtAuth, Result, ServerSpec, SharedKey, @@ -37,6 +42,10 @@ OAuthToken, TokenStoreUnavailable, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + AssertionStoreUnavailable, + SSOIdentityAssertion, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( ExchangedToken, ) @@ -71,6 +80,23 @@ def _with_inbound(token: str) -> Subject: return Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr(token)) +class _FakeAssertionStore: + """The SSO assertion read seam, canned per user_id and recording every lookup.""" + + def __init__(self, assertions: dict[str, SSOIdentityAssertion] | None = None) -> None: + self._assertions = dict(assertions or {}) + self.lookups: list[str] = [] + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + self.lookups.append(user_id) + return self._assertions.get(user_id) + + +def _assertion(id_token: str, expires_in: timedelta | None = timedelta(minutes=30)) -> SSOIdentityAssertion: + expires_at = datetime.now(timezone.utc) + expires_in if expires_in is not None else None + return SSOIdentityAssertion(id_token=SecretStr(id_token), expires_at=expires_at) + + def _spec(config): return ServerSpec(server_id="s", resource="https://upstream.example.com", config=config) @@ -471,9 +497,381 @@ async def test_id_jag_runs_both_legs_and_returns_the_leg2_bearer(): @pytest.mark.asyncio -async def test_id_jag_without_inbound_token_is_precondition_required_no_http(): +async def test_id_jag_without_inbound_token_or_stored_assertion_is_precondition_required_no_http(): + endpoint = _FakeTokenEndpoint([]) + store = _FakeAssertionStore() + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert endpoint.calls == [] + assert store.lookups == ["alice"] + + +@pytest.mark.asyncio +async def test_id_jag_exchanges_the_stored_sso_assertion_when_the_caller_presents_no_token(): + """The agent-triggered flow: a brokered LiteLLM credential carries no IdP token, so leg 1's + subject is the assertion captured for that user at SSO login.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + store = _FakeAssertionStore({"alice": _assertion("alice-id-token")}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Ok) + assert _emitted(result.ok)["Authorization"] == "Bearer final-access" + assert store.lookups == ["alice"] + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == "alice-id-token" + assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" + + +@pytest.mark.asyncio +async def test_id_jag_prefers_the_callers_own_token_over_the_stored_assertion(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + store = _FakeAssertionStore({"alice": _assertion("stored-id-token")}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials(_with_inbound("inbound-id-token"), _spec(_id_jag_config())) + + assert isinstance(result, Ok) + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == "inbound-id-token" + assert store.lookups == [] + + +@pytest.mark.asyncio +async def test_id_jag_refuses_an_expired_stored_assertion_without_calling_the_idp(): + endpoint = _FakeTokenEndpoint([]) + store = _FakeAssertionStore({"alice": _assertion("stale-id-token", expires_in=-timedelta(seconds=1))}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_id_jag_accepts_a_stored_assertion_that_declares_no_expiry(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + store = _FakeAssertionStore({"alice": _assertion("undated-id-token", expires_in=None)}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Ok) + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == "undated-id-token" + + +@pytest.mark.asyncio +async def test_id_jag_never_reads_the_store_for_an_unidentified_caller(): + """An empty subject_id must not select a credential; otherwise every anonymous caller would + share one store slot.""" + endpoint = _FakeTokenEndpoint([]) + store = _FakeAssertionStore({"": _assertion("anonymous-slot")}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id=""), _spec(_id_jag_config())) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert store.lookups == [] + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_id_jag_keeps_store_sourced_bearers_partitioned_per_user(): + endpoint = _FakeTokenEndpoint( + [ + Ok(ExchangedToken(access_token="alice-id-jag", expires_in=300)), + Ok(ExchangedToken(access_token="alice-bearer", expires_in=3600)), + Ok(ExchangedToken(access_token="bob-id-jag", expires_in=300)), + Ok(ExchangedToken(access_token="bob-bearer", expires_in=3600)), + ] + ) + store = _FakeAssertionStore( + {"alice": _assertion("alice-id-token"), "bob": _assertion("bob-id-token")} + ) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + alice = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + bob = await provider.resolve_credentials(Subject(tenant_id="", subject_id="bob"), _spec(_id_jag_config())) + + assert isinstance(alice, Ok) and isinstance(bob, Ok) + assert _emitted(alice.ok)["Authorization"] == "Bearer alice-bearer" + assert _emitted(bob.ok)["Authorization"] == "Bearer bob-bearer" + + +_DRIVER_DETAIL = "could not connect to host=pg-primary.internal port=5432 user=litellm" + + +class _OutageAssertionStore: + """A store whose backing DB is down, failing with a driver message full of internals.""" + + def __init__(self) -> None: + self.lookups: list[str] = [] + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + self.lookups.append(user_id) + raise AssertionStoreUnavailable(_DRIVER_DETAIL) + + +@pytest.mark.asyncio +async def test_id_jag_maps_an_assertion_store_outage_to_upstream_unavailable(): + """A store outage must not escape as an unhandled error, and must not be reported as a missing + assertion: telling the user to sign in again does not fix a database that is down.""" endpoint = _FakeTokenEndpoint([]) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=_OutageAssertionStore()) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_id_jag_store_outage_does_not_leak_driver_detail_to_the_caller(caplog): + """`upstream_unavailable` is rendered into the 503 body verbatim, so the driver's message, which + can name hosts, ports and users, must stay out of the summary and go to the log instead.""" + provider = UpstreamCredentialProvider( + token_endpoint=_FakeTokenEndpoint([]), sso_assertion_store=_OutageAssertionStore() + ) + + with caplog.at_level(logging.WARNING): + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert _DRIVER_DETAIL not in result.error.summary + assert "pg-primary.internal" not in result.error.summary + # The operator still needs it, so it must be in the log. + assert _DRIVER_DETAIL in caplog.text + + +@pytest.mark.asyncio +async def test_id_jag_invalidation_survives_an_assertion_store_outage(): + """invalidate_credentials runs on the upstream-401 retry path, so a store outage there must be + swallowed rather than turning a recoverable 401 into a 500.""" + provider = UpstreamCredentialProvider( + token_endpoint=_FakeTokenEndpoint([]), sso_assertion_store=_OutageAssertionStore() + ) + + await provider.invalidate_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + +class _FlakyAssertionStore: + """Serves an assertion, but fails while ``down`` is set.""" + + def __init__(self, assertion: SSOIdentityAssertion) -> None: + self._assertion = assertion + self.down = False + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + if self.down: + raise AssertionStoreUnavailable("connection refused") + return self._assertion + + +@pytest.mark.asyncio +async def test_id_jag_evicts_the_rejected_bearer_even_if_the_store_is_down_during_invalidation(): + """The upstream-401 recovery sequence with a transient store blip. + + Invalidation runs while the store is unreachable and the store recovers before the retry + resolves. Deriving the eviction key from a fresh lookup would evict nothing and then recompute + the identical key, handing the retry the very bearer the upstream just rejected. + """ + endpoint = _FakeTokenEndpoint(_two_leg_ok("rejected-bearer") + _two_leg_ok("reminted-bearer")) + store = _FlakyAssertionStore(_assertion("alice-id-token")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + subject = Subject(tenant_id="", subject_id="alice") + spec = _spec(_id_jag_config()) + + first = await provider.resolve_credentials(subject, spec) + assert isinstance(first, Ok) + assert _emitted(first.ok)["Authorization"] == "Bearer rejected-bearer" + + store.down = True + await provider.invalidate_credentials(subject, spec) + store.down = False + + second = await provider.resolve_credentials(subject, spec) + assert isinstance(second, Ok) + assert _emitted(second.ok)["Authorization"] == "Bearer reminted-bearer" + assert len(endpoint.calls) == 4 + + +class _SwitchableAssertionStore: + """Serves whichever assertion the test currently points it at, as a re-login would.""" + + def __init__(self, id_token: str) -> None: + self.id_token = id_token + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + return _assertion(self.id_token) + + +@pytest.mark.asyncio +async def test_id_jag_invalidation_clears_every_live_bearer_for_the_principal(): + """Overlapping store-sourced requests for one principal can hold different keys (a re-login + between them mints a different subject token). Invalidation must clear all of them: keeping + only the newest would let one request's 401 recovery evict the other's entry and leave its own + rejected bearer cached to be replayed on the retry.""" + endpoint = _FakeTokenEndpoint( + _two_leg_ok("bearer-from-first") + _two_leg_ok("bearer-from-second") + _two_leg_ok("reminted") + ) + store = _SwitchableAssertionStore("id-token-first") + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + subject = Subject(tenant_id="", subject_id="alice") + spec = _spec(_id_jag_config()) + + first = await provider.resolve_credentials(subject, spec) + store.id_token = "id-token-second" + second = await provider.resolve_credentials(subject, spec) + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(first.ok)["Authorization"] == "Bearer bearer-from-first" + assert _emitted(second.ok)["Authorization"] == "Bearer bearer-from-second" + + await provider.invalidate_credentials(subject, spec) + + # Point the store back at the first token. If that entry had survived the invalidation this + # would replay "bearer-from-first", which is the bearer an upstream may already have rejected. + store.id_token = "id-token-first" + third = await provider.resolve_credentials(subject, spec) + assert isinstance(third, Ok) + assert _emitted(third.ok)["Authorization"] == "Bearer reminted" + + +class _SequentialAssertionStore: + """Issues a distinct assertion per call unless pinned, so concurrent resolutions genuinely + mint distinct credentials rather than collapsing onto one through single-flight.""" + + def __init__(self) -> None: + self.pinned: str | None = None + self.issued: list[str] = [] + self._n = 0 + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + await asyncio.sleep(0) + if self.pinned is not None: + return _assertion(self.pinned) + self._n += 1 + token = f"id-token-{self._n}" + self.issued.append(token) + return _assertion(token) + + +class _CountingTokenEndpoint: + """Mints a unique bearer per exchange and yields, so exchanges interleave.""" + + def __init__(self) -> None: + self._n = 0 + + async def fetch(self, endpoint, client_id, grant_params, client_auth): + await asyncio.sleep(0) + self._n += 1 + return Ok(ExchangedToken(access_token=f"tok-{self._n}", expires_in=3600)) + + +@pytest.mark.asyncio +async def test_id_jag_invalidation_leaves_no_bearer_behind_under_concurrency(): + """After invalidation, no bearer minted before it may ever be served again. + + Drives many overlapping resolutions that each mint a distinct credential, invalidates once, + then replays every subject token that was issued. Any credential the eviction could not reach + would show up here as a replayed pre-invalidation bearer. + """ + concurrency = 20 + endpoint = _CountingTokenEndpoint() + store = _SequentialAssertionStore() + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + subject = Subject(tenant_id="t", subject_id="alice") + spec = _spec(_id_jag_config()) + + results = await asyncio.gather(*(provider.resolve_credentials(subject, spec) for _ in range(concurrency))) + before = {_emitted(r.ok)["Authorization"] for r in results if isinstance(r, Ok)} + issued = list(store.issued) + # Guard the guard: if these collapsed onto one credential the test would prove nothing. + assert len(before) > 1 + + await provider.invalidate_credentials(subject, spec) + + for token in issued: + store.pinned = token + replayed = await provider.resolve_credentials(subject, spec) + assert isinstance(replayed, Ok) + assert _emitted(replayed.ok)["Authorization"] not in before + + +@pytest.mark.asyncio +async def test_id_jag_never_serves_a_bearer_minted_for_a_different_caller(): + """Two unidentified-principal callers share a slot, so the fingerprint, not the key, is what + keeps them apart: a mismatch must read as a miss rather than hand over the other's bearer.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("first-callers-bearer") + _two_leg_ok("second-callers-bearer")) provider = UpstreamCredentialProvider(token_endpoint=endpoint) + spec = _spec(_id_jag_config()) + + first = await provider.resolve_credentials(_with_inbound("caller-one-token"), spec) + second = await provider.resolve_credentials(_with_inbound("caller-two-token"), spec) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(first.ok)["Authorization"] == "Bearer first-callers-bearer" + assert _emitted(second.ok)["Authorization"] == "Bearer second-callers-bearer" + + +@pytest.mark.asyncio +async def test_id_jag_rotating_the_signing_key_does_not_reuse_the_cached_bearer(): + """The cache key fingerprints the private-key-JWT client auth, so a rotated signing key + re-mints instead of serving a bearer authorized under the retired key.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("old-key-bearer") + _two_leg_ok("new-key-bearer")) + store = _FakeAssertionStore({"alice": _assertion("alice-id-token")}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + subject = Subject(tenant_id="", subject_id="alice") + + def _with_key(pem: str) -> IdJagConfig: + return _id_jag_config().model_copy( + update={"client_auth": PrivateKeyJwtAuth(private_key=SecretStr(pem), key_id="kid-1")} + ) + + first = await provider.resolve_credentials(subject, _spec(_with_key("-----OLD KEY-----"))) + second = await provider.resolve_credentials(subject, _spec(_with_key("-----NEW KEY-----"))) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(first.ok)["Authorization"] == "Bearer old-key-bearer" + assert _emitted(second.ok)["Authorization"] == "Bearer new-key-bearer" + assert len(endpoint.calls) == 4 + + +@pytest.mark.asyncio +async def test_id_jag_reads_a_naive_stored_expiry_as_utc(): + """A stored expires_at that lost its offset must still compare rather than raise: an aware/naive + comparison would be a TypeError on the egress path, turning a 412 into a 500.""" + endpoint = _FakeTokenEndpoint([]) + naive_past = datetime.now(timezone.utc).replace(tzinfo=None) - timedelta(hours=1) + store = _FakeAssertionStore( + {"alice": SSOIdentityAssertion(id_token=SecretStr("stale"), expires_at=naive_past)} + ) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + result = await provider.resolve_credentials( Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) ) @@ -483,6 +881,26 @@ async def test_id_jag_without_inbound_token_is_precondition_required_no_http(): assert endpoint.calls == [] +@pytest.mark.asyncio +async def test_invalidate_evicts_a_store_sourced_id_jag_bearer(): + """The upstream-401 recovery path. Keyed off the request alone the eviction would miss, and the + rejected bearer would be replayed until its TTL.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("first-bearer") + _two_leg_ok("second-bearer")) + store = _FakeAssertionStore({"alice": _assertion("alice-id-token")}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + subject = Subject(tenant_id="", subject_id="alice") + spec = _spec(_id_jag_config()) + + first = await provider.resolve_credentials(subject, spec) + await provider.invalidate_credentials(subject, spec) + second = await provider.resolve_credentials(subject, spec) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(first.ok)["Authorization"] == "Bearer first-bearer" + assert _emitted(second.ok)["Authorization"] == "Bearer second-bearer" + assert len(endpoint.calls) == 4 + + @pytest.mark.asyncio async def test_id_jag_propagates_a_leg1_error_without_calling_leg2(): endpoint = _FakeTokenEndpoint( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py index a3f46a49ba95..7b82e004f370 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py @@ -8,6 +8,7 @@ """ import json +import os import time from unittest.mock import AsyncMock, MagicMock, patch @@ -15,6 +16,8 @@ import pytest from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + AssertionStoreUnavailable, + DbSSOAssertionStore, assertion_from_sso_login, ema_assertion_retention_enabled, fetch_sso_identity_assertion, @@ -341,3 +344,23 @@ async def test_rotation_skips_unreadable_rows_but_rotates_readable_ones(): await rotate_sso_identity_assertions_master_key(prisma_client=prisma, new_master_key="another-new-salt-key-0000") assert stored["bad"] == "garbage-blob" assert stored["good"] != good_blob_before + + +@pytest.mark.asyncio +async def test_db_store_converts_a_driver_failure_into_assertion_store_unavailable(): + """The live store must not let a raw driver error escape: the resolver distinguishes an outage + from an absent assertion, and only a typed failure lets it do that.""" + prisma = MagicMock() + prisma.db.litellm_ssoidentityassertion.find_unique = AsyncMock(side_effect=RuntimeError("connection refused")) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with pytest.raises(AssertionStoreUnavailable): + await DbSSOAssertionStore().fetch("alice") + + +@pytest.mark.asyncio +async def test_db_store_returns_none_for_a_user_with_no_stored_assertion(): + """An absent row stays an absence, not an outage, so a user who never signed in still gets the + 412 that tells them to.""" + with patch.dict(os.environ, {"LITELLM_SALT_KEY": SALT_KEY}): + with patch("litellm.proxy.proxy_server.prisma_client", _make_prisma({})): + assert await DbSSOAssertionStore().fetch("nobody") is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 5f7f2267fc7a..bc285a6cd05c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -37,6 +37,7 @@ _deserialize_json_dict, _deserialize_json_list, _normalize_mcp_server_cost_info, + _obo_retry_applies, _should_strip_caller_authorization, _without_authorization, ) @@ -47,7 +48,7 @@ MCPEnvVarScope, MCPTransport, ) -from litellm.types.mcp import MCPAuth +from litellm.types.mcp import MCPAuth, MCPAuthType from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer @@ -8330,6 +8331,35 @@ def test_should_strip_caller_authorization_for_token_exchange(): assert _should_strip_caller_authorization(mcp_server=server, raw_headers=None, user_api_key_auth=None) is True +def _retry_gate_server(auth_type: MCPAuthType) -> MCPServer: + return MCPServer( + server_id="retry-gate", + name="retry-gate-server", + url="https://up.example.com", + transport=MCPTransport.http, + auth_type=auth_type, + ) + + +def test_obo_retry_applies_to_id_jag_without_an_inbound_subject_token(): + """ID-JAG can source its subject from the user's stored SSO assertion, so the upstream-401 + invalidate-and-retry path must engage even when the caller presented no token of its own; + otherwise a store-sourced bearer is replayed until its TTL after being rejected.""" + assert _obo_retry_applies(_retry_gate_server(MCPAuth.oauth2_id_jag), None) is True + assert _obo_retry_applies(_retry_gate_server(MCPAuth.oauth2_id_jag), "inbound-id-token") is True + + +def test_obo_retry_still_requires_a_subject_token_for_token_exchange(): + """token_exchange can only mint from an inbound token, so with none there is nothing to re-mint.""" + assert _obo_retry_applies(_retry_gate_server(MCPAuth.oauth2_token_exchange), None) is False + assert _obo_retry_applies(_retry_gate_server(MCPAuth.oauth2_token_exchange), "inbound-token") is True + + +def test_obo_retry_does_not_apply_to_other_auth_modes(): + for auth_type in (MCPAuth.none, MCPAuth.api_key, MCPAuth.oauth2, MCPAuth.true_passthrough): + assert _obo_retry_applies(_retry_gate_server(auth_type), "some-token") is False + + class _UpstreamAuthError(Exception): """Mimics a wrapped upstream 401 the way _extract_upstream_auth_failure detects it.""" diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 7686cc05fa68..f13a34bd9725 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -802,6 +802,11 @@ "count": 1 } }, + "src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/mcp-servers/_components/MCPLogoSelector.test.tsx": { "unused-imports/no-unused-imports": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx new file mode 100644 index 000000000000..e8730a5b9748 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx @@ -0,0 +1,158 @@ +import React from "react"; +import { Form, Input, Select, Tooltip } from "antd"; +import { InfoCircleOutlined } from "@ant-design/icons"; + +interface IdJagFormFieldsProps { + isEditing?: boolean; +} + +const fieldClassName = "rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"; + +const FieldLabel: React.FC<{ label: string; tooltip: string }> = ({ label, tooltip }) => ( + + {label} + + + + +); + +const IdJagFormFields: React.FC = ({ isEditing = false }) => { + const placeholderSuffix = isEditing ? " (leave blank to keep existing)" : ""; + + return ( + <> + + } + name="token_exchange_endpoint" + rules={[{ required: !isEditing, message: "The org token endpoint is required for ID-JAG" }]} + > + + + + } + name={["credentials", "id_jag_resource_token_endpoint"]} + rules={[{ required: !isEditing, message: "The resource token endpoint is required for ID-JAG" }]} + > + + + } + name={["credentials", "client_id"]} + rules={[{ required: !isEditing, message: "Client ID is required for ID-JAG" }]} + > + + + + } + name={["credentials", "client_secret"]} + dependencies={[["credentials", "client_private_key"]]} + rules={[ + ({ getFieldValue }) => ({ + validator: (_, value) => { + if (isEditing || value || getFieldValue(["credentials", "client_private_key"])) { + return Promise.resolve(); + } + return Promise.reject(new Error("Provide either a client secret or a client private key")); + }, + }), + ]} + > + + + + } + name={["credentials", "client_private_key"]} + > + + + + } + name={["credentials", "client_private_key_id"]} + > + + + + } + name={["credentials", "client_assertion_signing_alg"]} + > + + + + } + name="audience" + > + + + + } + name={["credentials", "id_jag_resource"]} + > + + + + } + name="subject_token_type" + > + + + } + name={["credentials", "scopes"]} + > + + +