From f04398e3fd502ea53b30b63d21b2571fd937822c Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 24 Jul 2026 18:27:55 -0700 Subject: [PATCH] feat(mcp): source the ID-JAG subject token from the stored SSO assertion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The oauth2_id_jag egress arm spends the user's IdP id_token as its RFC 8693 subject_token, but a front-door client presents only a gateway-minted llm_session_ bearer, so hop A had nothing usable and EMA was unreachable end to end for exactly the clients it exists for — while a gateway-issued credential presented via Authorization was forwarded verbatim to the IdP. One rule per concern, each with a single owner: - Provenance (credential_provenance.py, new): the inbound bearer is classified once at the edge (absent / gateway_credential / external_jwt / external_opaque) by the one property every gateway mint shares and no external token has: it verifies under this gateway's master key. A new mint format is refused as exchange material by construction, with no denylist to grow stale. id_jag exchanges the inbound bearer only when it is an external JWT; token_exchange (OBO) keeps exchange-what-was-presented but refuses gateway credentials. - Header selection (_subject_bearer_token): every egress surface (tools call/list, prompts, resources, templates, OpenAPI) selects the caller's bearer through one manager rule; for id_jag the bearer must also be the credential that admitted the caller, so a free-riding Authorization value bound to nobody is never exchanged. - Sourcing (_id_jag_subject_token): the caller's presented id_token when the inbound bearer is one, else the stored SSO assertion, else fail closed — shared verbatim by resolve and invalidate. - Storage (sso_assertion_store.py): the generic-SSO callback persists the IdP id_token + refresh token as one encrypted row per user (latest login wins), write-through to the DB so every pod sees it, gated on an id_jag server existing per the config and DB authorities (never the registry snapshot). Reads compose the shared per-user credential stack — Cached(Refreshing(db)) — for positive-only caching, expiry-skewed renewal and per-user single-flight; the TokenRefresher protocol now documents the raise-through contract the EMA refresher relies on. - Renewal lifecycle: expiry is judged by readers (30s buffer). A dead grant requires proof about the GRANT, and RFC 6749 has exactly one code that states it: invalid_grant. Every other outcome (invalid_client and friends are the gateway's own client config failing — ops-fixable), a 429, a 4xx without an error object, a 5xx, or a transport failure reads retryable and never touches the row. A proven verdict is recorded ON the row by stripping its refresh token — never in a cache beside it — so it holds on every pod, costs one POST total, and a re-login lifts it by construction. The verdict write re-reads the row and strips only a row still carrying the exact rejected token, preserving the row's other material, so a losing pod in a cross-pod rotation race cannot clobber the winner's renewed assertion. - Failure semantics, same story on every surface: absence -> 412 naming the gateway-SSO remedy; dead grant -> 401 + re-authenticate challenge; outage -> retryable 503, never a false re-login and never absence. Single-server routes surface these at the transport edge via preflight_id_jag (sibling of the OBO preflight), where an HTTP status and WWW-Authenticate still reach the client — in-session list handlers can only serialize a raise into a JSON-RPC error. The aggregate absorbs per-server faults, and tools/list outcomes now preserve the retryable-vs-terminal class (unavailable -> 503, precondition -> 412) instead of reporting a dependency outage as the gateway's own 500. Tests pin the rules, not the mechanisms: provenance totality over every current mint plus adversarial input, the one-POST-per-dead-grant bound, the re-signed-in caller never refused, absence-vs-outage by remedy, the cross-pod clobber no-op, the invalid_grant-only proof predicate, scope identity between login and refresh, winner-and-waiter verdict sharing in the shared stack, and the preflight wiring gated to single-server routes. Part of LIT-3856 (PR 2 of the EMA stack, on the #34072 assertion store). Co-Authored-By: Claude Fable 5 --- .../mcp_server/auth/user_api_key_auth_mcp.py | 23 + .../mcp_server/faults/list_outcomes.py | 24 +- .../mcp_server/mcp_server_manager.py | 155 +++- .../outbound_credentials/adapter.py | 11 +- .../credential_provenance.py | 180 ++++ .../outbound_credentials/oauth_token_store.py | 26 +- .../outbound_credentials/resolver.py | 111 ++- .../sso_assertion_store.py | 385 ++++++++- .../mcp_server/outbound_credentials/types.py | 11 + .../proxy/_experimental/mcp_server/server.py | 23 + litellm/proxy/management_endpoints/ui_sso.py | 5 +- .../mcp_server/faults/test_list_outcomes.py | 2 + .../outbound_credentials/test_adapter.py | 36 +- .../test_authz_code_refresher.py | 66 +- .../test_credential_provenance.py | 183 ++++ .../test_oauth_token_store.py | 6 + .../outbound_credentials/test_resolver.py | 302 ++++++- .../test_sso_assertion_store.py | 812 +++++++++++++++++- .../test_token_cache_codec.py | 8 +- .../test_token_endpoint.py | 5 +- .../outbound_credentials/test_types.py | 8 +- .../test_v2_token_store.py | 4 +- .../mcp_server/test_mcp_server.py | 61 ++ .../mcp_server/test_mcp_server_manager.py | 456 +++++++++- 24 files changed, 2695 insertions(+), 208 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/credential_provenance.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_credential_provenance.py diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index a27d6b92843..e8d31bf3431 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1208,6 +1208,29 @@ class MCPRequestHandler: ) return MCP_CLIENT_SIDE_AUTH_HEADER_NAME + @staticmethod + def authorization_is_free_rider(raw_headers: dict[str, str] | None) -> bool: + """Whether the request's Authorization bearer was NOT the admission credential. + + This is the inverse view of ``get_litellm_api_key_from_headers`` and must encode the + SAME header-preference rule: admission authenticates with the explicit primary header + only when its value is truthy (a missing or empty primary falls back to Authorization), + so the Authorization bearer was the admission credential exactly when the primary is + absent or empty. It lives in this module beside that accessor so one owner holds the + rule, and the agreement test pins the two against a shared header matrix. + """ + if not raw_headers: + return False + primary = next( + ( + value + for key, value in raw_headers.items() + if str(key).lower() == MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY.lower() + ), + None, + ) + return bool(primary) + @staticmethod def get_litellm_api_key_from_headers(headers: Headers) -> Optional[str]: """ diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py index 10463496409..4551617ceba 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -30,6 +30,8 @@ ListFaultCategory: TypeAlias = Literal[ "timeout", "unreachable", "upstream_error", + "unavailable", + "precondition", "internal", ] @@ -43,8 +45,11 @@ class ServerListOk(BaseModel): class ServerListFault(BaseModel): """Why a server contributed nothing to a listing: the caller must authenticate upstream (``auth_required``/``forbidden``), the upstream did not answer (``timeout``/``unreachable``), - the upstream answered outside its contract (``upstream_error``), or the gateway itself failed - (``internal``). ``status_code`` is the upstream HTTP status when one exists.""" + the upstream answered outside its contract (``upstream_error``), a dependency the gateway + needs to authorize egress did not answer so the caller should retry (``unavailable``), a + per-caller requirement is unmet and only the caller can meet it, e.g. a gateway SSO sign-in + (``precondition``), or the gateway itself failed (``internal``). ``status_code`` is the + upstream HTTP status when one exists.""" model_config = ConfigDict(frozen=True) tag: ListFaultCategory @@ -149,7 +154,16 @@ def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]: match outcome.tag: case "ok": return {"status": "ok", "tool_count": outcome.tool_count} - case "auth_required" | "forbidden" | "timeout" | "unreachable" | "upstream_error" | "internal": + case ( + "auth_required" + | "forbidden" + | "timeout" + | "unreachable" + | "upstream_error" + | "unavailable" + | "precondition" + | "internal" + ): return { "status": outcome.tag, **({"http_status": outcome.status_code} if outcome.status_code is not None else {}), @@ -171,6 +185,10 @@ def list_fault_http_status(fault: ServerListFault) -> int: return 504 case "unreachable" | "upstream_error": return 502 + case "unavailable": + return 503 + case "precondition": + return 412 case "internal": return 500 case _: diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0ee74960293..becdd5491b2 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -90,6 +90,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import ( LazyPerUserOAuthTokenStore, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + LiveSsoAssertionSource, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_provider import ( build_token_exchanger, ) @@ -809,6 +812,25 @@ def _caller_authorization_fans_out( ) +def _obo_retry_covers_caller( + auth_type: MCPAuthType | None, + subject_token: str | None, + user_api_key_auth: UserAPIKeyAuth | None, +) -> bool: + """Whether the invalidate-and-retry path has a subject to re-mint with after an upstream 401. + + ``token_exchange`` re-mints only from the caller's inbound token. ``id_jag`` also re-mints + from the stored SSO assertion, which the resolver looks up by the caller's user id, so a + caller with a user identity but no inbound assertion (an admission-key or session caller) + must route through the retry path too or its rejected cached bearer is never evicted. + """ + if auth_type not in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag): + return False + if subject_token: + return True + return auth_type == MCPAuth.oauth2_id_jag and user_api_key_auth is not None and bool(user_api_key_auth.user_id) + + def _extract_upstream_auth_failure( exc: BaseException, ) -> Optional[tuple[int, Optional[str]]]: @@ -1153,6 +1175,7 @@ class MCPServerManager: self._cred_provider = cred_provider or UpstreamCredentialProvider( oauth_token_store=self._per_user_oauth_token_store, token_exchanger=build_token_exchanger(), + sso_assertions=LiveSsoAssertionSource(), ) self.registry: dict[str, MCPServer] = {} self.config_mcp_servers: dict[str, MCPServer] = {} @@ -1251,6 +1274,7 @@ class MCPServerManager: mcp_auth_header=None, extra_headers=extra_headers, stdio_env=None, + user_api_key_auth=None, # internal path (cache warm-up / health check): no caller identity ) async def _noop(_session): @@ -2739,20 +2763,35 @@ class MCPServerManager: return auth_value return None - def _obo_subject_token( + def _subject_bearer_token( self, server: MCPServer, raw_headers: Optional[dict[str, str]], + oauth2_headers: dict[str, str] | None = None, ) -> Optional[str]: - """The caller's bearer as the token_exchange (OBO) subject token, for that mode only. + """The caller's bearer as the exchange subject material, resolved by the ONE rule every + egress surface shares: only the two exchange modes (token_exchange OBO and id_jag) read + the inbound bearer, and every surface (tools call and list, prompts, resources) resolves + it here so the RULE cannot diverge per surface. The INPUTS still can: the prompt/resource + helpers receive no ``oauth2_headers``, so a client-side oauth2 header participates only on + the tool surfaces until someone threads it through those signatures. Other modes return + None to avoid forwarding it. - Prompts/resources discovery and reads on a token_exchange server must exchange the caller's - token like the tools paths do, not connect with no credential. Other modes never read the - inbound bearer, so return None to avoid forwarding it. + For id_jag the bearer must additionally BE the credential that authenticated the caller: + admission prefers the explicit litellm key header, so when that header is present the + Authorization bearer is an unvalidated free rider bound to nobody, and exchanging it + would let a caller act upstream under any identity whose token they hold. Binding is + admission's job, so the rule is structural (was this the admission credential), never a + claims comparison re-deriving what admission already decided. Such callers resolve + through the stored assertion, which is bound to the authenticated user by construction. + The OBO mode keeps its documented exchange-what-was-presented semantics; its identical + free-rider shape predates this seam and is tracked as a follow-up. """ - if server.auth_type != MCPAuth.oauth2_token_exchange: + if server.auth_type not in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag): return None - return self._extract_bearer_token(None, raw_headers) + if server.auth_type == MCPAuth.oauth2_id_jag and MCPRequestHandler.authorization_is_free_rider(raw_headers): + return None + return self._extract_bearer_token(oauth2_headers, raw_headers) def _build_stdio_env( self, @@ -3036,6 +3075,39 @@ class MCPServerManager: ) raise_public(err) + async def preflight_id_jag( + self, + server: MCPServer, + oauth2_headers: dict[str, str] | None, + raw_headers: dict[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None, + ) -> None: + """Run the ID-JAG exchange at the transport edge, where an HTTP status and + ``WWW-Authenticate`` still reach the client. + + The list handlers cannot surface auth failures (the MCP session manager serializes a + raise into a JSON-RPC error, losing the header), so without this preflight a dead stored + assertion reads as an empty tool/prompt/resource catalog on single-server routes. Unlike + the OBO preflight there is no subject-not-yet-supplied case to wave through: id_jag + sources its subject from the caller's presented id_token or the stored SSO assertion, so + every resolution outcome is already the caller's answer — the re-login 401 (with the + resolver's challenge header), the retryable 503 of a store or IdP outage, or the 412 of a + user who never signed in. The subject bearer is selected by the same one rule egress uses, + so the preflight and the session it admits cannot disagree. A successful exchange is + cached by the resolver, so the session's list/call reuses it. + """ + if server.auth_type != MCPAuth.oauth2_id_jag: + return + spec = to_server_spec(server) + if spec is None or not isinstance(spec.config, IdJagConfig): + return + subject_token = self._subject_bearer_token(server, raw_headers, oauth2_headers) + match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): + case Ok(_): + return + case Error(err): + raise_public(err) + async def _create_mcp_client( self, server: MCPServer, @@ -3280,14 +3352,7 @@ class MCPServerManager: stdio_env = self._build_stdio_env(server, raw_headers) - # token_exchange (OBO) discovery needs the caller's token too: list it with the user's own - # token (mirrors the call path), not v1's deleted client_credentials fallback. Other modes - # never read the inbound bearer, so leave subject_token None to avoid forwarding it. - subject_token = ( - self._extract_bearer_token(oauth2_headers, raw_headers) - if server.auth_type == MCPAuth.oauth2_token_exchange - else None - ) + subject_token = self._subject_bearer_token(server, raw_headers, oauth2_headers=oauth2_headers) client = await self._create_mcp_client( server=server, @@ -3354,6 +3419,16 @@ class MCPServerManager: server_name=server.name, ) from e verbose_logger.warning(f"Failed to get tools from server {server.name}: {str(e)}") + if e.status_code == 503: + # The resolver's retryable verdict (credential store or IdP did not answer). Keep + # it retryable on the listing edge: "internal" would report the gateway itself + # failed and turn a transient outage into a 500. + raise MCPServerListError(ServerListFault(tag="unavailable", status_code=503), server.name) from e + if e.status_code == 412: + # The resolver's per-caller precondition (e.g. no stored SSO identity assertion, + # remedied only by the caller signing in to the gateway). Neither an upstream auth + # challenge nor a gateway fault. + raise MCPServerListError(ServerListFault(tag="precondition", status_code=412), server.name) from e raise MCPServerListError(ServerListFault(tag="internal", status_code=e.status_code), server.name) from e except MCPServerListError: raise @@ -3368,6 +3443,7 @@ class MCPServerManager: extra_headers: Optional[dict[str, str]] = None, add_prefix: bool = True, raw_headers: Optional[dict[str, str]] = None, + user_api_key_auth: UserAPIKeyAuth | None = None, ) -> list[Prompt]: """ Helper method to get prompts from a single MCP server with prefixed names. @@ -3392,7 +3468,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env = self._build_stdio_env(server, raw_headers) - subject_token = self._obo_subject_token(server, raw_headers) + subject_token = self._subject_bearer_token(server, raw_headers) client = await self._create_mcp_client( server=server, @@ -3400,6 +3476,7 @@ class MCPServerManager: extra_headers=extra_headers, stdio_env=stdio_env, subject_token=subject_token, + user_api_key_auth=user_api_key_auth, ) prompts = await client.list_prompts() @@ -3419,6 +3496,7 @@ class MCPServerManager: extra_headers: Optional[dict[str, str]] = None, add_prefix: bool = True, raw_headers: Optional[dict[str, str]] = None, + user_api_key_auth: UserAPIKeyAuth | None = None, ) -> list[Resource]: """Fetch available resources from a single MCP server.""" @@ -3434,7 +3512,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env = self._build_stdio_env(server, raw_headers) - subject_token = self._obo_subject_token(server, raw_headers) + subject_token = self._subject_bearer_token(server, raw_headers) client = await self._create_mcp_client( server=server, @@ -3442,6 +3520,7 @@ class MCPServerManager: extra_headers=extra_headers, stdio_env=stdio_env, subject_token=subject_token, + user_api_key_auth=user_api_key_auth, ) resources = await client.list_resources() @@ -3461,6 +3540,7 @@ class MCPServerManager: extra_headers: Optional[dict[str, str]] = None, add_prefix: bool = True, raw_headers: Optional[dict[str, str]] = None, + user_api_key_auth: UserAPIKeyAuth | None = None, ) -> list[ResourceTemplate]: """Fetch available resource templates from a single MCP server.""" @@ -3476,7 +3556,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env = self._build_stdio_env(server, raw_headers) - subject_token = self._obo_subject_token(server, raw_headers) + subject_token = self._subject_bearer_token(server, raw_headers) client = await self._create_mcp_client( server=server, @@ -3484,6 +3564,7 @@ class MCPServerManager: extra_headers=extra_headers, stdio_env=stdio_env, subject_token=subject_token, + user_api_key_auth=user_api_key_auth, ) resource_templates = await client.list_resource_templates() @@ -3505,6 +3586,7 @@ class MCPServerManager: mcp_auth_header: Optional[Union[str, dict[str, str]]] = None, extra_headers: Optional[dict[str, str]] = None, raw_headers: Optional[dict[str, str]] = None, + user_api_key_auth: UserAPIKeyAuth | None = None, ) -> ReadResourceResult: """Read resource contents from a specific MCP server.""" @@ -3517,7 +3599,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env = self._build_stdio_env(server, raw_headers) - subject_token = self._obo_subject_token(server, raw_headers) + subject_token = self._subject_bearer_token(server, raw_headers) client = await self._create_mcp_client( server=server, @@ -3525,6 +3607,7 @@ class MCPServerManager: extra_headers=extra_headers, stdio_env=stdio_env, subject_token=subject_token, + user_api_key_auth=user_api_key_auth, ) return await client.read_resource(url) @@ -3537,6 +3620,7 @@ class MCPServerManager: mcp_auth_header: Optional[Union[str, dict[str, str]]] = None, extra_headers: Optional[dict[str, str]] = None, raw_headers: Optional[dict[str, str]] = None, + user_api_key_auth: UserAPIKeyAuth | None = None, ) -> GetPromptResult: """Fetch a specific prompt definition from a single MCP server.""" @@ -3549,7 +3633,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env = self._build_stdio_env(server, raw_headers) - subject_token = self._obo_subject_token(server, raw_headers) + subject_token = self._subject_bearer_token(server, raw_headers) client = await self._create_mcp_client( server=server, @@ -3557,6 +3641,7 @@ class MCPServerManager: extra_headers=extra_headers, stdio_env=stdio_env, subject_token=subject_token, + user_api_key_auth=user_api_key_auth, ) get_prompt_request_params = GetPromptRequestParams( @@ -4593,12 +4678,14 @@ class MCPServerManager: subject_token: Optional[str], user_api_key_auth: Optional[UserAPIKeyAuth], ) -> CallToolResult: - """Call a token_exchange (OBO) tool; on an upstream 401/403 re-mint the token once and retry. + """Call an exchange-mode (token_exchange OBO or id_jag) tool; on an upstream 401/403 + re-mint the token once and retry. The exchanged token is baked into the client at build time, so the retry invalidates the - cached exchange and rebuilds the client (which re-exchanges). One retry only: a non-auth - failure or a second auth failure degrades to the normal ``isError`` result, and a re-exchange - that now fails surfaces its own 401 challenge from ``_create_mcp_client``. + cached exchange and rebuilds the client (which re-exchanges — for id_jag stored-assertion + callers that includes re-reading the assertion). One retry only: a non-auth failure or a + second auth failure degrades to the normal ``isError`` result, and a re-exchange that now + fails surfaces its own 401 challenge from ``_create_mcp_client``. """ try: return await client.call_tool( @@ -4633,7 +4720,7 @@ class MCPServerManager: proxy_logging_obj: Optional[ProxyLogging], host_progress_callback: Optional[Callable] = None, hook_extra_headers: Optional[dict[str, str]] = None, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, + user_api_key_auth: UserAPIKeyAuth | None = None, ) -> CallToolResult: """ Call a regular MCP tool using the MCP client. @@ -4680,15 +4767,10 @@ class MCPServerManager: if server_auth_header is None: server_auth_header = mcp_auth_header - # Extract subject token for OAuth2 Token Exchange (OBO) and ID-JAG flows - subject_token: Optional[str] = None + subject_token: str | None = self._subject_bearer_token(mcp_server, raw_headers, oauth2_headers=oauth2_headers) + # The exchange modes leave extra_headers None (the resolver injects the minted bearer). extra_headers: Optional[dict[str, str]] = None - if mcp_server.auth_type in ( - MCPAuth.oauth2_token_exchange, - MCPAuth.oauth2_id_jag, - ): - subject_token = self._extract_bearer_token(oauth2_headers, raw_headers) - elif mcp_server.auth_type == MCPAuth.oauth2: + if mcp_server.auth_type == MCPAuth.oauth2: if mcp_server.has_client_credentials: # For M2M OAuth servers, Authorization must come from token fetch. extra_headers = None @@ -4788,7 +4870,7 @@ class MCPServerManager: arguments=arguments, ) - if mcp_server.auth_type in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag) and subject_token: + if _obo_retry_covers_caller(mcp_server.auth_type, subject_token, user_api_key_auth): # 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. @@ -5041,7 +5123,7 @@ class MCPServerManager: subject_token: str | None = None if isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)): - subject_token = self._extract_bearer_token(oauth2_headers, raw_headers) + subject_token = self._subject_bearer_token(mcp_server, raw_headers, oauth2_headers=oauth2_headers) elif isinstance(spec.config, PassthroughConfig): inbound_token, forwarded_headers = _take_forwarded_authorization(forwarded_headers) per_server_token = _passthrough_token_from_mcp_auth_header(mcp_auth_header) @@ -5703,6 +5785,7 @@ class MCPServerManager: mcp_auth_header=None, extra_headers=extra_headers, stdio_env=None, + user_api_key_auth=None, # internal path (cache warm-up / health check): no caller identity ) try: @@ -5795,7 +5878,7 @@ class MCPServerManager: async def get_all_allowed_mcp_servers( self, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, + user_api_key_auth: UserAPIKeyAuth | None = None, ) -> list[LiteLLM_MCPServerTable]: """ Get all MCP servers that the user has access to. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 565c489e77c..9a4788e7d1a 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -18,6 +18,9 @@ from fastapi import HTTPException from pydantic import SecretStr from typing_extensions import assert_never +from litellm.proxy._experimental.mcp_server.outbound_credentials.credential_provenance import ( + classify_inbound_provenance, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, @@ -48,15 +51,19 @@ def to_subject(user_api_key_auth: Optional[UserAPIKeyAuth], subject_token: Optio """Map v1's authenticated principal onto the resolver's Subject. tenant_id / subject_id are empty for an unauthenticated caller; the per-user arms must reject - an empty subject rather than share one credential slot across callers. + an empty subject rather than share one credential slot across callers. The inbound token's + provenance (gateway-issued vs the caller's own external credential) is classified here, at the + edge, so the resolver core never has to recognize gateway credentials itself. """ inbound = SecretStr(subject_token) if subject_token else None + provenance = classify_inbound_provenance(subject_token) if user_api_key_auth is None: - return Subject(tenant_id="", subject_id="", inbound_token=inbound) + return Subject(tenant_id="", subject_id="", inbound_token=inbound, inbound_provenance=provenance) return Subject( tenant_id=user_api_key_auth.org_id or user_api_key_auth.team_id or "", subject_id=user_api_key_auth.user_id or "", inbound_token=inbound, + inbound_provenance=provenance, ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/credential_provenance.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/credential_provenance.py new file mode 100644 index 00000000000..531669e1884 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/credential_provenance.py @@ -0,0 +1,180 @@ +"""Provenance of the inbound subject token, decided at the resolver's v1 edge. + +The resolver must never forward a gateway-issued or gateway-honored credential to an external +token endpoint as exchange subject material. Recognizing "is this a credential this gateway +issued" by an enumerated denylist of prefixes is a leak generator: the set of minted formats +grows, and every format the list has not caught yet is disclosed. So this module recognizes a +gateway credential by the one property every one of them shares and no external token shares: +it is cryptographically ours, signed or encrypted under this gateway's master key. A new mint +format is caught by construction, with no edit here. + +``classify_inbound_provenance`` is the single owner both exchange arms read through +(``Subject.inbound_provenance``). id_jag's subject is an id_token (a JWT), so it exchanges the +inbound bearer only when it is ``external_jwt`` and otherwise sources the stored assertion (an +``external_opaque`` value is not an id_token and is never forwarded); token_exchange (OBO) keeps +its exchange-what-was-presented contract for both ``external_jwt`` and ``external_opaque`` tokens +but refuses a ``gateway_credential``. +""" + +from __future__ import annotations + +import hashlib +import secrets + +import jwt + +from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + ENVELOPE_ISSUER, + ENVELOPE_PREFIX, + REFRESH_ENVELOPE_PREFIX, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + SESSION_ISSUER, + SESSION_REFRESH_PREFIX, + SESSION_TOKEN_PREFIX, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import InboundTokenProvenance + +_VIRTUAL_KEY_PREFIX = "sk-" +_GATEWAY_MINT_PREFIXES = ( + _VIRTUAL_KEY_PREFIX, + SESSION_TOKEN_PREFIX, + SESSION_REFRESH_PREFIX, + ENVELOPE_PREFIX, + REFRESH_ENVELOPE_PREFIX, +) +_RESERVED_GATEWAY_ISSUERS = frozenset({SESSION_ISSUER, ENVELOPE_ISSUER}) + +# Recognition asks whether a token is THIS gateway's, not whether it is still usable, so every +# decode below disables claim validation (expiry, not-before, issued-at, audience, issuer). An +# expired or not-yet-valid gateway credential is still a gateway credential; letting a claim check +# reject it would drop it back to the exchangeable-external path and disclose it upstream. Only the +# signature (for the master-key check) speaks to provenance. +_RECOGNITION_CLAIM_CHECKS_DISABLED = { + "verify_exp": False, + "verify_nbf": False, + "verify_iat": False, + "verify_aud": False, + "verify_iss": False, +} + + +def _master_key() -> str | None: + from litellm.proxy.proxy_server import master_key + + return master_key + + +def _equals_master_key(token: str, master_key: str) -> bool: + """Constant-time equality with the master key, TOTAL over any input string. + + ``secrets.compare_digest`` raises ``TypeError`` on a ``str`` carrying non-ASCII characters, and + recognition must never raise into egress (an inbound bearer is fully attacker-controlled). Both + sides are hashed to fixed-length digests first, so any string compares without raising and the + master key's length is not leaked; a non-ASCII bearer reads "not the master key" rather than + 500-ing the request. + """ + token_digest = hashlib.sha256(token.encode("utf-8", "surrogatepass")).digest() + master_key_digest = hashlib.sha256(master_key.encode("utf-8", "surrogatepass")).digest() + return secrets.compare_digest(token_digest, master_key_digest) + + +def _verifies_under_master_key(token: str, master_key: str) -> bool: + """True when ``token`` is a JWT this gateway signed with the master key (HS256). + + Covers every master-key-signed mint (UI login session, onboarding, byok_session) and any + future one, since the check is the SIGNATURE, not the format and not the claims. Claim + validation is disabled: an expired or not-yet-valid gateway token is still this gateway's and + must be recognized, or it would fall through to the exchangeable-external path and leak. + """ + try: + jwt.decode( + token, + master_key, + algorithms=["HS256"], + options={"verify_signature": True, **_RECOGNITION_CLAIM_CHECKS_DISABLED}, + ) + return True + except Exception: # noqa: BLE001 # signature (not claims) failed -> not signed by us; other checks still run + return False + + +def _claims_a_reserved_issuer(token: str) -> bool: + """True when ``token`` is a JWT whose ``iss`` is one this gateway reserves (session/envelope). + + These are signed with master-key-DERIVED keys, so they do not verify under the master key + directly; they are already caught by their mint prefix, and this is the belt-and-suspenders + for an inner token presented without its prefix. The issuer is refused on the unverified + claim alone because the value is reserved to this gateway: no external token legitimately + carries it, so treating it as ours fails closed. Claim validation stays off for the same + reason as the master-key check: an expired reserved-issuer token is still ours. + """ + try: + claims = jwt.decode(token, options={"verify_signature": False, **_RECOGNITION_CLAIM_CHECKS_DISABLED}) + except Exception: # noqa: BLE001 # not a JWT + return False + return claims.get("iss") in _RESERVED_GATEWAY_ISSUERS + + +def _decrypts_under_gateway_key(token: str) -> bool: + """True when ``token`` decrypts under this gateway's salt key (the CLI / experimental UI + login blobs, which are encrypted rather than JWT-signed). Decryption succeeding is proof the + value is ours; an external token decrypts to nothing.""" + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + + try: + decrypted = decrypt_value_helper( + token, key="mcp_subject_token_provenance", exception_type="debug", return_original_value=False + ) + except Exception: # noqa: BLE001 # defensive: never let recognition raise into egress + return False + return decrypted is not None + + +def is_gateway_issued_credential(token: str) -> bool: + """Whether ``token`` is a credential this gateway issued or honors as its own. + + Complete by construction rather than by enumeration: a gateway credential is caught if it is + the master key, carries a gateway mint prefix, verifies under the master key, claims a + reserved gateway issuer, or decrypts under the gateway salt key. Every current mint (virtual + key, master key, MCP session/refresh, bridge envelope, UI login session, onboarding, + byok_session, CLI/experimental login) matches at least one, and a future mint signed or + encrypted with the master key matches without a change here. + """ + master_key = _master_key() + if master_key and _equals_master_key(token, master_key): + return True + if token.startswith(_GATEWAY_MINT_PREFIXES): + return True + if master_key and _verifies_under_master_key(token, master_key): + return True + if _claims_a_reserved_issuer(token): + return True + return _decrypts_under_gateway_key(token) + + +def _is_decodable_jwt(token: str) -> bool: + """Whether ``token`` is structurally a JWT (an id_token is a JWT). Signature and every + registered-claim check are disabled: this is a shape test to tell an id_token candidate from + an opaque bearer, not a validation. The gateway check runs first, so this only ever runs on a + non-gateway token; the org authorization server remains the authority on the id_token itself. + """ + try: + jwt.decode(token, options={"verify_signature": False, **_RECOGNITION_CLAIM_CHECKS_DISABLED}) + return True + except Exception: # noqa: BLE001 # not a JWT -> an opaque bearer, not an id_token candidate + return False + + +def classify_inbound_provenance(subject_token: str | None) -> InboundTokenProvenance: + """The one classification both exchange arms read. A gateway-issued credential is never + exchange subject material (any mode). Among the caller's own credentials, a JWT is an id_token + candidate (``external_jwt``) that id_jag can exchange, while an opaque value (``external_opaque``) + is not an id_token and id_jag falls to the stored assertion; OBO may forward either.""" + if not subject_token: + return "absent" + if is_gateway_issued_credential(subject_token): + return "gateway_credential" + if _is_decodable_jwt(subject_token): + return "external_jwt" + return "external_opaque" diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py index c1c70cf9050..11d1e8df2b3 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py @@ -91,6 +91,13 @@ class TokenRefresher(Protocol): ``server_id`` selects the upstream's config (token endpoint, client credentials, scopes) the grant runs against; ``(user_id, server_id)`` is the key the new token is persisted under. They are not derivable from ``token``, so the seam threads them alongside it. + + CONTRACT: an implementation MAY raise a domain exception instead of returning, to carry an + outcome richer than ``OAuthToken | None`` (the EMA assertion refresher raises its dead-grant + and store-outage verdicts, each mapping to a different caller remedy). ``RefreshingTokenStore`` + and the refresh coordinator MUST propagate such exceptions to the refresh winner and every + single-flight waiter alike — wrapping a refresher call in ``try/except -> None`` here would + silently collapse those verdicts into "no token" and change caller-visible statuses. """ async def refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None: ... @@ -159,6 +166,7 @@ class CachedOAuthTokenStore: *, default_ttl_seconds: float, expiry_skew_seconds: float = 60.0, + max_ttl_seconds: float | None = None, max_size: int = 4096, backend: TokenCacheBackend | None = None, clock: Callable[[], float] = time.time, @@ -166,13 +174,25 @@ class CachedOAuthTokenStore: self._inner = inner self._default_ttl_seconds = default_ttl_seconds self._expiry_skew_seconds = expiry_skew_seconds + self._max_ttl_seconds = max_ttl_seconds self._clock = clock self._backend: TokenCacheBackend = backend or InMemoryTokenCacheBackend(max_size=max_size, clock=clock) def _ttl(self, token: OAuthToken) -> float: - if token.expires_at is not None: - return max(0.0, token.expires_at - self._expiry_skew_seconds - self._clock()) - return self._default_ttl_seconds + """How long this token may be served from cache. + + A token that declares an expiry is held until then (minus the skew); one that does not is + held for the default. ``max_ttl_seconds`` caps both, for a credential whose holder must + converge on a replacement sooner than the token's own life: an identity assertion is + re-issued by a re-login that may REDUCE the user's claims, so serving the superseded one + for its full life would keep the old claims usable. + """ + ttl = ( + max(0.0, token.expires_at - self._expiry_skew_seconds - self._clock()) + if token.expires_at is not None + else self._default_ttl_seconds + ) + return ttl if self._max_ttl_seconds is None else min(ttl, self._max_ttl_seconds) async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: hit = await self._backend.get(user_id, server_id) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 69984a56311..bef5225eb2a 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -19,6 +19,7 @@ from __future__ import annotations import hashlib from functools import partial +from typing import Protocol import httpx from typing_extensions import assert_never @@ -41,6 +42,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( Ok, Result, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + SsoAssertionUnrenewable, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( ExchangedToken, ExchangedTokenCache, @@ -74,6 +78,29 @@ _JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer" _ID_JAG_REQUESTED_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:id-jag" +def _inbound_id_token(subject: Subject) -> str | None: + """The inbound bearer when it is the caller's own presented id_token, else None so id_jag + sourcing falls to the stored assertion. The ID-JAG exchange takes an id_token (a JWT), so the + inbound is used only when the edge classified it ``external_jwt``: a gateway-issued credential + and an opaque (non-JWT) bearer are both rejected here and resolve through the store instead.""" + if subject.inbound_token is None or subject.inbound_provenance != "external_jwt": + return None + return subject.inbound_token.get_secret_value() + + +class SsoAssertionSource(Protocol): + """Where the ID-JAG arm finds the caller's stored SSO identity assertion by user id.""" + + async def fetch_usable(self, user_id: str) -> OAuthToken | None: ... + + +class _NullSsoAssertionSource: + """Fail-closed default: with no assertion source wired, no stored assertion is ever found.""" + + async def fetch_usable(self, user_id: str) -> OAuthToken | None: + return None + + class _NullOAuthTokenStore: """Fail-closed default: with no token store wired, every user reads as not authorized.""" @@ -111,12 +138,14 @@ class UpstreamCredentialProvider: token_endpoint: TokenEndpointClient | None = None, exchanged_tokens: ExchangedTokenCache | None = None, client_credentials_source: ClientCredentialsTokenSource | None = None, + sso_assertions: SsoAssertionSource | 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_assertions: SsoAssertionSource = sso_assertions or _NullSsoAssertionSource() async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: @@ -170,15 +199,55 @@ class UpstreamCredentialProvider: return Error(CredError.of_not_implemented("api_key BYOK source not implemented yet")) 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: + async def _id_jag_subject_token(self, subject: Subject) -> Result[str, CredError]: + """The ONE sourcing rule for the ID-JAG subject token, shared by resolve and invalidate. + The ID-JAG exchange takes an id_token, so the subject is the caller's presented id_token + when the inbound bearer is one (the live, freshest identity evidence), else the stored SSO + assertion (itself an id_token, renewed when expired), else fail closed. Anything that is + not an id_token, a gateway-issued credential or an opaque bearer, is never forwarded: it + resolves through the store or not at all. An expired-and-unrenewable stored assertion is a + 401 re-login challenge, never a silent fall-through: the user HAS connected and the only + fix is theirs to perform.""" + inbound = _inbound_id_token(subject) + if inbound is not None: + return Ok(inbound) + try: + assertion = await self._sso_assertions.fetch_usable(subject.subject_id) + except SsoAssertionUnrenewable: 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." + CredError.of_unauthorized( + "The stored identity assertion for this user has expired and could not " + "be renewed; sign in to the gateway again to re-establish it.", + www_authenticate=( + 'Bearer error="invalid_token", error_description="identity assertion expired; re-authenticate"' + ), ) ) - token = subject.inbound_token.get_secret_value() + except TokenStoreUnavailable: + # Covers both causes with one honest, cause-neutral message: the assertion store (DB) + # not answering on the read and the IdP not answering on the renewal. Naming either + # one would misdiagnose the other. + return Error( + CredError.of_upstream_unavailable( + "the stored identity assertion could not be read or renewed; retry shortly" + ) + ) + if assertion is None: + return Error( + CredError.of_precondition_required( + "no stored identity assertion exists for this user; ID-JAG asserts the " + "calling user's identity upstream, so sign in to the gateway via SSO to " + "establish one (or present an IdP id_token as the bearer)." + ) + ) + return Ok(assertion.access_token) + + async def _id_jag(self, subject: Subject, server: ServerSpec, config: IdJagConfig) -> Result[httpx.Auth, CredError]: + match await self._id_jag_subject_token(subject): + case Error(err): + return Error(err) + case Ok(resolved): + token = resolved cache_key = _id_jag_cache_key(token, server.server_id, config) async def _exchange() -> Result[ExchangedToken, CredError]: @@ -248,6 +317,12 @@ class UpstreamCredentialProvider: No inbound token means there is nothing to exchange, so the arm fails closed with a 401 rather than falling through to a weaker source (§1.5); the exchanger handles the IdP round-trip and caching and returns the upstream token or a typed error. + + The arm keeps its exchange-what-was-presented contract for the caller's own external token, + but a gateway-issued credential (an admission key, session token, or gateway-signed JWT that + arrived on Authorization) is never disclosed to the external token endpoint; it fails closed + with the same 401 as a missing token so a misdirected gateway credential is refused, not + forwarded. """ inbound = subject.inbound_token if inbound is None: @@ -257,6 +332,14 @@ class UpstreamCredentialProvider: www_authenticate='Bearer error="invalid_request"', ) ) + if subject.inbound_provenance == "gateway_credential": + return Error( + CredError.of_unauthorized( + "Token exchange refuses a gateway-issued credential as the subject token; " + "present the caller's own identity token to exchange (OBO).", + www_authenticate='Bearer error="invalid_request"', + ) + ) match await self._token_exchanger.exchange( inbound.get_secret_value(), server, config, tenant_id=subject.tenant_id ): @@ -272,18 +355,22 @@ class UpstreamCredentialProvider: than serving the same rejected token until TTL. `token_exchange` and `id_jag` hold a 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. + other modes are a no-op. The `id_jag` eviction resolves its subject token through the same + rule that minted the cache entry, so a stored-assertion caller (no inbound token) evicts + the entry its resolve created; a failed resolution means no entry could exist, a no-op. """ - if subject.inbound_token is None: - return if isinstance(server.config, TokenExchangeConfig): + if subject.inbound_token is None: + return 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) - ) + match await self._id_jag_subject_token(subject): + case Ok(resolved): + self._exchanged_tokens.invalidate(_id_jag_cache_key(resolved, server.server_id, server.config)) + case Error(_): + return 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. 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 e0927cc4f64..8d13d7e30d4 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 @@ -5,8 +5,12 @@ The ``oauth2_id_jag`` egress arm needs the user's IdP ``id_token`` as its RFC 86 IdP assertion, so the assertion captured at the one SSO login is the only usable subject source for it. This module owns both sides of that state: the SSO callback persists here (write-through to the DB so a login on one pod is visible to every pod) and the resolver -seam reads back by ``user_id``. Retention is gated on an ``oauth2_id_jag`` server actually -being registered, so a gateway with no EMA upstream never stores bearer material. +seam reads back by ``user_id``. Capture currently hooks the generic SSO connector only (the +one whose ``GENERIC_*`` client also renews assertions below); Microsoft/Google SSO logins do +not feed this store, so EMA on those connectors answers 412 until they do. Retention is gated +on an ``oauth2_id_jag`` server actually being registered, so a gateway with no EMA upstream +never BEGINS storing bearer material; rows persisted while one existed are not purged by its +removal and remain until the user's next login replaces them. The row is one encrypted payload per user, latest login wins. ``expires_at`` mirrors the id_token ``exp`` claim and is judged by the reader, never enforced by deletion here: an @@ -17,6 +21,10 @@ truth, the same contract as the per-user OAuth credential store. from __future__ import annotations import json +import os +import time +import weakref +from collections.abc import Awaitable, Callable from datetime import datetime, timezone from typing import TYPE_CHECKING @@ -24,6 +32,12 @@ import jwt from pydantic import BaseModel, ConfigDict, SecretStr, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + CachedOAuthTokenStore, + OAuthToken, + RefreshingTokenStore, + TokenStoreUnavailable, +) if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -32,6 +46,20 @@ _ASSERTION_DECRYPT_LOG_KEY = "sso_identity_assertion" _STR_ADAPTER: TypeAdapter[str] = TypeAdapter(str) _MAYBE_STR_ADAPTER: TypeAdapter[str | None] = TypeAdapter(str | None) +# An assertion this close to expiry is treated as expired, so a subject token is never handed +# to the exchange with less lifetime than the two token-endpoint legs need to complete. +_ASSERTION_EXPIRY_BUFFER_SECONDS = 30 + +# The ceiling on how long any assertion is served from this pod's cache before the DB row is +# re-read: also the recovery bound for the one race the write chokepoint's invalidate cannot +# close (a read in flight across a re-login re-caching the superseded row). An assertion +# declaring an earlier expiry leaves the cache at that expiry (minus the buffer above) instead. +_ASSERTION_CACHE_TTL_SECONDS = 60.0 + +# The assertion is the user's SSO identity rather than a per-upstream credential, so it occupies a +# single per-user slot in the shared (user_id, server_id) keyed stack. +_ASSERTION_SERVER_KEY = "" + class SSOIdentityAssertion(BaseModel): """The IdP material an EMA exchange needs: ``id_token`` is the RFC 8693 subject token, @@ -127,31 +155,42 @@ async def persist_sso_identity_assertion(user_id: str, assertion: SSOIdentityAss "update": {"assertion_b64": encoded}, }, ) + # This is the one place the row is replaced, so it is the one place that must drop a cached + # predecessor: otherwise a re-login (notably one that reduces the user's IdP claims) would keep + # serving the superseded assertion until the old id_token expired. + await _drop_cached_assertion(user_id) async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | None: - """The stored assertion for ``user_id``, or ``None`` when absent, undecryptable (salt-key - rotation), or unparseable. Expiry is not judged here; the reader owns that policy.""" + """The stored assertion for ``user_id``, or ``None`` only when no row was ever written. + + ``None`` is reserved for that one determinate fact, so a fault can never be reported as "this + user has not signed in". A store that cannot be reached raises ``TokenStoreUnavailable``, which + is the contract ``OAuthTokenStore`` already defines for exactly this, and a row that cannot be + decrypted (salt-key rotation) or parsed raises ``SsoAssertionUnrenewable`` because the material + is unusable and only a fresh sign-in replaces it. Expiry is not judged here; the reader owns + that policy. + """ from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper # noqa: PLC0415 # runtime global from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime global if prisma_client is None: - return None - row = await prisma_client.db.litellm_ssoidentityassertion.find_unique(where={"user_id": user_id}) + raise TokenStoreUnavailable("the assertion store is not connected") + try: + row = await prisma_client.db.litellm_ssoidentityassertion.find_unique(where={"user_id": user_id}) + except Exception as exc: # noqa: BLE001 # an unreadable store is indeterminate, never an absence + raise TokenStoreUnavailable("the assertion store could not be read") from exc if row is None: return None raw = _MAYBE_STR_ADAPTER.validate_python( decrypt_value_helper(row.assertion_b64, _ASSERTION_DECRYPT_LOG_KEY, exception_type="debug") ) if raw is None: - return None + raise SsoAssertionUnrenewable("the stored identity assertion could not be decrypted") try: payload = _StoredAssertionPayload.model_validate_json(raw) - except ValidationError: - verbose_proxy_logger.warning( - "Stored SSO identity assertion for user_id=%s could not be parsed; treating as absent.", user_id - ) - return None + except ValidationError as exc: + raise SsoAssertionUnrenewable("the stored identity assertion could not be parsed") from exc return SSOIdentityAssertion( id_token=SecretStr(payload.id_token), refresh_token=SecretStr(payload.refresh_token) if payload.refresh_token else None, @@ -198,6 +237,328 @@ async def rotate_sso_identity_assertions_master_key(prisma_client: PrismaClient, ) +class SsoAssertionUnrenewable(Exception): + """The stored assertion is expired and its grant is definitively dead; only a fresh sign-in + produces a new one. Raised rather than returned so the verdict survives the shared refresh + stack, whose ``OAuthToken | None`` would otherwise collapse "the grant is dead" (re-login) + into "this user never connected" (a different remedy, and a different status).""" + + +# The renewal POST answers with the token body or raises the verdict its failure proves, so the +# outcome needs no intermediate representation: an unrenewable grant and an unreachable IdP are +# exactly the two terminal states the shared stack already propagates. +SsoTokenEndpointPost = Callable[[str, dict[str, str]], Awaitable[dict[str, object]]] + + +class _SsoRefreshClient(BaseModel): + """The gateway's own SSO client registration at the IdP. The stored refresh token was + issued to THIS client (the generic SSO app), never to an ``id_jag`` server's client, so + renewing the assertion must authenticate as it or the IdP answers ``invalid_grant``.""" + + model_config = ConfigDict(frozen=True) + client_id: str + client_secret: SecretStr | None + token_endpoint: str + + +def generic_sso_scopes(getenv: Callable[[str], str | None] = os.getenv) -> list[str]: + """The scopes the generic SSO client requests, shared by the login authorize request and the + EMA assertion refresh so the two can never diverge. ``openid`` must be among them or the IdP + returns no ``id_token`` on refresh, which would strand a renewable assertion as expired and + force a needless re-login; requesting exactly what login was granted keeps the refresh within + RFC 6749 scope while guaranteeing the id_token comes back. The refresh passes the source's + injected ``getenv`` so it reads the same environment as the rest of its client config.""" + raw = getenv("GENERIC_SCOPE") + return (raw if raw is not None else "openid email profile").split(" ") + + +def _sso_refresh_client_from_env(getenv: Callable[[str], str | None]) -> _SsoRefreshClient | None: + client_id = getenv("GENERIC_CLIENT_ID") + token_endpoint = getenv("GENERIC_TOKEN_ENDPOINT") + if not client_id or not token_endpoint: + return None + secret = getenv("GENERIC_CLIENT_SECRET") + return _SsoRefreshClient( + client_id=client_id, + client_secret=SecretStr(secret) if secret else None, + token_endpoint=token_endpoint, + ) + + +def _oauth_error_code(response: object) -> str | None: + """The RFC 6749 section 5.2 ``error`` code from a token-endpoint response body, or None + when the body is not a JSON object carrying one.""" + try: + body = response.json() # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType,reportUnknownVariableType] # httpx response is partially typed + except Exception: # noqa: BLE001 # an unparseable body simply carries no code + return None + if not isinstance(body, dict): + return None + error = body.get("error") # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # narrowed below + return error if isinstance(error, str) and error else None + + +async def _post_sso_token_endpoint(url: str, form: dict[str, str]) -> dict[str, object]: + """The renewal grant, answering with the token body or raising what its failure proves. + + Unrenewable requires PROOF of a verdict on the GRANT, and RFC 6749 section 5.2 has exactly + one code that states that fact: ``invalid_grant`` ("the provided ... refresh token is + invalid, expired, revoked"). Every other code describes the CLIENT or the REQUEST — + ``invalid_client`` is this gateway's own credentials failing, ``invalid_scope`` its scope + config — which an operator fixes with no help from the user, so recording them on the row + would permanently destroy renewable assertions org-wide over an ops mistake. Those, plus a + 429, a 4xx without the error object (an intermediary answering, not the token endpoint), a + 5xx, a transport failure, or a 2xx whose body is not a JSON object, all read unavailable, so + the caller retries instead of being told to sign in again. + """ + import httpx # noqa: PLC0415 + + from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # http_handler is untyped + ) + from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 + + try: + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) # pyright: ignore[reportUnknownVariableType] # http_handler is untyped + response = await client.post(url, headers={"Accept": "application/json"}, data=form) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # httpx handler partially typed + response.raise_for_status() # pyright: ignore[reportUnknownMemberType] # httpx handler partially typed + body: object = response.json() # pyright: ignore[reportUnknownMemberType] # shape-validated below + except httpx.HTTPStatusError as exc: + status_code = exc.response.status_code + error_code = _oauth_error_code(exc.response) + if 400 <= status_code < 500 and status_code != 429 and error_code == "invalid_grant": + raise SsoAssertionUnrenewable(f"the identity provider rejected the refresh grant ({status_code})") from exc + raise TokenStoreUnavailable( + f"the identity provider answered {status_code}" + (f" ({error_code})" if error_code else "") + ) from exc + except (SsoAssertionUnrenewable, TokenStoreUnavailable): + raise + except Exception as exc: # noqa: BLE001 # transport/parse failure carries no verdict on the grant + raise TokenStoreUnavailable(f"the identity provider could not be reached ({type(exc).__name__})") from exc + if not isinstance(body, dict): + raise TokenStoreUnavailable("the identity provider returned a non-object JSON body") + return body + + +def _to_oauth_token(assertion: SSOIdentityAssertion) -> OAuthToken: + """The assertion as the shared stack's credential: the id_token is the value EMA spends (the + RFC 8693 subject token), so it rides in ``access_token``. ``issuer`` has no reader.""" + return OAuthToken( + access_token=assertion.id_token.get_secret_value(), + expires_at=assertion.expires_at.timestamp() if assertion.expires_at is not None else None, + refresh_token=assertion.refresh_token.get_secret_value() if assertion.refresh_token else None, + ) + + +class _SsoAssertionDbStore: + """``OAuthTokenStore`` over the EMA assertion row. The assertion is the user's SSO identity, + not a per-upstream credential, so ``server_id`` is ignored: every id_jag server spends it.""" + + def __init__(self, fetch: Callable[[str], Awaitable[SSOIdentityAssertion | None]]) -> None: + self._fetch = fetch + + async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: + assertion = await self._fetch(user_id) + return None if assertion is None else _to_oauth_token(assertion) + + +class _SsoAssertionRefresher: + """``TokenRefresher`` that renews the assertion as the gateway's OWN generic SSO client, whose + registration issued the stored refresh token (an id_jag server's client would get + ``invalid_grant``). A renewed token is returned and persisted best-effort, so a rotated refresh + token is not lost while an in-hand one is never discarded; a dead grant raises + ``SsoAssertionUnrenewable`` (re-login) and anything proving nothing about the grant raises + ``TokenStoreUnavailable`` (retry), so a down IdP never tells a user to sign in again. + + A PROVEN dead grant is recorded where the assertion lives, by persisting the row without the + refresh token the IdP rejected, never in a cache beside it. Later reads then answer from the + row with no token-endpoint call, on every pod and for as long as the row stands, and a re-login + lifts it by replacing that row. A verdict in a side cache guarantees neither: it is per-pod, it + expires on its own clock, and it can refuse a caller who has just signed in again. + """ + + def __init__( + self, + fetch: Callable[[str], Awaitable[SSOIdentityAssertion | None]], + persist: Callable[[str, SSOIdentityAssertion], Awaitable[None]], + post: SsoTokenEndpointPost, + getenv: Callable[[str], str | None], + now: Callable[[], float] = time.time, + ) -> None: + self._fetch = fetch + self._persist = persist + self._post = post + self._getenv = getenv + self._now = now + + async def _record_dead_grant(self, user_id: str, rejected_refresh_token: str) -> None: + """Strip the rejected refresh token from the row — but only while the row still carries it. + + The verdict is a fact about ONE grant, so it is written by re-reading the row and rewriting + its own current material minus that grant, never by fabricating a row from this pod's stale + snapshot. When the re-read shows a different refresh token, another pod's renewal (the IdP + rotated the token this pod's POST lost the race to) or a fresh login already replaced the + grant this verdict describes, and writing it would clobber a live assertion into a forced + re-login. Best effort throughout: failing to record the verdict costs one more POST next + time, never a wrong answer.""" + try: + current = await self._fetch(user_id) + if ( + current is None + or current.refresh_token is None + or current.refresh_token.get_secret_value() != rejected_refresh_token + ): + return + await self._persist( + user_id, + SSOIdentityAssertion( + id_token=current.id_token, + refresh_token=None, + issuer=current.issuer, + expires_at=current.expires_at, + ), + ) + except Exception as exc: # noqa: BLE001 # recording the verdict must not mask it + verbose_proxy_logger.warning( + "Could not record the rejected SSO refresh grant for user_id=%s: %s", user_id, exc + ) + + async def refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None: + spent_refresh_token = token.refresh_token + if spent_refresh_token is None: + # No grant to spend, because the login never returned one or a proven rejection stripped + # it. Either way only a re-login helps, and answering costs no token-endpoint call. + raise SsoAssertionUnrenewable("the stored assertion carries no refresh token") + client = _sso_refresh_client_from_env(self._getenv) + if client is None: + verbose_proxy_logger.warning( + "Stored SSO assertion for user_id=%s is expired and the SSO client env " + "(GENERIC_CLIENT_ID/GENERIC_TOKEN_ENDPOINT) is not set; renewal is unavailable " + "until the deployment restores it.", + user_id, + ) + raise TokenStoreUnavailable("the SSO client environment is not configured") + scope = " ".join(generic_sso_scopes(self._getenv)) + form = { + "grant_type": "refresh_token", + "refresh_token": spent_refresh_token, + "client_id": client.client_id, + **({"client_secret": client.client_secret.get_secret_value()} if client.client_secret else {}), + **({"scope": scope} if scope else {}), + } + try: + body = await self._post(client.token_endpoint, form) + except SsoAssertionUnrenewable: + await self._record_dead_grant(user_id, spent_refresh_token) + raise + rotated = body.get("refresh_token") + carried_refresh = rotated if isinstance(rotated, str) and rotated else spent_refresh_token + refreshed = assertion_from_sso_login(body.get("id_token"), carried_refresh) + renewed = _to_oauth_token(refreshed) if refreshed is not None else None + if renewed is not None and ( + renewed.expires_at is None or self._now() < renewed.expires_at - _ASSERTION_EXPIRY_BUFFER_SECONDS + ): + try: + await self._persist(user_id, refreshed) + except Exception as exc: # noqa: BLE001 # a write-back failure must not discard the in-hand token + verbose_proxy_logger.warning( + "SSO assertion refresh for user_id=%s succeeded but the write-back failed; serving the " + "in-hand assertion (a rotated refresh token was not stored): %s", + user_id, + exc, + ) + return renewed + # A 2xx that produced no usable id_token, because the IdP omitted one or returned one already + # past expiry, spent the grant for nothing. It is recorded dead on the row exactly like a 4xx + # rejection, so the next read answers re-login without another POST; keeping the refresh token + # to retry would re-spend a one-time grant on every egress and never converge, and a fresh + # sign-in (which requests openid at the authorize leg) is what actually restores a usable one. + verbose_proxy_logger.warning( + "SSO assertion refresh for user_id=%s produced no usable id_token; treating as expired.", user_id + ) + await self._record_dead_grant(user_id, spent_refresh_token) + raise SsoAssertionUnrenewable("the refresh produced no usable id_token") + + +# Live sources, so the one write chokepoint can drop a superseded assertion from their caches the +# moment the row is replaced. Weak, so a discarded source never keeps an instance alive. +_LIVE_ASSERTION_SOURCES: weakref.WeakSet[LiveSsoAssertionSource] = weakref.WeakSet() + + +async def _drop_cached_assertion(user_id: str) -> None: + for source in tuple(_LIVE_ASSERTION_SOURCES): + try: + await source.invalidate(user_id) + except Exception as exc: # noqa: BLE001 # a cache drop must never fail the login write + verbose_proxy_logger.warning( + "Could not drop the cached SSO assertion for user_id=%s after its row was replaced: %s", user_id, exc + ) + + +class LiveSsoAssertionSource: + """The resolver's view of the assertion store, built on the shared per-user credential stack. + + ``Cached(Refreshing(db))`` is the composition ``authorization_code`` uses, so the read-through + cache (positive-only, never caching an absence, so a fresh login is seen at once), the + expiry-skewed renewal and the per-user single-flight all come from ``oauth_token_store`` + instead of being rebuilt here; passing a Redis cache and coordinator would extend that + single-flight across replicas without touching this class. + + That stack answers ``OAuthToken | None`` while EMA needs four outcomes, each with its own + remedy, so the refresher raises its two terminal verdicts and this adapter maps them with no + second read: a token is usable, ``SsoAssertionUnrenewable`` is re-login (401), + ``TokenStoreUnavailable`` is retry (503), and ``None`` then means no row was ever stored (412). + """ + + def __init__( + self, + fetch: Callable[[str], Awaitable[SSOIdentityAssertion | None]] = fetch_sso_identity_assertion, + persist: Callable[[str, SSOIdentityAssertion], Awaitable[None]] = persist_sso_identity_assertion, + post: SsoTokenEndpointPost = _post_sso_token_endpoint, + getenv: Callable[[str], str | None] = os.getenv, + cache_ttl_seconds: float = _ASSERTION_CACHE_TTL_SECONDS, + now: Callable[[], float] = time.time, + ) -> None: + self._store = CachedOAuthTokenStore( + RefreshingTokenStore( + _SsoAssertionDbStore(fetch), + _SsoAssertionRefresher(fetch, persist, post, getenv, now=now), + expiry_skew_seconds=_ASSERTION_EXPIRY_BUFFER_SECONDS, + clock=now, + ), + default_ttl_seconds=cache_ttl_seconds, + expiry_skew_seconds=_ASSERTION_EXPIRY_BUFFER_SECONDS, + max_ttl_seconds=cache_ttl_seconds, + clock=now, + ) + _LIVE_ASSERTION_SOURCES.add(self) + + async def invalidate(self, user_id: str) -> None: + """Drop this pod's cached assertion for ``user_id`` so the next read sees the replaced row. + + Called from the write chokepoint, so a re-login (in particular one that reduces the user's + claims) is honored immediately rather than when the superseded id_token finally expires. + """ + await self._store.invalidate(user_id, _ASSERTION_SERVER_KEY) + + async def fetch_usable(self, user_id: str) -> OAuthToken | None: + """The caller's usable assertion, ``None`` only when no row was ever stored. + + ``SsoAssertionUnrenewable`` (re-login) and ``TokenStoreUnavailable`` (retry) propagate, so + the arm maps three remedies without a parallel result union. ``None`` carries exactly one + meaning: anything indeterminate reads as unavailable, so a store outage is a retry rather + than a false "this user has never signed in", and never a 500. + """ + if not user_id: + return None + try: + return await self._store.fetch(user_id, _ASSERTION_SERVER_KEY) + except (SsoAssertionUnrenewable, TokenStoreUnavailable): + raise + except Exception as exc: # noqa: BLE001 # indeterminate, so never absence and never a 500 + raise TokenStoreUnavailable("the stored identity assertion could not be read") from exc + + async def retain_sso_identity_assertion_for_ema(user_id: str, assertion: SSOIdentityAssertion | None) -> None: """The SSO-callback hook: a no-op unless there is material AND an EMA server is registered. A store failure is logged and swallowed because the login itself must not fail on an diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 926d96c8868..036cfed9a59 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -383,6 +383,15 @@ AuthConfig = Annotated[ ] +# What the inbound bearer is, classified once at the edge so each exchange arm gates on the +# exact shape its mode requires. `gateway_credential` is one this gateway issued (never exchange +# material, any mode). `external_jwt` is a non-gateway JWT, i.e. an id_token candidate id_jag can +# exchange. `external_opaque` is a non-gateway non-JWT: id_jag rejects it (an id_token is a JWT) +# and falls to the stored assertion, while token_exchange (OBO) may still forward it per its +# exchange-what-was-presented contract. +InboundTokenProvenance = Literal["absent", "gateway_credential", "external_jwt", "external_opaque"] + + class Subject(BaseModel): """The validated inbound principal. NOT the v1 request object and NOT the LiteLLM key.""" @@ -392,6 +401,8 @@ class Subject(BaseModel): subject_id: str # Opaque, already-validated inbound identity. Only `token_exchange` / `passthrough` read it. inbound_token: SecretStr | None = None + # What `inbound_token` is, classified once at the edge (`to_subject`); each arm gates on it. + inbound_provenance: InboundTokenProvenance = "absent" class ServerSpec(BaseModel): diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 483f57f9139..0fad8696058 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2164,6 +2164,7 @@ if MCP_AVAILABLE: extra_headers=extra_headers, add_prefix=True, # Always add server prefix raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, ) all_prompts.extend(prompts) @@ -2217,6 +2218,7 @@ if MCP_AVAILABLE: extra_headers=extra_headers, add_prefix=True, # Always add server prefix raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, ) all_resources.extend(resources) @@ -2268,6 +2270,7 @@ if MCP_AVAILABLE: extra_headers=extra_headers, add_prefix=True, # Always add server prefix raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, ) all_resource_templates.extend(resource_templates) verbose_logger.debug( @@ -3113,6 +3116,7 @@ if MCP_AVAILABLE: mcp_auth_header=server_auth_header, extra_headers=extra_headers, raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, ) async def mcp_read_resource( @@ -3162,6 +3166,7 @@ if MCP_AVAILABLE: mcp_auth_header=server_auth_header, extra_headers=extra_headers, raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, ) def _get_standard_logging_mcp_tool_call( @@ -3634,6 +3639,7 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth], client_ip: Optional[str], allowed_server_ids: Optional[Set[str]] = None, + raw_headers: Optional[Dict[str, str]] = None, ) -> None: """Fail fast with HTTP 401 for MCP servers that need user auth but didn't receive it on this request. Covers both gateway-managed OAuth2 @@ -3748,6 +3754,21 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, ) + # id_jag (EMA): resolve the exchange here at the transport edge for the same reason as + # the OBO preflight above — the in-session list handlers can only serialize a raise + # into a JSON-RPC error, so a dead stored SSO assertion would read as an empty catalog + # instead of the re-login 401 the resolver minted for it. There is no subject-header + # gate: id_jag sources its subject from the caller's id_token or the stored assertion, + # so resolution always has an answer. Gated to single-server routes; the multi-server + # aggregate keeps absorbing per-server auth failures. + if server and server.auth_type == MCPAuth.oauth2_id_jag and len(mcp_servers or []) == 1: + await global_mcp_server_manager.preflight_id_jag( + server=server, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + ) + # Pass-through OAuth: when the admin has opted a server into # forwarding the client's bearer token (is_oauth_passthrough) and # the client hasn't supplied one, fail fast with 401 and point @@ -4071,6 +4092,7 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, client_ip=_client_ip, allowed_server_ids=toolset_allowed_server_ids, + raw_headers=raw_headers, ) # Pre-flight auth check for pass-through servers. Must run after @@ -4393,6 +4415,7 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, client_ip=_sse_client_ip, allowed_server_ids=toolset_allowed_server_ids, + raw_headers=raw_headers, ) # Pre-flight auth check for pass-through servers: surface upstream diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 8682b61f910..d0ab49243d0 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -66,6 +66,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( SSOIdentityAssertion, assertion_from_sso_login, + generic_sso_scopes, retain_sso_identity_assertion_for_ema, ) from litellm.proxy._types import ( @@ -1130,7 +1131,7 @@ def _setup_generic_sso_env_vars( ) -> Tuple[str, List[str], str, str, str, bool]: """Setup and validate Generic SSO environment variables.""" generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None) - generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ") + generic_scope = generic_sso_scopes() generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None) generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None) generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None) @@ -2708,7 +2709,7 @@ class SSOAuthenticationHandler: from fastapi_sso.sso.generic import create_provider generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None) - generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ") + generic_scope = generic_sso_scopes() generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None) generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None) generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index cb27e992ecb..210d1e6d780 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -94,6 +94,8 @@ def test_wire_value_carries_no_prose(): ("timeout", None, 504), ("unreachable", None, 502), ("upstream_error", 500, 502), + ("unavailable", 503, 503), + ("precondition", 412, 412), ("internal", None, 500), ], ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index bf757b64c9a..3682ae52980 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -358,6 +358,34 @@ def test_to_subject_maps_principal_fields(): assert subject.tenant_id == "org1" assert subject.subject_id == "user1" assert subject.inbound_token is None + assert subject.inbound_provenance == "absent" + + +def test_to_subject_classifies_an_external_jwt_as_id_token_candidate(): + import jwt as pyjwt + + external_jwt = pyjwt.encode({"iss": "https://idp.example.com", "sub": "alice"}, "idp-key", algorithm="HS256") + subject = to_subject(SimpleNamespace(org_id="", team_id="", user_id="alice"), external_jwt) + assert subject.inbound_provenance == "external_jwt" + + +def test_to_subject_classifies_an_opaque_inbound_token_as_external_opaque(): + subject = to_subject(SimpleNamespace(org_id="", team_id="", user_id="alice"), "opaque-not-a-jwt-token") + assert subject.inbound_provenance == "external_opaque" + + +def test_to_subject_classifies_a_gateway_credential_inbound_token(monkeypatch): + """The edge is where a gateway-issued bearer that rode Authorization is recognized, so the + resolver core never has to; a virtual key is stamped `gateway_credential`.""" + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "master_key", "sk-master-adapter-test") + subject = to_subject(SimpleNamespace(org_id="", team_id="", user_id="alice"), "sk-a-caller-virtual-key") + assert subject.inbound_provenance == "gateway_credential" + + +def test_to_subject_absent_token_classifies_absent(): + assert to_subject(SimpleNamespace(org_id="", team_id="", user_id="alice"), None).inbound_provenance == "absent" @pytest.mark.parametrize( @@ -512,9 +540,7 @@ def test_id_jag_client_secret_maps_to_config(): # ID-JAG asserts the user's id_token; the access_token default maps to id_token. assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:id_token" assert isinstance(spec.config.client_auth, ClientSecretAuth) - assert spec.config.client_auth.client_secret.get_secret_value() == ( - "litellm-client-secret" - ) + assert spec.config.client_auth.client_secret.get_secret_value() == ("litellm-client-secret") def test_id_jag_private_key_maps_to_private_key_jwt_auth(): @@ -540,9 +566,7 @@ def test_id_jag_private_key_wins_over_client_secret(): def test_id_jag_honors_explicit_subject_token_type(): - spec = to_server_spec( - _id_jag_server(subject_token_type="urn:ietf:params:oauth:token-type:saml2") - ) + spec = to_server_spec(_id_jag_server(subject_token_type="urn:ietf:params:oauth:token-type:saml2")) assert spec is not None and isinstance(spec.config, IdJagConfig) assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:saml2" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py index d0264319aab..67a047de8ea 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py @@ -38,19 +38,13 @@ def _endpoint(body, sink=None): def _recording_persist(sink): - async def persist( - user_id, server_id, access_token, refresh_token, expires_in, scopes - ): - sink.append( - (user_id, server_id, access_token, refresh_token, expires_in, scopes) - ) + async def persist(user_id, server_id, access_token, refresh_token, expires_in, scopes): + sink.append((user_id, server_id, access_token, refresh_token, expires_in, scopes)) return persist -def _refresher( - server=None, body=None, *, post_sink=None, persist_sink=None, clock=lambda: 1000.0 -): +def _refresher(server=None, body=None, *, post_sink=None, persist_sink=None, clock=lambda: 1000.0): return AuthorizationCodeRefresher( _lookup(server if server is not None else _Server()), _endpoint(body, post_sink), @@ -73,9 +67,7 @@ async def test_refreshes_persists_and_returns_typed_token(): post_sink=posted, persist_sink=persisted, ) - token = await refresher.refresh( - "alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt") - ) + token = await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")) assert token is not None assert token.access_token == "new-at" @@ -108,9 +100,7 @@ async def test_client_secret_basic_sends_authorization_header_not_body(): body={"access_token": "new-at"}, post_sink=posted, ) - token = await refresher.refresh( - "alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt") - ) + token = await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")) assert token is not None _url, form, headers = posted[0] @@ -136,38 +126,22 @@ async def test_client_secret_basic_without_secret_is_a_failed_refresh(): async def test_no_refresh_token_is_not_refreshable(): posted = [] refresher = _refresher(body={"access_token": "x"}, post_sink=posted) - assert ( - await refresher.refresh("alice", "srv", OAuthToken(access_token="old")) is None - ) + assert await refresher.refresh("alice", "srv", OAuthToken(access_token="old")) is None assert posted == [] # never hit the IdP @pytest.mark.asyncio async def test_unknown_server_or_no_token_url_yields_none(): - assert ( - await _refresher(server=None).refresh( - "a", "s", OAuthToken("old", refresh_token="rt") - ) - is None - ) + assert await _refresher(server=None).refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None no_url = _Server(token_url=None) - assert ( - await _refresher(server=no_url).refresh( - "a", "s", OAuthToken("old", refresh_token="rt") - ) - is None - ) + assert await _refresher(server=no_url).refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None @pytest.mark.asyncio async def test_grant_failure_does_not_persist(): persisted = [] - refresher = _refresher( - body=None, persist_sink=persisted - ) # token_endpoint signals failure - assert ( - await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None - ) + refresher = _refresher(body=None, persist_sink=persisted) # token_endpoint signals failure + assert await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None assert persisted == [] @@ -175,9 +149,7 @@ async def test_grant_failure_does_not_persist(): async def test_response_without_access_token_does_not_persist(): persisted = [] refresher = _refresher(body={"expires_in": 60}, persist_sink=persisted) - assert ( - await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None - ) + assert await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None assert persisted == [] @@ -185,13 +157,9 @@ async def test_response_without_access_token_does_not_persist(): async def test_unrotated_refresh_token_is_carried_forward(): persisted = [] refresher = _refresher(body={"access_token": "new-at"}, persist_sink=persisted) - token = await refresher.refresh( - "a", "s", OAuthToken("old", refresh_token="keep-rt") - ) + token = await refresher.refresh("a", "s", OAuthToken("old", refresh_token="keep-rt")) assert token is not None - assert ( - token.refresh_token == "keep-rt" - ) # response omitted refresh_token -> reuse the old one + assert token.refresh_token == "keep-rt" # response omitted refresh_token -> reuse the old one assert token.expires_at is None # no expires_in -> no known expiry assert persisted[0][3] == "keep-rt" @@ -200,9 +168,7 @@ async def test_unrotated_refresh_token_is_carried_forward(): async def test_unrecorded_scope_is_carried_forward(): persisted = [] refresher = _refresher(body={"access_token": "new-at"}, persist_sink=persisted) - token = await refresher.refresh( - "a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write")) - ) + token = await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write"))) assert token is not None # response omitted "scope" -> the user's recorded grant is preserved, not dropped assert token.scopes == ("read", "write") @@ -216,9 +182,7 @@ async def test_returned_scope_overrides_prior_when_present(): body={"access_token": "new-at", "scope": "read"}, persist_sink=persisted, ) - token = await refresher.refresh( - "a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write")) - ) + token = await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write"))) assert token is not None assert token.scopes == ("read",) # a present scope replaces the prior grant assert persisted[0][5] == ("read",) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_credential_provenance.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_credential_provenance.py new file mode 100644 index 00000000000..0f13da7fccd --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_credential_provenance.py @@ -0,0 +1,183 @@ +"""The gateway-credential recognizer must be complete by construction: every credential this +gateway mints classifies as `gateway_credential` (so it is never exchange subject material), and +a real external token classifies as `external`. These tests mint the real gateway formats so a +newly added mint that this recognizer fails to catch fails a test here, not in production.""" + +from datetime import datetime, timezone + +import jwt as pyjwt +import pytest + +import litellm.proxy.proxy_server as proxy_server +from litellm.proxy._experimental.mcp_server.outbound_credentials.credential_provenance import ( + classify_inbound_provenance, + is_gateway_issued_credential, +) + +_MASTER = "sk-master-key-for-provenance-tests" + + +@pytest.fixture(autouse=True) +def _master_key(monkeypatch): + monkeypatch.setattr(proxy_server, "master_key", _MASTER) + + +def _hs256(claims: dict, key: str) -> str: + return pyjwt.encode(claims, key, algorithm="HS256") + + +def _real_session_token() -> str: + from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + session_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + SessionPrincipal, + mint_session_token, + ) + + minted = mint_session_token( + SessionPrincipal(user_id="u", client_id="c"), + session_keys_from_master_key(_MASTER), + datetime.now(timezone.utc), + ) + return minted.token.get_secret_value() + + +def _gateway_credentials() -> dict[str, str]: + now = 2_000_000_000 + return { + "master_key": _MASTER, + "virtual_key": "sk-" + "a" * 40, + # The three the prefix denylist missed: master-key HS256 JWTs with no prefix and no iss. + "ui_login_session": _hs256({"user_id": "u", "key": "sk-embedded", "user_role": "proxy_admin"}, _MASTER), + "onboarding": _hs256({"token_type": "litellm_onboarding", "user_id": "u", "exp": now}, _MASTER), + "byok_session": _hs256({"user_id": "u", "server_id": "s", "type": "byok_session", "exp": now}, _MASTER), + # Prefixed gateway JWTs. + "mcp_session_real": _real_session_token(), + "envelope_prefixed": "llm_env_" + _hs256({"iss": "litellm-mcp-bridge"}, "derived-key"), + # Reserved-issuer inner tokens presented without their prefix (derived-key signed). + "session_inner_no_prefix": _hs256({"iss": "litellm-mcp-gateway", "exp": now}, "some-derived-key"), + "envelope_inner_no_prefix": _hs256({"iss": "litellm-mcp-bridge", "exp": now}, "some-derived-key"), + } + + +def _external_jwt_tokens() -> dict[str, str]: + """Non-gateway JWTs: id_token candidates id_jag can exchange (`external_jwt`).""" + now = 2_000_000_000 + return { + "external_idp_jwt": _hs256({"iss": "https://okta.example.com", "sub": "u", "aud": "gw", "exp": now}, "idp-key"), + "external_jwt_no_exp": _hs256({"iss": "https://idp.example.com", "sub": "u"}, "another-idp-key"), + } + + +def _external_opaque_tokens() -> dict[str, str]: + """Non-gateway, non-JWT bearers: not id_tokens, so id_jag rejects them (`external_opaque`).""" + return { + "opaque_access_token": "opaque-access-token-from-an-idp-1234567890", + # Three dot-segments but not a decodable JWT; the finding's example. + "jwt_shaped_but_undecodable": "abc.def.ghi", + } + + +# Evaluated once, and the case NAME is the test id. A minted session token embeds the current time +# and a random jti, so deriving the id from the token value instead would make it differ on every +# collection; under xdist the workers then disagree about which tests exist and the run errors out. +_GATEWAY_CREDENTIALS = _gateway_credentials() +_EXTERNAL_JWT_TOKENS = _external_jwt_tokens() +_EXTERNAL_OPAQUE_TOKENS = _external_opaque_tokens() + + +@pytest.mark.parametrize("name,token", list(_GATEWAY_CREDENTIALS.items()), ids=list(_GATEWAY_CREDENTIALS)) +def test_every_gateway_credential_is_recognized(name, token): + assert is_gateway_issued_credential(token) is True + assert classify_inbound_provenance(token) == "gateway_credential" + + +@pytest.mark.parametrize("name,token", list(_EXTERNAL_JWT_TOKENS.items()), ids=list(_EXTERNAL_JWT_TOKENS)) +def test_external_jwt_is_an_id_token_candidate(name, token): + assert is_gateway_issued_credential(token) is False + assert classify_inbound_provenance(token) == "external_jwt" + + +@pytest.mark.parametrize("name,token", list(_EXTERNAL_OPAQUE_TOKENS.items()), ids=list(_EXTERNAL_OPAQUE_TOKENS)) +def test_external_opaque_token_is_not_an_id_token(name, token): + """A non-gateway non-JWT bearer is not an id_token; id_jag must not forward it (the finding).""" + assert is_gateway_issued_credential(token) is False + assert classify_inbound_provenance(token) == "external_opaque" + + +@pytest.mark.parametrize("empty", [None, ""]) +def test_absent_inbound_token_classifies_absent(empty): + assert classify_inbound_provenance(empty) == "absent" + + +@pytest.mark.parametrize( + "temporal_claims", + [ + {"exp": 1}, # long expired + {"nbf": 9_999_999_999}, # not yet valid + {"iat": 9_999_999_999}, # issued in the future + ], +) +def test_temporally_invalid_gateway_jwt_is_still_recognized(temporal_claims): + """Recognition determines provenance, not usability: an expired or not-yet-valid gateway JWT + is still this gateway's and must classify as gateway_credential, or it drops to the + exchangeable-external path and is disclosed upstream.""" + token = _hs256({"user_id": "u", "key": "sk-embedded", **temporal_claims}, _MASTER) + assert is_gateway_issued_credential(token) is True + assert classify_inbound_provenance(token) == "gateway_credential" + + +def test_expired_external_jwt_stays_an_id_token_candidate(): + """An expired token that is NOT ours stays external_jwt; the org authorization server, not + this recognizer, is the authority on whether the caller's own token is still valid.""" + token = _hs256({"iss": "https://idp.example.com", "sub": "u", "exp": 1}, "external-idp-key") + assert is_gateway_issued_credential(token) is False + assert classify_inbound_provenance(token) == "external_jwt" + + +def test_encrypted_login_blob_is_recognized_by_decryption(): + """CLI / experimental UI login tokens are encrypted, not JWT-signed; decryption succeeding is + proof they are this gateway's, so they are recognized without a prefix.""" + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + blob = encrypt_value_helper('{"token": "cli-session-abc", "is_session_token": true}') + assert is_gateway_issued_credential(blob) is True + + +def test_recognizer_does_not_crash_without_a_master_key(monkeypatch): + """With no master key configured the crypto checks are skipped; an external token is still + classified external rather than raising into egress.""" + monkeypatch.setattr(proxy_server, "master_key", None) + assert is_gateway_issued_credential("some.external.jwt") is False + # A prefix-bearing gateway credential is still caught with no master key. + assert is_gateway_issued_credential("sk-still-caught-by-prefix") is True + + +@pytest.mark.parametrize( + "adversarial", + [ + "café-tökèn-with-non-ascii", # the reported crash: non-ASCII bytes hit the master-key compare + "bearer-\U0001f600-emoji", # supplementary-plane characters + "\x00\x01control-bytes", # control characters + "aaa.é.ccc", # non-ASCII inside a jwt-shaped value + ], +) +def test_non_ascii_inbound_bearer_never_raises_into_egress(adversarial): + """The inbound bearer is fully attacker-controlled, so recognition must be TOTAL: a non-ASCII + bearer used to reach ``secrets.compare_digest(token, master_key)`` and raise ``TypeError``, + 500-ing every egress. It must now classify as a non-gateway external token without raising. + Reverting the digest-based master-key compare makes this test raise.""" + assert is_gateway_issued_credential(adversarial) is False + assert classify_inbound_provenance(adversarial) in ("external_jwt", "external_opaque") + + +def test_master_key_compare_matches_only_the_exact_key(monkeypatch): + """The digest-based compare must still be exact equality, not a length match or near miss. A + master key with no gateway prefix isolates the compare from the other checks: the exact key is + recognized, a value that only shares its length or a prefix of it is not.""" + plain_master = "plain-master-key-no-gateway-prefix" + monkeypatch.setattr(proxy_server, "master_key", plain_master) + assert is_gateway_issued_credential(plain_master) is True + assert is_gateway_issued_credential(plain_master + "x") is False + assert is_gateway_issued_credential(plain_master[:-1] + "Z") is False diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_oauth_token_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_oauth_token_store.py index d72afaf46b6..a40f6e095eb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_oauth_token_store.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_oauth_token_store.py @@ -269,6 +269,12 @@ async def test_stale_read_after_refresh_rereads_before_starting_new_refresh(): async def test_refresh_failure_is_shared_by_joiners_not_re_run(): + """Pins the ``TokenRefresher`` contract: a refresher's raised exception MUST propagate to the + refresh winner and every single-flight waiter alike, and the failed attempt runs once. Domain + refreshers (the EMA assertion refresher) raise verdicts richer than ``OAuthToken | None``; + catching them here and returning None would silently collapse "the grant is dead" (re-login) + into "no token" (a different remedy and status).""" + class _FailingRefresher: def __init__(self) -> None: self.calls = 0 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 a710da81962..091fc149068 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 @@ -68,7 +68,21 @@ class _FakeTokenEndpoint: def _with_inbound(token: str) -> Subject: - return Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr(token)) + """A caller presenting their own id_token (the edge classified it `external_jwt`).""" + return Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr(token), inbound_provenance="external_jwt") + + +def _idp_jwt(sub: str = "user") -> str: + import jwt as pyjwt + + return pyjwt.encode( + {"sub": sub, "iss": "https://idp.example.com"}, "test-signing-key-32-bytes-long-xx", algorithm="HS256" + ) + + +_USER_IDP_JWT = _idp_jwt("user") +_ALICE_IDP_JWT = _idp_jwt("alice") +_BOB_IDP_JWT = _idp_jwt("bob") def _spec(config): @@ -247,9 +261,17 @@ _OBO = TokenExchangeConfig( @pytest.mark.asyncio -async def test_token_exchange_emits_the_exchanged_bearer(): +@pytest.mark.parametrize("external_provenance", ["external_jwt", "external_opaque"]) +async def test_token_exchange_emits_the_exchanged_bearer(external_provenance): + """OBO exchanges what was presented for the caller's own external token, JWT or opaque; the + id_token-only restriction is id_jag's, not OBO's.""" exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at"))) - subject = Subject(tenant_id="acme", subject_id="alice", inbound_token=SecretStr("caller-jwt")) + subject = Subject( + tenant_id="acme", + subject_id="alice", + inbound_token=SecretStr("caller-jwt"), + inbound_provenance=external_provenance, + ) result = await UpstreamCredentialProvider(token_exchanger=exchanger).resolve_credentials(subject, _spec(_OBO)) assert isinstance(result, Ok) assert _emitted(result.ok)["Authorization"] == "Bearer exchanged-at" @@ -257,6 +279,25 @@ async def test_token_exchange_emits_the_exchanged_bearer(): assert exchanger.calls == [("caller-jwt", "acme", "s")] +@pytest.mark.asyncio +async def test_token_exchange_refuses_a_gateway_credential_subject_token(): + """OBO keeps exchange-what-was-presented for the caller's own external token, but a + gateway-issued credential (the edge classified it `gateway_credential`) is refused with a + 401 and never disclosed to the external token endpoint.""" + exchanger = _FakeExchanger(Ok(OAuthToken(access_token="never"))) + subject = Subject( + tenant_id="acme", + subject_id="alice", + inbound_token=SecretStr("sk-a-gateway-admission-key"), + inbound_provenance="gateway_credential", + ) + result = await UpstreamCredentialProvider(token_exchanger=exchanger).resolve_credentials(subject, _spec(_OBO)) + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + # The gateway credential is never handed to the exchanger. + assert exchanger.calls == [] + + @pytest.mark.asyncio async def test_invalidate_credentials_drops_the_exchanged_token_for_the_subject_and_tenant(): exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at"))) @@ -446,22 +487,16 @@ async def test_id_jag_runs_both_legs_and_returns_the_leg2_bearer(): ] ) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - result = await provider.resolve_credentials( - _with_inbound("user-id-token"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(_id_jag_config())) assert isinstance(result, Ok) assert _emitted(result.ok)["Authorization"] == "Bearer final-access" leg1_endpoint, _, leg1_params = endpoint.calls[0] assert leg1_endpoint == "https://idp.example.com/token" - assert ( - leg1_params["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" - ) - assert ( - leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" - ) - assert leg1_params["subject_token"] == "user-id-token" + assert leg1_params["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" + assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" + assert leg1_params["subject_token"] == _USER_IDP_JWT leg2_endpoint, _, leg2_params = endpoint.calls[1] assert leg2_endpoint == "https://mcp-as.example.com/token" @@ -474,9 +509,7 @@ async def test_id_jag_runs_both_legs_and_returns_the_leg2_bearer(): async def test_id_jag_without_inbound_token_is_precondition_required_no_http(): endpoint = _FakeTokenEndpoint([]) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "precondition_required" @@ -485,13 +518,9 @@ async def test_id_jag_without_inbound_token_is_precondition_required_no_http(): @pytest.mark.asyncio async def test_id_jag_propagates_a_leg1_error_without_calling_leg2(): - endpoint = _FakeTokenEndpoint( - [Error(CredError.of_upstream_unavailable("leg1 down"))] - ) + endpoint = _FakeTokenEndpoint([Error(CredError.of_upstream_unavailable("leg1 down"))]) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - result = await provider.resolve_credentials( - _with_inbound("user-id-token"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "upstream_unavailable" @@ -508,9 +537,7 @@ async def test_id_jag_propagates_a_leg2_error(): ] ) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - result = await provider.resolve_credentials( - _with_inbound("user-id-token"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "upstream_unavailable" @@ -530,8 +557,8 @@ async def test_id_jag_reuses_the_cached_bearer_for_an_unchanged_config(): endpoint = _FakeTokenEndpoint(_two_leg_ok("first-bearer")) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - first = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) - second = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) + first = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(_id_jag_config())) + second = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(_id_jag_config())) assert isinstance(first, Ok) and isinstance(second, Ok) assert _emitted(second.ok)["Authorization"] == "Bearer first-bearer" @@ -566,8 +593,8 @@ async def test_id_jag_config_change_forces_a_fresh_exchange(changed): endpoint = _FakeTokenEndpoint(_two_leg_ok("old-policy-bearer") + _two_leg_ok("new-policy-bearer")) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - before = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) - after = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(changed)) + before = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(_id_jag_config())) + after = await provider.resolve_credentials(_with_inbound(_USER_IDP_JWT), _spec(changed)) assert isinstance(before, Ok) and isinstance(after, Ok) assert _emitted(after.ok)["Authorization"] == "Bearer new-policy-bearer" @@ -579,8 +606,8 @@ async def test_id_jag_does_not_share_the_cached_bearer_across_caller_tokens(): endpoint = _FakeTokenEndpoint(_two_leg_ok("alice-bearer") + _two_leg_ok("bob-bearer")) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - alice = await provider.resolve_credentials(_with_inbound("alice-id-token"), _spec(_id_jag_config())) - bob = await provider.resolve_credentials(_with_inbound("bob-id-token"), _spec(_id_jag_config())) + alice = await provider.resolve_credentials(_with_inbound(_ALICE_IDP_JWT), _spec(_id_jag_config())) + bob = await provider.resolve_credentials(_with_inbound(_BOB_IDP_JWT), _spec(_id_jag_config())) assert isinstance(alice, Ok) and isinstance(bob, Ok) assert _emitted(bob.ok)["Authorization"] == "Bearer bob-bearer" @@ -591,7 +618,7 @@ async def test_id_jag_does_not_share_the_cached_bearer_across_caller_tokens(): async def test_invalidate_credentials_evicts_the_id_jag_bearer_so_the_next_resolve_re_exchanges(): endpoint = _FakeTokenEndpoint(_two_leg_ok("rejected-bearer") + _two_leg_ok("fresh-bearer")) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - subject = _with_inbound("user-id-token") + subject = _with_inbound(_USER_IDP_JWT) first = await provider.resolve_credentials(subject, _spec(_id_jag_config())) await provider.invalidate_credentials(subject, _spec(_id_jag_config())) @@ -606,7 +633,7 @@ async def test_invalidate_credentials_evicts_the_id_jag_bearer_so_the_next_resol async def test_invalidate_credentials_for_id_jag_is_a_noop_without_a_caller_token(): endpoint = _FakeTokenEndpoint(_two_leg_ok("cached-bearer")) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - subject = _with_inbound("user-id-token") + subject = _with_inbound(_USER_IDP_JWT) first = await provider.resolve_credentials(subject, _spec(_id_jag_config())) await provider.invalidate_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) @@ -615,3 +642,212 @@ async def test_invalidate_credentials_for_id_jag_is_a_noop_without_a_caller_toke assert isinstance(first, Ok) and isinstance(second, Ok) assert _emitted(second.ok)["Authorization"] == "Bearer cached-bearer" assert len(endpoint.calls) == 2 + + +# -- ID-JAG subject sourcing: stored SSO assertion seam (EMA) -------------------------------- + + +class _FakeAssertionSource: + """Returns the canned lookup and records which user ids were asked for.""" + + def __init__(self, lookup): + self._lookup = lookup + self.asked: list[str] = [] + + async def fetch_usable(self, user_id: str): + self.asked.append(user_id) + if isinstance(self._lookup, Exception): + raise self._lookup + return self._lookup + + +def _usable_assertion(id_token: str = "hdr.stored-id-token.sig"): + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken + + return OAuthToken(access_token=id_token) + + +def _no_assertion(): + return None + + +def _expired_assertion(): + from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + SsoAssertionUnrenewable, + ) + + return SsoAssertionUnrenewable("dead grant") + + +def _no_inbound(subject_id: str = "alice") -> Subject: + return Subject(tenant_id="", subject_id=subject_id) + + +@pytest.mark.asyncio +async def test_id_jag_sources_the_stored_assertion_when_no_inbound_token(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + source = _FakeAssertionSource(_usable_assertion()) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertions=source) + + result = await provider.resolve_credentials(_no_inbound(), _spec(_id_jag_config())) + + assert isinstance(result, Ok) + assert source.asked == ["alice"] + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == "hdr.stored-id-token.sig" + + +@pytest.mark.asyncio +async def test_id_jag_prefers_the_presented_idp_jwt_over_the_stored_assertion(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + source = _FakeAssertionSource(_usable_assertion()) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertions=source) + + result = await provider.resolve_credentials(_with_inbound(_idp_jwt("live")), _spec(_id_jag_config())) + + assert isinstance(result, Ok) + assert source.asked == [] + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == _idp_jwt("live") + + +@pytest.mark.asyncio +async def test_id_jag_never_exchanges_a_gateway_credential_inbound_token(): + """A gateway-issued bearer that rode the Authorization header (the edge classified it + `gateway_credential`) must fall through to the stored assertion, never reach an external + IdP, no matter that a token happens to be present.""" + gateway_bearer = "sk-litellm-admission-key" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + source = _FakeAssertionSource(_usable_assertion()) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertions=source) + + subject = Subject( + tenant_id="", + subject_id="alice", + inbound_token=SecretStr(gateway_bearer), + inbound_provenance="gateway_credential", + ) + result = await provider.resolve_credentials(subject, _spec(_id_jag_config())) + + assert isinstance(result, Ok) + assert source.asked == ["alice"] + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == "hdr.stored-id-token.sig" + assert all(gateway_bearer not in str(call) for call in endpoint.calls) + + +@pytest.mark.asyncio +async def test_id_jag_never_forwards_an_opaque_inbound_token(): + """The ID-JAG exchange takes an id_token (a JWT). An opaque non-gateway bearer is not an + id_token, so id_jag sources the stored assertion and never forwards the opaque value.""" + opaque_bearer = "abc.def.ghi" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + source = _FakeAssertionSource(_usable_assertion()) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertions=source) + + subject = Subject( + tenant_id="", subject_id="alice", inbound_token=SecretStr(opaque_bearer), inbound_provenance="external_opaque" + ) + result = await provider.resolve_credentials(subject, _spec(_id_jag_config())) + + assert isinstance(result, Ok) + assert source.asked == ["alice"] + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == "hdr.stored-id-token.sig" + assert all(opaque_bearer not in str(call) for call in endpoint.calls) + + +@pytest.mark.asyncio +async def test_id_jag_expired_stored_assertion_is_a_401_challenge_not_a_412(): + endpoint = _FakeTokenEndpoint([]) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertions=_FakeAssertionSource(_expired_assertion()) + ) + + result = await provider.resolve_credentials(_no_inbound(), _spec(_id_jag_config())) + + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + unauthorized = result.error.unauthorized + assert unauthorized is not None and unauthorized.www_authenticate is not None + assert "invalid_token" in unauthorized.www_authenticate + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_id_jag_absent_stored_assertion_and_no_inbound_is_precondition_required(): + endpoint = _FakeTokenEndpoint([]) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertions=_FakeAssertionSource(_no_assertion())) + + result = await provider.resolve_credentials(_no_inbound(), _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_cache_key_is_the_resolved_token_so_a_new_login_re_exchanges(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("first-bearer") + _two_leg_ok("second-bearer")) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertions=_FakeAssertionSource(_usable_assertion()) + ) + + first = await provider.resolve_credentials(_no_inbound(), _spec(_id_jag_config())) + relogged = UpstreamCredentialProvider( + token_endpoint=endpoint, + exchanged_tokens=provider._exchanged_tokens, + sso_assertions=_FakeAssertionSource(_usable_assertion("hdr.new-login-id-token.sig")), + ) + second = await relogged.resolve_credentials(_no_inbound(), _spec(_id_jag_config())) + + 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_invalidate_for_id_jag_evicts_the_stored_assertion_caller_entry(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("rejected-bearer") + _two_leg_ok("fresh-bearer")) + source = _FakeAssertionSource(_usable_assertion()) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertions=source) + subject = _no_inbound() + + first = await provider.resolve_credentials(subject, _spec(_id_jag_config())) + await provider.invalidate_credentials(subject, _spec(_id_jag_config())) + second = await provider.resolve_credentials(subject, _spec(_id_jag_config())) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(second.ok)["Authorization"] == "Bearer fresh-bearer" + assert len(endpoint.calls) == 4 + + +@pytest.mark.asyncio +async def test_id_jag_empty_subject_id_with_null_source_stays_precondition_required(): + provider = UpstreamCredentialProvider(token_endpoint=_FakeTokenEndpoint([])) + + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id=""), _spec(_id_jag_config())) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + + +@pytest.mark.asyncio +async def test_id_jag_unavailable_renewal_is_503_not_a_relogin_challenge(): + """An unreachable IdP during renewal proves nothing about the stored assertion, so the + caller gets a retryable upstream_unavailable, never the re-login challenge.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + TokenStoreUnavailable, + ) + + endpoint = _FakeTokenEndpoint([]) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertions=_FakeAssertionSource(TokenStoreUnavailable("idp down")) + ) + + result = await provider.resolve_credentials(_no_inbound(), _spec(_id_jag_config())) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert endpoint.calls == [] 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 a3f46a49ba9..c2d1a51c007 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 @@ -1,10 +1,10 @@ """Tests for the SSO identity assertion store (EMA subject-token capture). -Pins the contract of the store that PR 2's ``_id_jag`` subject-sourcing seam will read: -the carrier validates untyped IdP token-response values at the boundary, retention is -gated on an ``oauth2_id_jag`` server being registered, the row is encrypted at rest and -round-trips exactly, a store failure never escapes into the login path, and a salt-key -rotation re-encrypts stored rows like the sibling per-user credential tables. +Pins the contract the resolver's ``_id_jag`` subject-sourcing seam reads: the carrier +validates untyped IdP token-response values at the boundary, retention is gated on an +``oauth2_id_jag`` server being registered, the row is encrypted at rest and round-trips +exactly, a store failure never escapes into the login path, and a salt-key rotation +re-encrypts stored rows like the sibling per-user credential tables. """ import json @@ -14,7 +14,11 @@ from unittest.mock import AsyncMock, MagicMock, patch import jwt as pyjwt import pytest +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + TokenStoreUnavailable, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + SsoAssertionUnrenewable, assertion_from_sso_login, ema_assertion_retention_enabled, fetch_sso_identity_assertion, @@ -49,9 +53,7 @@ def _make_prisma(stored: dict, db_has_id_jag_server: bool = False): ``db_has_id_jag_server`` drives the retention gate's authoritative DB fallback; it is wired explicitly so the gate never reads a truthy bare MagicMock.""" prisma = MagicMock() - prisma.db.litellm_mcpservertable.find_first = AsyncMock( - return_value=MagicMock() if db_has_id_jag_server else None - ) + prisma.db.litellm_mcpservertable.find_first = AsyncMock(return_value=MagicMock() if db_has_id_jag_server else None) async def _upsert(where, data): stored[where["user_id"]] = data["update"]["assertion_b64"] @@ -240,19 +242,35 @@ async def test_fetch_missing_row_returns_none(): @pytest.mark.asyncio -async def test_fetch_undecryptable_row_returns_none(): +async def test_fetch_undecryptable_row_is_unrenewable_not_absent(): + """A row that will not decrypt (salt-key rotation) is unusable material, not a user who never + signed in; the remedy is a fresh sign-in, so it must not read as absence.""" prisma = _make_prisma({"user-a": "not-an-encrypted-blob"}) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - assert await fetch_sso_identity_assertion("user-a") is None + with patch("litellm.proxy.proxy_server.prisma_client", prisma), pytest.raises(SsoAssertionUnrenewable): + await fetch_sso_identity_assertion("user-a") @pytest.mark.asyncio -async def test_fetch_unparseable_payload_returns_none(): +async def test_fetch_unparseable_payload_is_unrenewable_not_absent(): from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper prisma = _make_prisma({"user-a": encrypt_value_helper("]]not json")}) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - assert await fetch_sso_identity_assertion("user-a") is None + with patch("litellm.proxy.proxy_server.prisma_client", prisma), pytest.raises(SsoAssertionUnrenewable): + await fetch_sso_identity_assertion("user-a") + + +@pytest.mark.asyncio +async def test_fetch_reports_an_unreachable_store_as_unavailable_never_absent(): + """The reported defect: a transient store outage for a user who DOES have an assertion must be + retryable (503), never the precondition-required shape that means "you have never signed in". + ``None`` is reserved for the one determinate fact that no row was ever written.""" + with patch("litellm.proxy.proxy_server.prisma_client", None), pytest.raises(TokenStoreUnavailable): + await fetch_sso_identity_assertion("user-a") + + exploding = _make_prisma({}) + exploding.db.litellm_ssoidentityassertion.find_unique = AsyncMock(side_effect=RuntimeError("db down")) + with patch("litellm.proxy.proxy_server.prisma_client", exploding), pytest.raises(TokenStoreUnavailable): + await fetch_sso_identity_assertion("user-a") @pytest.mark.asyncio @@ -341,3 +359,769 @@ 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 + + +# -- LiveSsoAssertionSource: the resolver-facing usable-assertion lookup + refresh ------------- + + +from datetime import datetime, timedelta, timezone # noqa: E402 + +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( # noqa: E402 + LiveSsoAssertionSource, + SSOIdentityAssertion, + SsoAssertionUnrenewable, +) +from pydantic import SecretStr # noqa: E402 + +_SSO_ENV = { + "GENERIC_CLIENT_ID": "sso-client", + "GENERIC_CLIENT_SECRET": "sso-secret", + "GENERIC_TOKEN_ENDPOINT": "https://idp.example.com/token", +} + + +def _stored(expires_in_seconds: int | None, refresh_token: str | None = "rt_stored") -> SSOIdentityAssertion: + return SSOIdentityAssertion( + id_token=SecretStr(_make_id_token()), + refresh_token=SecretStr(refresh_token) if refresh_token else None, + expires_at=( + datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds) + if expires_in_seconds is not None + else None + ), + ) + + +def _source(stored: SSOIdentityAssertion | None, post_bodies: list, env: dict | None = None): + fetches: list[str] = [] + persisted: list[tuple[str, SSOIdentityAssertion]] = [] + posts: list[tuple[str, dict]] = [] + state = {"stored": stored} + + async def fetch(user_id: str): + fetches.append(user_id) + return state["stored"] + + async def persist(user_id: str, assertion: SSOIdentityAssertion): + persisted.append((user_id, assertion)) + state["stored"] = assertion + + async def post(url: str, form: dict): + posts.append((url, dict(form))) + body = post_bodies.pop(0) if post_bodies else None + if body is None: + raise TokenStoreUnavailable("idp unreachable") + return body + + environment = _SSO_ENV if env is None else env + source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=environment.get) + return source, fetches, persisted, posts + + +@pytest.mark.asyncio +async def test_source_empty_user_id_is_absent_without_a_db_read(): + source, fetches, _, _ = _source(_stored(3600), []) + assert await source.fetch_usable("") is None + assert fetches == [] + + +@pytest.mark.asyncio +async def test_source_missing_row_is_absent(): + source, _, _, posts = _source(None, []) + assert await source.fetch_usable("alice") is None + assert posts == [] + + +@pytest.mark.asyncio +async def test_source_unexpired_assertion_is_usable_without_refresh(): + stored = _stored(3600) + source, _, _, posts = _source(stored, []) + lookup = await source.fetch_usable("alice") + assert lookup is not None + assert lookup.access_token == stored.id_token.get_secret_value() + assert posts == [] + + +@pytest.mark.asyncio +async def test_source_near_expiry_assertion_counts_as_expired(): + """A token inside the buffer would die mid-exchange; it must refresh, not be served.""" + source, _, _, posts = _source(_stored(10, refresh_token=None), []) + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + assert posts == [] + + +@pytest.mark.asyncio +async def test_source_expired_with_refresh_renews_persists_and_returns_usable(): + new_id_token = _make_id_token(exp_offset=7200) + source, _, persisted, posts = _source( + _stored(-100), [{"id_token": new_id_token, "refresh_token": "rt_rotated", "access_token": "at"}] + ) + lookup = await source.fetch_usable("alice") + assert lookup is not None + assert lookup.access_token == new_id_token + assert lookup.refresh_token is not None + assert lookup.refresh_token == "rt_rotated" + assert len(persisted) == 1 and persisted[0][0] == "alice" + url, form = posts[0] + assert url == "https://idp.example.com/token" + assert form["grant_type"] == "refresh_token" + assert form["refresh_token"] == "rt_stored" + assert form["client_id"] == "sso-client" + assert form["client_secret"] == "sso-secret" + + +@pytest.mark.asyncio +async def test_source_refresh_requests_openid_scope_so_the_idp_returns_an_id_token(): + """The refresh grant must carry a scope containing ``openid`` or the IdP returns no id_token, + stranding a renewable assertion as expired and forcing a needless re-login (the round-9 + finding). With no GENERIC_SCOPE configured it falls to the same default the login uses, which + contains ``openid``. Dropping the scope from the refresh form makes this fail.""" + source, _, _, posts = _source(_stored(-100), [{"id_token": _make_id_token(exp_offset=7200)}]) + await source.fetch_usable("alice") + _, form = posts[0] + assert "openid" in form["scope"].split(" ") + + +@pytest.mark.asyncio +async def test_source_refresh_scope_matches_the_configured_login_scope(): + """Refresh requests exactly the scopes the login was granted, read through the same shared + owner off the injected env, so the two can never diverge and the refresh stays within RFC 6749 + granted scope. Sourcing the scope from the global environment instead of the injected getenv + makes this fail.""" + env = {**_SSO_ENV, "GENERIC_SCOPE": "openid email offline_access"} + source, _, _, posts = _source(_stored(-100), [{"id_token": _make_id_token(exp_offset=7200)}], env=env) + await source.fetch_usable("alice") + _, form = posts[0] + assert form["scope"] == "openid email offline_access" + + +@pytest.mark.asyncio +async def test_source_refresh_without_rotated_token_carries_the_stored_one_forward(): + source, _, _, _ = _source(_stored(-100), [{"id_token": _make_id_token(exp_offset=7200)}]) + lookup = await source.fetch_usable("alice") + assert lookup is not None + assert lookup.refresh_token is not None + assert lookup.refresh_token == "rt_stored" + + +@pytest.mark.asyncio +async def test_source_expired_without_refresh_token_is_expired(): + source, _, _, posts = _source(_stored(-100, refresh_token=None), []) + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + assert posts == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [{}, {"id_token": ""}, {"id_token": 123}, {"access_token": "at-only"}]) +async def test_source_refresh_granting_no_usable_id_token_records_a_dead_grant(body): + """A 2xx that carries no usable id_token spent the grant for nothing, so it is recorded dead on + the row (refresh token dropped) rather than left to re-POST. A second egress must not spend the + grant again: one-time refresh tokens and IdPs that omit id_token on refresh would otherwise + thrash the token endpoint on every call.""" + source, _, persisted, posts = _source(_stored(-100), [body, body]) + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + assert len(persisted) == 1 + assert persisted[0][1].refresh_token is None + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + assert len(posts) == 1 # the dead-grant row short-circuits the second attempt with no POST + + +@pytest.mark.asyncio +async def test_source_refresh_rejection_is_expired_and_unreachable_is_unavailable(): + """A 4xx is a verdict on the grant (re-login fixes it); an unreachable IdP proves nothing + about the grant, so the user must NOT be told to re-login for a transient outage.""" + + async def fetch(user_id: str): + return _stored(-100) + + async def persist(user_id, assertion): + return None + + async def rejected_post(url, form): + raise SsoAssertionUnrenewable("rejected") + + async def unreachable_post(url, form): + raise TokenStoreUnavailable("unreachable") + + rejected_source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=rejected_post, getenv=_SSO_ENV.get) + with pytest.raises(SsoAssertionUnrenewable): + await rejected_source.fetch_usable("alice") + + unreachable_source = LiveSsoAssertionSource( + fetch=fetch, persist=persist, post=unreachable_post, getenv=_SSO_ENV.get + ) + with pytest.raises(TokenStoreUnavailable): + await unreachable_source.fetch_usable("alice") + + +@pytest.mark.asyncio +async def test_source_missing_sso_env_is_unavailable_not_expired(): + """Env absence is a deployment problem; a re-login cannot fix it (SSO needs the same env), + so the state must be Unavailable, not Expired.""" + source, _, _, posts = _source(_stored(-100), [], env={}) + with pytest.raises(TokenStoreUnavailable): + await source.fetch_usable("alice") + assert posts == [] + + +@pytest.mark.asyncio +async def test_a_dead_grant_costs_one_post_because_the_verdict_lands_on_the_row(): + """A dead refresh grant costs one token-endpoint POST, not one per egress call, because the + rejection strips the refresh token from the row itself; later reads answer from the row. The + bound therefore holds on every pod and for as long as the row stands, and a re-login lifts it + by replacing the row rather than by waiting out a timer.""" + clock = {"now": time.time()} + posts: list[dict] = [] + state = {"stored": _stored(-100)} + + async def fetch(user_id: str): + return state["stored"] + + async def persist(user_id, assertion): + state["stored"] = assertion + + async def post(url, form): + posts.append(dict(form)) + raise SsoAssertionUnrenewable("rejected") + + source = LiveSsoAssertionSource( + fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get, now=lambda: clock["now"] + ) + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + assert len(posts) == 1 + + clock["now"] = clock["now"] + 31.0 + state["stored"] = _stored(3600) + lookup = await source.fetch_usable("alice") + assert lookup is not None + assert len(posts) == 1 + + +@pytest.mark.asyncio +async def test_a_late_dead_grant_verdict_never_clobbers_anothers_renewal(): + """The verdict is a fact about the ONE grant the IdP rejected. When another pod's renewal + already replaced that grant (rotating IdPs reject the loser of a cross-pod race with + ``invalid_grant``), the loser's verdict describes a token the row no longer holds, so + recording it must be a no-op — not an overwrite that destroys the winner's live assertion + and forces an org-visible re-login.""" + winner_row = _stored(7200, refresh_token="rt_winner") + state = {"stored": _stored(-100, refresh_token="rt_loser")} + persisted: list[tuple[str, SSOIdentityAssertion]] = [] + + async def fetch(user_id: str): + return state["stored"] + + async def persist(user_id, assertion): + persisted.append((user_id, assertion)) + state["stored"] = assertion + + async def post(url, form): + # The winner's write lands while this pod's POST is in flight. + state["stored"] = winner_row + raise SsoAssertionUnrenewable("invalid_grant: rt_loser was already spent") + + source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get) + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + assert persisted == [] + assert state["stored"] is winner_row + + +@pytest.mark.asyncio +async def test_dead_grant_recording_preserves_the_rows_current_material(): + """The verdict strips exactly the rejected grant; every other fact on the row (id_token, + issuer, expiry) is the row's own current material, re-read at recording time, never this + pod's stale snapshot of it.""" + stored = SSOIdentityAssertion( + id_token=SecretStr(_make_id_token(exp_offset=-100)), + refresh_token=SecretStr("rt_stored"), + issuer=ISSUER, + expires_at=datetime.now(timezone.utc) - timedelta(seconds=100), + ) + state = {"stored": stored} + persisted: list[tuple[str, SSOIdentityAssertion]] = [] + + async def fetch(user_id: str): + return state["stored"] + + async def persist(user_id, assertion): + persisted.append((user_id, assertion)) + state["stored"] = assertion + + async def post(url, form): + raise SsoAssertionUnrenewable("rejected") + + source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get) + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + assert len(persisted) == 1 + recorded = persisted[0][1] + assert recorded.refresh_token is None + assert recorded.id_token.get_secret_value() == stored.id_token.get_secret_value() + assert recorded.issuer == stored.issuer + assert recorded.expires_at == stored.expires_at + + +@pytest.mark.asyncio +async def test_recording_a_dead_grant_is_best_effort_when_the_store_is_unreachable(): + """Failing to record the verdict costs one more POST next time, never a wrong answer: the + caller still gets the re-login challenge (that verdict is proven), and nothing is written.""" + outage = {"armed": False} + persisted: list[tuple[str, SSOIdentityAssertion]] = [] + + async def fetch(user_id: str): + # The verdict recording's re-read is the first read after the rejected POST; killing + # exactly that one leaves every other layer's reads (initial, single-flight double-check, + # authority fall-through) working. + if outage["armed"]: + outage["armed"] = False + raise TokenStoreUnavailable("store went away") + return _stored(-100) + + async def persist(user_id, assertion): + persisted.append((user_id, assertion)) + + async def post(url, form): + outage["armed"] = True + raise SsoAssertionUnrenewable("rejected") + + source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get) + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + assert persisted == [] + + +@pytest.mark.asyncio +async def test_source_unavailable_refresh_keeps_the_grant_for_the_next_attempt(): + """A refresh failure that proves nothing about the grant (the gateway's own client config + rejected, the IdP down) must leave the stored refresh token in place, so the next attempt + after the operator's fix renews with no re-login.""" + new_id_token = _make_id_token(exp_offset=7200) + source, _, persisted, posts = _source( + _stored(-100), [None, {"id_token": new_id_token, "refresh_token": "rt_rotated"}] + ) + with pytest.raises(TokenStoreUnavailable): + await source.fetch_usable("alice") + lookup = await source.fetch_usable("alice") + assert lookup is not None + assert lookup.access_token == new_id_token + assert [form["refresh_token"] for _, form in posts] == ["rt_stored", "rt_stored"] + assert len(persisted) == 1 + + +@pytest.mark.asyncio +async def test_source_relogin_immediately_after_a_rejection_is_honored(): + """The 401 challenge tells the user to re-login; the fresh row must be served the moment it + lands (no clock advance, same pod). Nothing beside the row may hold the verdict: a rejection + is recorded ON the row, so a re-login lifts it by construction, and no per-pod state may + mask the authoritative read, or the recovery the challenge asks for silently fails.""" + posts: list[dict] = [] + state = {"stored": _stored(-100)} + + async def fetch(user_id: str): + return state["stored"] + + async def persist(user_id, assertion): + state["stored"] = assertion + + async def post(url, form): + posts.append(dict(form)) + raise SsoAssertionUnrenewable("rejected") + + source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get) + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + assert len(posts) == 1 + + state["stored"] = _stored(3600) # re-login lands right after the rejection, no clock advance + lookup = await source.fetch_usable("alice") + assert lookup is not None + assert len(posts) == 1 # the fresh usable row short-circuits before any refresh POST + + +@pytest.mark.asyncio +async def test_source_waiter_of_an_expired_winner_does_not_respend_the_grant(): + """The single-flight hand-off carries negative verdicts too: a waiter behind a winner whose + refresh was rejected observes the winner's raised verdict instead of re-POSTing the spent + token.""" + import asyncio + + posts: list[dict] = [] + + async def fetch(user_id: str): + await asyncio.sleep(0) + return _stored(-100) + + async def persist(user_id, assertion): + return None + + async def post(url, form): + posts.append(dict(form)) + await asyncio.sleep(0) + raise SsoAssertionUnrenewable("rejected") + + source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get) + results = await asyncio.gather(source.fetch_usable("alice"), source.fetch_usable("alice"), return_exceptions=True) + assert all(isinstance(r, SsoAssertionUnrenewable) for r in results) + assert len(posts) == 1 + + +@pytest.mark.asyncio +async def test_source_concurrent_expired_fetches_refresh_once(): + """Two flights hit the same expired assertion; the single-flight lock makes the second + re-read the row the first refreshed instead of spending the refresh token again. The fake + post yields to the loop so the flights genuinely interleave.""" + import asyncio + + fetches: list[str] = [] + posts: list[tuple[str, dict]] = [] + state = {"stored": _stored(-100)} + bodies = [{"id_token": _make_id_token(exp_offset=7200)}, {"id_token": _make_id_token(exp_offset=7200)}] + + async def fetch(user_id: str): + await asyncio.sleep(0) + fetches.append(user_id) + return state["stored"] + + async def persist(user_id: str, assertion: SSOIdentityAssertion): + await asyncio.sleep(0) + state["stored"] = assertion + + async def post(url: str, form: dict): + posts.append((url, dict(form))) + await asyncio.sleep(0) + return bodies.pop(0) + + source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get) + results = await asyncio.gather(source.fetch_usable("alice"), source.fetch_usable("alice")) + assert all(r is not None for r in results) + assert len(posts) == 1 + + +@pytest.mark.asyncio +async def test_source_memoizes_usable_lookups_within_ttl(): + """A usable assertion is served from the in-process memo, so repeated egress calls do not + pay a DB read per call; the memo expires with the TTL and the row is re-read.""" + clock = {"now": time.time()} + fetches: list[str] = [] + stored = _stored(3600) + + async def fetch(user_id: str): + fetches.append(user_id) + return stored + + async def persist(user_id, assertion): + return None + + async def post(url, form): + return None + + source = LiveSsoAssertionSource( + fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get, now=lambda: clock["now"], cache_ttl_seconds=60.0 + ) + first = await source.fetch_usable("alice") + second = await source.fetch_usable("alice") + assert first is not None and second is not None + assert fetches == ["alice"] + clock["now"] = clock["now"] + 61.0 + third = await source.fetch_usable("alice") + assert third is not None + assert fetches == ["alice", "alice"] + + +@pytest.mark.asyncio +async def test_a_dead_grant_verdict_never_refuses_a_caller_who_signed_in_again(): + """A rejected grant must not refuse the NEXT assertion, even one that is itself near expiry. + + The verdict is recorded on the row (its refresh token is stripped), so replacing the row at + re-login clears it by construction. A verdict parked in a side cache keyed by user would still + be inside its window here and would raise a false 401 at the moment the fresh assertion first + needs renewing, despite a perfectly good refresh token. + """ + state = {"stored": _stored(-100, refresh_token="rt_dead")} + outcomes = {"reject": True} + + async def fetch(user_id: str): + return state["stored"] + + async def persist(user_id, assertion): + state["stored"] = assertion + + async def post(url, form): + if outcomes["reject"]: + raise SsoAssertionUnrenewable("grant is dead") + return {"id_token": _make_id_token(exp_offset=3600), "refresh_token": "rt_fresh"} + + source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get) + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + # The verdict lives on the row now: the dead refresh token is gone, so a repeat costs no POST. + assert state["stored"].refresh_token is None + + # The user signs in again. The new assertion is ALSO near expiry, so renewal runs immediately; + # it must use the new grant rather than a remembered verdict. + outcomes["reject"] = False + state["stored"] = _stored(-5, refresh_token="rt_after_relogin") + renewed = await source.fetch_usable("alice") + assert renewed is not None + assert renewed.refresh_token == "rt_fresh" + + +@pytest.mark.asyncio +async def test_replacing_the_row_drops_a_cached_superseded_assertion(): + """A re-login must take effect at once, not when the superseded id_token finally expires. + + The assertion is identity material, so a user who re-authenticates after their IdP claims are + REDUCED must stop being able to spend the old one. ``persist_sso_identity_assertion`` is the + single place the row is replaced, so it drops the cached predecessor; without that hook the + warm entry below keeps serving the old id_token for its full remaining life. + """ + superseded = _make_id_token(exp_offset=3600, iss="https://old.example.com") + replacement = _make_id_token(exp_offset=3600, iss="https://new.example.com") + state = {"stored": SSOIdentityAssertion(id_token=SecretStr(superseded))} + + async def fetch(user_id: str): + return state["stored"] + + async def persist(user_id, assertion): + state["stored"] = assertion + + async def post(url, form): + raise TokenStoreUnavailable("unreachable") + + source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get) + warm = await source.fetch_usable("alice") + assert warm is not None + assert warm.access_token == superseded + + # The user signs in again; the row is replaced through the real write chokepoint. + state["stored"] = SSOIdentityAssertion(id_token=SecretStr(replacement)) + with patch("litellm.proxy.proxy_server.prisma_client", _make_prisma({})): + await persist_sso_identity_assertion("alice", state["stored"]) + + after = await source.fetch_usable("alice") + assert after is not None + assert after.access_token == replacement + + +@pytest.mark.asyncio +async def test_source_never_memoizes_absence_so_a_fresh_login_is_seen_immediately(): + fetches: list[str] = [] + state = {"stored": None} + + async def fetch(user_id: str): + fetches.append(user_id) + return state["stored"] + + async def persist(user_id, assertion): + return None + + async def post(url, form): + return None + + source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get) + assert await source.fetch_usable("alice") is None + state["stored"] = _stored(3600) + assert await source.fetch_usable("alice") is not None + assert fetches == ["alice", "alice"] + + +@pytest.mark.asyncio +async def test_source_store_outage_is_retryable_and_a_warm_cache_survives_it(): + """A warm cache keeps serving through an outage, and once cold the outage is RETRYABLE. + + It must never read as absence: this user has an assertion, so answering "you have never signed + in" would send them to a sign-in that fixes nothing while the real remedy is to retry. ``None`` + is reserved for the one case where no row was ever written. + """ + calls = {"n": 0} + stored = _stored(3600) + + async def flaky_fetch(user_id: str): + calls["n"] += 1 + if calls["n"] > 1: + raise RuntimeError("db down") + return stored + + async def persist(user_id, assertion): + return None + + async def post(url, form): + return None + + clock = {"now": time.time()} + source = LiveSsoAssertionSource( + fetch=flaky_fetch, + persist=persist, + post=post, + getenv=_SSO_ENV.get, + now=lambda: clock["now"], + cache_ttl_seconds=60.0, + ) + assert await source.fetch_usable("alice") is not None + assert await source.fetch_usable("alice") is not None + clock["now"] = clock["now"] + 61.0 + with pytest.raises(TokenStoreUnavailable): + await source.fetch_usable("alice") + + +@pytest.mark.asyncio +async def test_source_persist_failure_after_refresh_still_serves_the_in_hand_token(): + """A write-back failure is not a read failure: the freshly refreshed assertion is valid + and in hand, and discarding it would also strand a one-time rotated refresh token.""" + new_id_token = _make_id_token(exp_offset=7200) + posts: list[tuple[str, dict]] = [] + + async def fetch(user_id: str): + return _stored(-100) + + async def failing_persist(user_id, assertion): + raise RuntimeError("db write down") + + async def post(url, form): + posts.append((url, dict(form))) + return {"id_token": new_id_token, "refresh_token": "rt_rotated"} + + source = LiveSsoAssertionSource(fetch=fetch, persist=failing_persist, post=post, getenv=_SSO_ENV.get) + lookup = await source.fetch_usable("alice") + assert lookup is not None + assert lookup.access_token == new_id_token + assert len(posts) == 1 + + +@pytest.mark.asyncio +async def test_source_refresh_returning_an_already_expired_token_records_a_dead_grant(): + """Usable has ONE construction gate: a refreshed assertion passes the same expiry predicate a + stored one does, so an IdP handing back an already-dead token raises the re-login verdict + rather than being sent to an exchange. A short-lived token that lands expired every time would + thrash if its rotated refresh token were kept, so this is the same dead-grant outcome as a + missing id_token: the grant is recorded dead and re-login is the remedy, not another spend of + the rotation.""" + expired_new = _make_id_token(exp_offset=-10) + persisted: list[tuple[str, SSOIdentityAssertion]] = [] + + async def fetch(user_id: str): + return _stored(-100) + + async def persist(user_id, assertion): + persisted.append((user_id, assertion)) + + async def post(url, form): + return {"id_token": expired_new, "refresh_token": "rt_rotated"} + + source = LiveSsoAssertionSource(fetch=fetch, persist=persist, post=post, getenv=_SSO_ENV.get) + with pytest.raises(SsoAssertionUnrenewable): + await source.fetch_usable("alice") + assert len(persisted) == 1 + assert persisted[0][1].refresh_token is None + + +@pytest.mark.asyncio +async def test_source_waiter_gets_the_winners_token_when_persist_failed(): + """The memo is the single-flight hand-off: with the write-back failing and a one-time + refresh token, the waiter must receive the winner's in-hand assertion from the memo under + the lock instead of re-reading the stale row and re-spending the refresh grant.""" + import asyncio + + new_id_token = _make_id_token(exp_offset=7200) + posts: list[dict] = [] + + async def fetch(user_id: str): + await asyncio.sleep(0) + return _stored(-100) + + async def failing_persist(user_id, assertion): + raise RuntimeError("db write down") + + async def post(url, form): + posts.append(dict(form)) + await asyncio.sleep(0) + if form["refresh_token"] != "rt_stored" or len(posts) > 1: + raise SsoAssertionUnrenewable("rejected") + return {"id_token": new_id_token, "refresh_token": "rt_rotated"} + + source = LiveSsoAssertionSource(fetch=fetch, persist=failing_persist, post=post, getenv=_SSO_ENV.get) + results = await asyncio.gather(source.fetch_usable("alice"), source.fetch_usable("alice")) + assert all(r is not None for r in results) + assert len(posts) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status_code,body,expected_type", + [ + (400, {"error": "invalid_grant"}, SsoAssertionUnrenewable), + (401, {"error": "invalid_grant"}, SsoAssertionUnrenewable), + (401, {"error": "invalid_client"}, TokenStoreUnavailable), + (400, {"error": "invalid_scope"}, TokenStoreUnavailable), + (400, {"error": "invalid_request"}, TokenStoreUnavailable), + (429, {"error": "rate_limited"}, TokenStoreUnavailable), + (400, "not-json-object", TokenStoreUnavailable), + (400, None, TokenStoreUnavailable), + (502, {"error": "invalid_grant"}, TokenStoreUnavailable), + ], +) +async def test_post_helper_rejection_requires_a_proven_oauth_verdict(status_code, body, expected_type, monkeypatch): + """Rejected needs proof of a verdict ON THE GRANT, and only ``invalid_grant`` states one. + ``invalid_client``/``invalid_scope`` are the gateway's own client config failing — an ops + problem a re-login cannot fix, so treating them as grant death would destroy every renewable + assertion over a rotated GENERIC_CLIENT_SECRET. A rate limit, a 4xx without the error object + (an intermediary, not the token endpoint), or any 5xx likewise proves nothing about the + grant and must not send the user to re-login.""" + import json as _json + + import httpx + + from litellm.proxy._experimental.mcp_server.outbound_credentials import sso_assertion_store + + response = MagicMock() + response.status_code = status_code + if body is None: + response.json.side_effect = _json.JSONDecodeError("x", "y", 0) + else: + response.json.return_value = body + response.raise_for_status.side_effect = httpx.HTTPStatusError("err", request=MagicMock(), response=response) + client = MagicMock() + client.post = AsyncMock(return_value=response) + monkeypatch.setattr("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", lambda llm_provider: client) + with pytest.raises(expected_type): + await sso_assertion_store._post_sso_token_endpoint("https://idp.example.com/token", {"a": "b"}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [None, [], "string", 42]) +async def test_post_helper_2xx_non_object_body_is_unreachable_not_a_crash(body, monkeypatch): + """A 2xx whose body is not a JSON object must read Unreachable; letting it construct a + Granted outcome detonates downstream validation into the broad catch, which misreads the + state as Absent and misdirects the user to a 412 instead of a retryable failure.""" + client = MagicMock() + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + response.json.return_value = body + client.post = AsyncMock(return_value=response) + monkeypatch.setattr("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", lambda llm_provider: client) + from litellm.proxy._experimental.mcp_server.outbound_credentials import sso_assertion_store + + with pytest.raises(TokenStoreUnavailable): + await sso_assertion_store._post_sso_token_endpoint("https://idp.example.com/token", {"a": "b"}) + + +@pytest.mark.asyncio +async def test_post_helper_transport_failure_is_unreachable(monkeypatch): + client = MagicMock() + client.post = AsyncMock(side_effect=ConnectionError("down")) + monkeypatch.setattr("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", lambda llm_provider: client) + from litellm.proxy._experimental.mcp_server.outbound_credentials import sso_assertion_store + + with pytest.raises(TokenStoreUnavailable): + await sso_assertion_store._post_sso_token_endpoint("https://idp.example.com/token", {"a": "b"}) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py index 17a7e13af03..83e4af52e72 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py @@ -19,9 +19,7 @@ def _wrapping_codec(): def test_round_trips_the_access_token(): codec = _wrapping_codec() - token = codec.decode( - codec.encode(OAuthToken(access_token="at-123", expires_at=1234.5)) - ) + token = codec.decode(codec.encode(OAuthToken(access_token="at-123", expires_at=1234.5))) assert token is not None assert token.access_token == "at-123" @@ -30,9 +28,7 @@ def test_encode_encrypts_and_omits_the_refresh_token(): codec = _wrapping_codec() blob = codec.encode(OAuthToken(access_token="at", refresh_token="super-secret-rt")) assert blob.startswith("enc:") # encryption was applied - assert ( - "super-secret-rt" not in blob - ) # the long-lived secret never reaches the cache + assert "super-secret-rt" not in blob # the long-lived secret never reaches the cache decoded = codec.decode(blob) assert decoded is not None and decoded.refresh_token is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py index f100bd56f8f..3ea5e61e3a8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py @@ -35,10 +35,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ) from pydantic import SecretStr -_PATCH_TARGET = ( - "litellm.proxy._experimental.mcp_server.outbound_credentials." - "token_endpoint.get_async_httpx_client" -) +_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint.get_async_httpx_client" _ENDPOINT = "https://idp.example.com/oauth2/token" _CLIENT_ID = "litellm-client-id" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py index bb25ab6bd3c..cbdb59f3cd1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py @@ -98,9 +98,7 @@ def test_aws_sigv4_defaults_to_the_ambient_credential_chain(): def test_authconfig_discriminates_on_kind(): - api_key = _AUTH_CONFIG.validate_python( - {"kind": "api_key", "key_source": {"source": "shared", "value": "k"}} - ) + api_key = _AUTH_CONFIG.validate_python({"kind": "api_key", "key_source": {"source": "shared", "value": "k"}}) assert isinstance(api_key, ApiKeyConfig) assert isinstance(api_key.key_source, SharedKey) @@ -197,9 +195,7 @@ def test_id_jag_client_auth_discriminates_on_source(): def test_id_jag_client_auth_rejects_unknown_source(): with pytest.raises(ValidationError): - _AUTH_CONFIG.validate_python( - {**_ID_JAG_MINIMAL, "client_auth": {"source": "mystery"}} - ) + _AUTH_CONFIG.validate_python({**_ID_JAG_MINIMAL, "client_auth": {"source": "mystery"}}) def test_id_jag_config_defaults_id_token_subject_and_empty_optionals(): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_v2_token_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_v2_token_store.py index 97d14b65f5a..d00dd585133 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_v2_token_store.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_v2_token_store.py @@ -43,9 +43,7 @@ async def test_optional_fields_absent_yield_none(): @pytest.mark.asyncio async def test_unparseable_expiry_is_dropped_not_raised(): - store = V2PerUserTokenStore( - _reader({"access_token": "at", "expires_at": "not-a-date"}) - ) + store = V2PerUserTokenStore(_reader({"access_token": "at", "expires_at": "not-a-date"})) token = await store.fetch("alice", "srv") assert token is not None and token.expires_at is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index dff1f1d87c7..f9859b892da 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -805,6 +805,9 @@ async def test_mcp_get_prompt_success(): mcp_auth_header={"Authorization": "token"}, extra_headers={"X-Test": "1"}, raw_headers=None, + # An id_jag server resolves its stored SSO assertion by user id, so dropping the caller + # identity here would make prompts fail closed while tools/call still succeeded. + user_api_key_auth=user_api_key_auth, ) assert result is prompt_result @@ -866,6 +869,7 @@ async def test_mcp_read_resource_success(): mcp_auth_header={"Authorization": "token"}, extra_headers={"X-Test": "1"}, raw_headers=None, + user_api_key_auth=user_api_key_auth, ) assert result is read_result @@ -7726,3 +7730,60 @@ class TestPreemptive401ModeAware: await self._run(delegate, None, has_stored_token=False) assert exc.value.status_code == 401 await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False) + + +class TestPreemptiveIdJagPreflight: + """The transport-edge chokepoint runs the id_jag preflight on single-server routes only, + with the caller's raw headers, so a dead stored assertion surfaces its re-login 401 where a + WWW-Authenticate still reaches the client — and a multi-server aggregate keeps absorbing + per-server failures instead of letting one server 401 the whole connect.""" + + def _id_jag_server(self, alias: str): + return MCPServer( + server_id=f"id-{alias}", + name=alias, + alias=alias, + server_name=alias, + url=f"https://{alias}.test/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gw-client", + client_secret="gw-secret", + token_exchange_endpoint="https://org-idp.test/token", + id_jag_resource_token_endpoint="https://res-as.test/token", + mcp_info={"server_name": alias}, + ) + + async def _run(self, mcp_servers, server, raw_headers=None): + from litellm.proxy._experimental.mcp_server import server as server_module + + preflight = AsyncMock(return_value=None) + with ( + patch.object(server_module.global_mcp_server_manager, "get_mcp_server_by_name", return_value=server), + patch.object(server_module.global_mcp_server_manager, "preflight_id_jag", preflight), + ): + await server_module._raise_preemptive_401_for_unauthenticated_servers( + scope={"type": "http", "method": "POST", "path": f"/mcp/{server.alias}", "headers": []}, + mcp_servers=mcp_servers, + oauth2_headers=None, + mcp_server_auth_headers=None, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key"), + client_ip=None, + raw_headers=raw_headers, + ) + return preflight + + @pytest.mark.asyncio + async def test_single_server_route_runs_the_preflight_with_the_raw_headers(self): + server = self._id_jag_server("emasrv") + raw = {"x-litellm-api-key": "sk-litellm-virtual-key"} + preflight = await self._run([server.alias], server, raw_headers=raw) + preflight.assert_awaited_once() + assert preflight.await_args.kwargs["raw_headers"] == raw + assert preflight.await_args.kwargs["server"] is server + + @pytest.mark.asyncio + async def test_multi_server_aggregate_skips_the_preflight(self): + server = self._id_jag_server("emasrv") + preflight = await self._run([server.alias, "other-server"], server) + preflight.assert_not_awaited() 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 42b6cbee1c4..21c7b8bcdf5 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 @@ -1020,15 +1020,26 @@ class TestMCPServerManager: assert exc_info.value.www_authenticate == challenge @pytest.mark.asyncio - async def test_list_surfaces_non_auth_httpexception_as_internal_fault(self): - """A non-auth HTTP error (e.g. 412 no endpoint, 503 IdP down) now raises MCPServerListError - with an "internal" fault carrying the status code instead of absorbing to []: the silent - empty list made a misconfigured/unavailable server indistinguishable from a healthy server - with no tools. The aggregate absorbs it into that server's outcome, so one broken server - still does not blank the whole aggregate listing.""" + @pytest.mark.parametrize( + "status_code,expected_tag", + [ + (503, "unavailable"), # resolver's retryable verdict (store/IdP outage) stays retryable + (412, "precondition"), # resolver's per-caller precondition (e.g. gateway SSO sign-in) + (500, "internal"), # only the gateway's own failure reads as the gateway's fault + ], + ) + async def test_list_preserves_the_resolvers_verdict_class_on_non_auth_httpexceptions( + self, status_code, expected_tag + ): + """A non-auth HTTP error raises MCPServerListError with a fault that preserves the + resolver's retryable-vs-terminal verdict instead of absorbing to [] (the silent empty list + made a broken server indistinguishable from a healthy empty one) and instead of collapsing + everything to "internal" (which reported an IdP outage as the gateway's own 500). The + aggregate absorbs the fault into that server's outcome, so one broken server still does + not blank the whole aggregate listing.""" server = MCPServer( - server_id="te-412", - name="te-412-server", + server_id=f"te-{status_code}", + name=f"te-{status_code}-server", url="https://up.example.com", transport=MCPTransport.http, auth_type=MCPAuth.oauth2_token_exchange, @@ -1037,13 +1048,11 @@ class TestMCPServerManager: client_secret="csec", ) manager = MCPServerManager() - manager._create_mcp_client = AsyncMock( - side_effect=HTTPException(status_code=412, detail="token exchange endpoint is not configured") - ) + manager._create_mcp_client = AsyncMock(side_effect=HTTPException(status_code=status_code, detail="resolver")) with pytest.raises(MCPServerListError) as exc_info: await manager._get_tools_from_server(server=server, oauth2_headers={"Authorization": "Bearer subj-jwt"}) - assert exc_info.value.fault == ServerListFault(tag="internal", status_code=412) - assert exc_info.value.server_name == "te-412-server" + assert exc_info.value.fault == ServerListFault(tag=expected_tag, status_code=status_code) + assert exc_info.value.server_name == f"te-{status_code}-server" def _upstream_status_error(self, status_code: int, www_authenticate: Optional[str] = None) -> httpx.HTTPStatusError: """Build an httpx.HTTPStatusError shaped like the one the MCP SDK surfaces for an upstream @@ -1875,6 +1884,133 @@ class TestMCPServerManager: await manager.preflight_token_exchange(server=server, oauth2_headers=None, user_api_key_auth=None) assert resolved == ["good-subject"] + @staticmethod + def _id_jag_server(server_id: str) -> MCPServer: + return MCPServer( + server_id=server_id, + name=f"{server_id}-name", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://org-idp.example/oauth2/token", + id_jag_resource_token_endpoint="https://resource-as.example/oauth2/token", + ) + + @pytest.mark.asyncio + async def test_preflight_id_jag_surfaces_the_relogin_challenge_at_the_transport_edge(self): + """A dead stored assertion must raise the resolver's re-login 401 (challenge header + intact) from the preflight, so a single-server front-door client sees the remedy instead + of the session opening onto an empty tool/prompt/resource catalog.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + return Error( + CredError.of_unauthorized( + "the stored identity assertion has expired; sign in to the gateway again", + www_authenticate='Bearer error="invalid_token", error_description="re-authenticate"', + ) + ) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + with pytest.raises(HTTPException) as exc_info: + await manager.preflight_id_jag( + server=self._id_jag_server("idjag-preflight-401"), + oauth2_headers=None, + raw_headers={"x-litellm-api-key": "Bearer sk-admission"}, + user_api_key_auth=None, + ) + assert exc_info.value.status_code == 401 + headers = exc_info.value.headers or {} + www_authenticate = headers.get("WWW-Authenticate") or headers.get("www-authenticate") or "" + assert "invalid_token" in www_authenticate + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "error_factory,expected_status", + [ + (lambda CredError: CredError.of_upstream_unavailable("the assertion store could not be read"), 503), + (lambda CredError: CredError.of_precondition_required("no stored identity assertion"), 412), + ], + ) + async def test_preflight_id_jag_preserves_retryable_vs_terminal(self, error_factory, expected_status): + """A store/IdP outage must surface as the retryable 503 and a never-signed-in user as the + 412 — at the transport edge, with the same statuses the call surface uses, so the two + surfaces never tell the caller different stories.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + return Error(error_factory(CredError)) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + with pytest.raises(HTTPException) as exc_info: + await manager.preflight_id_jag( + server=self._id_jag_server("idjag-preflight-status"), + oauth2_headers=None, + raw_headers=None, + user_api_key_auth=None, + ) + assert exc_info.value.status_code == expected_status + + @pytest.mark.asyncio + async def test_preflight_id_jag_subject_follows_the_shared_bearer_rule(self): + """The preflight selects its subject with the same one rule egress uses + (`_subject_bearer_token`), including the free-rider guard: an Authorization bearer that + did not admit the caller is never handed to the resolver as exchange subject material.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + subjects = [] + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + subjects.append(subject.inbound_token.get_secret_value() if subject.inbound_token else None) + return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = self._id_jag_server("idjag-preflight-subject") + + # Admission consumed the explicit key header, so the Authorization bearer is a free + # rider: the preflight must resolve through the stored assertion (no inbound subject). + await manager.preflight_id_jag( + server=server, + oauth2_headers=None, + raw_headers={"x-litellm-api-key": "sk-admission", "authorization": "Bearer free-rider"}, + user_api_key_auth=None, + ) + # The bearer WAS the admission credential, so it is the subject. + await manager.preflight_id_jag( + server=server, + oauth2_headers=None, + raw_headers={"authorization": "Bearer caller-id-token"}, + user_api_key_auth=None, + ) + assert subjects == [None, "caller-id-token"] + + @pytest.mark.asyncio + async def test_preflight_id_jag_noop_for_other_auth_types(self): + """The preflight is mode-terminal: a non-id_jag server never reaches the resolver here + (its own preflight or challenge path owns it).""" + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + raise AssertionError("resolver must not run for non-id_jag servers") + + manager = MCPServerManager(cred_provider=_FakeProvider()) + await manager.preflight_id_jag( + server=self._token_exchange_server("te-not-id-jag"), + oauth2_headers={"Authorization": "Bearer subject"}, + raw_headers=None, + user_api_key_auth=None, + ) + @pytest.mark.asyncio async def test_call_regular_mcp_tool_passthrough_strips_authorization_when_admission_consumed_litellm_key( self, @@ -2670,6 +2806,7 @@ class TestMCPServerManager: extra_headers=None, stdio_env=None, subject_token=None, + user_api_key_auth=None, ) mock_client.list_resource_templates.assert_awaited_once() mock_prefix.assert_called_once_with(mock_templates, server, add_prefix=False) @@ -3443,7 +3580,7 @@ class TestMCPServerManager: # Capture the extra_headers passed to _create_mcp_client captured_extra_headers = None - async def capture_create_mcp_client(server, mcp_auth_header, extra_headers, stdio_env): + async def capture_create_mcp_client(server, mcp_auth_header, extra_headers, stdio_env, **kwargs): nonlocal captured_extra_headers captured_extra_headers = extra_headers return mock_client @@ -7573,7 +7710,7 @@ class TestHealthCheckInterpolatesGlobalEnvVars: mock_client = AsyncMock() mock_client.run_with_session = AsyncMock(return_value="ok") - async def _create(server, mcp_auth_header, extra_headers, stdio_env): + async def _create(server, mcp_auth_header, extra_headers, stdio_env, **kwargs): captured["extra_headers"] = extra_headers return mock_client @@ -9240,3 +9377,292 @@ class TestDiscoveryFailureLogging: assert "typo_row" in caplog.text assert "authorization_url, token_url" in caplog.text assert "unresolved" in caplog.text +@pytest.mark.parametrize( + "auth_type,subject_token,user_id,expected", + [ + (MCPAuth.oauth2_token_exchange, "hdr.jwt.sig", None, True), + (MCPAuth.oauth2_token_exchange, None, "alice", False), + (MCPAuth.oauth2_id_jag, "hdr.jwt.sig", None, True), + (MCPAuth.oauth2_id_jag, None, "alice", True), + (MCPAuth.oauth2_id_jag, None, None, False), + (MCPAuth.oauth2_id_jag, None, "", False), + (MCPAuth.oauth2, "hdr.jwt.sig", "alice", False), + (None, "hdr.jwt.sig", "alice", False), + ], +) +def test_obo_retry_covers_caller_matrix(auth_type, subject_token, user_id, expected): + """token_exchange re-mints only from the inbound token; id_jag also re-mints from the + stored SSO assertion keyed by user id, so an id_jag caller with a user identity but no + inbound assertion must still route through the invalidate-and-retry path.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _obo_retry_covers_caller + from litellm.proxy._types import UserAPIKeyAuth + + auth = UserAPIKeyAuth(api_key="sk-x", user_id=user_id) if user_id is not None else None + assert _obo_retry_covers_caller(auth_type, subject_token, auth) is expected + + +def test_obo_retry_requires_an_authenticated_caller_for_stored_assertion_retry(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _obo_retry_covers_caller + + assert _obo_retry_covers_caller(MCPAuth.oauth2_id_jag, None, None) is False + + +class TestSubjectBearerTokenSymmetry: + """One rule decides whether the caller's bearer becomes exchange subject material, and every + egress surface (tools call, tools list, prompts, resources) resolves through it; an id_jag + caller presenting a fresh IdP JWT must see identical sourcing on list and call.""" + + def _server(self, auth_type): + server = MagicMock() + server.auth_type = auth_type + return server + + @pytest.mark.parametrize( + "auth_type,expected", + [ + (MCPAuth.oauth2_token_exchange, "hdr.jwt.sig"), + (MCPAuth.oauth2_id_jag, "hdr.jwt.sig"), + (MCPAuth.oauth2, None), + (MCPAuth.api_key, None), + (None, None), + ], + ) + def test_subject_bearer_mode_matrix(self, auth_type, expected): + manager = MCPServerManager() + raw_headers = {"authorization": "Bearer hdr.jwt.sig"} + assert manager._subject_bearer_token(self._server(auth_type), raw_headers) == expected + + def test_oauth2_headers_win_over_raw(self): + manager = MCPServerManager() + token = manager._subject_bearer_token( + self._server(MCPAuth.oauth2_id_jag), + {"authorization": "Bearer raw.jwt.sig"}, + oauth2_headers={"Authorization": "Bearer oauth2.jwt.sig"}, + ) + assert token == "oauth2.jwt.sig" + + @pytest.mark.parametrize("key_header", ["x-litellm-api-key", "X-Litellm-Api-Key"]) + def test_id_jag_ignores_a_free_rider_jwt_when_admission_used_the_explicit_key(self, key_header): + """A caller admitted with the explicit litellm key can ride any valid IdP JWT in + Authorization; it was never validated as THIS caller's identity, so exchanging it would + let the caller act upstream as whoever the token belongs to. The bearer is subject + material only when it was the admission credential; such callers resolve through the + stored assertion, which is bound to the authenticated user.""" + manager = MCPServerManager() + raw_headers = {key_header: "sk-caller-a", "authorization": "Bearer stolen.user-b.jwt"} + assert manager._subject_bearer_token(self._server(MCPAuth.oauth2_id_jag), raw_headers) is None + assert ( + manager._subject_bearer_token( + self._server(MCPAuth.oauth2_id_jag), + raw_headers, + oauth2_headers={"Authorization": "Bearer stolen.user-b.jwt"}, + ) + is None + ) + + def test_id_jag_uses_the_bearer_when_it_was_the_admission_credential(self): + manager = MCPServerManager() + raw_headers = {"authorization": "Bearer admitted.idp.jwt"} + assert ( + manager._subject_bearer_token(self._server(MCPAuth.oauth2_id_jag), raw_headers) == "admitted.idp.jwt" + ) + + @pytest.mark.parametrize( + "raw_headers", + [ + {"x-litellm-api-key": "", "authorization": "Bearer admitted.idp.jwt"}, + {"X-Litellm-Api-Key": "", "authorization": "Bearer admitted.idp.jwt"}, + ], + ) + def test_id_jag_empty_explicit_key_is_not_a_free_rider(self, raw_headers): + """Admission skips a falsy primary header and authenticates with Authorization, so an + empty explicit key means the bearer IS the admission credential and stays usable.""" + manager = MCPServerManager() + assert ( + manager._subject_bearer_token(self._server(MCPAuth.oauth2_id_jag), raw_headers) == "admitted.idp.jwt" + ) + + @pytest.mark.parametrize( + "raw_headers", + [ + None, + {}, + {"authorization": "Bearer j.w.t"}, + {"x-litellm-api-key": "", "authorization": "Bearer j.w.t"}, + {"x-litellm-api-key": "sk-a", "authorization": "Bearer j.w.t"}, + {"X-LiteLLM-Api-Key": "sk-a", "authorization": "Bearer j.w.t"}, + {"x-litellm-api-key": " ", "authorization": "Bearer j.w.t"}, + {"x-litellm-api-key": "sk-a"}, + ], + ) + def test_free_rider_predicate_agrees_with_the_admission_accessor(self, raw_headers): + """The predicate and get_litellm_api_key_from_headers must encode ONE header-preference + rule: the bearer is a free rider exactly when admission's chosen credential is NOT the + Authorization value. A mutation to either side breaks this agreement.""" + from starlette.datastructures import Headers + + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + + headers = Headers(headers=dict(raw_headers or {})) + admission_credential = MCPRequestHandler.get_litellm_api_key_from_headers(headers) + authorization_value = headers.get("authorization") + expected_free_rider = ( + admission_credential != authorization_value if authorization_value is not None else bool(admission_credential) + ) + assert MCPRequestHandler.authorization_is_free_rider(raw_headers) == expected_free_rider + + @pytest.mark.asyncio + async def test_openapi_egress_uses_the_same_subject_bearer_rule(self): + """spec_path servers egress outside _create_mcp_client, so their extraction is the + eighth surface; it must resolve through the identical rule or the binding invariant + holds on MCP paths while the OpenAPI path exchanges a free-rider token.""" + manager = MCPServerManager() + server = MagicMock() + server.auth_type = MCPAuth.oauth2_id_jag + server.server_id = "idjag-openapi" + server.url = None + server.spec_path = "https://spec.example.com/openapi.yaml" + + captured: dict[str, object] = {} + + def fake_rule(mcp_server, raw_headers, oauth2_headers=None): + captured["called_with"] = (raw_headers, oauth2_headers) + return None + + raw = {"x-litellm-api-key": "sk-a", "authorization": "Bearer free.rider.jwt"} + with ( + patch.object(manager, "_subject_bearer_token", side_effect=fake_rule), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.to_server_spec" + ) as spec_mock, + ): + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientSecretAuth, + IdJagConfig, + ServerSpec, + ) + from pydantic import SecretStr + + spec_mock.return_value = ServerSpec( + server_id="idjag-openapi", + resource="https://up.example.com", + config=IdJagConfig( + org_token_endpoint="https://idp.example.com/token", + resource_token_endpoint="https://ras.example.com/token", + client_id="gw", + client_auth=ClientSecretAuth(client_secret=SecretStr("s")), + ), + ) + with patch.object( + manager._cred_provider, "resolve_credentials", new=AsyncMock(side_effect=RuntimeError("stop")) + ): + try: + await manager.resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers={"Authorization": "Bearer free.rider.jwt"}, + raw_headers=raw, + mcp_auth_header=None, + user_api_key_auth=None, + forwarded_headers=None, + ) + except Exception: + pass + + assert captured.get("called_with") == (raw, {"Authorization": "Bearer free.rider.jwt"}) + + def test_obo_keeps_exchange_what_was_presented_semantics(self): + manager = MCPServerManager() + raw_headers = {"x-litellm-api-key": "sk-caller-a", "authorization": "Bearer presented.obo.jwt"} + assert ( + manager._subject_bearer_token(self._server(MCPAuth.oauth2_token_exchange), raw_headers) + == "presented.obo.jwt" + ) + + @pytest.mark.asyncio + async def test_list_path_threads_the_id_jag_subject_token(self): + manager = MCPServerManager() + server = MagicMock() + server.auth_type = MCPAuth.oauth2_id_jag + server.server_id = "idjag-1" + server.name = "idjag_srv" + server.alias = "idjag_srv" + server.transport = MCPTransport.http + server.spec_path = None + + captured: dict[str, object] = {} + + async def fake_create_client(**kwargs): + captured.update(kwargs) + raise RuntimeError("stop after capture") + + from litellm.proxy._experimental.mcp_server.exceptions import MCPServerListError + + manager._create_mcp_client = fake_create_client + with patch.object(manager, "_build_stdio_env", return_value=None): + with pytest.raises(MCPServerListError): + await manager._get_tools_from_server( + server, + mcp_auth_header=None, + oauth2_headers={"Authorization": "Bearer caller.idp.jwt"}, + raw_headers={"authorization": "Bearer caller.idp.jwt"}, + user_api_key_auth=None, + ) + + assert captured.get("subject_token") == "caller.idp.jwt" + + +class TestEgressSurfacesThreadCallerIdentity: + """Every id_jag/OBO egress surface must forward the caller's ``user_api_key_auth`` into + ``_create_mcp_client`` so the stored SSO assertion resolves to the right user. A surface that + drops it silently breaks id_jag for virtual-key callers; prompts and resources did exactly + that. This pins the whole prompt/resource family so a surface cannot regress the threading.""" + + def _server(self): + server = MagicMock() + server.static_headers = None + server.transport = MCPTransport.http + server.url = "https://up.example.com/mcp" + server.name = "srv" + server.alias = "srv" + server.server_name = "srv" + server.server_id = "srv-1" + return server + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "invoke", + [ + lambda mgr, srv, uk: mgr.get_prompts_from_server(server=srv, raw_headers={}, user_api_key_auth=uk), + lambda mgr, srv, uk: mgr.get_resources_from_server(server=srv, raw_headers={}, user_api_key_auth=uk), + lambda mgr, srv, uk: mgr.get_resource_templates_from_server(server=srv, raw_headers={}, user_api_key_auth=uk), + lambda mgr, srv, uk: mgr.get_prompt_from_server( + server=srv, prompt_name="p", raw_headers={}, user_api_key_auth=uk + ), + lambda mgr, srv, uk: mgr.read_resource_from_server( + server=srv, url="https://up.example.com/r", raw_headers={}, user_api_key_auth=uk + ), + ], + ids=["list_prompts", "list_resources", "list_resource_templates", "get_prompt", "read_resource"], + ) + async def test_prompt_and_resource_surfaces_thread_user_api_key_auth(self, invoke): + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + caller = UserAPIKeyAuth(api_key="sk-caller", user_id="alice") + captured: dict = {} + + async def capture(**kwargs): + captured.update(kwargs) + return AsyncMock() + + with ( + patch.object(manager, "_create_mcp_client", side_effect=capture), + patch.object(manager, "_build_stdio_env", return_value=None), + patch.object(manager, "_subject_bearer_token", return_value=None), + ): + try: + await invoke(manager, self._server(), caller) + except Exception: + pass + + assert captured.get("user_api_key_auth") is caller