diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 6631e38f524..3bf9fd1c12e 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -21,8 +21,12 @@ from typing_extensions import assert_never from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, +<<<<<<< HEAD ClientAuth, ClientSecretAuth, +======= + ClientCredentialsConfig, +>>>>>>> 73df37ca23 (feat(mcp): migrate client_credentials (M2M) onto the v2 resolver arm) CredError, IdJagConfig, NoneConfig, @@ -70,10 +74,10 @@ 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), ``oauth2`` per-user tokens (``authorization_code``), ``oauth2_token_exchange`` - (OBO), and the client-forwarded token modes ``true_passthrough`` / ``oauth_delegate`` - (``PassthroughConfig``); client_credentials (M2M), delegated/passthrough oauth2, and SigV4 - return None and stay on v1. + all shared-key), ``oauth2`` per-user tokens (``authorization_code``), ``oauth2`` M2M + (``client_credentials``), ``oauth2_token_exchange`` (OBO), and the client-forwarded token + modes ``true_passthrough`` / ``oauth_delegate`` (``PassthroughConfig``); 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) @@ -95,13 +99,15 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: case MCPAuth.basic: return _shared_key_spec(server, resource, "Authorization", "Basic", encode=True) case MCPAuth.oauth2: + if server.has_client_credentials: + return _client_credentials_spec(server, resource) if server.needs_user_oauth_token and not server.delegate_auth_to_upstream: return ServerSpec( server_id=server.server_id, resource=resource, config=AuthorizationCodeConfig(), ) - # client_credentials (M2M) and delegate/passthrough oauth2 stay on v1 + # delegate/passthrough oauth2 stay on v1 return None case MCPAuth.oauth2_id_jag: return _id_jag_spec(server, resource) @@ -114,6 +120,29 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: assert_never(auth_type) +def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec: + """Build a client_credentials (M2M) spec; the explicit ``oauth2_flow`` opt-in owns the server. + + Missing grant fields (``client_id``/``client_secret``/``token_url``) are NOT a reason to defer: + v1 would connect unauthenticated and the upstream's 401 gets absorbed into an empty tool list, + so the arm fails closed with ``misconfigured`` instead, naming the missing fields (mirrors the + OBO ownership rule). ``audience`` is forwarded only when the operator set it; a missing one is + omitted, not derived, since a fabricated value risks the IdP rejecting the grant. + """ + return ServerSpec( + server_id=server.server_id, + resource=resource, + config=ClientCredentialsConfig( + client_id=server.client_id, + client_secret=SecretStr(server.client_secret) if server.client_secret else None, + token_url=server.token_url, + scopes=tuple(server.scopes or ()), + audience=server.audience, + token_endpoint_auth_method=server.token_endpoint_auth_method, + ), + ) + + def _token_exchange_spec(server: MCPServer, resource: str) -> Optional[ServerSpec]: """Build a token_exchange (OBO) spec, or defer (None) when it is not OBO-configured. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py new file mode 100644 index 00000000000..3f6c329118f --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -0,0 +1,334 @@ +"""The ``client_credentials`` (M2M) arm's token source and retrying bearer auth. + +Implements the client-credentials behavior contract for the v2 resolver: + +- **Acquisition**: POST ``grant_type=client_credentials`` to the configured token endpoint with + the configured scopes and (when set) the IdP's ``audience`` parameter, authenticating the + client per ``token_endpoint_auth_method`` (RFC 6749 section 2.3.1, shared helper). +- **Caching**: tokens are cached per ``(client identity, server)`` where the identity key hashes + ``token_url`` / ``client_id`` / ``client_secret`` / auth method / scopes / audience — rotating + or re-scoping the credentials changes the key, so a stale token can never be served for the + new identity (the contract's rotation-invalidation clause). +- **Expiry**: the cache TTL respects ``expires_in`` minus a skew so an entry lapses before the + real token does; a response with no ``expires_in`` is cached briefly + (``default_ttl_seconds``), not assumed long-lived. No refresh_token is ever expected. +- **401 recovery**: ``ClientCredentialsBearerAuth`` retries an upstream request exactly once + after a 401 — discard the cached token, mint a fresh one, resend; a second failure surfaces + the upstream's own auth error unchanged. +- **No user context**: nothing here reads a ``Subject``; every caller shares the one client + identity. + +The token-endpoint POST is injected (``M2MTokenEndpointPost``) so the grant orchestration is +testable without a live IdP; ``post_client_credentials_grant`` is the httpx edge and the one +place the untyped response boundary is contained. Failures are values: the source returns +``Result[OAuthToken, CredError]``; only the httpx edge touches exceptions. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import time +from collections.abc import AsyncGenerator, Awaitable, Callable, Generator +from dataclasses import dataclass +from typing import Annotated, Literal + +import httpx +from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + InMemoryTokenCacheBackend, + OAuthToken, + TokenCacheBackend, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientCredentialsConfig, + CredError, +) + + +class TokenEndpointSuccess(BaseModel): + """The endpoint returned a JSON object; field validation is the caller's job.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["success"] = "success" + body: dict[str, object] + + +class TokenEndpointDenied(BaseModel): + """The endpoint answered but did not grant a token (an HTTP error or a non-JSON body).""" + + model_config = ConfigDict(frozen=True) + tag: Literal["denied"] = "denied" + status_code: int + detail: str + + +class TokenEndpointUnreachable(BaseModel): + """The endpoint could not be reached (DNS, TLS, connect/read failure).""" + + model_config = ConfigDict(frozen=True) + tag: Literal["unreachable"] = "unreachable" + detail: str + + +TokenEndpointOutcome = Annotated[ + TokenEndpointSuccess | TokenEndpointDenied | TokenEndpointUnreachable, + Field(discriminator="tag"), +] + +M2MTokenEndpointPost = Callable[[str, "dict[str, str]", "dict[str, str]"], Awaitable[TokenEndpointOutcome]] + + +_TOKEN_BODY_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) + + +async def post_client_credentials_grant( + url: str, form: dict[str, str], headers: dict[str, str] +) -> TokenEndpointOutcome: + """POST the grant to the token endpoint and classify the transport outcome. + + The httpx edge: litellm's handler is partially typed (and raises ``HTTPStatusError`` itself on + a 4xx/5xx), so the untyped boundary is contained here and every field the caller reads comes + out of a validated ``TokenEndpointOutcome``. + """ + from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # handler is partially typed + ) + from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 + + try: + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) + response = await client.post( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # handler is partially typed + url, headers={"Accept": "application/json", **headers}, data=form + ) + except httpx.HTTPStatusError as status_err: + status_code = status_err.response.status_code + return TokenEndpointDenied(status_code=status_code, detail=f"token endpoint returned HTTP {status_code}") + except Exception as exc: # noqa: BLE001 # any transport failure is the same outcome: unreachable + return TokenEndpointUnreachable(detail=str(exc)) + if not isinstance(response, httpx.Response): + return TokenEndpointUnreachable(detail="token endpoint returned no response") + try: + body = _TOKEN_BODY_ADAPTER.validate_json(response.content) + except ValidationError: + return TokenEndpointDenied( + status_code=response.status_code, detail="token endpoint returned a non-JSON-object body" + ) + return TokenEndpointSuccess(body=body) + + +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 _parse_granted_scopes(raw: object) -> tuple[str, ...] | None: + return tuple(raw.split()) if isinstance(raw, str) and raw else None + + +@dataclass(frozen=True, slots=True) +class _PreparedGrant: + """A validated, ready-to-POST grant plus the identity key its token caches under.""" + + token_url: str + form: dict[str, str] + headers: dict[str, str] + identity_key: str + + +class ClientCredentialsTokenSource: + """Cached M2M access tokens, one per ``(client identity, server)``. + + ``get`` serves from the cache while the entry's TTL (derived from ``expires_in`` minus + ``expiry_skew_seconds``) holds, fetching under a per-server lock so concurrent misses + produce one grant. ``refetch`` is the 401-recovery path: it drops the failed token and + mints a fresh one, unless a concurrent caller already replaced it. + """ + + def __init__( + self, + post: M2MTokenEndpointPost = post_client_credentials_grant, + *, + backend: TokenCacheBackend | None = None, + default_ttl_seconds: float = 300.0, + expiry_skew_seconds: float = 60.0, + min_cache_seconds: float = 10.0, + clock: Callable[[], float] = time.time, + ) -> None: + self._post = post + self._backend: TokenCacheBackend = backend or InMemoryTokenCacheBackend(clock=clock) + self._default_ttl_seconds = default_ttl_seconds + self._expiry_skew_seconds = expiry_skew_seconds + self._min_cache_seconds = min_cache_seconds + self._clock = clock + self._locks: dict[str, asyncio.Lock] = {} + + def _lock(self, server_id: str) -> asyncio.Lock: + return self._locks.setdefault(server_id, asyncio.Lock()) + + async def get(self, server_id: str, config: ClientCredentialsConfig) -> Result[OAuthToken, CredError]: + match _prepare_grant(config): + case Error(err): + return Error(err) + case Ok(grant): + cached = await self._backend.get(grant.identity_key, server_id) + if cached is not None: + return Ok(cached) + async with self._lock(server_id): + cached = await self._backend.get(grant.identity_key, server_id) + if cached is not None: + return Ok(cached) + return await self._fetch_and_cache(server_id, grant) + + async def refetch(self, server_id: str, config: ClientCredentialsConfig, failed_access_token: str) -> str | None: + """Replace a token the upstream just 401'd; returns the fresh bearer value or ``None``. + + Runs under the same per-server lock as ``get``: if a concurrent caller already replaced + the failed token, that replacement is returned without another grant, so a burst of 401s + yields one fetch. A failed refetch returns ``None`` and the caller surfaces the + upstream's original auth error (the contract's retry-once-then-give-up clause). + """ + match _prepare_grant(config): + case Error(_): + return None + case Ok(grant): + async with self._lock(server_id): + cached = await self._backend.get(grant.identity_key, server_id) + if cached is not None and cached.access_token != failed_access_token: + return cached.access_token + await self._backend.delete(grant.identity_key, server_id) + match await self._fetch_and_cache(server_id, grant): + case Ok(token): + return token.access_token + case Error(_): + return None + + async def _fetch_and_cache(self, server_id: str, grant: _PreparedGrant) -> Result[OAuthToken, CredError]: + outcome = await self._post(grant.token_url, grant.form, grant.headers) + match outcome: + case TokenEndpointUnreachable(): + return Error(CredError.of_upstream_unavailable(f"OAuth2 token endpoint unreachable: {outcome.detail}")) + case TokenEndpointDenied(): + if outcome.status_code >= 500: + return Error(CredError.of_upstream_unavailable(f"OAuth2 token endpoint failed: {outcome.detail}")) + return Error(CredError.of_misconfigured(f"OAuth2 client_credentials grant rejected: {outcome.detail}")) + case TokenEndpointSuccess(): + return await self._cache_token(server_id, grant, outcome.body) + assert_never(outcome) + + async def _cache_token( + self, server_id: str, grant: _PreparedGrant, body: dict[str, object] + ) -> Result[OAuthToken, CredError]: + access_token = body.get("access_token") + if not isinstance(access_token, str) or not access_token: + return Error(CredError.of_misconfigured("OAuth2 token response is missing 'access_token'")) + expires_in = _parse_expires_in(body.get("expires_in")) + token = OAuthToken( + access_token=access_token, + expires_at=self._clock() + expires_in if expires_in is not None else None, + scopes=_parse_granted_scopes(body.get("scope")) or (), + ) + ttl = ( + max(expires_in - self._expiry_skew_seconds, self._min_cache_seconds) + if expires_in is not None + else self._default_ttl_seconds + ) + await self._backend.set(grant.identity_key, server_id, token, ttl) + return Ok(token) + + +def _prepare_grant(config: ClientCredentialsConfig) -> Result[_PreparedGrant, CredError]: + if not config.client_id or not config.client_secret or not config.token_url: + missing = ", ".join( + name + for name, present in ( + ("client_id", bool(config.client_id)), + ("client_secret", bool(config.client_secret)), + ("token_url", bool(config.token_url)), + ) + if not present + ) + return Error(CredError.of_misconfigured(f"client_credentials config is missing: {missing}")) + + from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( # noqa: PLC0415 + build_token_endpoint_client_auth, + ) + + client_auth = build_token_endpoint_client_auth( + auth_method=config.token_endpoint_auth_method, + client_id=config.client_id, + client_secret=config.client_secret.get_secret_value(), + ) + form = { + "grant_type": "client_credentials", + **client_auth.body, + **({"scope": " ".join(config.scopes)} if config.scopes else {}), + **({"audience": config.audience} if config.audience else {}), + } + return Ok( + _PreparedGrant( + token_url=config.token_url, + form=form, + headers=client_auth.headers, + identity_key=_identity_key(config), + ) + ) + + +def _identity_key(config: ClientCredentialsConfig) -> str: + """Hash of everything that names the client identity; any rotation yields a new key.""" + material = "\n".join( + ( + config.token_url or "", + config.client_id or "", + config.client_secret.get_secret_value() if config.client_secret else "", + config.token_endpoint_auth_method or "", + " ".join(config.scopes), + config.audience or "", + ) + ) + return hashlib.sha256(material.encode("utf-8")).hexdigest() + + +class ClientCredentialsBearerAuth(httpx.Auth): + """Bearer auth that retries an upstream 401 exactly once with a freshly minted token. + + The initial token was already resolved (so config/IdP failures surfaced as typed errors + before any upstream request); ``refetch`` is the source's 401-recovery callback. If the + refetch fails, or the retried request 401s again, the upstream's response stands. + """ + + def __init__(self, access_token: str, refetch: Callable[[str], Awaitable[str | None]]) -> None: + self.header_name = "Authorization" + self._access_token = SecretStr(access_token) + self._refetch = refetch + + async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]: + token = self._access_token.get_secret_value() + request.headers[self.header_name] = f"Bearer {token}" + response = yield request + if response.status_code != 401: + return + fresh = await self._refetch(token) + if fresh is None: + return + request.headers[self.header_name] = f"Bearer {fresh}" + yield request + + def sync_auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: + raise RuntimeError("ClientCredentialsBearerAuth only supports async httpx clients") diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 7e5c073870a..69984a56311 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -9,18 +9,24 @@ at runtime instead of returning `None`. `none`, `api_key` (shared-key source), and `passthrough` (forwards the caller's own inbound token) are live, as is `authorization_code`, which reads the user's token from the injected -`OAuthTokenStore`, 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. +`OAuthTokenStore`, `token_exchange`, which swaps the caller's inbound token through the injected +`TokenExchanger`, and `client_credentials`, which mints and caches the gateway's M2M token through +the injected `ClientCredentialsTokenSource`. The remaining arms are `not_implemented` stubs that +each land in a follow-up PR with their seam. Pure v2: no imports from v1. """ from __future__ import annotations import hashlib +from functools import partial import httpx from typing_extensions import assert_never +from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( + ClientCredentialsBearerAuth, + ClientCredentialsTokenSource, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( NoOpAuth, StaticHeaderAuth, @@ -104,11 +110,13 @@ class UpstreamCredentialProvider: token_exchanger: TokenExchanger | None = None, token_endpoint: TokenEndpointClient | None = None, exchanged_tokens: ExchangedTokenCache | None = None, + client_credentials_source: ClientCredentialsTokenSource | 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() async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: @@ -118,8 +126,8 @@ class UpstreamCredentialProvider: return self._api_key(config) case PassthroughConfig(): return self._passthrough(subject) - case ClientCredentialsConfig(): - return _not_implemented(AuthSpecKind.client_credentials) + case ClientCredentialsConfig() as config: + return await self._client_credentials(server.server_id, config) case TokenExchangeConfig() as config: return await self._token_exchange(subject, server, config) case IdJagConfig() as config: @@ -215,6 +223,23 @@ class UpstreamCredentialProvider: 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 _client_credentials( + self, server_id: str, config: ClientCredentialsConfig + ) -> Result[httpx.Auth, CredError]: + """The M2M arm: resolve a cached (or freshly minted) gateway token; no user context. + + The token is resolved here, before any upstream request, so a misconfigured grant or an + unreachable IdP surfaces as a typed ``CredError``. The returned auth carries the source's + ``refetch``, so an upstream 401 is retried exactly once with a freshly minted token (the + contract's invalid-token recovery); a second 401 surfaces the upstream's own error. + """ + match await self._client_credentials_source.get(server_id, config): + case Ok(token): + refetch = partial(self._client_credentials_source.refetch, server_id, config) + return Ok(ClientCredentialsBearerAuth(token.access_token, refetch)) + case Error(err): + return Error(err) + async def _token_exchange( self, subject: Subject, server: ServerSpec, config: TokenExchangeConfig ) -> Result[StaticHeaderAuth, CredError]: @@ -245,7 +270,9 @@ class UpstreamCredentialProvider: Used after an upstream rejects the injected credential, so the next resolve re-mints rather than serving the same rejected token until TTL. `token_exchange` and `id_jag` hold a - re-mintable cached credential here; other modes are a no-op. + 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. """ if subject.inbound_token is None: return diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 64a20255ab2..926d96c8868 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -184,7 +184,12 @@ class ClientCredentialsConfig(BaseModel): Fields are optional so the config can be built incomplete: a value may be supplied at runtime (`token_url` via RFC 8414 discovery, `client_id`/`secret` via DCR), and the - resolver arm raises `CredError.misconfigured` when a needed field is still absent. + resolver arm returns `CredError.misconfigured` when a needed field is still absent. + + `audience` is the IdP-specific audience parameter some authorization servers require on + the client_credentials grant (sent as `audience` in the token request when set). + `token_endpoint_auth_method` selects how the client authenticates to the token endpoint + (RFC 6749 section 2.3.1); `None` defaults to `client_secret_post`. """ model_config = ConfigDict(frozen=True) @@ -193,6 +198,8 @@ class ClientCredentialsConfig(BaseModel): client_secret: SecretStr | None = None token_url: str | None = None scopes: tuple[str, ...] = () + audience: str | None = None + token_endpoint_auth_method: Literal["client_secret_post", "client_secret_basic"] | None = None class TokenExchangeConfig(BaseModel): 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 707374e7061..bf757b64c9a 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 @@ -21,6 +21,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, + ClientCredentialsConfig, ClientSecretAuth, CredError, IdJagConfig, @@ -107,7 +108,6 @@ 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( @@ -124,6 +124,74 @@ def test_unmigrated_modes_defer_to_v1(server): assert to_server_spec(server) is None +def test_client_credentials_maps_full_config(): + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + url="https://up.example.com/mcp", + token_url="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + scopes=["read", "write"], + audience="https://up.example.com", + token_endpoint_auth_method="client_secret_basic", + ) + ) + assert spec is not None + config = spec.config + assert isinstance(config, ClientCredentialsConfig) + assert config.client_id == "cid" + assert config.client_secret is not None + assert config.client_secret.get_secret_value() == "csec" + assert config.token_url == "https://idp.example.com/token" + assert config.scopes == ("read", "write") + assert config.audience == "https://up.example.com" + assert config.token_endpoint_auth_method == "client_secret_basic" + + +def test_client_credentials_omits_audience_when_unset(): + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + token_url="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + ) + assert spec is not None + assert isinstance(spec.config, ClientCredentialsConfig) + assert spec.config.audience is None + + +def test_client_credentials_with_incomplete_grant_fields_is_owned_for_fail_closed(): + # An M2M server missing its grant fields is still owned by v2 (spec, not None) so it fails + # closed at the source (misconfigured, 500) rather than deferring to v1, which would connect + # unauthenticated and mask the upstream 401 as an empty tool list. + spec = to_server_spec(_server(auth_type=MCPAuth.oauth2, oauth2_flow="client_credentials", client_id="cid")) + assert spec is not None + assert isinstance(spec.config, ClientCredentialsConfig) + assert spec.config.token_url is None + assert spec.config.client_secret is None + + +def test_client_credentials_wins_over_delegate_flag(): + # v1 never delegates for M2M servers; the explicit oauth2_flow opt-in outranks the delegate flag. + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + delegate_auth_to_upstream=True, + token_url="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + ) + assert spec is not None + assert isinstance(spec.config, ClientCredentialsConfig) + + def test_token_exchange_maps_full_config(): spec = to_server_spec( _server( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py new file mode 100644 index 00000000000..bd00bb77abd --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py @@ -0,0 +1,320 @@ +"""Tests for the client_credentials (M2M) token source and its retrying bearer auth. + +These are the behavior-contract spec: grant shape (scopes / audience / client auth method), +rotation-aware cache keying, expires_in-driven expiry, error classification, and the +401 -> discard -> refetch -> retry-once recovery in ``ClientCredentialsBearerAuth``. +""" + +import httpx +import pytest +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( + ClientCredentialsBearerAuth, + ClientCredentialsTokenSource, + TokenEndpointDenied, + TokenEndpointOutcome, + TokenEndpointSuccess, + TokenEndpointUnreachable, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientCredentialsConfig, +) + + +class _Clock: + def __init__(self, t: float = 1000.0) -> None: + self.t = t + + def __call__(self) -> float: + return self.t + + +class _FakePoster: + """Records every grant POST and returns canned outcomes (last one repeats).""" + + def __init__(self, outcomes: "list[TokenEndpointOutcome]") -> None: + self._outcomes = outcomes + self.calls: "list[tuple[str, dict[str, str], dict[str, str]]]" = [] + + async def __call__(self, url: str, form: "dict[str, str]", headers: "dict[str, str]") -> TokenEndpointOutcome: + self.calls.append((url, dict(form), dict(headers))) + index = min(len(self.calls) - 1, len(self._outcomes) - 1) + return self._outcomes[index] + + +def _success(access_token: str = "m2m-token", **extra: object) -> TokenEndpointSuccess: + return TokenEndpointSuccess(body={"access_token": access_token, **extra}) + + +def _config(**overrides: object) -> ClientCredentialsConfig: + fields: "dict[str, object]" = { + "client_id": "cid", + "client_secret": SecretStr("csec"), + "token_url": "https://idp.example.com/token", + **overrides, + } + return ClientCredentialsConfig.model_validate(fields) + + +@pytest.mark.asyncio +async def test_grant_posts_client_credentials_with_scopes_and_audience(): + poster = _FakePoster([_success()]) + source = ClientCredentialsTokenSource(poster) + result = await source.get("s", _config(scopes=("read", "write"), audience="https://api.example.com")) + assert isinstance(result, Ok) + assert result.ok.access_token == "m2m-token" + url, form, _headers = poster.calls[0] + assert url == "https://idp.example.com/token" + assert form["grant_type"] == "client_credentials" + assert form["scope"] == "read write" + assert form["audience"] == "https://api.example.com" + assert form["client_id"] == "cid" + assert form["client_secret"] == "csec" + + +@pytest.mark.asyncio +async def test_grant_omits_scope_and_audience_when_not_configured(): + poster = _FakePoster([_success()]) + await ClientCredentialsTokenSource(poster).get("s", _config()) + _url, form, _headers = poster.calls[0] + assert "scope" not in form + assert "audience" not in form + + +@pytest.mark.asyncio +async def test_grant_honors_client_secret_basic(): + poster = _FakePoster([_success()]) + await ClientCredentialsTokenSource(poster).get("s", _config(token_endpoint_auth_method="client_secret_basic")) + _url, form, headers = poster.calls[0] + assert headers["Authorization"].startswith("Basic ") + assert "client_secret" not in form + assert "client_id" not in form + + +@pytest.mark.asyncio +async def test_missing_grant_fields_are_misconfigured_and_never_posted(): + poster = _FakePoster([_success()]) + result = await ClientCredentialsTokenSource(poster).get( + "s", ClientCredentialsConfig(client_id="cid", client_secret=SecretStr("csec")) + ) + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" + assert "token_url" in result.error.summary + assert poster.calls == [] + + +@pytest.mark.asyncio +async def test_token_is_cached_across_gets(): + poster = _FakePoster([_success(expires_in=3600)]) + source = ClientCredentialsTokenSource(poster) + first = await source.get("s", _config()) + second = await source.get("s", _config()) + assert isinstance(first, Ok) and isinstance(second, Ok) + assert second.ok.access_token == first.ok.access_token + assert len(poster.calls) == 1 + + +@pytest.mark.asyncio +async def test_expires_in_bounds_the_cache_lifetime(): + clock = _Clock(1000.0) + poster = _FakePoster([_success("t1", expires_in=120), _success("t2", expires_in=120)]) + source = ClientCredentialsTokenSource(poster, expiry_skew_seconds=60.0, clock=clock) + first = await source.get("s", _config()) + assert isinstance(first, Ok) + assert first.ok.expires_at == 1120.0 + clock.t = 1059.0 # within expires_in - skew + assert len(poster.calls) == 1 + within = await source.get("s", _config()) + assert isinstance(within, Ok) and within.ok.access_token == "t1" + clock.t = 1061.0 # past expires_in - skew: the entry lapsed before the real token does + lapsed = await source.get("s", _config()) + assert isinstance(lapsed, Ok) and lapsed.ok.access_token == "t2" + assert len(poster.calls) == 2 + + +@pytest.mark.asyncio +async def test_missing_expires_in_is_cached_briefly_not_an_hour(): + clock = _Clock(1000.0) + poster = _FakePoster([_success("t1"), _success("t2")]) + source = ClientCredentialsTokenSource(poster, default_ttl_seconds=300.0, clock=clock) + first = await source.get("s", _config()) + assert isinstance(first, Ok) + assert first.ok.expires_at is None + clock.t = 1301.0 # past the default TTL; v1 would still be serving its 3600s-cached token + second = await source.get("s", _config()) + assert isinstance(second, Ok) and second.ok.access_token == "t2" + assert len(poster.calls) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rotation", + [ + {"client_secret": SecretStr("rotated")}, + {"client_id": "cid-2"}, + {"scopes": ("admin",)}, + {"audience": "https://other.example.com"}, + {"token_url": "https://idp2.example.com/token"}, + ], +) +async def test_credential_rotation_invalidates_the_cached_token(rotation): + poster = _FakePoster([_success("old", expires_in=3600), _success("new", expires_in=3600)]) + source = ClientCredentialsTokenSource(poster) + before = await source.get("s", _config()) + after = await source.get("s", _config(**rotation)) + assert isinstance(before, Ok) and before.ok.access_token == "old" + assert isinstance(after, Ok) and after.ok.access_token == "new" + assert len(poster.calls) == 2 + + +@pytest.mark.asyncio +async def test_idp_4xx_is_misconfigured_and_5xx_is_unavailable(): + denied = await ClientCredentialsTokenSource( + _FakePoster([TokenEndpointDenied(status_code=401, detail="HTTP 401")]) + ).get("s", _config()) + assert isinstance(denied, Error) and denied.error.tag == "misconfigured" + down = await ClientCredentialsTokenSource( + _FakePoster([TokenEndpointDenied(status_code=503, detail="HTTP 503")]) + ).get("s", _config()) + assert isinstance(down, Error) and down.error.tag == "upstream_unavailable" + unreachable = await ClientCredentialsTokenSource(_FakePoster([TokenEndpointUnreachable(detail="dns")])).get( + "s", _config() + ) + assert isinstance(unreachable, Error) and unreachable.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_response_without_access_token_is_misconfigured(): + poster = _FakePoster([TokenEndpointSuccess(body={"token_type": "Bearer"})]) + result = await ClientCredentialsTokenSource(poster).get("s", _config()) + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" + + +@pytest.mark.asyncio +async def test_error_results_are_not_cached(): + poster = _FakePoster([TokenEndpointUnreachable(detail="down"), _success("recovered")]) + source = ClientCredentialsTokenSource(poster) + first = await source.get("s", _config()) + second = await source.get("s", _config()) + assert isinstance(first, Error) + assert isinstance(second, Ok) and second.ok.access_token == "recovered" + + +@pytest.mark.asyncio +async def test_refetch_discards_the_failed_token_and_mints_a_fresh_one(): + poster = _FakePoster([_success("stale", expires_in=3600), _success("fresh", expires_in=3600)]) + source = ClientCredentialsTokenSource(poster) + first = await source.get("s", _config()) + assert isinstance(first, Ok) + fresh = await source.refetch("s", _config(), failed_access_token="stale") + assert fresh == "fresh" + assert len(poster.calls) == 2 + after = await source.get("s", _config()) + assert isinstance(after, Ok) and after.ok.access_token == "fresh" + assert len(poster.calls) == 2 + + +@pytest.mark.asyncio +async def test_refetch_reuses_a_concurrent_replacement_without_a_second_grant(): + poster = _FakePoster([_success("replacement", expires_in=3600)]) + source = ClientCredentialsTokenSource(poster) + seeded = await source.get("s", _config()) + assert isinstance(seeded, Ok) + result = await source.refetch("s", _config(), failed_access_token="some-older-token") + assert result == "replacement" + assert len(poster.calls) == 1 + + +@pytest.mark.asyncio +async def test_refetch_returns_none_when_the_grant_fails(): + poster = _FakePoster([_success("stale"), TokenEndpointUnreachable(detail="down")]) + source = ClientCredentialsTokenSource(poster) + await source.get("s", _config()) + assert await source.refetch("s", _config(), failed_access_token="stale") is None + + +def _upstream(responses: "list[httpx.Response]") -> "tuple[httpx.MockTransport, list[str]]": + # The auth flow re-yields the same Request object on retry, so snapshot the Authorization + # value per send; holding the Request would show the post-retry mutation for both entries. + seen: "list[str]" = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request.headers.get("Authorization", "")) + return responses[min(len(seen) - 1, len(responses) - 1)] + + return httpx.MockTransport(handler), seen + + +@pytest.mark.asyncio +async def test_bearer_auth_sends_the_token_and_leaves_a_success_alone(): + transport, seen = _upstream([httpx.Response(200)]) + + async def refetch(failed: str) -> "str | None": + raise AssertionError("must not refetch on success") + + auth = ClientCredentialsBearerAuth("m2m-token", refetch) + async with httpx.AsyncClient(transport=transport, auth=auth) as client: + response = await client.get("https://upstream.example.com/mcp") + assert response.status_code == 200 + assert seen == ["Bearer m2m-token"] + + +@pytest.mark.asyncio +async def test_bearer_auth_retries_a_401_once_with_a_fresh_token(): + transport, seen = _upstream([httpx.Response(401), httpx.Response(200)]) + refetched: "list[str]" = [] + + async def refetch(failed: str) -> "str | None": + refetched.append(failed) + return "fresh-token" + + auth = ClientCredentialsBearerAuth("stale-token", refetch) + async with httpx.AsyncClient(transport=transport, auth=auth) as client: + response = await client.get("https://upstream.example.com/mcp") + assert response.status_code == 200 + assert refetched == ["stale-token"] + assert seen == ["Bearer stale-token", "Bearer fresh-token"] + + +@pytest.mark.asyncio +async def test_bearer_auth_surfaces_the_401_when_the_refetch_fails(): + transport, seen = _upstream([httpx.Response(401)]) + + async def refetch(failed: str) -> "str | None": + return None + + auth = ClientCredentialsBearerAuth("stale-token", refetch) + async with httpx.AsyncClient(transport=transport, auth=auth) as client: + response = await client.get("https://upstream.example.com/mcp") + assert response.status_code == 401 + assert len(seen) == 1 + + +@pytest.mark.asyncio +async def test_bearer_auth_gives_up_after_a_second_401(): + transport, seen = _upstream([httpx.Response(401), httpx.Response(401)]) + refetched: "list[str]" = [] + + async def refetch(failed: str) -> "str | None": + refetched.append(failed) + return "fresh-token" + + auth = ClientCredentialsBearerAuth("stale-token", refetch) + async with httpx.AsyncClient(transport=transport, auth=auth) as client: + response = await client.get("https://upstream.example.com/mcp") + assert response.status_code == 401 + assert len(seen) == 2 + assert refetched == ["stale-token"] + + +def test_bearer_auth_rejects_sync_clients(): + async def refetch(failed: str) -> "str | None": + return None + + auth = ClientCredentialsBearerAuth("token", refetch) + with httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200)), auth=auth) as client: + with pytest.raises(RuntimeError): + client.get("https://upstream.example.com/mcp") 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 ba7720ffd51..a710da81962 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,10 @@ """Tests for the resolver dispatch: live arms produce auth, stubbed arms fail closed. -`none`, `api_key` (shared-key source), `passthrough`, `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. +`none`, `api_key` (shared-key source), `passthrough`, `authorization_code`, `token_exchange`, and +`client_credentials` 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 @@ -324,9 +325,106 @@ async def test_passthrough_without_inbound_token_is_a_no_op(): assert isinstance(result.ok, NoOpAuth) +class _FakeM2MSource: + """A ClientCredentialsTokenSource returning a canned result and recording refetches.""" + + def __init__(self, result) -> None: + self._result = result + self.gets: list[str] = [] + self.refetches: list[tuple[str, str]] = [] + + async def get(self, server_id: str, config): + self.gets.append(server_id) + return self._result + + async def refetch(self, server_id: str, config, failed_access_token: str): + self.refetches.append((server_id, failed_access_token)) + return "fresh-m2m" + + +_M2M = ClientCredentialsConfig( + client_id="cid", + client_secret=SecretStr("csec"), + token_url="https://idp.example.com/token", +) + + +async def _emitted_async(auth: httpx.Auth, respond=None) -> tuple[httpx.Headers, list[httpx.Request]]: + """Drive the async auth flow one request at a time, replying via ``respond`` when given.""" + seen: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request) + return respond(request) if respond else httpx.Response(200) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler), auth=auth) as client: + await client.get("https://upstream.example.com/mcp") + return seen[-1].headers, seen + + +@pytest.mark.asyncio +async def test_client_credentials_emits_the_minted_bearer(): + source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at"))) + result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials( + _SUBJECT, _spec(_M2M) + ) + assert isinstance(result, Ok) + headers, _ = await _emitted_async(result.ok) + assert headers["Authorization"] == "Bearer m2m-at" + assert source.gets == ["s"] + + +@pytest.mark.asyncio +async def test_client_credentials_ignores_the_subject(): + # The contract's no-user-context clause: every caller shares the one client identity. + source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at"))) + provider = UpstreamCredentialProvider(client_credentials_source=source) + alice = await provider.resolve_credentials(Subject(tenant_id="t1", subject_id="alice"), _spec(_M2M)) + bob = await provider.resolve_credentials(Subject(tenant_id="t2", subject_id="bob"), _spec(_M2M)) + assert isinstance(alice, Ok) and isinstance(bob, Ok) + alice_headers, _ = await _emitted_async(alice.ok) + bob_headers, _ = await _emitted_async(bob.ok) + assert alice_headers["Authorization"] == bob_headers["Authorization"] == "Bearer m2m-at" + + +@pytest.mark.asyncio +async def test_client_credentials_auth_retries_a_401_through_the_source(): + source = _FakeM2MSource(Ok(OAuthToken(access_token="stale-at"))) + result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials( + _SUBJECT, _spec(_M2M) + ) + assert isinstance(result, Ok) + + def respond(request: httpx.Request) -> httpx.Response: + is_stale = request.headers["Authorization"] == "Bearer stale-at" + return httpx.Response(401) if is_stale else httpx.Response(200) + + headers, seen = await _emitted_async(result.ok, respond) + assert headers["Authorization"] == "Bearer fresh-m2m" + assert len(seen) == 2 + assert source.refetches == [("s", "stale-at")] + + +@pytest.mark.asyncio +async def test_client_credentials_propagates_the_source_error(): + source = _FakeM2MSource(Error(CredError.of_upstream_unavailable("idp down"))) + result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials( + _SUBJECT, _spec(_M2M) + ) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_client_credentials_with_no_source_wired_fails_closed_on_missing_config(): + # The default source validates the grant fields before any network is touched. + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, _spec(ClientCredentialsConfig())) + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" + + _STUBBED = [ ("api_key_byok", ApiKeyConfig(key_source=Byok())), - ("client_credentials", ClientCredentialsConfig()), ("aws_sigv4", AwsSigV4Config(region="us-east-1")), ]