mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
feat(mcp): migrate client_credentials (M2M) onto the v2 resolver arm
Replaces the not_implemented stub with a live arm: ClientCredentialsTokenSource mints and caches the M2M token (rotation-aware identity key, expires_in-driven TTL, audience and token_endpoint_auth_method support) and ClientCredentialsBearerAuth retries an upstream 401 exactly once with a freshly minted token. to_server_spec owns oauth2_flow=client_credentials servers and fails closed on incomplete grant config instead of connecting unauthenticated
This commit is contained in:
Tin Chi Lo
Tin
parent
ce88999ed1
commit
d4035a07c7
7 changed files with 901 additions and 18 deletions
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
@ -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")),
|
||||
]
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue