diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index baa4a5a9e0c1..8227a34e5927 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -71,6 +71,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import ( LazyPerUserOAuthTokenStore, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_provider import ( + build_token_exchanger, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( AuthorizationCodeConfig, ) @@ -531,7 +534,8 @@ def _resolve_oauth2_flow( def __init__(self, cred_provider: Optional[UpstreamCredentialProvider] = None): self._cred_provider = cred_provider or UpstreamCredentialProvider( - oauth_token_store=LazyPerUserOAuthTokenStore(self.get_mcp_server_by_id) + oauth_token_store=LazyPerUserOAuthTokenStore(self.get_mcp_server_by_id), + token_exchanger=build_token_exchanger(), ) self.registry: dict[str, MCPServer] = {} self.config_mcp_servers: dict[str, MCPServer] = {} @@ -1977,9 +1981,11 @@ async def _create_mcp_client( ): resolved_auth = None case Error(err): - if err.tag == "unauthorized": - # The arm signals a missing per-user token semantically; raise the - # per-server OAuth challenge here, where the full MCPServer is in hand. + if err.tag == "unauthorized" and isinstance(spec.config, AuthorizationCodeConfig): + # authorization_code's missing per-user token -> the per-server + # browser-OAuth challenge, built here where the full MCPServer is in + # hand. token_exchange and other modes carry their own 401 (e.g. OBO + # needs a caller token, not a browser flow), so they go via raise_public. raise_user_oauth_challenge(server) raise_public(err) return MCPClient( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 815fc2ba29dd..732c7185dbc9 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -26,6 +26,7 @@ ServerSpec, SharedKey, Subject, + TokenExchangeConfig, ) from litellm.types.mcp import MCPAuth @@ -61,8 +62,9 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is explicitly mapped or explicitly deferred, rather than silently falling through to v1. Live modes: ``none``, the static-header family (``api_key`` plus the Authorization schemes, - all shared-key), and ``oauth2`` per-user tokens (``authorization_code``); client_credentials - (M2M), delegated/passthrough oauth2, token exchange, and SigV4 return None and stay on v1. + all shared-key), ``oauth2`` per-user tokens (``authorization_code``), and + ``oauth2_token_exchange`` (RFC 8693 OBO); client_credentials (M2M), delegated/passthrough + oauth2, and SigV4 return None and stay on v1. """ if server.is_byok: return None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type) @@ -92,11 +94,39 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: ) # client_credentials (M2M) and delegate/passthrough oauth2 stay on v1 return None - case MCPAuth.oauth2_token_exchange | MCPAuth.aws_sigv4: - return None # token exchange and SigV4 are not migrated yet -> defer to v1 + case MCPAuth.oauth2_token_exchange: + return _token_exchange_spec(server, resource) + case MCPAuth.aws_sigv4: + return None # SigV4 is not migrated yet -> defer to v1 assert_never(auth_type) +def _token_exchange_spec(server: MCPServer, resource: str) -> Optional[ServerSpec]: + """Build a token_exchange (RFC 8693 OBO) spec, or defer (None) if the exchange config is absent. + + Mirrors v1's ``has_token_exchange_config`` precondition: an endpoint (``token_exchange_endpoint`` + or ``token_url``) plus ``client_id``/``client_secret`` must all be present, else there is nothing + to exchange against and the server stays on v1 (parity-safe). ``audience`` is forwarded only when + the operator set it; a missing one is omitted, not derived. + """ + endpoint = server.token_exchange_endpoint or server.token_url + if not endpoint or not server.client_id or not server.client_secret: + return None + return ServerSpec( + server_id=server.server_id, + resource=resource, + config=TokenExchangeConfig( + subject_token_type=server.subject_token_type, + token_exchange_endpoint=endpoint, + audience=server.audience, + client_id=server.client_id, + client_secret=SecretStr(server.client_secret), + token_endpoint_auth_method=server.token_endpoint_auth_method, + scopes=tuple(server.scopes or ()), + ), + ) + + def _shared_key_spec( server: MCPServer, resource: str, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index f9a9fa00b239..d94f442019cf 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -8,7 +8,8 @@ at runtime instead of returning `None`. `none` and `api_key` (shared-key source) are live, as is `authorization_code`, which reads the -user's token from the injected `OAuthTokenStore`. The remaining arms are `not_implemented` stubs +user's token from the injected `OAuthTokenStore`, and `token_exchange`, which swaps the caller's +inbound token through the injected `TokenExchanger`. The remaining arms are `not_implemented` stubs that each land in a follow-up PR with their seam. Pure v2: no imports from v1. """ @@ -31,6 +32,9 @@ Ok, Result, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( + TokenExchanger, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, @@ -55,16 +59,31 @@ async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: return None +class _NullTokenExchanger: + """Fail-closed default: with no exchanger wired, token_exchange cannot produce a credential.""" + + async def exchange( + self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig + ) -> Result[OAuthToken, CredError]: + return Error(CredError.of_misconfigured("token exchange collaborator not wired")) + + class UpstreamCredentialProvider: """Produces the one `httpx.Auth` for a `(subject, upstream)` pair, per declared mode. Collaborators (the per-mode credential stores and token fetchers) are injected as each arm is built; the live `none` and `api_key`-shared arms read from the config and need none, while - `authorization_code` reads the user's token from the injected `OAuthTokenStore`. + `authorization_code` reads the user's token from the injected `OAuthTokenStore` and + `token_exchange` swaps the caller's token through the injected `TokenExchanger`. """ - def __init__(self, oauth_token_store: OAuthTokenStore | None = None) -> None: + def __init__( + self, + oauth_token_store: OAuthTokenStore | None = None, + token_exchanger: TokenExchanger | None = None, + ) -> None: self._oauth_token_store: OAuthTokenStore = oauth_token_store or _NullOAuthTokenStore() + self._token_exchanger: TokenExchanger = token_exchanger or _NullTokenExchanger() async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: @@ -76,8 +95,8 @@ async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Res return _not_implemented(AuthSpecKind.passthrough) case ClientCredentialsConfig(): return _not_implemented(AuthSpecKind.client_credentials) - case TokenExchangeConfig(): - return _not_implemented(AuthSpecKind.token_exchange) + case TokenExchangeConfig() as config: + return await self._token_exchange(subject, server, config) case AuthorizationCodeConfig(): return await self._authorization_code(subject, server) case AwsSigV4Config(): @@ -110,6 +129,29 @@ async def _authorization_code(self, subject: Subject, server: ServerSpec) -> Res return Error(CredError.of_unauthorized("Authorization required: complete the OAuth flow for this server.")) return Ok(StaticHeaderAuth(f"Bearer {token.access_token}", header_name="Authorization")) + async def _token_exchange( + self, subject: Subject, server: ServerSpec, config: TokenExchangeConfig + ) -> Result[StaticHeaderAuth, CredError]: + """RFC 8693 OBO: exchange the caller's inbound token for an upstream-bound bearer. + + 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. + """ + inbound = subject.inbound_token + if inbound is None: + return Error( + CredError.of_unauthorized( + "Token exchange requires a caller token to exchange (OBO).", + www_authenticate='Bearer error="invalid_request"', + ) + ) + match await self._token_exchanger.exchange(inbound.get_secret_value(), server, config): + case Ok(token): + return Ok(StaticHeaderAuth(f"Bearer {token.access_token}", header_name="Authorization")) + case Error(err): + return Error(err) + 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/token_exchange_provider.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py new file mode 100644 index 000000000000..f370cba702a7 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py @@ -0,0 +1,62 @@ +"""Composition root for the v2-native token_exchange (OBO) exchanger. + +Wires the pure ``Rfc8693TokenExchanger`` to its runtime edges: the real httpx POST against the IdP and +the configured cache sizing/TTL constants. ``build_token_exchanger`` is built once at egress +construction and reused, so the in-process exchanged-token cache survives across requests. Unlike the +per-user store, nothing here reads a runtime global at build time (the httpx client is acquired per +call), so it needs no lazy wrapper. +""" + +from __future__ import annotations + +from litellm._logging import verbose_logger +from litellm.constants import ( + MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, + MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, + MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + InMemoryTokenCacheBackend, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( + Rfc8693TokenExchanger, +) + + +async def _post_exchange_endpoint( + url: str, form: dict[str, str], client_auth_headers: dict[str, str] +) -> dict[str, object] | None: + from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 + get_async_httpx_client, # pyright: ignore + ) + from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 + + # litellm's httpx handler and httpx.Response are only partially typed; the IdP returns a JSON + # object and the exchanger validates each field, so the untyped boundary is contained here. + # A failed exchange is a miss, not a 500 (matches v1), so any error becomes None. + headers = {"Accept": "application/json", **client_auth_headers} + try: + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) # pyright: ignore + response = await client.post(url, headers=headers, data=form) # pyright: ignore + response.raise_for_status() # pyright: ignore + parsed: object = response.json() # pyright: ignore + except Exception as exc: # noqa: BLE001 + verbose_logger.warning("MCP token exchange request failed: %s", exc) + return None + if not isinstance(parsed, dict): + # A valid-but-non-object JSON body (list/string/number) would crash the field parsing; map it + # to a miss so it surfaces as a typed upstream_unavailable, not a 500. + verbose_logger.warning("MCP token exchange returned non-object JSON (%s)", type(parsed).__name__) + return None + return parsed # pyright: ignore + + +def build_token_exchanger() -> Rfc8693TokenExchanger: + return Rfc8693TokenExchanger( + _post_exchange_endpoint, + cache=InMemoryTokenCacheBackend(max_size=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE), + default_ttl_seconds=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, + min_ttl_seconds=MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, + expiry_buffer_seconds=MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchanger.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchanger.py new file mode 100644 index 000000000000..9b7795b918db --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchanger.py @@ -0,0 +1,239 @@ +"""v2-native RFC 8693 token exchange (OBO): swap the caller's token for an upstream one. + +The pure core of the ``token_exchange`` mode. Given the caller's ``subject_token`` and the server's +``TokenExchangeConfig``, ``Rfc8693TokenExchanger.exchange`` POSTs the RFC 8693 token-exchange grant to +the configured endpoint and returns the upstream-bound ``access_token`` as a typed ``OAuthToken``, or a +typed ``CredError`` - never a raise (the HTTP edge is the injected ``ExchangeHttpPost``, whose adapter +contains the I/O). The exchanged token is cached and single-flighted per ``(subject_token, server)`` so +a repeated caller token skips the IdP round-trip and concurrent calls collapse to one exchange, reusing +the shared in-process cache + coordinator foundation. A rotated caller token hashes to a new key and +re-exchanges. Pure v2 apart from the shared RFC 6749 client-auth helper, which carries no v1 state. + +A missing/expired exchange is an error, never a fall-through to a weaker source (§1.5): the caller +presenting no token is the resolver arm's 401, and an IdP that does not return a usable token is an +``upstream_unavailable`` here. +""" + +from __future__ import annotations + +import hashlib +import time +from collections.abc import Awaitable, Callable +from typing import Protocol + +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + InMemoryTokenCacheBackend, + InProcessRefreshCoordinator, + OAuthToken, + RefreshCoordinator, + TokenCacheBackend, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( + build_token_endpoint_client_auth, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + CredError, + ServerSpec, + TokenExchangeConfig, +) + +# A token with no declared expiry is cached for this long; one with an expiry is cached until then +# minus the skew buffer, floored at the minimum. Values mirror v1's MCP_OAUTH2_* constants; the +# composition root injects the configured ones. +_DEFAULT_TTL_SECONDS = 3600.0 +_MIN_TTL_SECONDS = 10.0 +_EXPIRY_BUFFER_SECONDS = 60.0 + +_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" + +# RFC 8693 3 token-type URNs that are not usable as an upstream Bearer access token. token_type +# already rejects the common non-access case (N_A); this catches a malformed STS that mints one of +# these but still labels it Bearer. An access_token / jwt / absent / unknown type is accepted (lenient). +_NON_ACCESS_ISSUED_TOKEN_TYPES = frozenset( + { + "urn:ietf:params:oauth:token-type:refresh_token", + "urn:ietf:params:oauth:token-type:id_token", + "urn:ietf:params:oauth:token-type:saml1", + "urn:ietf:params:oauth:token-type:saml2", + } +) + +# The IdP returns an opaque JSON object; the post adapter hands it over untyped and the exchanger +# validates each field, so no Any leaks past this seam (None == any transport/HTTP failure). The +# second dict is the form body; the third is the client-auth headers (HTTP Basic for +# client_secret_basic, empty for client_secret_post). +ExchangeHttpPost = Callable[[str, "dict[str, str]", "dict[str, str]"], Awaitable["dict[str, object] | None"]] + + +class TokenExchanger(Protocol): + """Exchanges a caller token for an upstream-bound one, per the server's token_exchange config.""" + + async def exchange( + self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig + ) -> Result[OAuthToken, CredError]: ... + + +def _cache_key(subject_token: str, config: TokenExchangeConfig) -> str: + """Bind the cache entry to the caller token AND the exchange config that minted it. + + A rotated caller token, endpoint, audience, scope, client_id, secret, auth method, or + subject_token_type all change the key, so a config change forces a fresh exchange instead of + serving a token minted for the old config until TTL. Everything is hashed, so no secret is held + in the key. + """ + secret = config.client_secret.get_secret_value() if config.client_secret else "" + material = "\x00".join( + ( + subject_token, + config.token_exchange_endpoint or "", + config.audience or "", + config.subject_token_type, + config.client_id or "", + secret, + config.token_endpoint_auth_method or "", + " ".join(config.scopes), + ) + ) + return hashlib.sha256(material.encode()).hexdigest() + + +def _parse_expires_in(raw: object) -> int | None: + if isinstance(raw, bool): + return None + if isinstance(raw, int): + return raw + if isinstance(raw, str): + try: + return int(raw) + except ValueError: + return None + return None + + +def _build_exchange_form( + *, + subject_token: str, + subject_token_type: str, + audience: str | None, + scopes: tuple[str, ...], +) -> dict[str, str]: + return { + "grant_type": _GRANT_TYPE, + "subject_token": subject_token, + "subject_token_type": subject_token_type, + **({"audience": audience} if audience else {}), + **({"scope": " ".join(scopes)} if scopes else {}), + } + + +class Rfc8693TokenExchanger: + """``TokenExchanger`` that runs the RFC 8693 grant once per caller token, then caches the result. + + The HTTP post is injected (``None`` on any IdP failure, mirroring v1: a failed exchange is a miss, + not a 500). The cache and single-flight coordinator default to the in-process foundation; a + deployment with no shared state needs nothing more (v1's exchanged-token cache is per-process too). + The clock is injected so TTL/expiry is deterministic in tests. + """ + + def __init__( + self, + http_post: ExchangeHttpPost, + *, + cache: TokenCacheBackend | None = None, + coordinator: RefreshCoordinator | None = None, + clock: Callable[[], float] = time.time, + default_ttl_seconds: float = _DEFAULT_TTL_SECONDS, + min_ttl_seconds: float = _MIN_TTL_SECONDS, + expiry_buffer_seconds: float = _EXPIRY_BUFFER_SECONDS, + ) -> None: + self._http_post = http_post + self._cache: TokenCacheBackend = cache or InMemoryTokenCacheBackend(clock=clock) + self._coordinator: RefreshCoordinator = coordinator or InProcessRefreshCoordinator() + self._clock = clock + self._default_ttl_seconds = default_ttl_seconds + self._min_ttl_seconds = min_ttl_seconds + self._expiry_buffer_seconds = expiry_buffer_seconds + + async def exchange( + self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig + ) -> Result[OAuthToken, CredError]: + endpoint = config.token_exchange_endpoint + client_id = config.client_id + client_secret = config.client_secret + if not endpoint or not client_id or client_secret is None: + return Error( + CredError.of_misconfigured( + "token_exchange requires token_exchange_endpoint, client_id and client_secret" + ) + ) + + cache_key = _cache_key(subject_token, config) + server_id = server.server_id + cached = await self._cache.get(cache_key, server_id) + if cached is not None: + return Ok(cached) + + client_auth = build_token_endpoint_client_auth( + auth_method=config.token_endpoint_auth_method, + client_id=client_id, + client_secret=client_secret.get_secret_value(), + ) + form = { + **_build_exchange_form( + subject_token=subject_token, + subject_token_type=config.subject_token_type, + audience=config.audience, + scopes=config.scopes, + ), + **client_auth.body, + } + + async def run_exchange() -> OAuthToken | None: + fresh = await self._cache.get(cache_key, server_id) + if fresh is not None: + return fresh + body = await self._http_post(endpoint, form, client_auth.headers) + if body is None: + return None + token = self._token_from_body(body) + if token is None: + return None + await self._cache.set(cache_key, server_id, token, self._ttl_seconds(token)) + return token + + async def reread() -> OAuthToken | None: + return await self._cache.get(cache_key, server_id) + + token = await self._coordinator.run(cache_key, server_id, refresh=run_exchange, reread=reread) + if token is None: + return Error(CredError.of_upstream_unavailable("token exchange did not return a usable access token")) + return Ok(token) + + def _token_from_body(self, body: dict[str, object]) -> OAuthToken | None: + access_token = body.get("access_token") + if not isinstance(access_token, str) or not access_token: + return None + # token_type is forwarded downstream as Bearer, so a present-but-non-Bearer type (e.g. N_A) + # must fail closed rather than be minted as a bogus Bearer; an absent type defaults to Bearer. + token_type = body.get("token_type") + if isinstance(token_type, str) and token_type.strip().lower() != "bearer": + return None + # issued_token_type says what representation was minted; reject a clearly-non-access type + # (refresh/id/saml) even if token_type claimed Bearer. access_token / jwt / absent / unknown pass. + issued_token_type = body.get("issued_token_type") + if isinstance(issued_token_type, str) and issued_token_type in _NON_ACCESS_ISSUED_TOKEN_TYPES: + return None + expires_in = _parse_expires_in(body.get("expires_in")) + expires_at = self._clock() + expires_in if expires_in is not None else None + return OAuthToken(access_token=access_token, expires_at=expires_at) + + def _ttl_seconds(self, token: OAuthToken) -> float: + if token.expires_at is None: + return self._default_ttl_seconds + lifetime = token.expires_at - self._clock() + return max(lifetime - self._expiry_buffer_seconds, self._min_ttl_seconds) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 671de63eabe1..d475f931a259 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -183,17 +183,23 @@ class ClientCredentialsConfig(BaseModel): class TokenExchangeConfig(BaseModel): """RFC 8693 OBO; swap the caller's live subject_token for a token bound to the upstream's - audience (`server.resource`, RFC 8707). The gateway authenticates to the exchange endpoint - as an OAuth client (`client_id`/`client_secret`); the inbound token is sent only to that - endpoint, never to the upstream. + audience. The gateway authenticates to the exchange endpoint as an OAuth client + (`client_id`/`client_secret`); the inbound token is sent only to that endpoint, never to the + upstream. + + `audience` is the RFC 8693 target; it is optional and sent only when the operator configured + one, since both `audience` and `resource` are optional in the spec and the authorization server + applies its own default when neither is sent (fabricating one risks `invalid_target`). """ model_config = ConfigDict(frozen=True) kind: Literal[AuthSpecKind.token_exchange] = AuthSpecKind.token_exchange subject_token_type: str = "urn:ietf:params:oauth:token-type:access_token" token_exchange_endpoint: str | None = None + audience: str | None = None client_id: str | None = None client_secret: SecretStr | None = None + token_endpoint_auth_method: Literal["client_secret_basic", "client_secret_post"] | None = None scopes: tuple[str, ...] = () 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 383dc2556074..a448af165154 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 @@ -24,6 +24,7 @@ CredError, NoneConfig, SharedKey, + TokenExchangeConfig, ) from litellm.types.mcp import MCPAuth, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -65,9 +66,7 @@ def test_authorization_schemes_map_with_their_prefix(auth_type, prefix): def test_basic_scheme_base64_encodes_the_token(): - spec = to_server_spec( - _server(auth_type=MCPAuth.basic, authentication_token="user:pass") - ) + spec = to_server_spec(_server(auth_type=MCPAuth.basic, authentication_token="user:pass")) assert spec is not None and isinstance(spec.config, ApiKeyConfig) assert spec.config.value_prefix == "Basic" expected = base64.b64encode(b"user:pass").decode() @@ -89,17 +88,16 @@ def test_oauth2_user_token_maps_to_authorization_code(oauth2_flow): [ _server(auth_type=MCPAuth.api_key), # no token configured _server(auth_type=MCPAuth.bearer_token), # no token configured + _server(auth_type=MCPAuth.oauth2, oauth2_flow="client_credentials"), # M2M -> v1 + _server(auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True), # delegated upstream OAuth -> v1 + _server(auth_type=MCPAuth.oauth2_token_exchange), # no endpoint/client creds -> incomplete -> v1 _server( - auth_type=MCPAuth.oauth2, oauth2_flow="client_credentials" - ), # M2M -> v1 - _server( - auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True - ), # delegated upstream OAuth -> v1 - _server(auth_type=MCPAuth.oauth2_token_exchange), + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp/token", + client_id="cid", + ), # missing client_secret -> incomplete -> v1 _server(auth_type=MCPAuth.aws_sigv4), - _server( - auth_type=None, oauth_passthrough=True, extra_headers=["Authorization"] - ), + _server(auth_type=None, oauth_passthrough=True, extra_headers=["Authorization"]), ], ) def test_unmigrated_modes_defer_to_v1(server): @@ -107,6 +105,59 @@ def test_unmigrated_modes_defer_to_v1(server): assert to_server_spec(server) is None +def test_token_exchange_maps_full_config(): + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2_token_exchange, + url="https://up.example.com/mcp", + token_exchange_endpoint="https://idp.example.com/token", + audience="https://up.example.com", + client_id="cid", + client_secret="csec", + subject_token_type="urn:ietf:params:oauth:token-type:jwt", + scopes=["a", "b"], + ) + ) + assert spec is not None + config = spec.config + assert isinstance(config, TokenExchangeConfig) + assert config.token_exchange_endpoint == "https://idp.example.com/token" + assert config.audience == "https://up.example.com" + assert config.client_id == "cid" + assert config.client_secret is not None + assert config.client_secret.get_secret_value() == "csec" + assert config.subject_token_type == "urn:ietf:params:oauth:token-type:jwt" + assert config.scopes == ("a", "b") + + +def test_token_exchange_falls_back_to_token_url_when_no_exchange_endpoint(): + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2_token_exchange, + token_url="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + ) + assert spec is not None + assert isinstance(spec.config, TokenExchangeConfig) + assert spec.config.token_exchange_endpoint == "https://idp.example.com/token" + + +def test_token_exchange_omits_audience_when_unset(): + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + ) + assert spec is not None + assert isinstance(spec.config, TokenExchangeConfig) + assert spec.config.audience is None + + @pytest.mark.parametrize( "server", [ @@ -115,9 +166,7 @@ def test_unmigrated_modes_defer_to_v1(server): # static token must not route a BYOK server to a v2 shared-key spec with the wrong value. _server(auth_type=MCPAuth.bearer_token, is_byok=True, authentication_token="x"), _server(auth_type=MCPAuth.basic, is_byok=True, authentication_token="x"), - _server( - auth_type=MCPAuth.authorization, is_byok=True, authentication_token="x" - ), + _server(auth_type=MCPAuth.authorization, is_byok=True, authentication_token="x"), _server(auth_type=MCPAuth.token, is_byok=True, authentication_token="x"), _server(auth_type=None, is_byok=True), ], @@ -161,9 +210,7 @@ def test_raise_public_maps_each_error_to_its_status(error, status): def test_raise_public_emits_unauthorized_challenge(): body = {"error": "byok_auth_required", "server_id": "s1"} - error = CredError.of_unauthorized( - "needs key", www_authenticate='Bearer resource_metadata="/x"', body=body - ) + error = CredError.of_unauthorized("needs key", www_authenticate='Bearer resource_metadata="/x"', body=body) with pytest.raises(HTTPException) as exc_info: raise_public(error) exc = exc_info.value @@ -191,8 +238,7 @@ def test_raise_user_oauth_challenge_points_at_per_server_prm(): exc = exc_info.value assert exc.status_code == 401 assert ( - exc.headers["WWW-Authenticate"] - == 'Bearer resource_metadata="/.well-known/oauth-protected-resource/mcp/my-srv"' + exc.headers["WWW-Authenticate"] == 'Bearer resource_metadata="/.well-known/oauth-protected-resource/mcp/my-srv"' ) 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 73e9a52b9370..23d9af756902 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 @@ -1,9 +1,9 @@ """Tests for the resolver dispatch: live arms produce auth, stubbed arms fail closed. -`none`, `api_key` (shared-key source), and `authorization_code` are implemented; every other arm, -plus the `api_key` BYOK source, returns a typed `not_implemented` error until its mode lands. -Parametrizing the stubs over one config each also guards reachability: a dropped `case` would hit -`assert_never` and raise instead of returning the stub. +`none`, `api_key` (shared-key source), `authorization_code`, and `token_exchange` are implemented; +every other arm, plus the `api_key` BYOK source, returns a typed `not_implemented` error until its +mode lands. Parametrizing the stubs over one config each also guards reachability: a dropped `case` +would hit `assert_never` and raise instead of returning the stub. """ import httpx @@ -16,11 +16,13 @@ AwsSigV4Config, Byok, ClientCredentialsConfig, + CredError, Error, NoneConfig, NoOpAuth, Ok, PassthroughConfig, + Result, ServerSpec, SharedKey, StaticHeaderAuth, @@ -37,9 +39,7 @@ def _spec(config): - return ServerSpec( - server_id="s", resource="https://upstream.example.com", config=config - ) + return ServerSpec(server_id="s", resource="https://upstream.example.com", config=config) def _emitted(auth: httpx.Auth) -> httpx.Headers: @@ -52,9 +52,7 @@ def _emitted(auth: httpx.Auth) -> httpx.Headers: @pytest.mark.asyncio async def test_none_mode_yields_a_no_op_auth(): - result = await UpstreamCredentialProvider().resolve_credentials( - _SUBJECT, _spec(NoneConfig()) - ) + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, _spec(NoneConfig())) assert isinstance(result, Ok) assert isinstance(result.ok, NoOpAuth) @@ -66,9 +64,7 @@ async def test_api_key_shared_emits_the_configured_header(): value_prefix="", key_source=SharedKey(value=SecretStr("secret-key")), ) - result = await UpstreamCredentialProvider().resolve_credentials( - _SUBJECT, _spec(config) - ) + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, _spec(config)) assert isinstance(result, Ok) assert isinstance(result.ok, StaticHeaderAuth) assert _emitted(result.ok)["X-API-Key"] == "secret-key" @@ -81,9 +77,7 @@ async def test_api_key_shared_honors_authorization_scheme(): value_prefix="Bearer", key_source=SharedKey(value=SecretStr("tok")), ) - result = await UpstreamCredentialProvider().resolve_credentials( - _SUBJECT, _spec(config) - ) + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, _spec(config)) assert isinstance(result, Ok) assert _emitted(result.ok)["Authorization"] == "Bearer tok" @@ -101,9 +95,7 @@ async def fetch(self, user_id: str, server_id: str): @pytest.mark.asyncio async def test_authorization_code_emits_bearer_for_a_stored_token(): store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="at-alice")}) - result = await UpstreamCredentialProvider( - oauth_token_store=store - ).resolve_credentials( + result = await UpstreamCredentialProvider(oauth_token_store=store).resolve_credentials( Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) ) assert isinstance(result, Ok) @@ -112,9 +104,7 @@ async def test_authorization_code_emits_bearer_for_a_stored_token(): @pytest.mark.asyncio async def test_authorization_code_without_token_is_semantically_unauthorized(): - result = await UpstreamCredentialProvider( - oauth_token_store=_FakeTokenStore({}) - ).resolve_credentials( + result = await UpstreamCredentialProvider(oauth_token_store=_FakeTokenStore({})).resolve_credentials( Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) ) assert isinstance(result, Error) @@ -131,9 +121,7 @@ class _Unavailable: async def fetch(self, user_id: str, server_id: str): raise TokenStoreUnavailable("down") - result = await UpstreamCredentialProvider( - oauth_token_store=_Unavailable() - ).resolve_credentials( + result = await UpstreamCredentialProvider(oauth_token_store=_Unavailable()).resolve_credentials( Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) ) assert isinstance(result, Error) @@ -156,22 +144,15 @@ async def test_authorization_code_isolates_by_subject(): alice = await provider.resolve_credentials( Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) ) - bob = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="bob"), _spec(AuthorizationCodeConfig()) - ) - assert ( - isinstance(alice, Ok) - and _emitted(alice.ok)["Authorization"] == "Bearer at-alice" - ) + bob = await provider.resolve_credentials(Subject(tenant_id="", subject_id="bob"), _spec(AuthorizationCodeConfig())) + assert isinstance(alice, Ok) and _emitted(alice.ok)["Authorization"] == "Bearer at-alice" assert isinstance(bob, Error) and bob.error.tag == "unauthorized" @pytest.mark.asyncio async def test_has_user_token_reflects_the_stored_token(): present = UpstreamCredentialProvider( - oauth_token_store=_FakeTokenStore( - {("alice", "s"): OAuthToken(access_token="at")} - ) + oauth_token_store=_FakeTokenStore({("alice", "s"): OAuthToken(access_token="at")}) ) absent = UpstreamCredentialProvider(oauth_token_store=_FakeTokenStore({})) spec = _spec(AuthorizationCodeConfig()) @@ -185,17 +166,76 @@ async def test_has_user_token_false_for_a_non_per_user_mode(): # A none-mode server has no per-user token to check. provider = UpstreamCredentialProvider() spec = _spec(NoneConfig()) - assert ( - await provider.has_user_token(Subject(tenant_id="", subject_id="a"), spec) - is False + assert await provider.has_user_token(Subject(tenant_id="", subject_id="a"), spec) is False + + +class _FakeExchanger: + def __init__(self, result: Result[OAuthToken, CredError]) -> None: + self._result = result + self.calls: list[tuple[str, str]] = [] + + async def exchange(self, subject_token, server, config): + self.calls.append((subject_token, server.server_id)) + return self._result + + +_OBO = TokenExchangeConfig( + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret=SecretStr("csec"), +) + + +@pytest.mark.asyncio +async def test_token_exchange_emits_the_exchanged_bearer(): + exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at"))) + subject = Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr("caller-jwt")) + result = await UpstreamCredentialProvider(token_exchanger=exchanger).resolve_credentials(subject, _spec(_OBO)) + assert isinstance(result, Ok) + assert _emitted(result.ok)["Authorization"] == "Bearer exchanged-at" + # The arm hands the unwrapped caller token to the exchanger, never the upstream. + assert exchanger.calls == [("caller-jwt", "s")] + + +@pytest.mark.asyncio +async def test_token_exchange_without_caller_token_is_unauthorized(): + exchanger = _FakeExchanger(Ok(OAuthToken(access_token="never"))) + result = await UpstreamCredentialProvider(token_exchanger=exchanger).resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_OBO) + ) + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + assert result.error.unauthorized.www_authenticate == 'Bearer error="invalid_request"' + # No caller token means nothing to exchange: the IdP is never hit. + assert exchanger.calls == [] + + +@pytest.mark.asyncio +async def test_token_exchange_propagates_the_exchanger_error(): + err = CredError.of_upstream_unavailable("idp down") + result = await UpstreamCredentialProvider(token_exchanger=_FakeExchanger(Error(err))).resolve_credentials( + Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr("jwt")), + _spec(_OBO), ) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_token_exchange_without_an_exchanger_fails_closed(): + # The fail-closed default (no exchanger wired) must not produce a credential. + result = await UpstreamCredentialProvider().resolve_credentials( + Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr("jwt")), + _spec(_OBO), + ) + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" _STUBBED = [ ("api_key_byok", ApiKeyConfig(key_source=Byok())), ("passthrough", PassthroughConfig()), ("client_credentials", ClientCredentialsConfig()), - ("token_exchange", TokenExchangeConfig()), ("aws_sigv4", AwsSigV4Config(region="us-east-1")), ] @@ -203,8 +243,6 @@ async def test_has_user_token_false_for_a_non_per_user_mode(): @pytest.mark.asyncio @pytest.mark.parametrize("label, config", _STUBBED) async def test_unbuilt_arms_fail_closed_with_not_implemented(label, config): - result = await UpstreamCredentialProvider().resolve_credentials( - _SUBJECT, _spec(config) - ) + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, _spec(config)) assert isinstance(result, Error) assert result.error.tag == "not_implemented" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchange_provider.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchange_provider.py new file mode 100644 index 000000000000..645afd31fbe5 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchange_provider.py @@ -0,0 +1,73 @@ +"""Tests for the token_exchange composition root: the built exchanger and the HTTP edge contract. + +`build_token_exchanger` wires the pure exchanger to its runtime edges; `_post_exchange_endpoint` is +the I/O edge that maps any transport/HTTP failure to None and parses a JSON body on success. +""" + +from unittest.mock import patch + +import pytest + +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_provider import ( + _post_exchange_endpoint, + build_token_exchanger, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( + Rfc8693TokenExchanger, +) + +_HTTP_CLIENT = "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" + + +def test_build_token_exchanger_returns_an_exchanger(): + assert isinstance(build_token_exchanger(), Rfc8693TokenExchanger) + + +def test_build_gives_each_caller_an_independent_cache(): + # Separate builds must not share a cache, so one egress instance cannot serve another's tokens. + assert build_token_exchanger() is not build_token_exchanger() + + +@pytest.mark.asyncio +async def test_post_returns_none_on_transport_error(): + with patch(_HTTP_CLIENT, side_effect=RuntimeError("boom")): + result = await _post_exchange_endpoint("https://idp/token", {"grant_type": "x"}, {}) + assert result is None + + +@pytest.mark.asyncio +async def test_post_parses_json_body_on_success(): + class _Resp: + def raise_for_status(self) -> None: + return None + + def json(self) -> dict[str, object]: + return {"access_token": "x", "expires_in": 60} + + class _Client: + async def post(self, url, headers, data): + return _Resp() + + with patch(_HTTP_CLIENT, return_value=_Client()): + result = await _post_exchange_endpoint("https://idp/token", {"grant_type": "x"}, {}) + assert result == {"access_token": "x", "expires_in": 60} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("payload", [["a", "b"], "a-string", 42], ids=["list", "str", "int"]) +async def test_post_returns_none_on_non_object_json(payload): + # A valid-but-non-object JSON body must become a miss, not crash field parsing downstream. + class _Resp: + def raise_for_status(self) -> None: + return None + + def json(self) -> object: + return payload + + class _Client: + async def post(self, url, headers, data): + return _Resp() + + with patch(_HTTP_CLIENT, return_value=_Client()): + result = await _post_exchange_endpoint("https://idp/token", {"grant_type": "x"}, {}) + assert result is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchanger.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchanger.py new file mode 100644 index 000000000000..24ea3bcac672 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchanger.py @@ -0,0 +1,358 @@ +"""Tests for the pure RFC 8693 token exchanger: the OBO swap, caching, and single-flight. + +Drives `Rfc8693TokenExchanger` through an injected fake HTTP post and clock, so the exchange, the +form it sends, the per-caller-token cache, and the failure mapping are pinned without a live IdP. +""" + +import asyncio + +import pytest +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials import ( + Error, + Ok, + ServerSpec, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( + Rfc8693TokenExchanger, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + TokenExchangeConfig, +) + +_GRANT = "urn:ietf:params:oauth:grant-type:token-exchange" + +_CONFIG = TokenExchangeConfig( + token_exchange_endpoint="https://idp.example.com/token", + audience="https://up.example.com", + client_id="cid", + client_secret=SecretStr("csec"), + scopes=("s1", "s2"), +) +_SERVER = ServerSpec(server_id="srv", resource="https://up.example.com", config=_CONFIG) + + +class _Clock: + def __init__(self, now: float = 1000.0) -> None: + self.now = now + + def __call__(self) -> float: + return self.now + + +class _RecordingPost: + def __init__(self, body: dict[str, object] | None) -> None: + self._body = body + self.calls: list[tuple[str, dict[str, str]]] = [] + self.headers: list[dict[str, str]] = [] + + async def __call__(self, url: str, form: dict[str, str], headers: dict[str, str]) -> dict[str, object] | None: + self.calls.append((url, dict(form))) + self.headers.append(dict(headers)) + return self._body + + +def _spec(config: TokenExchangeConfig) -> ServerSpec: + return ServerSpec(server_id="srv", resource="https://up.example.com", config=config) + + +@pytest.mark.asyncio +async def test_exchange_emits_token_and_sends_rfc8693_form(): + post = _RecordingPost({"access_token": "exchanged", "expires_in": 3600}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("caller-jwt", _SERVER, _CONFIG) + assert isinstance(result, Ok) + assert result.ok.access_token == "exchanged" + url, form = post.calls[0] + assert url == "https://idp.example.com/token" + assert form == { + "grant_type": _GRANT, + "subject_token": "caller-jwt", + "subject_token_type": "urn:ietf:params:oauth:token-type:access_token", + "client_id": "cid", + "client_secret": "csec", + "audience": "https://up.example.com", + "scope": "s1 s2", + } + # client_secret_post is the default: creds in the body, no client-auth header + assert post.headers[0] == {} + + +@pytest.mark.asyncio +async def test_client_secret_basic_sends_authorization_header_and_omits_body_creds(): + import base64 + + config = TokenExchangeConfig( + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret=SecretStr("csec"), + token_endpoint_auth_method="client_secret_basic", + ) + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _spec(config), config) + assert isinstance(result, Ok) + _, form = post.calls[0] + assert "client_id" not in form and "client_secret" not in form + assert post.headers[0]["Authorization"] == "Basic " + base64.b64encode(b"cid:csec").decode() + + +@pytest.mark.asyncio +async def test_client_secret_post_keeps_creds_in_body_with_no_auth_header(): + config = TokenExchangeConfig( + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret=SecretStr("csec"), + token_endpoint_auth_method="client_secret_post", + ) + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _spec(config), config) + _, form = post.calls[0] + assert form["client_id"] == "cid" and form["client_secret"] == "csec" + assert "Authorization" not in post.headers[0] + + +@pytest.mark.asyncio +async def test_exchange_caches_per_caller_token(): + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + await exchanger.exchange("jwt", _SERVER, _CONFIG) + second = await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert isinstance(second, Ok) and second.ok.access_token == "x" + assert len(post.calls) == 1 + + +@pytest.mark.asyncio +async def test_rotated_caller_token_re_exchanges(): + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + await exchanger.exchange("jwt-1", _SERVER, _CONFIG) + await exchanger.exchange("jwt-2", _SERVER, _CONFIG) + assert len(post.calls) == 2 + + +@pytest.mark.asyncio +async def test_rotated_config_re_exchanges_before_ttl(): + # Same caller token + server, but the operator rotated the audience/scope: the cached token was + # minted for the old config, so it must re-exchange (not serve the stale token) before TTL. + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + rotated = _CONFIG.model_copy(update={"audience": "https://new.example.com", "scopes": ("s3",)}) + await exchanger.exchange("jwt", _SERVER, _CONFIG) + await exchanger.exchange("jwt", _SERVER, rotated) + assert len(post.calls) == 2 + _, second_form = post.calls[1] + assert second_form["audience"] == "https://new.example.com" + assert second_form["scope"] == "s3" + + +@pytest.mark.asyncio +async def test_rotated_token_endpoint_auth_method_re_exchanges_before_ttl(): + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + rotated = _CONFIG.model_copy(update={"token_endpoint_auth_method": "client_secret_basic"}) + await exchanger.exchange("jwt", _SERVER, _CONFIG) + await exchanger.exchange("jwt", _SERVER, rotated) + assert len(post.calls) == 2 + _, second_form = post.calls[1] + assert "client_id" not in second_form + assert "client_secret" not in second_form + assert "Authorization" in post.headers[1] + + +@pytest.mark.asyncio +async def test_concurrent_callers_single_flight_one_exchange(): + release = asyncio.Event() + + class _Blocking: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, url, form, headers): + self.calls += 1 + await release.wait() + return {"access_token": "x", "expires_in": 3600} + + post = _Blocking() + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + first = asyncio.create_task(exchanger.exchange("jwt", _SERVER, _CONFIG)) + second = asyncio.create_task(exchanger.exchange("jwt", _SERVER, _CONFIG)) + await asyncio.sleep(0.02) + release.set() + r1, r2 = await asyncio.gather(first, second) + assert post.calls == 1 + assert isinstance(r1, Ok) and isinstance(r2, Ok) + + +@pytest.mark.asyncio +async def test_idp_failure_is_upstream_unavailable(): + result = await Rfc8693TokenExchanger(_RecordingPost(None), clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_missing_access_token_is_upstream_unavailable(): + post = _RecordingPost({"token_type": "Bearer"}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_type", ["N_A", "n_a", "DPoP", "mac"]) +async def test_non_bearer_token_type_is_refused(token_type): + # RFC 8693 2.2.1: the resolver forwards the exchanged token as `Bearer`. A non-Bearer token_type + # (e.g. N_A = not a standalone access token) must fail closed, not be minted as a bogus Bearer. + post = _RecordingPost({"access_token": "x", "token_type": token_type, "expires_in": 3600}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_type", ["Bearer", "bearer", "BEARER"]) +async def test_bearer_token_type_is_accepted_case_insensitively(token_type): + post = _RecordingPost({"access_token": "x", "token_type": token_type, "expires_in": 3600}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Ok) and result.ok.access_token == "x" + + +@pytest.mark.asyncio +async def test_absent_token_type_defaults_to_bearer(): + # Many IdPs omit token_type; absence must not fail the exchange (RFC 6750 default is Bearer). + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Ok) and result.ok.access_token == "x" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "issued_token_type", + [ + "urn:ietf:params:oauth:token-type:refresh_token", + "urn:ietf:params:oauth:token-type:id_token", + "urn:ietf:params:oauth:token-type:saml2", + ], +) +async def test_non_access_issued_token_type_is_refused_even_when_bearer(issued_token_type): + # A malformed STS could mint a refresh/id/saml token but label it Bearer; issued_token_type must + # still fail it closed rather than forward a non-access token as an upstream access credential. + post = _RecordingPost( + {"access_token": "x", "token_type": "Bearer", "issued_token_type": issued_token_type, "expires_in": 3600} + ) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "issued_token_type", + ["urn:ietf:params:oauth:token-type:access_token", "urn:ietf:params:oauth:token-type:jwt", "custom-unknown", None], +) +async def test_access_or_unknown_issued_token_type_is_accepted(issued_token_type): + # access_token / jwt are usable; an absent or unrecognized type is accepted (lenient), so real + # IdPs that omit issued_token_type or use a custom URN keep working. + body: dict[str, object] = {"access_token": "x", "token_type": "Bearer", "expires_in": 3600} + if issued_token_type is not None: + body["issued_token_type"] = issued_token_type + result = await Rfc8693TokenExchanger(_RecordingPost(body), clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Ok) and result.ok.access_token == "x" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "config", + [ + TokenExchangeConfig(client_id="c", client_secret=SecretStr("s")), + TokenExchangeConfig(token_exchange_endpoint="https://idp/token", client_secret=SecretStr("s")), + TokenExchangeConfig(token_exchange_endpoint="https://idp/token", client_id="c"), + ], +) +async def test_incomplete_config_is_misconfigured_without_hitting_idp(config): + post = _RecordingPost({"access_token": "x"}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _spec(config), config) + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" + assert post.calls == [] + + +@pytest.mark.asyncio +async def test_cached_token_expires_after_its_ttl(): + clock = _Clock(1000.0) + # expires_in=120, buffer=60 -> ttl 60 -> cached until t=1060. + post = _RecordingPost({"access_token": "x", "expires_in": 120}) + exchanger = Rfc8693TokenExchanger(post, clock=clock) + await exchanger.exchange("jwt", _SERVER, _CONFIG) + clock.now = 1059.0 + await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert len(post.calls) == 1 + clock.now = 1061.0 + await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert len(post.calls) == 2 + + +@pytest.mark.asyncio +async def test_audience_and_scope_omitted_when_unset(): + config = TokenExchangeConfig( + token_exchange_endpoint="https://idp/token", + client_id="cid", + client_secret=SecretStr("csec"), + ) + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _spec(config), config) + _, form = post.calls[0] + assert "audience" not in form + assert "scope" not in form + + +@pytest.mark.asyncio +async def test_string_expires_in_is_honored(): + clock = _Clock(1000.0) + # "120" parsed as int -> ttl max(120-60, 10) = 60 -> cached until 1060. + post = _RecordingPost({"access_token": "x", "expires_in": "120"}) + exchanger = Rfc8693TokenExchanger(post, clock=clock) + await exchanger.exchange("jwt", _SERVER, _CONFIG) + clock.now = 1061.0 + await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert len(post.calls) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [ + {"access_token": "x"}, + {"access_token": "x", "expires_in": "not-a-number"}, + {"access_token": "x", "expires_in": True}, + ], + ids=["missing", "unparseable", "bool"], +) +async def test_unusable_expires_in_falls_back_to_default_ttl(body): + clock = _Clock(1000.0) + post = _RecordingPost(body) + exchanger = Rfc8693TokenExchanger(post, clock=clock, default_ttl_seconds=300.0) + await exchanger.exchange("jwt", _SERVER, _CONFIG) + clock.now = 1299.0 + await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert len(post.calls) == 1 + + +@pytest.mark.asyncio +async def test_distributed_coordinator_refresh_and_reread_use_the_cache(): + # Mimics the cross-replica coordinator contract: the winner's refresh populates the cache, a + # re-entrant refresh sees the fresh entry, and a loser reads it back via reread, all without a + # second IdP call. + class _ReplayCoordinator: + async def run(self, user_id, server_id, refresh, reread): + first = await refresh() + second = await refresh() + via_reread = await reread() + assert first is not None and second is not None and via_reread is not None + return via_reread + + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, coordinator=_ReplayCoordinator(), clock=_Clock()) + result = await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Ok) and result.ok.access_token == "x" + assert len(post.calls) == 1