diff --git a/docker/build_from_pip/Dockerfile.build_from_pip b/docker/build_from_pip/Dockerfile.build_from_pip index a5733f0e1a09..fe2ebb9e72af 100644 --- a/docker/build_from_pip/Dockerfile.build_from_pip +++ b/docker/build_from_pip/Dockerfile.build_from_pip @@ -29,7 +29,7 @@ RUN uv venv --python python && \ "litellm[proxy,proxy-runtime]==${LITELLM_VERSION}" \ "google-cloud-aiplatform==1.133.0" \ "google-genai==1.37.0" \ - "anthropic[vertex]==0.84.0" \ + "anthropic[vertex]==1.3.0" \ "grpcio==1.78.0" \ "prometheus-client==0.20.0" \ "langfuse==2.59.7" \ diff --git a/litellm/llms/anthropic/wif.py b/litellm/llms/anthropic/wif.py index e9ef47da59c5..a965c74710b7 100644 --- a/litellm/llms/anthropic/wif.py +++ b/litellm/llms/anthropic/wif.py @@ -1,11 +1,11 @@ """Anthropic workload identity federation: exchanges an external OIDC identity -token for a short-lived ``sk-ant-oat01`` token via the shared RFC 7523 engine.""" +token for a short-lived ``sk-ant-oat01`` token through the Anthropic SDK's credentials helpers.""" import os from collections.abc import Callable, Mapping from itertools import chain from types import MappingProxyType -from typing import Final, NoReturn, TypeVar +from typing import TYPE_CHECKING, Final, NoReturn, TypeVar from urllib.parse import urlsplit, urlunsplit from pydantic import BaseModel, ConfigDict, ValidationError @@ -20,24 +20,19 @@ identity_source_ref, ) from litellm.llms.base_llm.auth.internal_issuer import internal_issuer_assertion_source -from litellm.llms.base_llm.auth.token_exchange import ( - JwtBearerTokenExchangeEngine, - default_token_exchange_engine, -) from litellm.llms.base_llm.auth.types import ( AssertionSourceError, ExchangeError, ExchangeResult, InsecureTokenUrl, MalformedTokenResponse, - MintedToken, TokenEndpointError, - TokenExchangeSpec, TokenTransportError, ) -from litellm.types.llms.anthropic import ANTHROPIC_TOKEN_EXCHANGE_PATH -_JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer" +if TYPE_CHECKING: + from litellm.llms.anthropic.wif_exchange import AnthropicWifTokenExchange + _DEFAULT_API_BASE: Final = "https://api.anthropic.com" _INLINE_ENV_VAR: Final = "ANTHROPIC_IDENTITY_TOKEN" _DISABLE_WIF_PARAM: Final = "anthropic_disable_workload_identity_federation" @@ -262,48 +257,18 @@ def _build_variant( ) from e -def build_anthropic_wif_spec(params: AnthropicWifParams, api_base: str) -> TokenExchangeSpec: - return TokenExchangeSpec( - token_url=api_base.rstrip("/") + ANTHROPIC_TOKEN_EXCHANGE_PATH, - assertion_ref=params.assertion_ref, - assertion_field="assertion", - static_body=MappingProxyType( - { - name: value - for name, value in ( - ("grant_type", _JWT_BEARER_GRANT_TYPE), - ("federation_rule_id", params.federation_rule_id), - ("organization_id", params.organization_id), - ("service_account_id", params.service_account_id), - ("workspace_id", params.workspace_id), - ) - if value is not None - } - ), - body_encoding="json", - request_headers=MappingProxyType({}), - assertion_source=params.assertion_source, - cache_key_identity=( - params.federation_rule_id, - params.organization_id, - params.service_account_id or "", - params.workspace_id or "", - ), - ) - - def get_anthropic_wif_token( litellm_params: Mapping[str, object] | None, api_base: str | None, model: str, - engine: JwtBearerTokenExchangeEngine = default_token_exchange_engine, + exchange: "AnthropicWifTokenExchange | None" = None, ) -> str | None: params: Final = resolve_anthropic_wif_params(litellm_params) if params is None: return None exchange_base: Final = resolve_anthropic_base(api_base) _raise_if_exchange_host_untrusted(exchange_base, model) - result: Final = engine.get_token(build_anthropic_wif_spec(params, exchange_base)) + result: Final = _exchange_or_default(exchange).get_token(params, exchange_base) return _token_from_result(result, model, params) @@ -311,21 +276,36 @@ async def aget_anthropic_wif_token( litellm_params: Mapping[str, object] | None, api_base: str | None, model: str, - engine: JwtBearerTokenExchangeEngine = default_token_exchange_engine, + exchange: "AnthropicWifTokenExchange | None" = None, ) -> str | None: params: Final = resolve_anthropic_wif_params(litellm_params) if params is None: return None exchange_base: Final = resolve_anthropic_base(api_base) _raise_if_exchange_host_untrusted(exchange_base, model) - result: Final = await engine.aget_token(build_anthropic_wif_spec(params, exchange_base)) + result: Final = await _exchange_or_default(exchange).aget_token(params, exchange_base) return _token_from_result(result, model, params) +def _exchange_or_default(exchange: "AnthropicWifTokenExchange | None") -> "AnthropicWifTokenExchange": + """The SDK is a proxy extra, not a core dependency, so it is imported only once a federated + deployment actually needs an exchange, and a missing install says what to add.""" + if exchange is not None: + return exchange + try: + from litellm.llms.anthropic.wif_exchange import default_anthropic_wif_exchange + except ImportError as e: + raise ImportError( + "Anthropic workload identity federation needs the anthropic SDK: " + "install anthropic>=1.3.0 (part of the litellm[proxy] extra)." + ) from e + return default_anthropic_wif_exchange + + def _token_from_result(result: ExchangeResult, model: str, params: AnthropicWifParams) -> str: match result: - case MintedToken(): - return result.access_token.get_secret_value() + case str(): + return result case _: _raise_anthropic_wif_error( result, diff --git a/litellm/llms/anthropic/wif_exchange.py b/litellm/llms/anthropic/wif_exchange.py new file mode 100644 index 000000000000..0bad0f958f1c --- /dev/null +++ b/litellm/llms/anthropic/wif_exchange.py @@ -0,0 +1,471 @@ +"""Anthropic workload identity federation on the Anthropic SDK. ``WorkloadIdentityCredentials`` +performs the RFC 7523 exchange and ``TokenCache`` owns refresh timing, single flight and the 401 +retry. LiteLLM keeps what the SDK cannot know: where the identity token comes from (secret refs, the +credential-dir allowlist, identity sources), one bounded cache per deployment, service metrics, and +the redacted error values ``wif.py`` maps onto the public exception contract. + +The SDK speaks ``httpx2``, a separate HTTP stack from the ``httpx`` LiteLLM's handlers use, so the +exchange client is built here with LiteLLM's SSL settings instead of being borrowed from a handler.""" + +import asyncio +import json +import os +import threading +import time +from collections.abc import Callable, Coroutine, Mapping +from concurrent.futures import Executor, ThreadPoolExecutor +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias + +import httpx2 +from anthropic.lib.credentials import AccessToken, TokenCache, WorkloadIdentityCredentials, WorkloadIdentityError +from pydantic import SecretStr, TypeAdapter +from typing_extensions import assert_never + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.base_llm.auth.oauth_endpoint import ( + drop_reflected_credential, + redact_oauth_error_body, + validate_token_endpoint_url, +) +from litellm.llms.base_llm.auth.types import ( + AssertionReader, + AssertionSource, + AssertionSourceError, + ExchangeCallType, + ExchangeError, + ExchangeResult, + InsecureTokenUrl, + MalformedTokenResponse, + TokenEndpointError, + TokenExchangeMetricsSink, + TokenTransportError, +) +from litellm.types.services import ServiceTypes + +if TYPE_CHECKING: + from litellm.llms.anthropic.wif import AnthropicWifParams + +CALL_TYPE_COLD_MINT: Final[ExchangeCallType] = "cold_mint" +CALL_TYPE_REFRESH: Final[ExchangeCallType] = "refresh" +CALL_TYPE_CACHE_HIT: Final = "cache_hit" +MAX_ASSERTION_BYTES: Final = 16 * 1024 +EXCHANGE_TIMEOUT_SECONDS: Final = 30.0 +EXCHANGE_CONNECT_TIMEOUT_SECONDS: Final = 5.0 +_DETAIL_CAP: Final = 256 +_METRICS_QUEUE_LIMIT: Final = 1000 + + +def _default_assertion_reader(ref: str) -> str | None: + from litellm.secret_managers.main import get_secret_str + + return get_secret_str(ref) + + +def _read_assertion(fetch: AssertionSource, ref: str) -> SecretStr | AssertionSourceError: + from litellm.secret_managers.main import OidcPathNotAllowedError + + try: + raw: Final = fetch() + except OidcPathNotAllowedError: + return AssertionSourceError(kind="disallowed_path", source_ref=ref) + except (ValueError, ImportError) as e: + return AssertionSourceError(kind="unreadable", source_ref=ref, detail=str(e)[:_DETAIL_CAP]) + except Exception: # noqa: BLE001 # injected readers (secret managers) raise arbitrarily; all failures become values + return AssertionSourceError(kind="unreadable", source_ref=ref) + if raw is None: + return AssertionSourceError(kind="missing", source_ref=ref) + stripped: Final = raw.strip() + if not stripped: + return AssertionSourceError(kind="empty", source_ref=ref) + if len(stripped.encode("utf-8")) > MAX_ASSERTION_BYTES: + return AssertionSourceError(kind="oversized", source_ref=ref) + return SecretStr(stripped) + + +class _AssertionUnavailable(WorkloadIdentityError): + """Raised out of the identity token provider so ``TokenCache`` treats a source failure like a + failed exchange: a still-valid cached token keeps serving through the advisory window.""" + + def __init__(self, error: AssertionSourceError) -> None: + super().__init__(f"identity token {error.kind} from {error.source_ref}") + self.error: Final = error + + +class _IdentityTokenSource: + """Reads the identity token fresh for every exchange, so a rotated file or a re-minted + internal-issuer assertion is what the next exchange carries, and remembers the assertion on + the wire so an endpoint that echoes it is scrubbed out of the error.""" + + def __init__(self, fetch: AssertionSource, ref: str) -> None: + self._fetch: Final = fetch + self._ref: Final = ref + self.last_sent: SecretStr | None = None + + def __call__(self) -> str: + match _read_assertion(self._fetch, self._ref): + case SecretStr() as assertion: + self.last_sent = assertion + return assertion.get_secret_value() + case AssertionSourceError() as error: + raise _AssertionUnavailable(error) + + +class _SdkErrorFields(Protocol): + """``WorkloadIdentityError`` declares ``body: Any``; reading it through this protocol keeps + ``Any`` out of the error mapping.""" + + status_code: int | None + body: object + + +_SdkBody: TypeAlias = Mapping[str, object] | str | None +_SDK_BODY_ADAPTER: Final = TypeAdapter[_SdkBody](_SdkBody) + + +def exchange_error(error: WorkloadIdentityError, assertion: SecretStr | None) -> ExchangeError: + if isinstance(error, _AssertionUnavailable): + return error.error + return _endpoint_error(error, str(error).removesuffix("."), assertion) + + +def _endpoint_error(fields: _SdkErrorFields, message: str, assertion: SecretStr | None) -> ExchangeError: + """The SDK folds every failure into one exception class; ``status_code`` and ``body`` tell + them apart. A missing status is a transport failure, a 2xx is a response the SDK could not + use (its message quotes the body, so a reflected assertion is scrubbed from it), and a 4xx/5xx + without a body is one the SDK refused to read for size.""" + if fields.status_code is None: + return TokenTransportError(detail=message) + if fields.status_code < 400: + return MalformedTokenResponse(detail=drop_reflected_credential(message, assertion)[:_DETAIL_CAP]) + body: Final = _SDK_BODY_ADAPTER.validate_python(fields.body) + if body is None: + return TokenEndpointError(status_code=fields.status_code, redacted_body=message) + return redact_oauth_error_body(fields.status_code, body if isinstance(body, str) else json.dumps(body), assertion) + + +def _emit(event: Callable[[], None]) -> None: + try: + event() + except Exception as e: # noqa: BLE001 # metrics are best-effort; a sink failure must never surface to the mint + verbose_logger.debug("token exchange metrics sink raised: %s", e) + + +class _MeteredExchange: + """Times each SDK exchange for the metrics sink and counts them, which is how a ``get_token`` + that never reached the provider is recognised as a cache hit.""" + + def __init__( + self, credentials: WorkloadIdentityCredentials, source: _IdentityTokenSource, sink: TokenExchangeMetricsSink + ) -> None: + self._credentials: Final = credentials + self._source: Final = source + self._sink: Final = sink + self.exchanges: int = 0 + self._minted: bool = False + + def __call__(self, *, force_refresh: bool = False) -> AccessToken: + call_type: Final = CALL_TYPE_REFRESH if self._minted else CALL_TYPE_COLD_MINT + self.exchanges += 1 + started: Final = time.monotonic() + try: + token: Final = self._credentials(force_refresh=force_refresh) + except WorkloadIdentityError as e: + error: Final = exchange_error(e, self._source.last_sent) + _emit( + lambda: self._sink.exchange_failure( + call_type=call_type, duration_seconds=time.monotonic() - started, error=error + ) + ) + raise + self._minted = True + _emit(lambda: self._sink.exchange_success(call_type=call_type, duration_seconds=time.monotonic() - started)) + return token + + +@dataclass(frozen=True, slots=True) +class _CacheKey: + exchange_base: str + assertion_ref: str + federation_rule_id: str + organization_id: str + service_account_id: str | None + workspace_id: str | None + + +@dataclass(frozen=True, slots=True) +class _Deployment: + cache: TokenCache + exchange: _MeteredExchange + source: _IdentityTokenSource + + +def new_exchange_client() -> httpx2.Client: + """The client the SDK would build for itself ignores LiteLLM's SSL settings, so the exchange + gets one built the way ``HTTPHandler`` builds its own: the same CA bundle, verification switch + and client certificate. Redirects stay off: only the bound base URL passed the host allowlist, + and a 3xx must not replay the assertion elsewhere.""" + from litellm.llms.custom_httpx.http_handler import get_ssl_configuration + + return httpx2.Client( + verify=get_ssl_configuration(), + cert=os.getenv("SSL_CERTIFICATE", litellm.ssl_certificate), + timeout=httpx2.Timeout(EXCHANGE_TIMEOUT_SECONDS, connect=EXCHANGE_CONNECT_TIMEOUT_SECONDS), + follow_redirects=False, + ) + + +class AnthropicWifTokenExchange: + """One ``TokenCache`` per (endpoint, identity, federation target), bounded so a deployment + churn cannot grow the process without limit. Everything credential-shaped stays inside the + SDK objects; callers get a token or a typed, redacted ``ExchangeError``.""" + + def __init__( + self, + http_client: httpx2.Client | None = None, + assertion_reader: AssertionReader = _default_assertion_reader, + max_entries: int = 64, + metrics_sink: TokenExchangeMetricsSink | None = None, + clock: Callable[[], float] = time.time, + ) -> None: + self._lock: Final = threading.Lock() + self._http_client: httpx2.Client | None = http_client + self._assertion_reader: Final = assertion_reader + self._max_entries: Final = max_entries + self._clock: Final = clock + self._metrics_sink: Final = metrics_sink if metrics_sink is not None else ServiceLoggingMetricsSink() + self._deployments: Final[dict[_CacheKey, _Deployment]] = {} # mutable-ok: bounded cache, guarded by _lock + + def get_token(self, params: "AnthropicWifParams", exchange_base: str) -> ExchangeResult: + match validate_token_endpoint_url(exchange_base): + case InsecureTokenUrl() as insecure: + _emit( + lambda: self._metrics_sink.exchange_failure( + call_type=CALL_TYPE_COLD_MINT, duration_seconds=0.0, error=insecure + ) + ) + return insecure + case str(): + pass + deployment: Final = self._deployment(params, exchange_base) + exchanges_before: Final = deployment.exchange.exchanges + try: + token: Final = deployment.cache.get_token() + except WorkloadIdentityError as e: + return exchange_error(e, deployment.source.last_sent) + if not token.strip(): + deployment.cache.invalidate() + return MalformedTokenResponse(detail="empty access_token") + if deployment.exchange.exchanges == exchanges_before: + _emit(self._metrics_sink.cache_hit) + return token + + async def aget_token(self, params: "AnthropicWifParams", exchange_base: str) -> ExchangeResult: + return await asyncio.to_thread(self.get_token, params, exchange_base) + + def _deployment(self, params: "AnthropicWifParams", exchange_base: str) -> _Deployment: + key: Final = _CacheKey( + exchange_base=exchange_base, + assertion_ref=params.assertion_ref, + federation_rule_id=params.federation_rule_id, + organization_id=params.organization_id, + service_account_id=params.service_account_id, + workspace_id=params.workspace_id, + ) + with self._lock: + existing: Final = self._deployments.get(key) + if existing is not None: + return existing + if len(self._deployments) >= self._max_entries: + del self._deployments[next(iter(self._deployments))] + created: Final = self._new_deployment(params, exchange_base) + self._deployments[key] = created + return created + + def _new_deployment(self, params: "AnthropicWifParams", exchange_base: str) -> _Deployment: + if self._http_client is None: + self._http_client = new_exchange_client() + fetch: Final[AssertionSource] = ( + params.assertion_source + if params.assertion_source is not None + else lambda: self._assertion_reader(params.assertion_ref) + ) + source: Final = _IdentityTokenSource(fetch, params.assertion_ref) + credentials: Final = WorkloadIdentityCredentials( + identity_token_provider=source, + federation_rule_id=params.federation_rule_id, + organization_id=params.organization_id, + service_account_id=params.service_account_id, + workspace_id=params.workspace_id, + http_client=self._http_client, + ) + credentials.bind_base_url(exchange_base) + exchange: Final = _MeteredExchange(credentials, source, self._metrics_sink) + return _Deployment(cache=TokenCache(exchange, time_source=self._clock), exchange=exchange, source=source) + + +def _error_summary(error: ExchangeError) -> str: + match error: + case AssertionSourceError(): + return f"AssertionSourceError: assertion {error.kind} from {error.source_ref}" + case InsecureTokenUrl(): + return f"InsecureTokenUrl: insecure token endpoint host {error.host}" + case TokenEndpointError(): + return f"TokenEndpointError: HTTP {error.status_code}: {error.redacted_body}" + case TokenTransportError(): + return f"TokenTransportError: {error.detail}" + case MalformedTokenResponse(): + return f"MalformedTokenResponse: {error.detail}" + case _: + assert_never(error) + + +class _MetricsFailure(Exception): + """Never raised: typed carriers handed to the service failure hook so the prometheus + ``error_class`` label names the ``ExchangeError`` variant; the message is the redacted + ``_error_summary`` and carries no credential material.""" + + +class TokenExchangeAssertionSourceFailure(_MetricsFailure): ... + + +class TokenExchangeInsecureUrlFailure(_MetricsFailure): ... + + +class TokenExchangeEndpointFailure(_MetricsFailure): ... + + +class TokenExchangeTransportFailure(_MetricsFailure): ... + + +class TokenExchangeMalformedResponseFailure(_MetricsFailure): ... + + +def _failure_exception(error: ExchangeError) -> _MetricsFailure: + summary: Final = _error_summary(error) + match error: + case AssertionSourceError(): + return TokenExchangeAssertionSourceFailure(summary) + case InsecureTokenUrl(): + return TokenExchangeInsecureUrlFailure(summary) + case TokenEndpointError(): + return TokenExchangeEndpointFailure(summary) + case TokenTransportError(): + return TokenExchangeTransportFailure(summary) + case MalformedTokenResponse(): + return TokenExchangeMalformedResponseFailure(summary) + case _: + assert_never(error) + + +class _ServiceLoggingHooks(Protocol): + """The slice of ``litellm._service_logger.ServiceLogging`` the metrics sink calls; a protocol + so tests inject a recorder instead of monkeypatching.""" + + async def async_service_success_hook(self, service: ServiceTypes, call_type: str, duration: float) -> None: ... + + async def async_service_failure_hook( + self, service: ServiceTypes, duration: float, error: str | Exception, call_type: str + ) -> None: ... + + +_HooksCoroFactory: TypeAlias = Callable[ + [_ServiceLoggingHooks], # mutable-ok: Callable param-list syntax, not a list + Coroutine[object, object, None], +] + + +def _default_service_logging() -> _ServiceLoggingHooks: + from litellm._service_logger import ServiceLogging + + return ServiceLogging() + + +class ServiceLoggingMetricsSink: + """Default sink: bridges exchange metrics onto litellm's ServiceTypes pattern + (prometheus ``litellm_anthropic_wif_*`` via ``service_callback``). Exchanges run on sync + threads with no event loop, and the service hooks are async, so every emission is + fire-and-forget on a dedicated single worker thread that owns its own short-lived loop -- + the mint path only ever pays for an executor queue put.""" + + def __init__( + self, + service_logging_factory: Callable[[], _ServiceLoggingHooks] = _default_service_logging, + executor: Executor | None = None, + ) -> None: + self._lock: Final = threading.Lock() + self._service_logging_factory: Final = service_logging_factory + self._service_logging: _ServiceLoggingHooks | None = None + self._executor: Executor | None = executor + self._queued: int = 0 # rebind-ok: backlog depth, guarded by _lock + + def _service_logging_instance(self) -> _ServiceLoggingHooks: + with self._lock: + if self._service_logging is None: + self._service_logging = self._service_logging_factory() + return self._service_logging + + def _executor_instance(self) -> Executor: + with self._lock: + if self._executor is None: + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="litellm-token-exchange-metrics") + return self._executor + + def _emit(self, coro_factory: _HooksCoroFactory) -> None: + try: + asyncio.run(coro_factory(self._service_logging_instance())) + except Exception as e: # noqa: BLE001 # metrics are best-effort; emission failures must never surface + verbose_logger.debug("token exchange metrics emission failed: %s", e) + + def _submit(self, coro_factory: _HooksCoroFactory) -> None: + """Drop the event rather than queue it once the backlog is full. A stalled telemetry + backend must not let request volume grow an unbounded queue in the proxy: losing a + metric sample is always cheaper than losing the process.""" + with self._lock: + if self._queued >= _METRICS_QUEUE_LIMIT: + verbose_logger.debug("token exchange metrics queue full, dropping event") + return + self._queued += 1 + try: + self._executor_instance().submit(self._emit_and_release, coro_factory) + except Exception as e: # noqa: BLE001 # a rejected submit must not surface to the mint + with self._lock: + self._queued -= 1 + verbose_logger.debug("token exchange metrics submit failed: %s", e) + + def _emit_and_release(self, coro_factory: _HooksCoroFactory) -> None: + try: + self._emit(coro_factory) + finally: + with self._lock: + self._queued -= 1 + + def exchange_success(self, *, call_type: ExchangeCallType, duration_seconds: float) -> None: + def start(hooks: _ServiceLoggingHooks) -> Coroutine[object, object, None]: + return hooks.async_service_success_hook( + service=ServiceTypes.ANTHROPIC_WIF, call_type=call_type, duration=duration_seconds + ) + + self._submit(start) + + def exchange_failure(self, *, call_type: ExchangeCallType, duration_seconds: float, error: ExchangeError) -> None: + failure: Final = _failure_exception(error) + + def start(hooks: _ServiceLoggingHooks) -> Coroutine[object, object, None]: + return hooks.async_service_failure_hook( + service=ServiceTypes.ANTHROPIC_WIF, duration=duration_seconds, error=failure, call_type=call_type + ) + + self._submit(start) + + def cache_hit(self) -> None: + def start(hooks: _ServiceLoggingHooks) -> Coroutine[object, object, None]: + return hooks.async_service_success_hook( + service=ServiceTypes.ANTHROPIC_WIF_CACHE, call_type=CALL_TYPE_CACHE_HIT, duration=0.0 + ) + + self._submit(start) + + +default_anthropic_wif_exchange: Final = AnthropicWifTokenExchange() diff --git a/litellm/llms/base_llm/auth/__init__.py b/litellm/llms/base_llm/auth/__init__.py index 291e1492c989..281ec252d0da 100644 --- a/litellm/llms/base_llm/auth/__init__.py +++ b/litellm/llms/base_llm/auth/__init__.py @@ -26,14 +26,8 @@ rfc7638_thumbprint, sign_es256_jwt, ) -from litellm.llms.base_llm.auth.token_exchange import ( - ADVISORY_REFRESH_BACKOFF_SECONDS, - ADVISORY_REFRESH_SECONDS, - MANDATORY_REFRESH_SECONDS, - MAX_ASSERTION_BYTES, +from litellm.llms.base_llm.auth.oauth_endpoint import ( MAX_RESPONSE_BYTES, - JwtBearerTokenExchangeEngine, - default_token_exchange_engine, redact_oauth_error_body, validate_token_endpoint_url, ) @@ -41,48 +35,36 @@ AssertionReader, AssertionSource, AssertionSourceError, - BodyEncoding, ExchangeError, ExchangeResult, InsecureTokenUrl, MalformedTokenResponse, - MintedToken, SyncTokenPoster, TokenEndpointError, - TokenExchangeSpec, TokenTransportError, ) __all__ = ( - "ADVISORY_REFRESH_BACKOFF_SECONDS", - "ADVISORY_REFRESH_SECONDS", "ALG", - "MANDATORY_REFRESH_SECONDS", - "MAX_ASSERTION_BYTES", "MAX_RESPONSE_BYTES", "AnthropicIdentitySourceConfig", "AnthropicIdentitySourceKind", "AssertionReader", "AssertionSource", "AssertionSourceError", - "BodyEncoding", "ExchangeError", "ExchangeResult", "InsecureTokenUrl", "InternalIssuerSource", - "JwtBearerTokenExchangeEngine", "KeycloakSource", "MalformedTokenResponse", - "MintedToken", "SecretReader", "SigningKeyReader", "SyncTokenPoster", "TokenEndpointError", - "TokenExchangeSpec", "TokenTransportError", "build_jwk", "build_jwks", - "default_token_exchange_engine", "fetch_keycloak_assertion", "identity_source_config_adapter", "identity_source_ref", diff --git a/litellm/llms/base_llm/auth/client_credentials.py b/litellm/llms/base_llm/auth/client_credentials.py index b2e683a9ed40..0f430bc3229a 100644 --- a/litellm/llms/base_llm/auth/client_credentials.py +++ b/litellm/llms/base_llm/auth/client_credentials.py @@ -22,7 +22,7 @@ from typing_extensions import assert_never from litellm.llms.base_llm.auth.identity_source import KeycloakSource, ref_for_error_message -from litellm.llms.base_llm.auth.token_exchange import ( +from litellm.llms.base_llm.auth.oauth_endpoint import ( MAX_RESPONSE_BYTES, endpoint_url_for_error_message, redact_oauth_error_body, @@ -222,6 +222,6 @@ def keycloak_assertion_source( secret_reader: SecretReader = _default_secret_reader, ) -> Callable[[], str]: """A zero-arg closure that fetches fresh on every call: the shape an ``oidc/keycloak/...`` - ref dispatches to once wired into ``TokenExchangeSpec.assertion_source`` (Phase 1 decision 7) + ref dispatches to once wired into ``AnthropicWifParams.assertion_source`` (Phase 1 decision 7) -- the caller parses the config and closes this function over it, with no registry involved.""" return lambda: fetch_keycloak_assertion(config, poster=poster, secret_reader=secret_reader) diff --git a/litellm/llms/base_llm/auth/internal_issuer.py b/litellm/llms/base_llm/auth/internal_issuer.py index 444ca7fd5b28..ccbacaacdb56 100644 --- a/litellm/llms/base_llm/auth/internal_issuer.py +++ b/litellm/llms/base_llm/auth/internal_issuer.py @@ -71,7 +71,7 @@ def internal_issuer_assertion_source( clock: Callable[[], float] = time.time, ) -> Callable[[], str]: """A zero-arg closure that mints fresh on every call: the shape an ``oidc/internal_issuer/...`` - ref dispatches to once wired into ``TokenExchangeSpec.assertion_source`` (Phase 1 decision 7) + ref dispatches to once wired into ``AnthropicWifParams.assertion_source`` (Phase 1 decision 7) -- the caller parses the config and closes this function over it, with no registry involved.""" return lambda: mint_internal_issuer_assertion(config, key_reader=key_reader, clock=clock) diff --git a/litellm/llms/base_llm/auth/oauth_endpoint.py b/litellm/llms/base_llm/auth/oauth_endpoint.py new file mode 100644 index 000000000000..405a1878bed2 --- /dev/null +++ b/litellm/llms/base_llm/auth/oauth_endpoint.py @@ -0,0 +1,181 @@ +"""OAuth 2.0 token endpoint hygiene shared by every grant LiteLLM posts itself: HTTPS pinning, +RFC 6749 5.2 error redaction with the sent credential scrubbed out of an echoing response, and the +response guard a raw ``httpx`` poster needs. Providers map the resulting typed values onto their +own public exception contract.""" + +import re +from collections.abc import Mapping, Sequence +from typing import Final, TypeAlias +from urllib.parse import unquote, unquote_plus, urlsplit, urlunsplit + +import httpx +from pydantic import SecretStr, TypeAdapter, ValidationError + +from litellm.llms.base_llm.auth.types import InsecureTokenUrl, TokenEndpointError + +MAX_RESPONSE_BYTES: Final = 1024 * 1024 + +_REDACTION_CAP: Final = 256 +_LOCAL_HOSTS: Final = frozenset({"localhost", "127.0.0.1", "::1"}) +_OAUTH_ERROR_FIELDS: Final = ("error", "error_description", "error_uri") +_NESTED_ERROR_FIELDS: Final = ("type", "message") +_OVERSIZED_BODY_MESSAGE: Final = "oversized error response omitted" +_NON_OBJECT_BODY_MESSAGE: Final = "non-object error response omitted" +_NO_OAUTH_FIELDS_MESSAGE: Final = "error response carried no RFC 6749 fields" +_UNSTRUCTURED_BODY_MESSAGE: Final = "non-JSON error response omitted" +_REFLECTED_VALUE_MESSAGE: Final = "" +# A credential fragment shorter than this is not worth the false positives; longer, and a run +# shared with the assertion is reflection rather than coincidence. +_REFLECTION_MIN_RUN: Final = 8 +# Everything a base64url credential is NOT made of, stripped so a fragment split by delimiters +# still lines up against the assertion. +_CREDENTIAL_CHARS: Final = re.compile(r"[^A-Za-z0-9._~+/=-]") +_SENTINEL_BODY_MESSAGES: Final = frozenset({_OVERSIZED_BODY_MESSAGE, _NON_OBJECT_BODY_MESSAGE}) + + +_RedactableBody: TypeAlias = Mapping[str, object] | list[object] | str | int | float | bool | None +_REDACTABLE_BODY_ADAPTER: Final = TypeAdapter[_RedactableBody](_RedactableBody) + + +def endpoint_url_for_error_message(url: str) -> str: + """``url`` reduced to scheme, host and path for operator-facing errors. + + A token endpoint is configuration, not a secret, and naming it is what makes these errors + actionable. But nothing stops an operator writing a credential into it, as a query parameter + or as userinfo, and these errors reach model callers, so neither part is echoed. + """ + parsed: Final = urlsplit(url) + host: Final = parsed.hostname or "" + authority: Final = f"{host}:{parsed.port}" if parsed.port is not None else host + return urlunsplit((parsed.scheme, authority, parsed.path, "", "")) + + +def validate_token_endpoint_url(url: str) -> str | InsecureTokenUrl: + parsed: Final = urlsplit(url) + if parsed.scheme == "https": + return url + if parsed.scheme == "http" and (parsed.hostname or "") in _LOCAL_HOSTS: + return url + return InsecureTokenUrl(host=parsed.hostname or "") + + +def redact_oauth_error_body( + status_code: int, + body_text: str, + assertion: SecretStr | Sequence[SecretStr] | None = None, +) -> TokenEndpointError: + """``assertion`` may be every form of the credential that went out on the wire. + + A grant that encodes its credential before sending it (``client_secret_basic`` base64s + ``id:secret``) can have that encoded form echoed back, and it decodes straight to the secret, + so checking only the raw value lets reversible material through. + """ + rendered: Final = _redact_body_text(body_text) + secrets: Final = () if assertion is None else (assertion,) if isinstance(assertion, SecretStr) else tuple(assertion) + redacted: Final = next( + ( + _REFLECTED_VALUE_MESSAGE + for secret in secrets + if drop_reflected_credential(rendered, secret) is _REFLECTED_VALUE_MESSAGE + ), + rendered, + ) + return TokenEndpointError(status_code=status_code, redacted_body=redacted) + + +def drop_reflected_credential(rendered: str, assertion: SecretStr | None) -> str: + """Catches an endpoint that echoes the submitted credential back, verbatim or in fragments, + however it split or percent-encoded it. + + Both sides are reduced to the characters a credential is made of before comparison. Stripping + only the rendered side would stop matching a secret that carries spaces or punctuation of its + own, which is exactly the hand-set passphrase most at risk of being echoed. + + This stops an accidental or naive echo. It cannot stop an endpoint that deliberately re-encodes + or interleaves the credential, and it is not what keeps the credential from the endpoint, which + already holds it. What it protects is blast radius: keeping the value out of the caller's error + and out of third-party log sinks. + """ + if assertion is None: + return rendered + secret: Final = assertion.get_secret_value() + if not secret: + return rendered + if secret in rendered: + return _REFLECTED_VALUE_MESSAGE + compacted_secret: Final = _CREDENTIAL_CHARS.sub("", secret) + if not compacted_secret: + return rendered + return _REFLECTED_VALUE_MESSAGE if _shares_a_credential_run(rendered, compacted_secret) else rendered + + +def _shares_a_credential_run(rendered: str, compacted_secret: str) -> bool: + """``unquote`` covers a credential sent form-encoded, without every caller enumerating that + shape for itself: percent-escaping is reversible and applies to any field, query string + included. + + A secret shorter than the probe run is compared whole: a window longer than the secret can + never be found inside it, which would leave a short client secret unprotected in every shape + but the verbatim one. + """ + # unquote covers %XX; unquote_plus additionally covers the "+" a form-encoded body uses for a + # space. Both are kept rather than only the wider one, because "+" is a base64 character and + # decoding it away would lose a run that the undecoded candidate still matches on. + run: Final = min(_REFLECTION_MIN_RUN, len(compacted_secret)) + compacted_candidates: Final = tuple( + _CREDENTIAL_CHARS.sub("", candidate) for candidate in (rendered, unquote(rendered), unquote_plus(rendered)) + ) + return any( + compacted[start : start + run] in compacted_secret + for compacted in compacted_candidates + for start in range(len(compacted) - run + 1) + ) + + +def _redact_body_text(body_text: str) -> str: + if body_text in _SENTINEL_BODY_MESSAGES: + return body_text + if len(body_text) > MAX_RESPONSE_BYTES: + return _OVERSIZED_BODY_MESSAGE + try: + parsed: Final = _REDACTABLE_BODY_ADAPTER.validate_json(body_text) + except ValidationError: + return _UNSTRUCTURED_BODY_MESSAGE + match parsed: + case Mapping(): + return _format_oauth_error_fields(parsed) + case _: + return _NON_OBJECT_BODY_MESSAGE + + +def _format_oauth_error_fields(body: Mapping[str, object]) -> str: + fields: Final = tuple( + f"{name}: {_format_oauth_error_value(value)}" + for name in _OAUTH_ERROR_FIELDS + for value in (body.get(name),) + if value is not None + ) + return "; ".join(fields) if fields else _NO_OAUTH_FIELDS_MESSAGE + + +def _format_oauth_error_value(value: object) -> str: + """RFC 6749 types ``error`` as a string, but Anthropic (and other providers) nest their + own ``{"type": ..., "message": ...}`` envelope there; render that rather than a dict repr.""" + if isinstance(value, Mapping): + nested: Final = tuple( + f"{str(part)[:_REDACTION_CAP]}" + for key in _NESTED_ERROR_FIELDS + for part in (value.get(key),) + if part is not None + ) + if nested: + return " - ".join(nested) + return str(value)[:_REDACTION_CAP] + + +def require_posted_response(response: httpx.Response | None, endpoint_label: str) -> httpx.Response: + """The legacy ``HTTPHandler`` carries no return annotation, so a patched or stubbed client can + hand a poster ``None`` back; a transport error beats dereferencing it.""" + if response is None: + raise httpx.TransportError(f"{endpoint_label} returned no response") + return response diff --git a/litellm/llms/base_llm/auth/token_exchange.py b/litellm/llms/base_llm/auth/token_exchange.py deleted file mode 100644 index be2fb46d1b35..000000000000 --- a/litellm/llms/base_llm/auth/token_exchange.py +++ /dev/null @@ -1,861 +0,0 @@ -"""RFC 7523 JWT-bearer token exchange engine, shared across providers. - -One sync state machine per process: bounded engine-owned entry map, two-tier -refresh (advisory background refresh + mandatory single-flight), HTTPS pinning, -response caps, and RFC 6749 5.2 redaction. Providers describe a grant profile as -a ``TokenExchangeSpec`` and map the typed ``ExchangeError`` union to their own -public exception contract. -""" - -import asyncio -import hashlib -import json -import re -import threading -import time -from collections.abc import Callable, Coroutine, Mapping, Sequence -from concurrent.futures import Executor, ThreadPoolExecutor -from dataclasses import dataclass -from math import inf -from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol, TypeAlias -from urllib.parse import unquote, unquote_plus, urlencode, urlsplit, urlunsplit - -import httpx -from pydantic import BaseModel, SecretStr, TypeAdapter, ValidationError -from typing_extensions import assert_never - -from litellm._logging import verbose_logger -from litellm.llms.base_llm.auth.types import ( - AssertionReader, - AssertionSource, - AssertionSourceError, - ExchangeCallType, - ExchangeError, - ExchangeResult, - InsecureTokenUrl, - MalformedTokenResponse, - MintedToken, - SyncTokenPoster, - TokenEndpointError, - TokenExchangeMetricsSink, - TokenExchangeSpec, - TokenTransportError, -) -from litellm.types.services import ServiceTypes - -if TYPE_CHECKING: - from litellm.llms.custom_httpx.http_handler import HTTPHandler - -CALL_TYPE_COLD_MINT: Final[ExchangeCallType] = "cold_mint" -CALL_TYPE_MANDATORY_REFRESH: Final[ExchangeCallType] = "mandatory_refresh" -CALL_TYPE_ADVISORY_REFRESH: Final[ExchangeCallType] = "advisory_refresh" -CALL_TYPE_CACHE_HIT: Final = "cache_hit" - -ADVISORY_REFRESH_SECONDS: Final = 120.0 -MANDATORY_REFRESH_SECONDS: Final = 30.0 -ADVISORY_REFRESH_LIFETIME_FRACTION: Final = 0.5 -MANDATORY_REFRESH_LIFETIME_FRACTION: Final = 0.125 -ADVISORY_REFRESH_BACKOFF_SECONDS: Final = 5.0 -FALLBACK_TOKEN_TTL_SECONDS: Final = 60.0 -# Metrics are best-effort, so the backlog is capped and further events are dropped. Request volume -# must not be able to grow this queue without bound when a telemetry backend stalls. -_METRICS_QUEUE_LIMIT: Final = 1000 -MAX_ASSERTION_BYTES: Final = 16 * 1024 -MAX_RESPONSE_BYTES: Final = 1024 * 1024 - -_REDACTION_CAP: Final = 256 -_FOLLOWER_WAIT_GRACE_SECONDS: Final = 5.0 -_LOCAL_HOSTS: Final = frozenset({"localhost", "127.0.0.1", "::1"}) -_OAUTH_ERROR_FIELDS: Final = ("error", "error_description", "error_uri") -_NESTED_ERROR_FIELDS: Final = ("type", "message") -_CONTENT_TYPES: Final = MappingProxyType({"json": "application/json", "form": "application/x-www-form-urlencoded"}) -_OVERSIZED_BODY_MESSAGE: Final = "oversized error response omitted" -_NON_OBJECT_BODY_MESSAGE: Final = "non-object error response omitted" -_NO_OAUTH_FIELDS_MESSAGE: Final = "error response carried no RFC 6749 fields" -_UNSTRUCTURED_BODY_MESSAGE: Final = "non-JSON error response omitted" -_REFLECTED_VALUE_MESSAGE: Final = "" -# A credential fragment shorter than this is not worth the false positives; longer, and a run -# shared with the assertion is reflection rather than coincidence. -_REFLECTION_MIN_RUN: Final = 8 -# Everything a base64url credential is NOT made of, stripped so a fragment split by delimiters -# still lines up against the assertion. -_CREDENTIAL_CHARS: Final = re.compile(r"[^A-Za-z0-9._~+/=-]") -_SENTINEL_BODY_MESSAGES: Final = frozenset({_OVERSIZED_BODY_MESSAGE, _NON_OBJECT_BODY_MESSAGE}) - - -class _TokenExchangeResponse(BaseModel): - access_token: str - expires_in: int | None = None - token_type: str | None = None - - -_RedactableBody: TypeAlias = Mapping[str, object] | list[object] | str | int | float | bool | None -_REDACTABLE_BODY_ADAPTER: Final = TypeAdapter[_RedactableBody](_RedactableBody) - - -def endpoint_url_for_error_message(url: str) -> str: - """``url`` reduced to scheme, host and path for operator-facing errors. - - A token endpoint is configuration, not a secret, and naming it is what makes these errors - actionable. But nothing stops an operator writing a credential into it, as a query parameter - or as userinfo, and these errors reach model callers, so neither part is echoed. - """ - parsed: Final = urlsplit(url) - host: Final = parsed.hostname or "" - authority: Final = f"{host}:{parsed.port}" if parsed.port is not None else host - return urlunsplit((parsed.scheme, authority, parsed.path, "", "")) - - -def validate_token_endpoint_url(url: str) -> str | InsecureTokenUrl: - parsed: Final = urlsplit(url) - if parsed.scheme == "https": - return url - if parsed.scheme == "http" and (parsed.hostname or "") in _LOCAL_HOSTS: - return url - return InsecureTokenUrl(host=parsed.hostname or "") - - -def redact_oauth_error_body( - status_code: int, - body_text: str, - assertion: SecretStr | Sequence[SecretStr] | None = None, -) -> TokenEndpointError: - """``assertion`` may be every form of the credential that went out on the wire. - - A grant that encodes its credential before sending it (``client_secret_basic`` base64s - ``id:secret``) can have that encoded form echoed back, and it decodes straight to the secret, - so checking only the raw value lets reversible material through. - """ - rendered: Final = _redact_body_text(body_text) - secrets: Final = () if assertion is None else (assertion,) if isinstance(assertion, SecretStr) else tuple(assertion) - redacted: Final = next( - ( - _REFLECTED_VALUE_MESSAGE - for secret in secrets - if _drop_reflected_assertion(rendered, secret) is _REFLECTED_VALUE_MESSAGE - ), - rendered, - ) - return TokenEndpointError(status_code=status_code, redacted_body=redacted) - - -def _drop_reflected_assertion(rendered: str, assertion: SecretStr | None) -> str: - """Catches an endpoint that echoes the submitted credential back, verbatim or in fragments, - however it split or percent-encoded it. - - Both sides are reduced to the characters a credential is made of before comparison. Stripping - only the rendered side would stop matching a secret that carries spaces or punctuation of its - own, which is exactly the hand-set passphrase most at risk of being echoed. - - This stops an accidental or naive echo. It cannot stop an endpoint that deliberately re-encodes - or interleaves the credential, and it is not what keeps the credential from the endpoint, which - already holds it. What it protects is blast radius: keeping the value out of the caller's error - and out of third-party log sinks. - """ - if assertion is None: - return rendered - secret: Final = assertion.get_secret_value() - if not secret: - return rendered - if secret in rendered: - return _REFLECTED_VALUE_MESSAGE - compacted_secret: Final = _CREDENTIAL_CHARS.sub("", secret) - if not compacted_secret: - return rendered - return _REFLECTED_VALUE_MESSAGE if _shares_a_credential_run(rendered, compacted_secret) else rendered - - -def _shares_a_credential_run(rendered: str, compacted_secret: str) -> bool: - """``unquote`` covers a credential sent form-encoded, without every caller enumerating that - shape for itself: percent-escaping is reversible and applies to any field, query string - included. - - A secret shorter than the probe run is compared whole: a window longer than the secret can - never be found inside it, which would leave a short client secret unprotected in every shape - but the verbatim one. - """ - # unquote covers %XX; unquote_plus additionally covers the "+" a form-encoded body uses for a - # space. Both are kept rather than only the wider one, because "+" is a base64 character and - # decoding it away would lose a run that the undecoded candidate still matches on. - run: Final = min(_REFLECTION_MIN_RUN, len(compacted_secret)) - compacted_candidates: Final = tuple( - _CREDENTIAL_CHARS.sub("", candidate) for candidate in (rendered, unquote(rendered), unquote_plus(rendered)) - ) - return any( - compacted[start : start + run] in compacted_secret - for compacted in compacted_candidates - for start in range(len(compacted) - run + 1) - ) - - -def _redact_body_text(body_text: str) -> str: - if body_text in _SENTINEL_BODY_MESSAGES: - return body_text - if len(body_text) > MAX_RESPONSE_BYTES: - return _OVERSIZED_BODY_MESSAGE - try: - parsed: Final = _REDACTABLE_BODY_ADAPTER.validate_json(body_text) - except ValidationError: - return _UNSTRUCTURED_BODY_MESSAGE - match parsed: - case Mapping(): - return _format_oauth_error_fields(parsed) - case _: - return _NON_OBJECT_BODY_MESSAGE - - -def _format_oauth_error_fields(body: Mapping[str, object]) -> str: - fields: Final = tuple( - f"{name}: {_format_oauth_error_value(value)}" - for name in _OAUTH_ERROR_FIELDS - for value in (body.get(name),) - if value is not None - ) - return "; ".join(fields) if fields else _NO_OAUTH_FIELDS_MESSAGE - - -def _format_oauth_error_value(value: object) -> str: - """RFC 6749 types ``error`` as a string, but Anthropic (and other providers) nest their - own ``{"type": ..., "message": ...}`` envelope there; render that rather than a dict repr.""" - if isinstance(value, Mapping): - nested: Final = tuple( - f"{str(part)[:_REDACTION_CAP]}" - for key in _NESTED_ERROR_FIELDS - for part in (value.get(key),) - if part is not None - ) - if nested: - return " - ".join(nested) - return str(value)[:_REDACTION_CAP] - - -def _error_summary(error: ExchangeError) -> str: - match error: - case AssertionSourceError(): - return f"AssertionSourceError: assertion {error.kind} from {error.source_ref}" - case InsecureTokenUrl(): - return f"InsecureTokenUrl: insecure token endpoint host {error.host}" - case TokenEndpointError(): - return f"TokenEndpointError: HTTP {error.status_code}: {error.redacted_body}" - case TokenTransportError(): - return f"TokenTransportError: {error.detail}" - case MalformedTokenResponse(): - return f"MalformedTokenResponse: {error.detail}" - case _: - assert_never(error) - - -class _MetricsFailure(Exception): - """Never raised: typed carriers handed to the service failure hook so the prometheus - ``error_class`` label names the ``ExchangeError`` variant; the message is the redacted - ``_error_summary`` and carries no credential material.""" - - -class TokenExchangeAssertionSourceFailure(_MetricsFailure): ... - - -class TokenExchangeInsecureUrlFailure(_MetricsFailure): ... - - -class TokenExchangeEndpointFailure(_MetricsFailure): ... - - -class TokenExchangeTransportFailure(_MetricsFailure): ... - - -class TokenExchangeMalformedResponseFailure(_MetricsFailure): ... - - -def _failure_exception(error: ExchangeError) -> _MetricsFailure: - summary: Final = _error_summary(error) - match error: - case AssertionSourceError(): - return TokenExchangeAssertionSourceFailure(summary) - case InsecureTokenUrl(): - return TokenExchangeInsecureUrlFailure(summary) - case TokenEndpointError(): - return TokenExchangeEndpointFailure(summary) - case TokenTransportError(): - return TokenExchangeTransportFailure(summary) - case MalformedTokenResponse(): - return TokenExchangeMalformedResponseFailure(summary) - case _: - assert_never(error) - - -def _cache_key(spec: TokenExchangeSpec) -> str: - return hashlib.sha256( - "\x1f".join((spec.token_url, spec.assertion_ref, *spec.cache_key_identity)).encode() - ).hexdigest() - - -def _assertion_fetch(reader: AssertionReader, spec: TokenExchangeSpec) -> AssertionSource: - """``spec.assertion_source`` (an identity source's own fetch/mint closure) takes priority over - the engine-level reader when set; either way, failures are reported against ``spec.assertion_ref``.""" - if spec.assertion_source is not None: - return spec.assertion_source - return lambda: reader(spec.assertion_ref) - - -def _read_assertion(fetch: AssertionSource, ref: str) -> SecretStr | AssertionSourceError: - from litellm.secret_managers.main import OidcPathNotAllowedError - - try: - raw: Final = fetch() - except OidcPathNotAllowedError: - return AssertionSourceError(kind="disallowed_path", source_ref=ref) - except (ValueError, ImportError) as e: - return AssertionSourceError(kind="unreadable", source_ref=ref, detail=str(e)[:_REDACTION_CAP]) - except Exception: # noqa: BLE001 # injected readers (secret managers) raise arbitrarily; all failures become values - return AssertionSourceError(kind="unreadable", source_ref=ref) - if raw is None: - return AssertionSourceError(kind="missing", source_ref=ref) - stripped: Final = raw.strip() - if not stripped: - return AssertionSourceError(kind="empty", source_ref=ref) - if len(stripped.encode("utf-8")) > MAX_ASSERTION_BYTES: - return AssertionSourceError(kind="oversized", source_ref=ref) - return SecretStr(stripped) - - -def _serialize_body(spec: TokenExchangeSpec, assertion: SecretStr) -> bytes: - if spec.body_encoding == "json": - return json.dumps( - { # mutable-ok: transient body dict consumed inline by the serializer - **spec.static_body, - spec.assertion_field: assertion.get_secret_value(), - } - ).encode() - return urlencode( - { # mutable-ok: transient body dict consumed inline by the serializer - **spec.static_body, - spec.assertion_field: assertion.get_secret_value(), - } - ).encode() - - -def _sanitize_expires_in(expires_in: int | None) -> float: - if expires_in is None or expires_in <= 0: - return FALLBACK_TOKEN_TTL_SECONDS - return float(expires_in) - - -@dataclass(frozen=True, slots=True) -class _RefreshWindows: - advisory: float - mandatory: float - - -def _refresh_windows(lifetime_seconds: float | None) -> _RefreshWindows: - """A token whose whole life is shorter than the flat windows sits inside them from the moment it - is minted, so every request would arm another background exchange against the token endpoint. - Scaling each window by a fraction of the observed lifetime makes a 60s token refresh around its - half life instead; at a lifetime of 240s and above both fractions reach the flat windows, so - ordinary long-lived tokens keep exactly the 120s/30s behaviour.""" - if lifetime_seconds is None or lifetime_seconds <= 0.0: - return _RefreshWindows(advisory=ADVISORY_REFRESH_SECONDS, mandatory=MANDATORY_REFRESH_SECONDS) - return _RefreshWindows( - advisory=min(ADVISORY_REFRESH_SECONDS, lifetime_seconds * ADVISORY_REFRESH_LIFETIME_FRACTION), - mandatory=min(MANDATORY_REFRESH_SECONDS, lifetime_seconds * MANDATORY_REFRESH_LIFETIME_FRACTION), - ) - - -def _capped_body_text(response: httpx.Response) -> str: - if len(response.content) > MAX_RESPONSE_BYTES: - return _OVERSIZED_BODY_MESSAGE - return response.text - - -def _default_assertion_reader(ref: str) -> str | None: - from litellm.secret_managers.main import get_secret_str - - return get_secret_str(ref) - - -def _new_exchange_handler() -> "HTTPHandler": - from litellm.llms.custom_httpx.http_handler import HTTPHandler - - handler: Final = HTTPHandler(timeout=httpx.Timeout(timeout=30.0, connect=5.0)) - handler.client.follow_redirects = False - return handler - - -def require_posted_response(response: httpx.Response | None, endpoint_label: str) -> httpx.Response: - """The legacy ``HTTPHandler`` carries no return annotation, so a patched or stubbed client can - hand a poster ``None`` back; a transport error beats dereferencing it.""" - if response is None: - raise httpx.TransportError(f"{endpoint_label} returned no response") - return response - - -class _HttpxSyncTokenPoster: - """Default poster: a dedicated HTTPHandler (no logging_obj, so litellm's - pre/post-call body logging never sees the exchange POST); returns the - response for any status.""" - - def __init__(self, handler_factory: Callable[[], "HTTPHandler"] = _new_exchange_handler) -> None: - self._lock: Final = threading.Lock() - self._handler_factory: Final = handler_factory - self._handler: HTTPHandler | None = None - - def _handler_instance(self) -> "HTTPHandler": - with self._lock: - if self._handler is None: - self._handler = self._handler_factory() - return self._handler - - def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: - try: - response: Final[httpx.Response | None] = self._handler_instance().post( # pyright: ignore[reportUnknownMemberType] # HTTPHandler.post is legacy-untyped; the result is validated below - url, - content=content, - headers=dict(headers), # mutable-ok: HTTPHandler.post requires a concrete dict - timeout=timeout, - ) - except httpx.HTTPStatusError as e: - return e.response - return require_posted_response(response, "token endpoint") - - -class _ServiceLoggingHooks(Protocol): - """The slice of ``litellm._service_logger.ServiceLogging`` the metrics sink calls; a protocol - so tests inject a recorder instead of monkeypatching.""" - - async def async_service_success_hook(self, service: ServiceTypes, call_type: str, duration: float) -> None: ... - - async def async_service_failure_hook( - self, service: ServiceTypes, duration: float, error: str | Exception, call_type: str - ) -> None: ... - - -_HooksCoroFactory: TypeAlias = Callable[ - [_ServiceLoggingHooks], # mutable-ok: Callable param-list syntax, not a list - Coroutine[object, object, None], -] - - -def _default_service_logging() -> _ServiceLoggingHooks: - from litellm._service_logger import ServiceLogging - - return ServiceLogging() - - -class ServiceLoggingMetricsSink: - """Default sink: bridges engine metrics onto litellm's ServiceTypes pattern - (prometheus ``litellm_anthropic_wif_*`` via ``service_callback``). The engine's entry points - are sync threads with no event loop, and the service hooks are async, so every emission is - fire-and-forget on a dedicated single worker thread that owns its own short-lived loop -- - the mint path only ever pays for an executor queue put.""" - - def __init__( - self, - service_logging_factory: Callable[[], _ServiceLoggingHooks] = _default_service_logging, - executor: Executor | None = None, - ) -> None: - self._lock: Final = threading.Lock() - self._service_logging_factory: Final = service_logging_factory - self._service_logging: _ServiceLoggingHooks | None = None - self._executor: Executor | None = executor - self._queued: int = 0 # rebind-ok: backlog depth, guarded by _lock - - def _service_logging_instance(self) -> _ServiceLoggingHooks: - with self._lock: - if self._service_logging is None: - self._service_logging = self._service_logging_factory() - return self._service_logging - - def _executor_instance(self) -> Executor: - with self._lock: - if self._executor is None: - self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="litellm-token-exchange-metrics") - return self._executor - - def _emit(self, coro_factory: _HooksCoroFactory) -> None: - try: - asyncio.run(coro_factory(self._service_logging_instance())) - except Exception as e: # noqa: BLE001 # metrics are best-effort; emission failures must never surface - verbose_logger.debug("token exchange metrics emission failed: %s", e) - - def _submit(self, coro_factory: _HooksCoroFactory) -> None: - """Drop the event rather than queue it once the backlog is full. A stalled telemetry - backend must not let request volume grow an unbounded queue in the proxy: losing a - metric sample is always cheaper than losing the process.""" - with self._lock: - if self._queued >= _METRICS_QUEUE_LIMIT: - verbose_logger.debug("token exchange metrics queue full, dropping event") - return - self._queued += 1 - try: - self._executor_instance().submit(self._emit_and_release, coro_factory) - except Exception as e: # noqa: BLE001 # a rejected submit must not surface to the mint - with self._lock: - self._queued -= 1 - verbose_logger.debug("token exchange metrics submit failed: %s", e) - - def _emit_and_release(self, coro_factory: _HooksCoroFactory) -> None: - try: - self._emit(coro_factory) - finally: - with self._lock: - self._queued -= 1 - - def exchange_success(self, *, call_type: ExchangeCallType, duration_seconds: float) -> None: - def start(hooks: _ServiceLoggingHooks) -> Coroutine[object, object, None]: - return hooks.async_service_success_hook( - service=ServiceTypes.ANTHROPIC_WIF, call_type=call_type, duration=duration_seconds - ) - - self._submit(start) - - def exchange_failure(self, *, call_type: ExchangeCallType, duration_seconds: float, error: ExchangeError) -> None: - failure: Final = _failure_exception(error) - - def start(hooks: _ServiceLoggingHooks) -> Coroutine[object, object, None]: - return hooks.async_service_failure_hook( - service=ServiceTypes.ANTHROPIC_WIF, duration=duration_seconds, error=failure, call_type=call_type - ) - - self._submit(start) - - def cache_hit(self) -> None: - def start(hooks: _ServiceLoggingHooks) -> Coroutine[object, object, None]: - return hooks.async_service_success_hook( - service=ServiceTypes.ANTHROPIC_WIF_CACHE, call_type=CALL_TYPE_CACHE_HIT, duration=0.0 - ) - - self._submit(start) - - -class _Entry: - """Single-flight state for one cache key; mutable by design, confined to the - engine, and only ever mutated under the engine lock.""" - - __slots__ = ("backoff_until", "done", "force_refresh", "in_flight", "last_error", "lifetime_seconds", "token") - - def __init__(self, force_refresh: bool = False) -> None: - self.token: MintedToken | None = None - self.lifetime_seconds: float | None = None - self.in_flight: bool = False - self.done: Final = threading.Event() - self.backoff_until: float = float("-inf") - self.force_refresh: bool = force_refresh - self.last_error: ExchangeError | None = None - - def arm(self) -> None: - self.in_flight = True - self.last_error = None - self.done.clear() - - def _store(self, token: MintedToken, now: float) -> None: - self.token = token - self.lifetime_seconds = None if token.expires_at is None else max(token.expires_at - now, 0.0) - self.last_error = None - - def publish(self, result: ExchangeResult, now: float) -> None: - match result: - case MintedToken(): - self._store(result, now) - case _: - self.last_error = result - self.backoff_until = now + ADVISORY_REFRESH_BACKOFF_SECONDS - self.force_refresh = False - self.in_flight = False - self.done.set() - - def publish_advisory(self, result: ExchangeResult, now: float) -> None: - """A failed advisory refresh records only the backoff, never ``last_error``: a follower whose - cached token expires while this runs must be free to re-lead a fresh mint and recover.""" - match result: - case MintedToken(): - self._store(result, now) - case _: - self.backoff_until = now + ADVISORY_REFRESH_BACKOFF_SECONDS - self.in_flight = False - self.done.set() - - -@dataclass(frozen=True, slots=True) -class _Serve: - token: MintedToken - - -@dataclass(frozen=True, slots=True) -class _ServeAndRefresh: - token: MintedToken - - -@dataclass(frozen=True, slots=True) -class _Lead: - call_type: ExchangeCallType - - -@dataclass(frozen=True, slots=True) -class _Follow: - pass - - -@dataclass(frozen=True, slots=True) -class _Fail: - error: ExchangeError - - -_Decision: TypeAlias = _Serve | _ServeAndRefresh | _Lead | _Follow | _Fail - - -@dataclass(frozen=True, slots=True) -class _Unauthorized: - response: httpx.Response - assertion: SecretStr - - -class JwtBearerTokenExchangeEngine: - def __init__( - self, - poster: SyncTokenPoster | None = None, - assertion_reader: AssertionReader | None = None, - clock: Callable[[], float] = time.monotonic, - refresh_executor: Executor | None = None, - max_entries: int = 64, - metrics_sink: TokenExchangeMetricsSink | None = None, - ) -> None: - self._poster: Final[SyncTokenPoster] = poster if poster is not None else _HttpxSyncTokenPoster() - self._assertion_reader: Final[AssertionReader] = ( - assertion_reader if assertion_reader is not None else _default_assertion_reader - ) - self._clock: Final = clock - self._refresh_executor: Executor | None = refresh_executor - self._max_entries: Final = max_entries - self._metrics_sink: Final[TokenExchangeMetricsSink] = ( - metrics_sink if metrics_sink is not None else ServiceLoggingMetricsSink() - ) - self._lock: Final = threading.Lock() - self._entries: Final[dict[str, _Entry]] = {} # mutable-ok: engine-owned map guarded by _lock - - def get_token(self, spec: TokenExchangeSpec) -> ExchangeResult: - """A follower whose leader published nothing re-classifies rather than recursing, so a - contended entry cannot grow the stack one frame per failed leader.""" - while True: - with self._lock: - entry = self._get_or_create_entry_locked(spec) # rebind-ok: re-read per follower round - decision = self._classify_and_arm_locked(entry) # rebind-ok: re-read per follower round - match decision: - case _Serve(token=token): - self._report_cache_hit() - return token - case _ServeAndRefresh(token=token): - self._report_cache_hit() - self._executor_instance().submit(self._advisory_refresh, spec, entry) - return token - case _Fail(error=error): - return error - case _Lead(call_type=call_type): - return self._lead(spec, entry, call_type) - case _Follow(): - followed = self._await_leader(spec, entry) # rebind-ok: one leader wait per round - if followed is not None: - return followed - case _: - assert_never(decision) - - async def aget_token(self, spec: TokenExchangeSpec) -> ExchangeResult: - return await asyncio.to_thread(self.get_token, spec) - - def invalidate(self, spec: TokenExchangeSpec) -> None: - key: Final = _cache_key(spec) - with self._lock: - if key in self._entries: - self._entries[key] = _Entry(force_refresh=True) - - def _get_or_create_entry_locked(self, spec: TokenExchangeSpec) -> _Entry: - key: Final = _cache_key(spec) - existing: Final = self._entries.get(key) - if existing is not None: - return existing - if len(self._entries) >= self._max_entries: - self._evict_locked() - created: Final = _Entry() - self._entries[key] = created - return created - - def _evict_locked(self) -> None: - now: Final = self._clock() - stale: Final = tuple( - key - for key, entry in self._entries.items() - if not entry.in_flight - and (entry.token is None or (entry.token.expires_at is not None and entry.token.expires_at <= now)) - ) - for key in stale: - del self._entries[key] - if len(self._entries) < self._max_entries: - return - # Evict soonest-to-expire first, and take as many as the overshoot needs rather than one, so a - # burst of distinct identities does not leave the map permanently above max_entries. An entry - # a leader owns or a follower waits on is never a candidate, so a moment where every entry is - # in flight still over-inserts; that residue is bounded by the concurrent mints themselves. - evictable: Final = sorted( - ( - entry.token.expires_at if entry.token is not None and entry.token.expires_at is not None else -inf, - key, - ) - for key, entry in self._entries.items() - if not entry.in_flight - ) - for _, key in evictable[: len(self._entries) - self._max_entries + 1]: - del self._entries[key] - - def _classify_and_arm_locked(self, entry: _Entry) -> _Decision: - token: Final = entry.token - if token is not None and not entry.force_refresh: - if token.expires_at is None: - return _Serve(token=token) - windows: Final = _refresh_windows(entry.lifetime_seconds) - remaining: Final = token.expires_at - self._clock() - if remaining > windows.advisory: - return _Serve(token=token) - if remaining > windows.mandatory: - if entry.in_flight or self._clock() < entry.backoff_until: - return _Serve(token=token) - entry.arm() - return _ServeAndRefresh(token=token) - if entry.in_flight: - return _Follow() - if entry.last_error is not None and self._clock() < entry.backoff_until: - return _Fail(error=entry.last_error) - entry.arm() - return _Lead(call_type=CALL_TYPE_COLD_MINT if token is None else CALL_TYPE_MANDATORY_REFRESH) - - def _executor_instance(self) -> Executor: - with self._lock: - if self._refresh_executor is None: - self._refresh_executor = ThreadPoolExecutor( - max_workers=2, thread_name_prefix="litellm-token-exchange-refresh" - ) - return self._refresh_executor - - def _lead(self, spec: TokenExchangeSpec, entry: _Entry, call_type: ExchangeCallType) -> ExchangeResult: - started: Final = self._clock() - result: Final = self._exchange_never_raises(spec) - duration: Final = self._clock() - started - with self._lock: - entry.publish(result, now=self._clock()) - self._report_exchange(call_type, duration, result) - return result - - def _await_leader(self, spec: TokenExchangeSpec, entry: _Entry) -> "ExchangeResult | None": - """None means the finished round left neither a valid token nor an error - (a failed advisory refresh); the caller re-enters and leads a fresh exchange.""" - leader_finished: Final = entry.done.wait(2 * spec.timeout_seconds + _FOLLOWER_WAIT_GRACE_SECONDS) - with self._lock: - token: Final = entry.token - if token is not None and (token.expires_at is None or token.expires_at > self._clock()): - return token - if entry.last_error is not None: - return entry.last_error - if leader_finished: - return None - return TokenTransportError(detail="timed out waiting for the token exchange leader") - - def _advisory_refresh(self, spec: TokenExchangeSpec, entry: _Entry) -> None: - started: Final = self._clock() - result: Final = self._exchange_never_raises(spec) - duration: Final = self._clock() - started - with self._lock: - now: Final = self._clock() - entry.publish_advisory(result, now=now) - stale_expires_at: Final = entry.token.expires_at if entry.token is not None else None - stale_mandatory: Final = _refresh_windows(entry.lifetime_seconds).mandatory - self._report_exchange(CALL_TYPE_ADVISORY_REFRESH, duration, result) - if isinstance(result, MintedToken): - return - seconds_to_mandatory_wall: Final = ( - max(stale_expires_at - now - stale_mandatory, 0.0) if stale_expires_at is not None else 0.0 - ) - verbose_logger.warning( - "Advisory token refresh against %s failed (%s); serving the cached token for up to " - "%.0fs before the mandatory refresh wall; next attempt after %.0fs backoff", - urlsplit(spec.token_url).hostname or "", - _error_summary(result), - seconds_to_mandatory_wall, - ADVISORY_REFRESH_BACKOFF_SECONDS, - ) - - def _report_exchange(self, call_type: ExchangeCallType, duration_seconds: float, result: ExchangeResult) -> None: - try: - match result: - case MintedToken(): - self._metrics_sink.exchange_success(call_type=call_type, duration_seconds=duration_seconds) - case _: - self._metrics_sink.exchange_failure( - call_type=call_type, duration_seconds=duration_seconds, error=result - ) - except Exception as e: # noqa: BLE001 # metrics are best-effort; a sink failure must never fail a mint - verbose_logger.debug("token exchange metrics emission failed: %s", e) - - def _report_cache_hit(self) -> None: - try: - self._metrics_sink.cache_hit() - except Exception as e: # noqa: BLE001 # metrics are best-effort; a sink failure must never fail a serve - verbose_logger.debug("token exchange cache-hit metric emission failed: %s", e) - - def _exchange_never_raises(self, spec: TokenExchangeSpec) -> ExchangeResult: - """The single-flight leader and the advisory refresher must always publish a result: an - unhandled exception here would leave the entry armed (in_flight, cleared event) forever, so - every subsequent caller for this key would follow a leader that never finishes.""" - try: - return self._exchange(spec) - except Exception as e: # noqa: BLE001 # a leader must resolve its entry; any failure becomes a value - return TokenTransportError(detail=f"{type(e).__name__}: {e}"[:_REDACTION_CAP]) - - def _exchange(self, spec: TokenExchangeSpec) -> ExchangeResult: - first: Final = self._attempt_exchange(spec) - if not isinstance(first, _Unauthorized): - return first - second: Final = self._attempt_exchange(spec) - if isinstance(second, _Unauthorized): - return redact_oauth_error_body( - second.response.status_code, _capped_body_text(second.response), second.assertion - ) - return second - - def _attempt_exchange(self, spec: TokenExchangeSpec) -> "ExchangeResult | _Unauthorized": - assertion: Final = _read_assertion(_assertion_fetch(self._assertion_reader, spec), spec.assertion_ref) - if isinstance(assertion, AssertionSourceError): - return assertion - url_check: Final = validate_token_endpoint_url(spec.token_url) - if isinstance(url_check, InsecureTokenUrl): - return url_check - try: - response: Final = self._poster.post( - spec.token_url, - content=_serialize_body(spec, assertion), - headers=MappingProxyType({"content-type": _CONTENT_TYPES[spec.body_encoding], **spec.request_headers}), - timeout=spec.timeout_seconds, - ) - except Exception as e: # noqa: BLE001 # injected posters may raise beyond httpx; transport failures become values - return TokenTransportError(detail=f"{type(e).__name__}: {e}"[:_REDACTION_CAP]) - if response.status_code == 401: - return _Unauthorized(response=response, assertion=assertion) - return self._parse_response(response, assertion) - - def _parse_response(self, response: httpx.Response, assertion: SecretStr | None = None) -> ExchangeResult: - if not 200 <= response.status_code < 300: - return redact_oauth_error_body(response.status_code, _capped_body_text(response), assertion) - if len(response.content) > MAX_RESPONSE_BYTES: - return MalformedTokenResponse(detail="token response body exceeds the 1 MiB cap") - try: - parsed: Final = _TokenExchangeResponse.model_validate_json(response.content) - except ValidationError: - return MalformedTokenResponse(detail="token response failed RFC 6749 5.1 schema validation") - if parsed.token_type is not None and parsed.token_type.lower() != "bearer": - return MalformedTokenResponse(detail="token response carried a non-bearer token_type") - if not parsed.access_token.strip(): - return MalformedTokenResponse(detail="token response carried an empty access_token") - return MintedToken( - access_token=SecretStr(parsed.access_token), - expires_at=self._clock() + _sanitize_expires_in(parsed.expires_in), - ) - - -default_token_exchange_engine: Final = JwtBearerTokenExchangeEngine() diff --git a/litellm/llms/base_llm/auth/types.py b/litellm/llms/base_llm/auth/types.py index 9d3a7a5012bc..10aaba8e1816 100644 --- a/litellm/llms/base_llm/auth/types.py +++ b/litellm/llms/base_llm/auth/types.py @@ -1,48 +1,16 @@ -"""Provider-agnostic types for the RFC 7523 JWT-bearer token exchange engine.""" +"""Provider-agnostic value types for OAuth token exchanges: where an assertion comes from, the typed +failure union every grant maps onto its public exception contract, and the observability seam.""" from collections.abc import Callable, Mapping from dataclasses import dataclass from typing import Literal, Protocol, TypeAlias import httpx -from pydantic import SecretStr -BodyEncoding: TypeAlias = Literal["json", "form"] AssertionReader: TypeAlias = Callable[[str], str | None] # mutable-ok: Callable param-list syntax, not a list AssertionSource: TypeAlias = Callable[[], str | None] # mutable-ok: Callable param-list syntax, not a list -@dataclass(frozen=True, slots=True) -class TokenExchangeSpec: - """One grant profile as pure data: one instance per (provider, deployment, identity). - - ``token_url`` must be derived from deployment config/env only, never per-request caller - input. ``assertion_ref`` is a ``oidc/...`` get_secret ref resolved fresh on every exchange. - - ``assertion_source``, when set, is a zero-arg per-config fetch/mint closure that the engine - prefers over its own engine-level ``AssertionReader`` -- the dispatch mechanism identity - sources beyond token_file/env (e.g. ``internal_issuer``, ``keycloak``) use to plug into the - shared engine without a global registry. ``assertion_ref`` still names the cache-key - discriminator and the ref echoed into operator-facing errors either way. - """ - - token_url: str - assertion_ref: str - assertion_field: str - static_body: Mapping[str, str] - body_encoding: BodyEncoding - request_headers: Mapping[str, str] - cache_key_identity: tuple[str, ...] - timeout_seconds: float = 30.0 - assertion_source: AssertionSource | None = None - - -@dataclass(frozen=True, slots=True) -class MintedToken: - access_token: SecretStr - expires_at: float | None - - @dataclass(frozen=True, slots=True) class AssertionSourceError: kind: Literal["missing", "empty", "oversized", "unreadable", "disallowed_path"] @@ -74,13 +42,13 @@ class MalformedTokenResponse: ExchangeError: TypeAlias = ( AssertionSourceError | InsecureTokenUrl | TokenEndpointError | TokenTransportError | MalformedTokenResponse ) -ExchangeResult: TypeAlias = MintedToken | ExchangeError +ExchangeResult: TypeAlias = str | ExchangeError -ExchangeCallType: TypeAlias = Literal["cold_mint", "mandatory_refresh", "advisory_refresh"] +ExchangeCallType: TypeAlias = Literal["cold_mint", "refresh"] class TokenExchangeMetricsSink(Protocol): - """Observability seam for the exchange engine. Implementations must be best-effort: never raise + """Observability seam for a token exchange. Implementations must be best-effort: never raise into the mint path, never block the calling thread, and never receive credential material -- ``ExchangeError`` values are redacted by construction.""" diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 0f77ae1d4681..b3462203c4b7 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -750,6 +750,5 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum): # OAuth constants ANTHROPIC_OAUTH_TOKEN_PREFIX: Final = "sk-ant-oat" ANTHROPIC_OAUTH_BETA_HEADER: Final = "oauth-2025-04-20" -ANTHROPIC_TOKEN_EXCHANGE_PATH: Final = "/v1/oauth/token" ANTHROPIC_PROMPT_CACHING_SCOPE_BETA_HEADER: Final = "prompt-caching-scope-2026-01-05" diff --git a/pyproject.toml b/pyproject.toml index b889a3a0e600..f489cfb3fb72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -151,7 +151,7 @@ proxy-runtime = [ # feature surface without forcing the base SDK install to grow. "google-cloud-aiplatform>=1.133.0,<2.0", "google-genai>=1.37.0,<2.0", - "anthropic[vertex]>=0.84.0,<1.0", + "anthropic[vertex]>=1.3.0,<2", "grpcio==1.78.0", "prometheus-client>=0.20.0,<1.0", "langfuse>=2.59.7,<3.0", @@ -233,6 +233,7 @@ proxy-dev = [ "opentelemetry-instrumentation-fastapi==0.49b0", "azure-identity==1.25.2", "a2a-sdk==1.1.0", + "anthropic[vertex]==1.3.0", ] ci = [ # These are lazily imported at call sites; keep them out of core deps to @@ -257,7 +258,7 @@ ci = [ "argon2-cffi==25.1.0", "assemblyai==0.52.4", "jsonlines==4.0.0", - "anthropic==0.84.0", + "anthropic==1.3.0", "blockbuster==1.5.26", "beautifulsoup4==4.14.3", "pylint==4.0.5", diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index fd7b30bc3140..184558ef508b 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -168,7 +168,7 @@ "limit": 3 }, "RET504": { - "limit": 173 + "limit": 172 }, "RUF012": { "limit": 239 @@ -198,7 +198,7 @@ "limit": 56 }, "SIM102": { - "limit": 310 + "limit": 309 }, "SIM103": { "limit": 119 diff --git a/tests/test_litellm/llms/__init__.py b/tests/test_litellm/llms/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/test_litellm/llms/anthropic/batches/test_handler.py b/tests/test_litellm/llms/anthropic/batches/test_handler.py index 398c94a67980..7a6309c0ed11 100644 --- a/tests/test_litellm/llms/anthropic/batches/test_handler.py +++ b/tests/test_litellm/llms/anthropic/batches/test_handler.py @@ -19,6 +19,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx +import httpx2 import pytest @@ -272,18 +273,18 @@ def test_retrieve_batch_sync_runs_to_result(handler, patched_client): } -class _BlockingPoster: - """A token-endpoint poster that blocks until released, so the test can prove +class _BlockingTokenEndpoint: + """A token endpoint that blocks until released, so the test can prove the exchange ran off the event loop's own thread instead of freezing it.""" def __init__(self): self.release = threading.Event() self.thread_ids = [] - def post(self, url, *, content, headers, timeout): + def __call__(self, request: httpx2.Request) -> httpx2.Response: self.thread_ids.append(threading.get_ident()) self.release.wait(timeout=5) - return httpx.Response( + return httpx2.Response( 200, json={ "access_token": "sk-ant-oat01-batches-seam", @@ -300,7 +301,7 @@ async def test_aretrieve_batch_wif_exchange_does_not_block_event_loop(handler, p other concurrent coroutine until the exchange finished.""" from litellm.llms.anthropic import common_utils as anthropic_common_utils from litellm.llms.anthropic.wif import get_anthropic_wif_token - from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + from litellm.llms.anthropic.wif_exchange import AnthropicWifTokenExchange fake_client, _ = patched_client for name in ( @@ -313,11 +314,11 @@ async def test_aretrieve_batch_wif_exchange_does_not_block_event_loop(handler, p for name, value in _WIF_ENV.items(): monkeypatch.setenv(name, value) - poster = _BlockingPoster() - engine = JwtBearerTokenExchangeEngine(poster=poster) + token_endpoint = _BlockingTokenEndpoint() + exchange = AnthropicWifTokenExchange(http_client=httpx2.Client(transport=httpx2.MockTransport(token_endpoint))) def routed_through_injected_engine(litellm_params, api_base, model): - return get_anthropic_wif_token(litellm_params, api_base, model, engine) + return get_anthropic_wif_token(litellm_params, api_base, model, exchange) monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", routed_through_injected_engine) @@ -342,16 +343,16 @@ async def ticker(): ) await asyncio.sleep(0.05) # The ticker kept advancing while the token exchange was still blocked on - # poster.release, proving the exchange did not run on the event loop. + # token_endpoint.release, proving the exchange did not run on the event loop. assert len(ticks) > 0 assert not retrieve_task.done() - poster.release.set() + token_endpoint.release.set() batch = await retrieve_task await ticker_task assert batch.id == "msgbatch_abc" - assert poster.thread_ids - assert poster.thread_ids[0] != threading.get_ident() + assert token_endpoint.thread_ids + assert token_endpoint.thread_ids[0] != threading.get_ident() sent_headers = fake_client.get.call_args.kwargs["headers"] assert sent_headers["authorization"] == "Bearer sk-ant-oat01-batches-seam" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py index 768a2834fbe6..9fdd26e1a6ae 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py @@ -11,8 +11,8 @@ from litellm.llms.minimax.messages.transformation import MinimaxMessagesConfig from litellm.llms.tencent.messages.transformation import TencentAnthropicMessagesConfig from tests.test_litellm.llms.anthropic.test_anthropic_wif import ( - ScriptedPoster, - make_engine, + ScriptedTokenEndpoint, + make_exchange, token_response, write_token_file, ) @@ -90,8 +90,8 @@ def test_tencent_validate_environment_never_attaches_anthropic_wif_credential( def test_wif_token_exchange_reaches_only_anthropic_not_minimax_or_tencent( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - """get_anthropic_wif_token's engine parameter is the only DI seam in the WIF minting chain; - validate_anthropic_messages_environment always uses the module's default engine, so this + """get_anthropic_wif_token's exchange parameter is the only DI seam in the WIF minting chain; + validate_anthropic_messages_environment always uses the module's default exchange, so this drives that seam directly with the exact litellm_params AnthropicModelInfo.get_auth_header would receive from each config, proving MiniMax/Tencent never reach the token endpoint even when a mint would otherwise succeed.""" @@ -100,19 +100,19 @@ def test_wif_token_exchange_reaches_only_anthropic_not_minimax_or_tencent( monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) token_file = write_token_file(tmp_path, "jwt-assertion-value") litellm_params = {"anthropic_identity_token_file": str(token_file)} - poster = ScriptedPoster([token_response("sk-ant-oat01-canary")]) - engine = make_engine(poster) + token_endpoint = ScriptedTokenEndpoint([token_response("sk-ant-oat01-canary")]) + exchange = make_exchange(token_endpoint) minted: Final = get_anthropic_wif_token( litellm_params, "https://api.anthropic.com", "claude-sonnet-4-5", - engine, + exchange, ) assert minted == "sk-ant-oat01-canary" - assert len(poster.requests) == 1 + assert len(token_endpoint.requests) == 1 for config in (MinimaxMessagesConfig(), TencentAnthropicMessagesConfig()): assert config._allows_workload_identity is False - assert len(poster.requests) == 1 + assert len(token_endpoint.requests) == 1 diff --git a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py b/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py index f385c2f22110..bbe3d4e89407 100644 --- a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py +++ b/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py @@ -11,6 +11,7 @@ import time import httpx +import httpx2 import pytest from unittest.mock import Mock, patch @@ -449,18 +450,18 @@ def test_get_error_class(self): } -class _BlockingPoster: - """A token-endpoint poster that blocks until released, so the test can prove +class _BlockingTokenEndpoint: + """A token endpoint that blocks until released, so the test can prove the exchange ran off the event loop's own thread instead of freezing it.""" def __init__(self): self.release = threading.Event() self.thread_ids = [] - def post(self, url, *, content, headers, timeout): + def __call__(self, request: httpx2.Request) -> httpx2.Response: self.thread_ids.append(threading.get_ident()) self.release.wait(timeout=5) - return httpx.Response( + return httpx2.Response( 200, json={ "access_token": "sk-ant-oat01-files-seam", @@ -482,7 +483,7 @@ def setup_method(self): async def test_avalidate_environment_wif_exchange_does_not_block_event_loop(self, monkeypatch): from litellm.llms.anthropic import common_utils as anthropic_common_utils from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token - from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + from litellm.llms.anthropic.wif_exchange import AnthropicWifTokenExchange for name in ( "ANTHROPIC_API_KEY", @@ -494,16 +495,16 @@ async def test_avalidate_environment_wif_exchange_does_not_block_event_loop(self for name, value in _WIF_ENV.items(): monkeypatch.setenv(name, value) - poster = _BlockingPoster() - engine = JwtBearerTokenExchangeEngine(poster=poster) + token_endpoint = _BlockingTokenEndpoint() + exchange = AnthropicWifTokenExchange(http_client=httpx2.Client(transport=httpx2.MockTransport(token_endpoint))) sync_calls = [] def sync_shim(litellm_params, api_base, model): sync_calls.append(model) - return get_anthropic_wif_token(litellm_params, api_base, model, engine) + return get_anthropic_wif_token(litellm_params, api_base, model, exchange) async def async_shim(litellm_params, api_base, model): - return await aget_anthropic_wif_token(litellm_params, api_base, model, engine) + return await aget_anthropic_wif_token(litellm_params, api_base, model, exchange) monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim) monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim) @@ -530,18 +531,18 @@ async def ticker(): ) await asyncio.sleep(0.05) # The ticker kept advancing while the exchange was still blocked on - # poster.release, proving avalidate_environment did not run it inline. + # token_endpoint.release, proving avalidate_environment did not run it inline. assert len(ticks) > 0 assert not validate_task.done() - poster.release.set() + token_endpoint.release.set() headers = await validate_task await ticker_task assert headers["authorization"] == "Bearer sk-ant-oat01-files-seam" assert sync_calls == [] - assert poster.thread_ids - assert poster.thread_ids[0] != threading.get_ident() + assert token_endpoint.thread_ids + assert token_endpoint.thread_ids[0] != threading.get_ident() class TestProviderConfigRegistration: diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index c2dfb91591b4..a5f1737ec0f5 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -19,11 +19,13 @@ from unittest.mock import patch import httpx +import httpx2 import pytest sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) from litellm.proxy._types import SpecialHeaders # noqa: E402 # sys.path must be patched before importing litellm +from litellm.llms.anthropic.wif_exchange import AnthropicWifTokenExchange # noqa: E402 # sys.path must be patched before importing litellm # Fake tokens for testing (not real secrets) FAKE_OAUTH_TOKEN = "sk-ant-oat01-fake-token-for-testing-123456789abcdef" @@ -2219,18 +2221,29 @@ def test_create_anthropic_model_list_response_empty(): PROXY_CREDENTIAL_HEADER_NAMES = sorted(SpecialHeaders.litellm_credential_header_names()) -class RecordingPoster: - def __init__(self, response): +def minted_token_response() -> httpx2.Response: + return httpx2.Response( + 200, + json={"access_token": FAKE_MINTED_TOKEN, "token_type": "Bearer", "expires_in": 3600}, + ) + + +class RecordingTokenEndpoint: + def __init__(self, response: httpx2.Response) -> None: self.requests = [] self.thread_ids = [] self._response = response - def post(self, url, *, content, headers, timeout): - self.requests.append((url, content, dict(headers))) + def __call__(self, request: httpx2.Request) -> httpx2.Response: + self.requests.append((str(request.url), request.content, dict(request.headers))) self.thread_ids.append(threading.get_ident()) return self._response +def exchange_against(token_endpoint: RecordingTokenEndpoint) -> AnthropicWifTokenExchange: + return AnthropicWifTokenExchange(http_client=httpx2.Client(transport=httpx2.MockTransport(token_endpoint))) + + @pytest.fixture def clean_anthropic_env(monkeypatch): for name in ANTHROPIC_ENV_VARS: @@ -2238,62 +2251,46 @@ def clean_anthropic_env(monkeypatch): @pytest.fixture -def wif_engine(monkeypatch, clean_anthropic_env): - """Route the wiring's WIF tier through a fresh engine (never the module +def wif_exchange(monkeypatch, clean_anthropic_env): + """Route the wiring's WIF tier through a fresh exchange (never the module singleton, to avoid cross-test cache pollution) and count its consultations.""" - import httpx - from litellm.llms.anthropic import common_utils as anthropic_common_utils from litellm.llms.anthropic.wif import get_anthropic_wif_token - from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine - poster = RecordingPoster( - httpx.Response( - 200, - json={"access_token": FAKE_MINTED_TOKEN, "token_type": "Bearer", "expires_in": 3600}, - ) - ) - engine = JwtBearerTokenExchangeEngine(poster=poster) + token_endpoint = RecordingTokenEndpoint(minted_token_response()) + exchange = exchange_against(token_endpoint) calls = [] def with_injected_engine(litellm_params, api_base, model): calls.append(model) - return get_anthropic_wif_token(litellm_params, api_base, model, engine) + return get_anthropic_wif_token(litellm_params, api_base, model, exchange) monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", with_injected_engine) - return poster, calls + return token_endpoint, calls @pytest.fixture -def wif_async_engine(monkeypatch, clean_anthropic_env): - """Route both WIF facades through one fresh engine; the poster records the +def wif_async_exchange(monkeypatch, clean_anthropic_env): + """Route both WIF facades through one fresh exchange; the token_endpoint records the thread each exchange ran on and sync-facade consultations are counted so async tests can prove the mint went through the async seam, off the loop.""" - import httpx - from litellm.llms.anthropic import common_utils as anthropic_common_utils from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token - from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine - poster = RecordingPoster( - httpx.Response( - 200, - json={"access_token": FAKE_MINTED_TOKEN, "token_type": "Bearer", "expires_in": 3600}, - ) - ) - engine = JwtBearerTokenExchangeEngine(poster=poster) + token_endpoint = RecordingTokenEndpoint(minted_token_response()) + exchange = exchange_against(token_endpoint) sync_calls = [] def sync_shim(litellm_params, api_base, model): sync_calls.append(model) - return get_anthropic_wif_token(litellm_params, api_base, model, engine) + return get_anthropic_wif_token(litellm_params, api_base, model, exchange) async def async_shim(litellm_params, api_base, model): - return await aget_anthropic_wif_token(litellm_params, api_base, model, engine) + return await aget_anthropic_wif_token(litellm_params, api_base, model, exchange) monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim) monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim) - return poster, sync_calls + return token_endpoint, sync_calls def _validate_chat_environment(api_key=None): @@ -2312,24 +2309,24 @@ def _validate_chat_environment(api_key=None): class TestWifTierPrecedence: """WIF is the LOWEST credential tier: any api_key / auth_token source must - win without the engine ever being consulted.""" + win without the exchange ever being consulted.""" def _set_wif_env(self, monkeypatch): for name, value in WIF_ENV.items(): monkeypatch.setenv(name, value) - def test_explicit_api_key_beats_wif(self, monkeypatch, wif_engine): - poster, calls = wif_engine + def test_explicit_api_key_beats_wif(self, monkeypatch, wif_exchange): + token_endpoint, calls = wif_exchange self._set_wif_env(monkeypatch) headers = _validate_chat_environment(api_key=FAKE_REGULAR_KEY) assert headers["x-api-key"] == FAKE_REGULAR_KEY assert calls == [] - assert poster.requests == [] + assert token_endpoint.requests == [] - def test_api_key_env_beats_wif(self, monkeypatch, wif_engine): - poster, calls = wif_engine + def test_api_key_env_beats_wif(self, monkeypatch, wif_exchange): + token_endpoint, calls = wif_exchange self._set_wif_env(monkeypatch) monkeypatch.setenv("ANTHROPIC_API_KEY", FAKE_REGULAR_KEY) @@ -2337,10 +2334,10 @@ def test_api_key_env_beats_wif(self, monkeypatch, wif_engine): assert headers["x-api-key"] == FAKE_REGULAR_KEY assert calls == [] - assert poster.requests == [] + assert token_endpoint.requests == [] - def test_auth_token_env_beats_wif(self, monkeypatch, wif_engine): - poster, calls = wif_engine + def test_auth_token_env_beats_wif(self, monkeypatch, wif_exchange): + token_endpoint, calls = wif_exchange self._set_wif_env(monkeypatch) monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", FAKE_AUTH_TOKEN) @@ -2348,18 +2345,18 @@ def test_auth_token_env_beats_wif(self, monkeypatch, wif_engine): assert headers["authorization"] == f"Bearer {FAKE_AUTH_TOKEN}" assert calls == [] - assert poster.requests == [] + assert token_endpoint.requests == [] - def test_wif_alone_mints_once(self, monkeypatch, wif_engine): - poster, calls = wif_engine + def test_wif_alone_mints_once(self, monkeypatch, wif_exchange): + token_endpoint, calls = wif_exchange self._set_wif_env(monkeypatch) headers = _validate_chat_environment() assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" assert calls == ["claude-sonnet-4-5"] - assert len(poster.requests) == 1 - assert poster.requests[0][0] == "https://api.anthropic.com/v1/oauth/token" + assert len(token_endpoint.requests) == 1 + assert token_endpoint.requests[0][0] == "https://api.anthropic.com/v1/oauth/token" class TestWifZeroBehaviorChange: @@ -2381,7 +2378,7 @@ def test_unconfigured_raises_same_authentication_error(self, clean_anthropic_env class TestWifHeaderContract: - def test_minted_token_headers(self, monkeypatch, wif_engine): + def test_minted_token_headers(self, monkeypatch, wif_exchange): for name, value in WIF_ENV.items(): monkeypatch.setenv(name, value) @@ -2438,12 +2435,12 @@ class TestWifServerOwnedAuthHeaderStrip: """A WIF-minted token must never ride alongside a caller-supplied credential header, but that stripping must fire only when a mint actually happened.""" - def test_mint_strips_caller_supplied_x_api_key(self, monkeypatch, wif_engine): + def test_mint_strips_caller_supplied_x_api_key(self, monkeypatch, wif_exchange): """Security regression: without the strip, a caller-forwarded x-api-key would sit next to the server-minted Authorization on the outgoing request.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo - poster, _ = wif_engine + token_endpoint, _ = wif_exchange for name, value in WIF_ENV.items(): monkeypatch.setenv(name, value) caller_key = "sk-ant-CALLER-SUPPLIED" @@ -2461,7 +2458,7 @@ def test_mint_strips_caller_supplied_x_api_key(self, monkeypatch, wif_engine): assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" assert "x-api-key" not in headers assert caller_key not in headers.values() - assert len(poster.requests) == 1 + assert len(token_endpoint.requests) == 1 def test_server_owned_set_is_every_proxy_credential_header(self): """The strip list must track the proxy's own key-header list, not a hand-rolled @@ -2472,7 +2469,7 @@ def test_server_owned_set_is_every_proxy_credential_header(self): assert {"x-litellm-api-key", "api-key", "x-goog-api-key"} < _SERVER_OWNED_AUTH_HEADERS @pytest.mark.parametrize("header_name", PROXY_CREDENTIAL_HEADER_NAMES) - def test_mint_strips_every_proxy_credential_header(self, monkeypatch, wif_engine, header_name): + def test_mint_strips_every_proxy_credential_header(self, monkeypatch, wif_exchange, header_name): """A LiteLLM virtual key arrives in any of the proxy's accepted key headers; once a mint happened none of them may reach Anthropic in any header slot.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -2497,7 +2494,7 @@ def test_mint_strips_every_proxy_credential_header(self, monkeypatch, wif_engine assert headers["user-agent"] == "caller/1.0" @pytest.mark.parametrize("header_name", PROXY_CREDENTIAL_HEADER_NAMES) - def test_skills_surface_strips_caller_credentials_too(self, monkeypatch, wif_engine, header_name): + def test_skills_surface_strips_caller_credentials_too(self, monkeypatch, wif_exchange, header_name): """Skills builds its own headers as well; every minting surface needs the same strip.""" from litellm.llms.anthropic.skills.transformation import AnthropicSkillsConfig @@ -2515,7 +2512,7 @@ def test_skills_surface_strips_caller_credentials_too(self, monkeypatch, wif_eng assert all(caller_key not in value for value in headers.values()) assert headers["user-agent"] == "caller/1.0" - def test_passthrough_honors_a_case_variant_caller_key_instead_of_minting(self, monkeypatch, wif_engine): + def test_passthrough_honors_a_case_variant_caller_key_instead_of_minting(self, monkeypatch, wif_exchange): """The passthrough surface hands the caller's own credential upstream rather than minting. That check was case-sensitive, so X-Api-Key slipped past it and the caller's key would have travelled beside a minted Bearer.""" @@ -2525,7 +2522,7 @@ def test_passthrough_honors_a_case_variant_caller_key_instead_of_minting(self, m for name, value in WIF_ENV.items(): monkeypatch.setenv(name, value) - poster, _ = wif_engine + token_endpoint, _ = wif_exchange caller_key = "sk-ant-CALLER-SUPPLIED" headers, _ = AnthropicMessagesConfig().validate_anthropic_messages_environment( @@ -2540,10 +2537,10 @@ def test_passthrough_honors_a_case_variant_caller_key_instead_of_minting(self, m assert headers["X-Api-Key"] == caller_key assert "authorization" not in {name.lower() for name in headers} - assert len(poster.requests) == 0 + assert len(token_endpoint.requests) == 0 @pytest.mark.parametrize("header_name", PROXY_CREDENTIAL_HEADER_NAMES) - def test_batches_surface_strips_caller_credentials_too(self, monkeypatch, wif_engine, header_name): + def test_batches_surface_strips_caller_credentials_too(self, monkeypatch, wif_exchange, header_name): """Batches builds its own headers on the create path, so it needs the same strip: the handler's retrieve path passes none, but this entry point takes the caller's.""" from litellm.llms.anthropic.batches.transformation import AnthropicBatchesConfig @@ -2568,7 +2565,7 @@ def test_batches_surface_strips_caller_credentials_too(self, monkeypatch, wif_en assert headers["user-agent"] == "caller/1.0" @pytest.mark.parametrize("header_name", PROXY_CREDENTIAL_HEADER_NAMES) - def test_files_surface_strips_caller_credentials_too(self, monkeypatch, wif_engine, header_name): + def test_files_surface_strips_caller_credentials_too(self, monkeypatch, wif_exchange, header_name): """The files surface builds its own headers, so it needs the same strip the chat surface has: without it a minted federation Bearer travels beside the caller's own credential.""" from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig @@ -2622,7 +2619,7 @@ class TestWifResolvedApiKeyThreading: validate_environment path and the async aget_auth_header path, never a stale None left over from the original unresolved parameter.""" - def test_validate_environment_carries_minted_token(self, monkeypatch, wif_engine): + def test_validate_environment_carries_minted_token(self, monkeypatch, wif_exchange): for name, value in WIF_ENV.items(): monkeypatch.setenv(name, value) @@ -2633,7 +2630,7 @@ def test_validate_environment_carries_minted_token(self, monkeypatch, wif_engine assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" @pytest.mark.asyncio - async def test_aget_auth_header_carries_minted_token(self, monkeypatch, wif_async_engine): + async def test_aget_auth_header_carries_minted_token(self, monkeypatch, wif_async_exchange): from litellm.llms.anthropic.common_utils import AnthropicModelInfo for name, value in WIF_ENV.items(): @@ -2659,10 +2656,10 @@ def test_oat_branch_carries_oauth_beta(self, clean_anthropic_env): "anthropic-beta": "oauth-2025-04-20", } - def test_wif_fallback_returns_bearer_and_beta(self, monkeypatch, wif_engine): + def test_wif_fallback_returns_bearer_and_beta(self, monkeypatch, wif_exchange): from litellm.llms.anthropic.common_utils import AnthropicModelInfo - poster, calls = wif_engine + token_endpoint, calls = wif_exchange for name, value in WIF_ENV.items(): monkeypatch.setenv(name, value) @@ -2672,7 +2669,7 @@ def test_wif_fallback_returns_bearer_and_beta(self, monkeypatch, wif_engine): "authorization": f"Bearer {FAKE_MINTED_TOKEN}", "anthropic-beta": "oauth-2025-04-20", } - assert len(poster.requests) == 1 + assert len(token_endpoint.requests) == 1 def test_no_credentials_still_returns_none(self, clean_anthropic_env): from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -2772,10 +2769,10 @@ class TestWifLitellmParamsPlumbing: def _inline_identity_token(self, monkeypatch): monkeypatch.setenv("WIF_PARAMS_TEST_TOKEN", "params-jwt") - def test_files_mints_from_litellm_params(self, wif_engine): + def test_files_mints_from_litellm_params(self, wif_exchange): from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig - poster, _ = wif_engine + token_endpoint, _ = wif_exchange headers = AnthropicFilesConfig().validate_environment( headers={}, model="", @@ -2785,12 +2782,12 @@ def test_files_mints_from_litellm_params(self, wif_engine): ) assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" - assert len(poster.requests) == 1 + assert len(token_endpoint.requests) == 1 - def test_batches_mints_from_litellm_params(self, wif_engine): + def test_batches_mints_from_litellm_params(self, wif_exchange): from litellm.llms.anthropic.batches.transformation import AnthropicBatchesConfig - poster, _ = wif_engine + token_endpoint, _ = wif_exchange headers = AnthropicBatchesConfig().validate_environment( headers={}, model="", @@ -2800,13 +2797,13 @@ def test_batches_mints_from_litellm_params(self, wif_engine): ) assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" - assert len(poster.requests) == 1 + assert len(token_endpoint.requests) == 1 - def test_skills_mints_from_litellm_params(self, wif_engine): + def test_skills_mints_from_litellm_params(self, wif_exchange): from litellm.llms.anthropic.skills.transformation import AnthropicSkillsConfig from litellm.types.router import GenericLiteLLMParams - poster, _ = wif_engine + token_endpoint, _ = wif_exchange headers = AnthropicSkillsConfig().validate_environment( headers={}, litellm_params=GenericLiteLLMParams( @@ -2817,14 +2814,14 @@ def test_skills_mints_from_litellm_params(self, wif_engine): ) assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" - assert len(poster.requests) == 1 + assert len(token_endpoint.requests) == 1 - def test_messages_mints_from_litellm_params(self, wif_engine): + def test_messages_mints_from_litellm_params(self, wif_exchange): from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, ) - poster, _ = wif_engine + token_endpoint, _ = wif_exchange headers, _ = AnthropicMessagesConfig().validate_anthropic_messages_environment( headers={}, model="claude-sonnet-4-5", @@ -2834,13 +2831,13 @@ def test_messages_mints_from_litellm_params(self, wif_engine): ) assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" - assert len(poster.requests) == 1 + assert len(token_endpoint.requests) == 1 class TestWifTokenUrlParity: """Both credential tiers must derive the SAME clean token URL from any form of the deployment base; a mismatch also duplicates mints because token_url is in - the engine cache key.""" + the exchange cache key.""" @pytest.mark.parametrize( "configured_base", @@ -2851,13 +2848,13 @@ class TestWifTokenUrlParity: "https://gw.example.com/v1/messages/", ], ) - def test_both_tiers_share_one_clean_token_url(self, monkeypatch, wif_engine, configured_base): + def test_both_tiers_share_one_clean_token_url(self, monkeypatch, wif_exchange, configured_base): # This is about deriving one URL from many spellings of the same base, not about which # hosts an operator trusts with org-scoped credentials, so the private host is allowlisted. monkeypatch.setenv("LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", "gw.example.com") from litellm.llms.anthropic.common_utils import AnthropicModelInfo - poster, _ = wif_engine + token_endpoint, _ = wif_exchange for name, value in WIF_ENV.items(): monkeypatch.setenv(name, value) @@ -2872,7 +2869,7 @@ def test_both_tiers_share_one_clean_token_url(self, monkeypatch, wif_engine, con ) AnthropicModelInfo.get_auth_header(api_base=configured_base, allow_workload_identity=True) - assert [url for (url, _, _) in poster.requests] == ["https://gw.example.com/v1/oauth/token"] + assert [url for (url, _, _) in token_endpoint.requests] == ["https://gw.example.com/v1/oauth/token"] class TestWifAsyncSeam: @@ -2880,10 +2877,10 @@ class TestWifAsyncSeam: mint never blocks the event loop.""" @pytest.mark.asyncio - async def test_aget_auth_header_runs_exchange_off_event_loop(self, monkeypatch, wif_async_engine): + async def test_aget_auth_header_runs_exchange_off_event_loop(self, monkeypatch, wif_async_exchange): from litellm.llms.anthropic.common_utils import AnthropicModelInfo - poster, sync_calls = wif_async_engine + token_endpoint, sync_calls = wif_async_exchange for name, value in WIF_ENV.items(): monkeypatch.setenv(name, value) @@ -2894,16 +2891,16 @@ async def test_aget_auth_header_runs_exchange_off_event_loop(self, monkeypatch, "anthropic-beta": "oauth-2025-04-20", } assert sync_calls == [] - assert poster.thread_ids == [poster.thread_ids[0]] - assert poster.thread_ids[0] != threading.get_ident() + assert token_endpoint.thread_ids == [token_endpoint.thread_ids[0]] + assert token_endpoint.thread_ids[0] != threading.get_ident() @pytest.mark.asyncio - async def test_avalidate_messages_environment_mints_off_loop(self, wif_async_engine, monkeypatch): + async def test_avalidate_messages_environment_mints_off_loop(self, wif_async_exchange, monkeypatch): from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, ) - poster, sync_calls = wif_async_engine + token_endpoint, sync_calls = wif_async_exchange monkeypatch.setenv("WIF_PARAMS_TEST_TOKEN", "params-jwt") headers, _ = await AnthropicMessagesConfig().avalidate_anthropic_messages_environment( @@ -2916,7 +2913,7 @@ async def test_avalidate_messages_environment_mints_off_loop(self, wif_async_eng assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" assert sync_calls == [] - assert poster.thread_ids[0] != threading.get_ident() + assert token_endpoint.thread_ids[0] != threading.get_ident() @pytest.mark.asyncio async def test_avalidate_delegates_to_subclass_sync_override(self): @@ -3006,7 +3003,6 @@ def test_completion_mints_and_never_leaks_config(self, monkeypatch, tmp_path, cl import litellm from litellm.llms.anthropic import common_utils as anthropic_common_utils from litellm.llms.anthropic.wif import get_anthropic_wif_token - from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine monkeypatch.setattr(litellm, "api_key", None) monkeypatch.setattr(litellm, "anthropic_key", None) @@ -3014,11 +3010,12 @@ def test_completion_mints_and_never_leaks_config(self, monkeypatch, tmp_path, cl token_file = tmp_path / "identity-token" token_file.write_text("e2e-oidc-assertion", encoding="utf-8") - engine = JwtBearerTokenExchangeEngine() + token_endpoint = RecordingTokenEndpoint(minted_token_response()) + exchange = exchange_against(token_endpoint) monkeypatch.setattr( anthropic_common_utils, "get_anthropic_wif_token", - lambda litellm_params, api_base, model: get_anthropic_wif_token(litellm_params, api_base, model, engine), + lambda litellm_params, api_base, model: get_anthropic_wif_token(litellm_params, api_base, model, exchange), ) wif_kwarg_names: Final = ( @@ -3040,12 +3037,6 @@ def test_completion_mints_and_never_leaks_config(self, monkeypatch, tmp_path, cl } with respx.mock: - token_route = respx.post("https://api.anthropic.com/v1/oauth/token").mock( - return_value=httpx.Response( - 200, - json={"access_token": FAKE_MINTED_TOKEN, "token_type": "Bearer", "expires_in": 3600}, - ) - ) messages_route = respx.post("https://api.anthropic.com/v1/messages").mock( return_value=httpx.Response(200, json=anthropic_response) ) @@ -3061,8 +3052,8 @@ def test_completion_mints_and_never_leaks_config(self, monkeypatch, tmp_path, cl ) assert response.choices[0].message.content == "Hello from WIF" - assert token_route.call_count == 1 - exchange_body = json.loads(token_route.calls[0].request.content) + assert len(token_endpoint.requests) == 1 + exchange_body = json.loads(token_endpoint.requests[0][1]) assert exchange_body["assertion"] == "e2e-oidc-assertion" assert exchange_body["federation_rule_id"] == "fdrl_e2e" @@ -3087,7 +3078,6 @@ def test_completion_with_trailing_slash_api_base_mints_at_clean_token_url( import litellm from litellm.llms.anthropic import common_utils as anthropic_common_utils from litellm.llms.anthropic.wif import get_anthropic_wif_token - from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine monkeypatch.setattr(litellm, "api_key", None) monkeypatch.setattr(litellm, "anthropic_key", None) @@ -3095,11 +3085,12 @@ def test_completion_with_trailing_slash_api_base_mints_at_clean_token_url( token_file = tmp_path / "identity-token" token_file.write_text("e2e-oidc-assertion", encoding="utf-8") - engine = JwtBearerTokenExchangeEngine() + token_endpoint = RecordingTokenEndpoint(minted_token_response()) + exchange = exchange_against(token_endpoint) monkeypatch.setattr( anthropic_common_utils, "get_anthropic_wif_token", - lambda litellm_params, api_base, model: get_anthropic_wif_token(litellm_params, api_base, model, engine), + lambda litellm_params, api_base, model: get_anthropic_wif_token(litellm_params, api_base, model, exchange), ) anthropic_response = { @@ -3113,12 +3104,6 @@ def test_completion_with_trailing_slash_api_base_mints_at_clean_token_url( } with respx.mock: - token_route = respx.post("https://api.anthropic.com/v1/oauth/token").mock( - return_value=httpx.Response( - 200, - json={"access_token": FAKE_MINTED_TOKEN, "token_type": "Bearer", "expires_in": 3600}, - ) - ) messages_route = respx.post(url__regex=r"https://api\.anthropic\.com/v1/messages.*").mock( return_value=httpx.Response(200, json=anthropic_response) ) @@ -3132,7 +3117,7 @@ def test_completion_with_trailing_slash_api_base_mints_at_clean_token_url( ) assert response.choices[0].message.content == "Hello from WIF" - assert token_route.call_count == 1 + assert len(token_endpoint.requests) == 1 assert messages_route.calls[0].request.headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" def test_get_auth_header_with_litellm_params_mints_via_real_engine( @@ -3144,26 +3129,20 @@ def test_get_auth_header_with_litellm_params_mints_via_real_engine( from litellm.llms.anthropic import common_utils as anthropic_common_utils from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.anthropic.wif import get_anthropic_wif_token - from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) token_file = tmp_path / "identity-token" token_file.write_text("e2e-oidc-assertion", encoding="utf-8") - engine = JwtBearerTokenExchangeEngine() + token_endpoint = RecordingTokenEndpoint(minted_token_response()) + exchange = exchange_against(token_endpoint) monkeypatch.setattr( anthropic_common_utils, "get_anthropic_wif_token", - lambda litellm_params, api_base, model: get_anthropic_wif_token(litellm_params, api_base, model, engine), + lambda litellm_params, api_base, model: get_anthropic_wif_token(litellm_params, api_base, model, exchange), ) with respx.mock: - token_route = respx.post("https://api.anthropic.com/v1/oauth/token").mock( - return_value=httpx.Response( - 200, - json={"access_token": FAKE_MINTED_TOKEN, "token_type": "Bearer", "expires_in": 3600}, - ) - ) result = AnthropicModelInfo.get_auth_header( allow_workload_identity=True, litellm_params={ @@ -3177,8 +3156,8 @@ def test_get_auth_header_with_litellm_params_mints_via_real_engine( "authorization": f"Bearer {FAKE_MINTED_TOKEN}", "anthropic-beta": "oauth-2025-04-20", } - assert token_route.call_count == 1 - exchange_body = json.loads(token_route.calls[0].request.content) + assert len(token_endpoint.requests) == 1 + exchange_body = json.loads(token_endpoint.requests[0][1]) assert exchange_body["federation_rule_id"] == "fdrl_e2e" @@ -3195,14 +3174,14 @@ def _env_only_wif(monkeypatch) -> None: # noqa: D401 monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "oidc/env/WIF_TEST_JWT") monkeypatch.setenv("WIF_TEST_JWT", "jwt-assertion-value") - def test_vertex_anthropic_never_mints_or_sends_the_assertion(self, monkeypatch, wif_engine): + def test_vertex_anthropic_never_mints_or_sends_the_assertion(self, monkeypatch, wif_exchange): from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( VertexAIAnthropicConfig, ) import litellm - poster, calls = wif_engine + token_endpoint, calls = wif_exchange self._env_only_wif(monkeypatch) with pytest.raises(litellm.AuthenticationError): @@ -3217,12 +3196,12 @@ def test_vertex_anthropic_never_mints_or_sends_the_assertion(self, monkeypatch, ) assert calls == [] - assert poster.requests == [] + assert token_endpoint.requests == [] - def test_anthropic_itself_still_mints(self, monkeypatch, wif_engine): + def test_anthropic_itself_still_mints(self, monkeypatch, wif_exchange): from litellm.llms.anthropic.chat.transformation import AnthropicConfig - poster, calls = wif_engine + token_endpoint, calls = wif_exchange self._env_only_wif(monkeypatch) headers = AnthropicConfig().validate_environment( @@ -3236,7 +3215,7 @@ def test_anthropic_itself_still_mints(self, monkeypatch, wif_engine): ) assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" - assert len(poster.requests) == 1 + assert len(token_endpoint.requests) == 1 def test_auth_header_facade_defaults_to_refusing_to_mint(self, monkeypatch, clean_anthropic_env): """The facade is reachable from provider code that has nothing to do with Anthropic, so a @@ -3260,13 +3239,13 @@ class NewCompatibleProvider(AnthropicConfig): assert config_allows_workload_identity(AnthropicConfig()) is True assert config_allows_workload_identity(NewCompatibleProvider()) is False - def test_model_discovery_gates_on_the_instance(self, monkeypatch, wif_engine): + def test_model_discovery_gates_on_the_instance(self, monkeypatch, wif_exchange): """get_models is inherited, so it must consult the instance rather than trusting its caller.""" from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( VertexAIAnthropicConfig, ) - poster, calls = wif_engine + token_endpoint, calls = wif_exchange self._env_only_wif(monkeypatch) with pytest.raises(ValueError, match="ANTHROPIC_API_KEY"): @@ -3274,7 +3253,7 @@ def test_model_discovery_gates_on_the_instance(self, monkeypatch, wif_engine): api_base="https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" ) - assert poster.requests == [] + assert token_endpoint.requests == [] def _models_page_response(page: dict, status_code: int = 200): @@ -3433,12 +3412,12 @@ def test_get_models_error_is_sanitized_not_raw_response_text(self, monkeypatch, assert "invalid x-api-key" in str(exc_info.value) assert reflected_payload not in str(exc_info.value) - def test_discover_models_threads_litellm_params_into_wif(self, monkeypatch, wif_engine): + def test_discover_models_threads_litellm_params_into_wif(self, monkeypatch, wif_exchange): """The gap this phase fixes: get_models only ever saw api_key/api_base, so a WIF source configured in litellm_params (rather than ANTHROPIC_* env vars) could not discover.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo - poster, calls = wif_engine + token_endpoint, calls = wif_exchange client = RecordingModelsClient([{"data": [{"id": "claude-wif"}], "has_more": False, "last_id": None}]) monkeypatch.setattr("litellm.module_level_client", client) monkeypatch.setenv("DISC_JWT", "jwt-assertion-value") @@ -3452,7 +3431,7 @@ def test_discover_models_threads_litellm_params_into_wif(self, monkeypatch, wif_ ) assert models == ["anthropic/claude-wif"] - assert len(poster.requests) == 1 + assert len(token_endpoint.requests) == 1 assert client.calls[0].headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" def test_discover_models_without_litellm_params_behaves_like_get_models(self, monkeypatch, clean_anthropic_env): @@ -3469,12 +3448,12 @@ def test_discover_models_without_litellm_params_behaves_like_get_models(self, mo assert models == ["anthropic/claude-env"] assert client.calls[0].headers["x-api-key"] == FAKE_REGULAR_KEY - def test_discover_models_explicit_api_key_beats_wif(self, monkeypatch, wif_engine): + def test_discover_models_explicit_api_key_beats_wif(self, monkeypatch, wif_exchange): """Same precedence discover_models must honor as every other Anthropic auth surface: WIF is the lowest tier.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo - poster, calls = wif_engine + token_endpoint, calls = wif_exchange client = RecordingModelsClient([{"data": [], "has_more": False, "last_id": None}]) monkeypatch.setattr("litellm.module_level_client", client) @@ -3489,18 +3468,16 @@ def test_discover_models_explicit_api_key_beats_wif(self, monkeypatch, wif_engin assert client.calls[0].headers["x-api-key"] == FAKE_REGULAR_KEY assert calls == [] - assert poster.requests == [] + assert token_endpoint.requests == [] class TestWifExchangeTransportHardening: def test_token_exchange_client_does_not_follow_redirects(self): """Only the initial token URL is validated, so a 3xx must not be allowed to replay the assertion to an origin that was never checked.""" - from litellm.llms.base_llm.auth.token_exchange import _HttpxSyncTokenPoster - - handler = _HttpxSyncTokenPoster()._handler_instance() + from litellm.llms.anthropic.wif_exchange import new_exchange_client - assert handler.client.follow_redirects is False + assert new_exchange_client().follow_redirects is False class TestWifParamsAreNotClientSettable: diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py b/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py index fe7ea23042f9..17c036eb608c 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py @@ -12,6 +12,7 @@ import httpx +import httpx2 import pytest from litellm.llms.anthropic.batches.transformation import AnthropicBatchesConfig @@ -438,7 +439,7 @@ async def test_afile_content_resolves_wif_via_async_facade( from litellm.llms.anthropic import common_utils as anthropic_common_utils from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token - from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + from litellm.llms.anthropic.wif_exchange import AnthropicWifTokenExchange for name in ( "ANTHROPIC_API_KEY", @@ -454,23 +455,25 @@ async def test_afile_content_resolves_wif_via_async_facade( minted = "sk-ant-oat01-files-minted" thread_ids = [] - class ThreadRecordingPoster: - def post(self, url, *, content, headers, timeout): + class ThreadRecordingTokenEndpoint: + def __call__(self, request: httpx2.Request) -> httpx2.Response: thread_ids.append(threading.get_ident()) - return httpx.Response( + return httpx2.Response( 200, json={"access_token": minted, "token_type": "Bearer", "expires_in": 3600}, ) - engine = JwtBearerTokenExchangeEngine(poster=ThreadRecordingPoster()) + exchange = AnthropicWifTokenExchange( + http_client=httpx2.Client(transport=httpx2.MockTransport(ThreadRecordingTokenEndpoint())) + ) sync_calls = [] def sync_shim(litellm_params, api_base, model): sync_calls.append(model) - return get_anthropic_wif_token(litellm_params, api_base, model, engine) + return get_anthropic_wif_token(litellm_params, api_base, model, exchange) async def async_shim(litellm_params, api_base, model): - return await aget_anthropic_wif_token(litellm_params, api_base, model, engine) + return await aget_anthropic_wif_token(litellm_params, api_base, model, exchange) monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim) monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim) @@ -517,7 +520,7 @@ async def test_afile_content_mints_from_the_deployment_litellm_params( process-wide env vars could ever mint on a batch-result download.""" from litellm.llms.anthropic import common_utils as anthropic_common_utils from litellm.llms.anthropic.wif import aget_anthropic_wif_token - from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + from litellm.llms.anthropic.wif_exchange import AnthropicWifTokenExchange for name in ( "ANTHROPIC_API_KEY", @@ -533,17 +536,17 @@ async def test_afile_content_mints_from_the_deployment_litellm_params( minted = "sk-ant-oat01-credential-minted" - class Poster: - def post(self, url, *, content, headers, timeout): - return httpx.Response( + class TokenEndpoint: + def __call__(self, request: httpx2.Request) -> httpx2.Response: + return httpx2.Response( 200, json={"access_token": minted, "token_type": "Bearer", "expires_in": 3600}, ) - engine = JwtBearerTokenExchangeEngine(poster=Poster()) + exchange = AnthropicWifTokenExchange(http_client=httpx2.Client(transport=httpx2.MockTransport(TokenEndpoint()))) async def async_shim(litellm_params, api_base, model): - return await aget_anthropic_wif_token(litellm_params, api_base, model, engine) + return await aget_anthropic_wif_token(litellm_params, api_base, model, exchange) monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim) diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_wif.py b/tests/test_litellm/llms/anthropic/test_anthropic_wif.py index 515d0c5ede60..e8db51c50fe5 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_wif.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_wif.py @@ -1,10 +1,10 @@ -import concurrent.futures import json +import time from collections.abc import Callable, Mapping from pathlib import Path from typing import Final -import httpx +import httpx2 import jwt import pytest from cryptography.hazmat.primitives import serialization @@ -14,18 +14,16 @@ from litellm.llms.anthropic.wif import ( AnthropicWifParams, _raise_anthropic_wif_error, - build_anthropic_wif_spec, get_anthropic_wif_token, resolve_anthropic_wif_params, ) +from litellm.llms.anthropic.wif_exchange import AnthropicWifTokenExchange from litellm.llms.base_llm.auth.identity_source import ( InternalIssuerSource, KeycloakSource, identity_source_ref, ) from litellm.llms.base_llm.auth.jwt_signing import build_jwks, rfc7638_thumbprint -from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine -from litellm.types.router import GenericLiteLLMParams from litellm.llms.base_llm.auth.types import ( AssertionSourceError, ExchangeError, @@ -34,6 +32,7 @@ TokenEndpointError, TokenTransportError, ) +from litellm.types.router import GenericLiteLLMParams WIF_ENV_VARS: Final = ( "ANTHROPIC_FEDERATION_RULE_ID", @@ -58,9 +57,12 @@ def _clean_wif_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv(name, raising=False) +SDK_BETA_HEADER: Final = "oauth-2025-04-20,oidc-federation-2026-04-01" + + class FakeClock: - def __init__(self, start: float = 1_000.0) -> None: - self.now = start + def __init__(self) -> None: + self.now = time.time() def __call__(self) -> float: return self.now @@ -70,52 +72,40 @@ def advance(self, seconds: float) -> None: class RecordedRequest: - def __init__(self, url: str, content: bytes, headers: Mapping[str, str], timeout: float) -> None: - self.url = url - self.content = content - self.headers = dict(headers) - self.timeout = timeout + def __init__(self, request: httpx2.Request) -> None: + self.url = str(request.url) + self.content = request.content + self.headers = dict(request.headers) def json_body(self) -> dict: return json.loads(self.content) -class ScriptedPoster: - def __init__(self, responses: list[httpx.Response]) -> None: +class ScriptedTokenEndpoint: + def __init__(self, responses: list[httpx2.Response]) -> None: self.requests: list[RecordedRequest] = [] self._responses = list(responses) - def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: - self.requests.append(RecordedRequest(url, content, headers, timeout)) + def __call__(self, request: httpx2.Request) -> httpx2.Response: + self.requests.append(RecordedRequest(request)) if len(self._responses) > 1: return self._responses.pop(0) return self._responses[0] -class ManualExecutor(concurrent.futures.Executor): - def __init__(self) -> None: - self.pending: list[Callable[[], None]] = [] - - def submit(self, fn, /, *args, **kwargs): - future: concurrent.futures.Future = concurrent.futures.Future() - self.pending.append(lambda: fn(*args, **kwargs)) - return future - - -def token_response(token: str = "sk-ant-oat01-minted", expires_in: int | None = 3600) -> httpx.Response: +def token_response(token: str = "sk-ant-oat01-minted", expires_in: int | None = 3600) -> httpx2.Response: body: Final[dict[str, str | int]] = { "access_token": token, "token_type": "Bearer", **({} if expires_in is None else {"expires_in": expires_in}), } - return httpx.Response(200, json=body) + return httpx2.Response(200, json=body) -def make_engine(poster: ScriptedPoster, clock: FakeClock | None = None) -> JwtBearerTokenExchangeEngine: - return JwtBearerTokenExchangeEngine( - poster=poster, - clock=clock if clock is not None else FakeClock(), - refresh_executor=ManualExecutor(), +def make_exchange(endpoint: ScriptedTokenEndpoint, clock: FakeClock | None = None) -> AnthropicWifTokenExchange: + return AnthropicWifTokenExchange( + http_client=httpx2.Client(transport=httpx2.MockTransport(endpoint)), + clock=clock if clock is not None else time.time, ) @@ -131,8 +121,8 @@ def test_minimal_body_and_headers(self, tmp_path: Path, monkeypatch: pytest.Monk monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) monkeypatch.setenv("ANTHROPIC_SCOPE", "user:inference") token_file = write_token_file(tmp_path, "jwt-assertion-value\n") - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster) + endpoint = ScriptedTokenEndpoint([token_response()]) + exchange = make_exchange(endpoint) token = get_anthropic_wif_token( { @@ -142,14 +132,14 @@ def test_minimal_body_and_headers(self, tmp_path: Path, monkeypatch: pytest.Monk }, "https://api.anthropic.com", "claude-sonnet-4-5", - engine, + exchange, ) assert token == "sk-ant-oat01-minted" - assert len(poster.requests) == 1 - request = poster.requests[0] + assert len(endpoint.requests) == 1 + request = endpoint.requests[0] assert request.url == "https://api.anthropic.com/v1/oauth/token" - assert "anthropic-beta" not in request.headers + assert request.headers["anthropic-beta"] == SDK_BETA_HEADER assert request.headers["content-type"] == "application/json" assert request.json_body() == { "grant_type": GRANT_TYPE, @@ -161,8 +151,8 @@ def test_minimal_body_and_headers(self, tmp_path: Path, monkeypatch: pytest.Monk def test_optional_fields_present_when_set(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) token_file = write_token_file(tmp_path, "jwt-assertion-value") - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster) + endpoint = ScriptedTokenEndpoint([token_response()]) + exchange = make_exchange(endpoint) get_anthropic_wif_token( { @@ -174,11 +164,11 @@ def test_optional_fields_present_when_set(self, tmp_path: Path, monkeypatch: pyt }, "https://api.anthropic.com", "claude-sonnet-4-5", - engine, + exchange, ) - request = poster.requests[0] - assert "anthropic-beta" not in request.headers + request = endpoint.requests[0] + assert request.headers["anthropic-beta"] == SDK_BETA_HEADER assert request.headers["content-type"] == "application/json" assert request.json_body() == { "grant_type": GRANT_TYPE, @@ -189,31 +179,6 @@ def test_optional_fields_present_when_set(self, tmp_path: Path, monkeypatch: pyt "assertion": "jwt-assertion-value", } - def test_spec_cache_key_identity(self): - params = AnthropicWifParams( - federation_rule_id="fdrl_1", - organization_id="org-1", - assertion_ref="oidc/env/ANTHROPIC_IDENTITY_TOKEN", - ) - spec = build_anthropic_wif_spec(params, "https://api.anthropic.com") - assert spec.cache_key_identity == ("fdrl_1", "org-1", "", "") - assert spec.body_encoding == "json" - assert spec.assertion_field == "assertion" - - def test_full_params_spec_has_no_request_headers(self): - """The token exchange sends no anthropic-beta header at all (verified against the - live endpoint); this must hold even for a fully populated params set, so a future - edit cannot reintroduce the header gated on service_account_id or workspace_id.""" - params = AnthropicWifParams( - federation_rule_id="fdrl_1", - organization_id="org-1", - service_account_id="svcacct_1", - workspace_id="wrkspc_1", - assertion_ref="oidc/env/ANTHROPIC_IDENTITY_TOKEN", - ) - spec = build_anthropic_wif_spec(params, "https://api.anthropic.com") - assert dict(spec.request_headers) == {} - class TestExchangeHostTrust: """A federated exchange sends the workload's identity token to api_base and presents the minted @@ -223,14 +188,14 @@ class TestExchangeHostTrust: def _mint(self, api_base: str | None, monkeypatch: pytest.MonkeyPatch) -> str: monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") - poster = ScriptedPoster([token_response()]) + endpoint = ScriptedTokenEndpoint([token_response()]) get_anthropic_wif_token( {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"}, api_base, "claude-sonnet-4-5", - make_engine(poster), + make_exchange(endpoint), ) - return poster.requests[0].url + return endpoint.requests[0].url def test_anthropic_is_trusted_without_configuration(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", raising=False) @@ -239,17 +204,17 @@ def test_anthropic_is_trusted_without_configuration(self, monkeypatch: pytest.Mo def test_an_unlisted_host_never_receives_the_identity_token(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", raising=False) monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") - poster = ScriptedPoster([token_response()]) + endpoint = ScriptedTokenEndpoint([token_response()]) with pytest.raises(litellm.AuthenticationError) as exc_info: get_anthropic_wif_token( {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"}, "https://attacker.example", "claude-sonnet-4-5", - make_engine(poster), + make_exchange(endpoint), ) - assert poster.requests == [], "the exchange must be refused before anything is sent" + assert endpoint.requests == [], "the exchange must be refused before anything is sent" assert "attacker.example" in str(exc_info.value) assert "LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS" in str(exc_info.value), ( "an operator running a private gateway has to be told how to allow it" @@ -258,17 +223,17 @@ def test_an_unlisted_host_never_receives_the_identity_token(self, monkeypatch: p def test_a_lookalike_host_does_not_pass_on_a_substring(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", raising=False) monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") - poster = ScriptedPoster([token_response()]) + endpoint = ScriptedTokenEndpoint([token_response()]) with pytest.raises(litellm.AuthenticationError): get_anthropic_wif_token( {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"}, "https://api.anthropic.com.evil.test", "claude-sonnet-4-5", - make_engine(poster), + make_exchange(endpoint), ) - assert poster.requests == [] + assert endpoint.requests == [] def test_an_operator_can_allow_a_private_gateway(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", "gateway.internal") @@ -281,7 +246,7 @@ def test_a_gateway_listed_with_its_port_is_trusted(self, monkeypatch: pytest.Mon def test_allowlist_matching_ignores_hostname_case(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", "Gateway.Internal:8443") assert self._mint("https://gateway.internal:8443", monkeypatch) == "https://gateway.internal:8443/v1/oauth/token" - assert self._mint("https://GATEWAY.internal:8443", monkeypatch) == "https://GATEWAY.internal:8443/v1/oauth/token" + assert self._mint("https://GATEWAY.internal:8443", monkeypatch) == "https://gateway.internal:8443/v1/oauth/token" class TestBaseUrlDerivation: @@ -294,15 +259,15 @@ def _mint(self, api_base: str | None, monkeypatch: pytest.MonkeyPatch) -> str: "LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", "gw.example.com,env.example.com,base.example.com,model.example.com", ) - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster) + endpoint = ScriptedTokenEndpoint([token_response()]) + exchange = make_exchange(endpoint) get_anthropic_wif_token( {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"}, api_base, "claude-sonnet-4-5", - engine, + exchange, ) - return poster.requests[0].url + return endpoint.requests[0].url def test_explicit_api_base_wins(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("ANTHROPIC_API_BASE", "https://env.example.com") @@ -510,8 +475,6 @@ def test_empty_workspace_env_coerced_to_none(self, monkeypatch: pytest.MonkeyPat ) assert params is not None assert params.workspace_id is None - spec = build_anthropic_wif_spec(params, "https://api.anthropic.com") - assert "workspace_id" not in spec.static_body @pytest.mark.parametrize( "litellm_params", @@ -528,10 +491,10 @@ def test_gate_unmet_returns_none(self, litellm_params: dict): assert resolve_anthropic_wif_params(litellm_params) is None def test_gate_unmet_facade_returns_none_without_engine_call(self): - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster) - assert get_anthropic_wif_token({}, None, "claude-sonnet-4-5", engine) is None - assert poster.requests == [] + endpoint = ScriptedTokenEndpoint([token_response()]) + exchange = make_exchange(endpoint) + assert get_anthropic_wif_token({}, None, "claude-sonnet-4-5", exchange) is None + assert endpoint.requests == [] class TestServiceAccountIdIsOptional: @@ -553,12 +516,12 @@ def test_activates_and_omits_service_account_id_when_unset(self, tmp_path: Path, assert params is not None assert params.service_account_id is None - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster) - token = get_anthropic_wif_token(litellm_params, "https://api.anthropic.com", "claude-sonnet-4-5", engine) + endpoint = ScriptedTokenEndpoint([token_response()]) + exchange = make_exchange(endpoint) + token = get_anthropic_wif_token(litellm_params, "https://api.anthropic.com", "claude-sonnet-4-5", exchange) assert token == "sk-ant-oat01-minted" - assert "service_account_id" not in poster.requests[0].json_body() + assert "service_account_id" not in endpoint.requests[0].json_body() class TestInlineRefRestrictions: @@ -566,8 +529,8 @@ class TestInlineRefRestrictions: @pytest.mark.parametrize("bad_ref", [RAW_JWT, "oidc/env_path/ANTHROPIC_TOKEN_PATH"]) def test_rejected_inline_refs(self, bad_ref: str): - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster) + endpoint = ScriptedTokenEndpoint([token_response()]) + exchange = make_exchange(endpoint) with pytest.raises(litellm.AuthenticationError) as exc_info: get_anthropic_wif_token( @@ -578,20 +541,20 @@ def test_rejected_inline_refs(self, bad_ref: str): }, None, "claude-sonnet-4-5", - engine, + exchange, ) assert "oidc/env/" in exc_info.value.message assert "oidc/file/" in exc_info.value.message assert self.RAW_JWT not in exc_info.value.message - assert poster.requests == [] + assert endpoint.requests == [] class TestFileAllowlistAndSymlink: SECRET_CONTENT: Final = "super-secret-jwt-content" - def _call(self, token_file: Path, poster: ScriptedPoster) -> str | None: - engine = make_engine(poster) + def _call(self, token_file: Path, endpoint: ScriptedTokenEndpoint) -> str | None: + exchange = make_exchange(endpoint) return get_anthropic_wif_token( { "anthropic_federation_rule_id": "fdrl_1", @@ -600,30 +563,30 @@ def _call(self, token_file: Path, poster: ScriptedPoster) -> str | None: }, "https://api.anthropic.com", "claude-sonnet-4-5", - engine, + exchange, ) def test_file_outside_allowlist_rejected(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path / "allowed")) token_file = write_token_file(tmp_path / "outside", self.SECRET_CONTENT) - poster = ScriptedPoster([token_response()]) + endpoint = ScriptedTokenEndpoint([token_response()]) with pytest.raises(litellm.AuthenticationError) as exc_info: - self._call(token_file, poster) + self._call(token_file, endpoint) assert str(token_file) in exc_info.value.message assert self.SECRET_CONTENT not in exc_info.value.message - assert poster.requests == [] + assert endpoint.requests == [] def test_disallowed_path_message_names_allowlist_and_env_var(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): """The disallowed_path error must explain the allowlist and name the env var an operator would set, not surface as a bare '(disallowed_path)' code dump.""" monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path / "allowed")) token_file = write_token_file(tmp_path / "outside", self.SECRET_CONTENT) - poster = ScriptedPoster([token_response()]) + endpoint = ScriptedTokenEndpoint([token_response()]) with pytest.raises(litellm.AuthenticationError) as exc_info: - self._call(token_file, poster) + self._call(token_file, endpoint) message = exc_info.value.message assert "(disallowed_path)" not in message @@ -637,21 +600,21 @@ def test_symlink_escape_rejected(self, tmp_path: Path, monkeypatch: pytest.Monke outside_file = write_token_file(tmp_path / "outside", self.SECRET_CONTENT) link = allowed / "identity-token" link.symlink_to(outside_file) - poster = ScriptedPoster([token_response()]) + endpoint = ScriptedTokenEndpoint([token_response()]) with pytest.raises(litellm.AuthenticationError) as exc_info: - self._call(link, poster) + self._call(link, endpoint) assert self.SECRET_CONTENT not in exc_info.value.message - assert poster.requests == [] + assert endpoint.requests == [] def test_file_inside_allowlist_succeeds(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) token_file = write_token_file(tmp_path, self.SECRET_CONTENT) - poster = ScriptedPoster([token_response()]) + endpoint = ScriptedTokenEndpoint([token_response()]) - assert self._call(token_file, poster) == "sk-ant-oat01-minted" - assert poster.requests[0].json_body()["assertion"] == self.SECRET_CONTENT + assert self._call(token_file, endpoint) == "sk-ant-oat01-minted" + assert endpoint.requests[0].json_body()["assertion"] == self.SECRET_CONTENT class TestErrorMappingExhaustive: @@ -704,15 +667,15 @@ def test_assertion_source_error_without_detail_is_unchanged(self): def test_endpoint_error_raised_through_facade(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") - poster = ScriptedPoster([httpx.Response(500, json={"error": "server_error"})]) - engine = make_engine(poster) + endpoint = ScriptedTokenEndpoint([httpx2.Response(500, json={"error": "server_error"})]) + exchange = make_exchange(endpoint) with pytest.raises(litellm.AuthenticationError) as exc_info: get_anthropic_wif_token( {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"}, None, "claude-sonnet-4-5", - engine, + exchange, ) assert exc_info.value.llm_provider == "anthropic" @@ -742,11 +705,11 @@ def test_token_endpoint_error_message_has_no_doubled_period( self, litellm_params: dict, status_code: int, body: dict, monkeypatch: pytest.MonkeyPatch ): monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") - poster = ScriptedPoster([httpx.Response(status_code, json=body)]) - engine = make_engine(poster) + endpoint = ScriptedTokenEndpoint([httpx2.Response(status_code, json=body)]) + exchange = make_exchange(endpoint) with pytest.raises(litellm.AuthenticationError) as exc_info: - get_anthropic_wif_token(litellm_params, None, "claude-sonnet-4-5", engine) + get_anthropic_wif_token(litellm_params, None, "claude-sonnet-4-5", exchange) assert ".." not in exc_info.value.message @@ -760,10 +723,10 @@ class TestDenialHints: def _raise(self, litellm_params: dict, status_code: int, monkeypatch: pytest.MonkeyPatch) -> str: monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") - poster = ScriptedPoster([httpx.Response(status_code, json={"error": "invalid_grant"})]) - engine = make_engine(poster) + endpoint = ScriptedTokenEndpoint([httpx2.Response(status_code, json={"error": "invalid_grant"})]) + exchange = make_exchange(endpoint) with pytest.raises(litellm.AuthenticationError) as exc_info: - get_anthropic_wif_token(litellm_params, None, "claude-sonnet-4-5", engine) + get_anthropic_wif_token(litellm_params, None, "claude-sonnet-4-5", exchange) return exc_info.value.message def test_401_points_at_console_authentication_history(self, monkeypatch: pytest.MonkeyPatch): @@ -810,26 +773,26 @@ class TestFileRereadOnRefresh: def test_mandatory_refresh_carries_rotated_assertion(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) token_file = write_token_file(tmp_path, "first-assertion") - clock = FakeClock(start=1_000.0) - poster = ScriptedPoster( + clock = FakeClock() + endpoint = ScriptedTokenEndpoint( [token_response("sk-ant-oat01-first", 3600), token_response("sk-ant-oat01-second", 3600)] ) - engine = make_engine(poster, clock=clock) + exchange = make_exchange(endpoint, clock=clock) litellm_params = { "anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1", "anthropic_identity_token_file": str(token_file), } - first = get_anthropic_wif_token(litellm_params, "https://api.anthropic.com", "claude-sonnet-4-5", engine) + first = get_anthropic_wif_token(litellm_params, "https://api.anthropic.com", "claude-sonnet-4-5", exchange) token_file.write_text("second-assertion", encoding="utf-8") clock.advance(3600 - 10) - second = get_anthropic_wif_token(litellm_params, "https://api.anthropic.com", "claude-sonnet-4-5", engine) + second = get_anthropic_wif_token(litellm_params, "https://api.anthropic.com", "claude-sonnet-4-5", exchange) assert first == "sk-ant-oat01-first" assert second == "sk-ant-oat01-second" - assert len(poster.requests) == 2 - assert poster.requests[1].json_body()["assertion"] == "second-assertion" + assert len(endpoint.requests) == 2 + assert endpoint.requests[1].json_body()["assertion"] == "second-assertion" _ISSUER_PRIVATE_VALUE: Final = 55566677788899900011122233344455566677788899900011122233344455 @@ -862,7 +825,7 @@ def fake_get_secret_str(secret_name: str, default_value: str | None = None) -> s class TestIdentitySourceDiscriminatorAbsentIsByteIdenticalToLegacy: """anthropic_identity_source unset must resolve exactly like today: no new dispatch code - runs, and no assertion_source closure is attached, so the engine falls back to its own + runs, and no assertion_source closure is attached, so the exchange falls back to its own reader precisely as it always has.""" def test_file_config_carries_no_assertion_source(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): @@ -959,13 +922,13 @@ def test_full_exchange_sends_the_minted_assertion(self, monkeypatch: pytest.Monk "litellm.secret_managers.main.get_secret_str", _get_secret_str_returning(pem, ISSUER_SIGNING_KEY_REF), ) - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster) + endpoint = ScriptedTokenEndpoint([token_response()]) + exchange = make_exchange(endpoint) - token = get_anthropic_wif_token(self.LITELLM_PARAMS, "https://api.anthropic.com", "claude-sonnet-4-5", engine) + token = get_anthropic_wif_token(self.LITELLM_PARAMS, "https://api.anthropic.com", "claude-sonnet-4-5", exchange) assert token == "sk-ant-oat01-minted" - sent_assertion = poster.requests[0].json_body()["assertion"] + sent_assertion = endpoint.requests[0].json_body()["assertion"] jwt.decode( sent_assertion, _issuer_signing_key().public_key(), algorithms=["ES256"], options={"verify_aud": False} ) @@ -973,7 +936,7 @@ def test_full_exchange_sends_the_minted_assertion(self, monkeypatch: pytest.Monk class TestKeycloakIdentitySourceDispatch: """A config.yaml-shaped litellm_params block for the keycloak identity source. The minted - closure's own network behavior is covered by test_client_credentials.py's DI-poster tests; + closure's own network behavior is covered by test_client_credentials.py's DI-endpoint tests; this only proves wif.py threads the fields into the right config and hash.""" LITELLM_PARAMS: Final = { diff --git a/tests/test_litellm/llms/anthropic/test_wif_exchange.py b/tests/test_litellm/llms/anthropic/test_wif_exchange.py new file mode 100644 index 000000000000..cbe60c7ba1df --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_wif_exchange.py @@ -0,0 +1,1068 @@ +import asyncio +import concurrent.futures +import json +import logging +import time +from collections.abc import Callable, Mapping +from typing import Final + +import httpx2 +import pytest + +from litellm.llms.anthropic.wif import AnthropicWifParams +from litellm.llms.anthropic.wif_exchange import ( + _DETAIL_CAP, + _METRICS_QUEUE_LIMIT, + CALL_TYPE_CACHE_HIT, + EXCHANGE_CONNECT_TIMEOUT_SECONDS, + EXCHANGE_TIMEOUT_SECONDS, + MAX_ASSERTION_BYTES, + AnthropicWifTokenExchange, + ServiceLoggingMetricsSink, + TokenExchangeEndpointFailure, + TokenExchangeTransportFailure, + _default_assertion_reader, + _error_summary, + new_exchange_client, +) +from litellm.llms.base_llm.auth.oauth_endpoint import MAX_RESPONSE_BYTES +from litellm.llms.base_llm.auth.types import ( + AssertionSourceError, + ExchangeError, + InsecureTokenUrl, + MalformedTokenResponse, + TokenEndpointError, + TokenTransportError, +) +from litellm.secret_managers.main import OidcPathNotAllowedError, _resolve_oidc_file_path +from litellm.types.services import ServiceTypes + +DEFAULT_REF: Final = "oidc/env/TEST_ASSERTION" +DEFAULT_ASSERTION: Final = "test-jwt-assertion" +EXCHANGE_BASE: Final = "https://token.example" +EXCHANGE_URL: Final = "https://token.example/v1/oauth/token" +GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer" +SDK_BETA_HEADER: Final = "oauth-2025-04-20,oidc-federation-2026-04-01" +TOKEN_TTL: Final = 3600 + + +class FakeClock: + """Starts at the real time: the SDK stamps ``expires_at`` from ``time.time()`` and only its + cache's refresh decisions read the injected clock.""" + + def __init__(self) -> None: + self.now = time.time() + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +class RecordedRequest: + def __init__(self, request: httpx2.Request) -> None: + self.url = str(request.url) + self.content = request.content + self.headers = dict(request.headers) + + def json_body(self) -> dict: + return json.loads(self.content) + + +class ScriptedTokenEndpoint: + """A token endpoint behind ``httpx2.MockTransport``: answers with the scripted responses in + order (repeating the last) and records every request the SDK put on the wire.""" + + def __init__( + self, + responses: list[httpx2.Response], + on_request: Callable[[RecordedRequest], None] | None = None, + ) -> None: + self.requests: list[RecordedRequest] = [] + self._responses = list(responses) + self._on_request = on_request + + def __call__(self, request: httpx2.Request) -> httpx2.Response: + recorded = RecordedRequest(request) + self.requests.append(recorded) + if self._on_request is not None: + self._on_request(recorded) + if len(self._responses) > 1: + return self._responses.pop(0) + return self._responses[0] + + +class RaisingTokenEndpoint: + def __init__(self, error: Exception) -> None: + self.calls = 0 + self._error = error + + def __call__(self, request: httpx2.Request) -> httpx2.Response: + self.calls += 1 + raise self._error + + +class EchoingUnauthorizedEndpoint: + """401s every attempt, echoing the submitted assertion back into the error body: a token + endpoint that reflects the request.""" + + def __init__(self) -> None: + self.requests: list[RecordedRequest] = [] + + def __call__(self, request: httpx2.Request) -> httpx2.Response: + recorded = RecordedRequest(request) + self.requests.append(recorded) + submitted = recorded.json_body()["assertion"] + return httpx2.Response(401, json={"error": "invalid_grant", "error_description": f"bad assertion {submitted}"}) + + +class RotatingAssertionSource: + """A per-call assertion source that mints a fresh value on every read, the shape the + internal_issuer and keycloak identity sources take.""" + + def __init__(self, values: list[str]) -> None: + self._values = iter(values) + self.calls = 0 + + def __call__(self) -> str: + self.calls += 1 + return next(self._values) + + +class RecordingMetricsSink: + def __init__(self) -> None: + self.successes: list[tuple[str, float]] = [] + self.failures: list[tuple[str, float, ExchangeError]] = [] + self.cache_hits = 0 + + def exchange_success(self, *, call_type: str, duration_seconds: float) -> None: + self.successes.append((call_type, duration_seconds)) + + def exchange_failure(self, *, call_type: str, duration_seconds: float, error: ExchangeError) -> None: + self.failures.append((call_type, duration_seconds, error)) + + def cache_hit(self) -> None: + self.cache_hits += 1 + + +class RaisingMetricsSink: + def exchange_success(self, *, call_type: str, duration_seconds: float) -> None: + raise RuntimeError("metrics sink down") + + def exchange_failure(self, *, call_type: str, duration_seconds: float, error: ExchangeError) -> None: + raise RuntimeError("metrics sink down") + + def cache_hit(self) -> None: + raise RuntimeError("metrics sink down") + + +def token_response(token: str = "sk-ant-oat01-minted", expires_in: int | None = TOKEN_TTL) -> httpx2.Response: + body: Final[dict[str, str | int]] = { + "access_token": token, + "token_type": "Bearer", + **({} if expires_in is None else {"expires_in": expires_in}), + } + return httpx2.Response(200, json=body) + + +def make_params( + *, + federation_rule_id: str = "fdrl_1", + organization_id: str = "org-1", + service_account_id: str | None = None, + workspace_id: str | None = None, + assertion_ref: str = DEFAULT_REF, + assertion_source: Callable[[], str | None] | None = None, +) -> AnthropicWifParams: + return AnthropicWifParams( + federation_rule_id=federation_rule_id, + organization_id=organization_id, + service_account_id=service_account_id, + workspace_id=workspace_id, + assertion_ref=assertion_ref, + assertion_source=assertion_source, + ) + + +def make_exchange( + endpoint: Callable[[httpx2.Request], httpx2.Response], + reader: Mapping[str, str] | Callable[[str], str | None] | None = None, + clock: FakeClock | None = None, + max_entries: int = 64, + metrics_sink=None, +) -> AnthropicWifTokenExchange: + resolved_reader = reader if callable(reader) else (reader or {DEFAULT_REF: DEFAULT_ASSERTION}).get + return AnthropicWifTokenExchange( + http_client=httpx2.Client(transport=httpx2.MockTransport(endpoint)), + assertion_reader=resolved_reader, + max_entries=max_entries, + metrics_sink=metrics_sink if metrics_sink is not None else RecordingMetricsSink(), + clock=clock if clock is not None else time.time, + ) + + +def mint( + exchange: AnthropicWifTokenExchange, params: AnthropicWifParams | None = None, exchange_base: str = EXCHANGE_BASE +) -> str: + result = exchange.get_token(params if params is not None else make_params(), exchange_base) + assert isinstance(result, str), result + return result + + +class TestFreshMintWireExact: + def test_url_headers_and_full_body(self): + endpoint = ScriptedTokenEndpoint([token_response()]) + exchange = make_exchange(endpoint) + + token = mint(exchange, make_params(service_account_id="svcacct_1", workspace_id="wrkspc_1")) + + assert token == "sk-ant-oat01-minted" + (request,) = endpoint.requests + assert request.url == EXCHANGE_URL + assert request.headers["anthropic-beta"] == SDK_BETA_HEADER + assert request.headers["content-type"] == "application/json" + assert request.headers["user-agent"].startswith("anthropic-python/") + assert request.json_body() == { + "grant_type": GRANT_TYPE, + "assertion": DEFAULT_ASSERTION, + "federation_rule_id": "fdrl_1", + "organization_id": "org-1", + "service_account_id": "svcacct_1", + "workspace_id": "wrkspc_1", + } + + def test_unset_optional_ids_are_absent_from_the_body(self): + endpoint = ScriptedTokenEndpoint([token_response()]) + + mint(make_exchange(endpoint)) + + assert set(endpoint.requests[0].json_body()) == { + "grant_type", + "assertion", + "federation_rule_id", + "organization_id", + } + + def test_a_second_call_is_served_from_cache_without_a_post(self): + endpoint = ScriptedTokenEndpoint([token_response()]) + exchange = make_exchange(endpoint) + + first = mint(exchange) + second = mint(exchange) + + assert first == second == "sk-ant-oat01-minted" + assert len(endpoint.requests) == 1 + + +class TestUnauthorizedRetry: + def test_a_401_is_retried_once_with_a_freshly_read_assertion(self): + assertions = {DEFAULT_REF: "assertion-v1"} + endpoint = ScriptedTokenEndpoint([httpx2.Response(401, json={"error": "invalid_grant"}), token_response()]) + exchange = make_exchange(endpoint, reader=assertions.get) + + def rotate(_request: RecordedRequest) -> None: + assertions[DEFAULT_REF] = "assertion-v2" + + endpoint._on_request = rotate + + assert mint(exchange) == "sk-ant-oat01-minted" + assert [request.json_body()["assertion"] for request in endpoint.requests] == ["assertion-v1", "assertion-v2"] + + def test_401_twice_is_endpoint_error(self): + endpoint = ScriptedTokenEndpoint([httpx2.Response(401, json={"error": "invalid_grant"})]) + + result = make_exchange(endpoint).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, TokenEndpointError) + assert result.status_code == 401 + assert "invalid_grant" in result.redacted_body + assert len(endpoint.requests) == 2 + + def test_401_retry_redacts_the_assertion_actually_sent_not_a_fresh_reread(self): + """Regression: with a rotating identity source, the reflection-drop check must match the + assertion the failing (second) attempt actually sent. Re-reading for the check would mint a + THIRD value that was never sent, so the reflection probe would miss and the actually-sent, + actually-reflected second assertion would leak into the error.""" + endpoint = EchoingUnauthorizedEndpoint() + source = RotatingAssertionSource(["assertion-v1", "assertion-v2", "assertion-v3"]) + + result = make_exchange(endpoint).get_token(make_params(assertion_source=source), EXCHANGE_BASE) + + assert isinstance(result, TokenEndpointError) + assert [request.json_body()["assertion"] for request in endpoint.requests] == ["assertion-v1", "assertion-v2"] + assert source.calls == 2, "the failing attempt's own assertion must be reused, never re-read a third time" + assert "assertion-v1" not in result.redacted_body + assert "assertion-v2" not in result.redacted_body + assert "assertion-v3" not in result.redacted_body + + +class TestSdkOutcomeMapping: + """The SDK folds every failure into one exception class; each shape must land on the typed + error ``wif.py`` maps onto the public exception contract, with nothing echoed through.""" + + def test_error_object_body_is_reduced_to_rfc6749_fields_and_capped(self): + endpoint = ScriptedTokenEndpoint( + [ + httpx2.Response( + 400, + json={ + "error": "invalid_grant", + "error_description": "d" * 500, + "error_uri": "https://errors.example/e1", + "assertion_echo": "LEAKED-ASSERTION", + }, + ) + ] + ) + + result = make_exchange(endpoint).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, TokenEndpointError) + assert result.status_code == 400 + assert "invalid_grant" in result.redacted_body + assert "d" * 256 in result.redacted_body + assert "d" * 257 not in result.redacted_body + assert "https://errors.example/e1" in result.redacted_body + assert "LEAKED-ASSERTION" not in result.redacted_body + + def test_reflected_assertion_in_an_error_body_is_dropped(self): + result = make_exchange(EchoingUnauthorizedEndpoint()).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, TokenEndpointError) + assert DEFAULT_ASSERTION not in result.redacted_body + + def test_plain_text_error_body_is_not_echoed(self): + endpoint = ScriptedTokenEndpoint([httpx2.Response(502, text="t" * 500)]) + + result = make_exchange(endpoint).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, TokenEndpointError) + assert result.status_code == 502 + assert result.redacted_body == "non-JSON error response omitted" + + def test_oversized_error_body_is_never_parsed(self): + endpoint = ScriptedTokenEndpoint( + [httpx2.Response(400, content=b'{"error": "' + b"x" * MAX_RESPONSE_BYTES + b'"}')] + ) + + result = make_exchange(endpoint).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, TokenEndpointError) + assert result.status_code == 400 + assert "x" * 10 not in result.redacted_body + + def test_oversized_success_body_is_malformed(self): + endpoint = ScriptedTokenEndpoint( + [httpx2.Response(200, content=b'{"access_token": "' + b"x" * MAX_RESPONSE_BYTES + b'"}')] + ) + + result = make_exchange(endpoint).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, MalformedTokenResponse) + assert "x" * 10 not in result.detail + + def test_non_json_success_body_that_echoes_the_assertion_is_scrubbed(self): + """The SDK quotes up to 256 characters of an unparseable 2xx body in its message, so a + reflected assertion has to be dropped before that message becomes an error detail.""" + endpoint = ScriptedTokenEndpoint([httpx2.Response(200, text=f"{DEFAULT_ASSERTION}")]) + + result = make_exchange(endpoint).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, MalformedTokenResponse) + assert DEFAULT_ASSERTION not in result.detail + + def test_non_json_success_body_without_an_echo_keeps_the_diagnosis(self): + endpoint = ScriptedTokenEndpoint([httpx2.Response(200, text="upstream proxy page")]) + + result = make_exchange(endpoint).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, MalformedTokenResponse) + assert "non-JSON" in result.detail + + def test_json_array_success_body_is_malformed(self): + endpoint = ScriptedTokenEndpoint([httpx2.Response(200, json=["a", "b"])]) + + result = make_exchange(endpoint).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, MalformedTokenResponse) + + def test_a_non_bearer_token_type_is_refused(self): + endpoint = ScriptedTokenEndpoint( + [httpx2.Response(200, json={"access_token": "tok", "token_type": "mac", "expires_in": 300})] + ) + + result = make_exchange(endpoint).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, MalformedTokenResponse) + assert "mac" in result.detail + + def test_missing_expires_in_is_malformed(self): + endpoint = ScriptedTokenEndpoint([token_response(expires_in=None)]) + + result = make_exchange(endpoint).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, MalformedTokenResponse) + assert "expires_in" in result.detail + + @pytest.mark.parametrize("access_token", ["", " "]) + def test_empty_access_token_is_malformed_and_not_cached(self, access_token: str): + endpoint = ScriptedTokenEndpoint( + [ + httpx2.Response(200, json={"access_token": access_token, "token_type": "Bearer", "expires_in": 3600}), + token_response("reminted"), + ] + ) + exchange = make_exchange(endpoint) + + first = exchange.get_token(make_params(), EXCHANGE_BASE) + second = exchange.get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(first, MalformedTokenResponse) + assert "empty access_token" in first.detail + assert second == "reminted" + assert len(endpoint.requests) == 2 + + @pytest.mark.parametrize("expires_in", [0, -5]) + def test_a_token_expired_on_arrival_is_never_cached(self, expires_in: int): + endpoint = ScriptedTokenEndpoint( + [token_response("short-lived", expires_in=expires_in), token_response("reminted")] + ) + exchange = make_exchange(endpoint) + + assert mint(exchange) == "short-lived" + assert mint(exchange) == "reminted" + assert len(endpoint.requests) == 2 + + def test_transport_failure_names_the_endpoint(self): + endpoint = RaisingTokenEndpoint(httpx2.ConnectError("connection refused")) + + result = make_exchange(endpoint).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, TokenTransportError) + assert EXCHANGE_URL in result.detail + assert "connection refused" in result.detail + assert endpoint.calls == 1 + + def test_a_redirect_is_an_error_not_a_second_post(self): + """Only the bound base URL passed the host allowlist, so a 3xx must never carry the + assertion to the location it names.""" + endpoint = ScriptedTokenEndpoint( + [httpx2.Response(302, headers={"location": "https://elsewhere.example/v1/oauth/token"})] + ) + + result = make_exchange(endpoint).get_token(make_params(), EXCHANGE_BASE) + + assert not isinstance(result, str) + assert len(endpoint.requests) == 1 + + +def test_sentinel_leak_audit(caplog: pytest.LogCaptureFixture): + jwt_sentinel = "JWT-SENTINEL-2c9f1e7ab4" + token_sentinel = "sk-ant-oat01-TOKEN-SENTINEL-90d4c3aa17" + ref = "oidc/env/SENTINEL_ASSERTION" + params = make_params(assertion_ref=ref) + + def exchange_with(endpoint, reader: Mapping[str, str] | None = None, clock: FakeClock | None = None): + return make_exchange(endpoint, reader=reader if reader is not None else {ref: jwt_sentinel}, clock=clock) + + with caplog.at_level(logging.DEBUG): + clock = FakeClock() + serving = exchange_with( + ScriptedTokenEndpoint( + [token_response(token_sentinel), httpx2.Response(500, json={"error": "server_error"})] + ), + clock=clock, + ) + minted = serving.get_token(params, EXCHANGE_BASE) + endpoint_error = exchange_with(ScriptedTokenEndpoint([httpx2.Response(400, json={"error": "invalid_grant"})])) + endpoint_error_result = endpoint_error.get_token(params, EXCHANGE_BASE) + transport_error = exchange_with(RaisingTokenEndpoint(httpx2.ConnectError("boom"))).get_token( + params, EXCHANGE_BASE + ) + malformed_error = exchange_with(ScriptedTokenEndpoint([httpx2.Response(200, json={"unexpected": "shape"})])) + malformed_result = malformed_error.get_token(params, EXCHANGE_BASE) + echoed_result = exchange_with(EchoingUnauthorizedEndpoint()).get_token(params, EXCHANGE_BASE) + oversized_error = exchange_with( + ScriptedTokenEndpoint([token_response()]), reader={ref: jwt_sentinel + "x" * MAX_ASSERTION_BYTES} + ).get_token(params, EXCHANGE_BASE) + insecure_error = exchange_with(ScriptedTokenEndpoint([token_response()])).get_token( + params, "http://token.example" + ) + clock.advance(TOKEN_TTL - 100) + stale = serving.get_token(params, EXCHANGE_BASE) + + assert minted == token_sentinel + assert stale == token_sentinel, "an advisory refresh failure keeps serving the cached token" + assert isinstance(oversized_error, AssertionSourceError) + assert oversized_error.kind == "oversized" + audited_values = [ + str(endpoint_error_result), + repr(endpoint_error_result), + str(transport_error), + repr(transport_error), + str(malformed_result), + repr(malformed_result), + str(echoed_result), + repr(echoed_result), + str(oversized_error), + repr(oversized_error), + str(insecure_error), + repr(insecure_error), + caplog.text, + ] + for value in audited_values: + assert jwt_sentinel not in value + assert token_sentinel not in value + + +class TestAssertionGuards: + @pytest.mark.parametrize( + "assertion_value,expected_kind", + [ + ("x" * (MAX_ASSERTION_BYTES + 1), "oversized"), + (" \n\t ", "empty"), + (None, "missing"), + ], + ) + def test_bad_assertion_values(self, assertion_value: str | None, expected_kind: str): + endpoint = ScriptedTokenEndpoint([token_response()]) + + result = make_exchange(endpoint, reader=lambda ref: assertion_value).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, AssertionSourceError) + assert result.kind == expected_kind + assert result.source_ref == DEFAULT_REF + assert len(endpoint.requests) == 0 + + @pytest.mark.parametrize( + "raised,expected_kind", + [ + (OidcPathNotAllowedError("path outside allowed credential directories"), "disallowed_path"), + (ValueError("Environment variable ANTHROPIC_IDENTITY_TOKEN not found"), "unreadable"), + (ImportError("needs PyJWT and cryptography: pip install 'litellm[proxy]'"), "unreadable"), + (OSError("permission denied"), "unreadable"), + ], + ) + def test_raising_reader(self, raised: Exception, expected_kind: str): + endpoint = ScriptedTokenEndpoint([token_response()]) + + def reader(ref: str) -> str | None: + raise raised + + result = make_exchange(endpoint, reader=reader).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, AssertionSourceError) + assert result.kind == expected_kind + assert len(endpoint.requests) == 0 + + @pytest.mark.parametrize( + "raised,expected_detail", + [ + ( + ValueError("Keycloak token endpoint returned invalid_client"), + "Keycloak token endpoint returned invalid_client", + ), + (ValueError("x" * (_DETAIL_CAP + 100)), "x" * _DETAIL_CAP), + (ImportError("the internal_issuer identity source needs PyJWT: pip install 'litellm[proxy]'"), None), + ], + ) + def test_operator_diagnosable_messages_are_captured_as_detail(self, raised: Exception, expected_detail: str | None): + def reader(ref: str) -> str | None: + raise raised + + result = make_exchange(ScriptedTokenEndpoint([token_response()]), reader=reader).get_token( + make_params(), EXCHANGE_BASE + ) + + assert isinstance(result, AssertionSourceError) + assert result.detail is not None + assert result.detail == (expected_detail if expected_detail is not None else str(raised)) + + @pytest.mark.parametrize( + "raised", + [OidcPathNotAllowedError("path outside allowed credential directories"), OSError("permission denied")], + ) + def test_non_value_error_never_populates_detail(self, raised: Exception): + """Only the ValueError and ImportError branches carry operator-diagnosable text; every other + reader failure stays detail=None.""" + + def reader(ref: str) -> str | None: + raise raised + + result = make_exchange(ScriptedTokenEndpoint([token_response()]), reader=reader).get_token( + make_params(), EXCHANGE_BASE + ) + + assert isinstance(result, AssertionSourceError) + assert result.detail is None + + +class TestAssertionSourceOverridesReader: + """``AnthropicWifParams.assertion_source`` is how a per-config identity source (internal_issuer, + keycloak) plugs in; it must win over the exchange-level reader, and failures must still be + reported against ``assertion_ref``.""" + + def test_assertion_source_is_used_instead_of_the_reader(self): + endpoint = ScriptedTokenEndpoint([token_response()]) + calls: list[str] = [] + + def reader(ref: str) -> str | None: + calls.append(ref) + return "from-reader" + + token = mint(make_exchange(endpoint, reader=reader), make_params(assertion_source=lambda: "from-source")) + + assert token == "sk-ant-oat01-minted" + assert endpoint.requests[0].json_body()["assertion"] == "from-source" + assert calls == [] + + def test_assertion_source_failure_is_reported_against_assertion_ref(self): + endpoint = ScriptedTokenEndpoint([token_response()]) + + def raising_source() -> str | None: + raise ValueError("keycloak token endpoint returned invalid_client") + + result = make_exchange(endpoint, reader=lambda ref: "from-reader").get_token( + make_params(assertion_source=raising_source, assertion_ref="oidc/keycloak/abc123"), EXCHANGE_BASE + ) + + assert isinstance(result, AssertionSourceError) + assert result.source_ref == "oidc/keycloak/abc123" + assert result.detail == "keycloak token endpoint returned invalid_client" + assert len(endpoint.requests) == 0 + + def test_assertion_source_is_re_invoked_on_401_retry(self): + values = iter(["assertion-v1", "assertion-v2"]) + endpoint = ScriptedTokenEndpoint([httpx2.Response(401, json={"error": "invalid_grant"}), token_response()]) + + token = mint( + make_exchange(endpoint, reader=lambda ref: "from-reader"), + make_params(assertion_source=lambda: next(values)), + ) + + assert token == "sk-ant-oat01-minted" + assert [request.json_body()["assertion"] for request in endpoint.requests] == ["assertion-v1", "assertion-v2"] + + +class TestOidcFilePathAllowlistRaisesTypedError: + """Assertion-source failures are classified by exception type (see TestAssertionGuards); that + only works if the real oidc/file allowlist raises OidcPathNotAllowedError, not a bare ValueError.""" + + def test_out_of_allowlist_absolute_path(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", raising=False) + + with pytest.raises(OidcPathNotAllowedError): + _resolve_oidc_file_path("/etc/not-a-credential-dir/token") + + def test_relative_path(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", raising=False) + + with pytest.raises(OidcPathNotAllowedError): + _resolve_oidc_file_path("relative/token/path") + + +class TestHttpsEnforcement: + def test_plain_http_rejected_host_only_zero_posts(self): + endpoint = ScriptedTokenEndpoint([token_response()]) + sink = RecordingMetricsSink() + + result = make_exchange(endpoint, metrics_sink=sink).get_token(make_params(), "http://token.example") + + assert result == InsecureTokenUrl(host="token.example") + assert "/v1/oauth/token" not in str(result) + assert len(endpoint.requests) == 0 + assert [(call_type, error) for call_type, _duration, error in sink.failures] == [("cold_mint", result)] + + @pytest.mark.parametrize("base", ["http://localhost:8080", "http://127.0.0.1", "http://[::1]"]) + def test_localhost_http_allowed(self, base: str): + endpoint = ScriptedTokenEndpoint([token_response()]) + + token = mint(make_exchange(endpoint), exchange_base=base) + + assert token == "sk-ant-oat01-minted" + assert endpoint.requests[0].url == f"{base}/v1/oauth/token" + + +class TestCacheIdentity: + def test_cache_key_semantics(self): + endpoint = ScriptedTokenEndpoint([token_response()]) + assertions = {DEFAULT_REF: DEFAULT_ASSERTION, "oidc/env/OTHER": "other-assertion"} + exchange = make_exchange(endpoint, reader=assertions.get) + + mint(exchange) + mint(exchange, make_params(service_account_id="svc-2")) + mint(exchange, make_params(workspace_id="wrkspc-2")) + mint(exchange, make_params(federation_rule_id="fdrl_2")) + mint(exchange, make_params(organization_id="org-2")) + mint(exchange, exchange_base="https://other.example") + mint(exchange, make_params(assertion_ref="oidc/env/OTHER")) + assert len(endpoint.requests) == 7 + + assertions[DEFAULT_REF] = "rotated-assertion" + assert mint(exchange) == "sk-ant-oat01-minted" + assert len(endpoint.requests) == 7, "a still-valid token is served without re-reading the assertion" + + def test_bounded_eviction_drops_the_oldest_deployment(self): + endpoint = ScriptedTokenEndpoint([token_response()]) + exchange = make_exchange(endpoint, max_entries=4) + + def mint_org(index: int) -> None: + mint(exchange, make_params(organization_id=f"org-{index}")) + + for index in range(5): + mint_org(index) + assert len(endpoint.requests) == 5 + + mint_org(0) + assert len(endpoint.requests) == 6, "the oldest entry (index 0) should have been evicted" + + mint_org(2) + assert len(endpoint.requests) == 6, "a younger entry should still be cached" + + mint_org(1) + assert len(endpoint.requests) == 7, "re-inserting index 0 should have evicted the next oldest entry" + + +async def test_aget_token_loop_responsive(): + def sleeping_endpoint(request: httpx2.Request) -> httpx2.Response: + time.sleep(0.3) + return token_response() + + exchange = make_exchange(sleeping_endpoint) + ticks = {"count": 0} + stop = asyncio.Event() + + async def ticker() -> None: + while not stop.is_set(): + ticks["count"] += 1 + await asyncio.sleep(0.01) + + ticker_task = asyncio.create_task(ticker()) + result = await exchange.aget_token(make_params(), EXCHANGE_BASE) + stop.set() + await ticker_task + + assert result == "sk-ant-oat01-minted" + assert ticks["count"] >= 5, "the event loop was blocked during aget_token" + assert exchange.get_token(make_params(), EXCHANGE_BASE) == result + + +class TestRefreshWindows: + def test_well_before_expiry_the_cache_serves_without_posting(self): + clock = FakeClock() + sink = RecordingMetricsSink() + endpoint = ScriptedTokenEndpoint([token_response("old"), token_response("new")]) + exchange = make_exchange(endpoint, clock=clock, metrics_sink=sink) + + mint(exchange) + clock.advance(TOKEN_TTL - 121) + assert mint(exchange) == "old" + + assert len(endpoint.requests) == 1 + assert sink.cache_hits == 1 + + def test_inside_the_advisory_window_the_token_is_refreshed_inline(self): + clock = FakeClock() + sink = RecordingMetricsSink() + endpoint = ScriptedTokenEndpoint([token_response("old"), token_response("new")]) + exchange = make_exchange(endpoint, clock=clock, metrics_sink=sink) + + mint(exchange) + clock.advance(TOKEN_TTL - 119) + assert mint(exchange) == "new" + + assert len(endpoint.requests) == 2 + assert [call_type for call_type, _ in sink.successes] == ["cold_mint", "refresh"] + assert sink.cache_hits == 0 + + def test_an_advisory_refresh_failure_serves_the_cached_token_and_backs_off(self): + clock = FakeClock() + sink = RecordingMetricsSink() + endpoint = ScriptedTokenEndpoint([token_response("old"), httpx2.Response(500, json={"error": "server_error"})]) + exchange = make_exchange(endpoint, clock=clock, metrics_sink=sink) + + mint(exchange) + clock.advance(TOKEN_TTL - 119) + assert mint(exchange) == "old" + assert mint(exchange) == "old" + + assert len(endpoint.requests) == 2, "a refresh that just failed is not retried on the very next call" + ((call_type, _duration, error),) = sink.failures + assert call_type == "refresh" + assert isinstance(error, TokenEndpointError) + assert error.status_code == 500 + + def test_inside_the_mandatory_window_a_failed_refresh_is_the_error(self): + clock = FakeClock() + sink = RecordingMetricsSink() + endpoint = ScriptedTokenEndpoint([token_response("old"), httpx2.Response(503, json={"error": "unavailable"})]) + exchange = make_exchange(endpoint, clock=clock, metrics_sink=sink) + + mint(exchange) + clock.advance(TOKEN_TTL - 29) + result = exchange.get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, TokenEndpointError) + assert result.status_code == 503 + assert [call_type for call_type, _duration, _error in sink.failures] == ["refresh"] + + def test_a_refresh_reads_the_assertion_again(self): + clock = FakeClock() + assertions = {DEFAULT_REF: "first-assertion"} + endpoint = ScriptedTokenEndpoint([token_response("old"), token_response("new")]) + exchange = make_exchange(endpoint, reader=assertions.get, clock=clock) + + mint(exchange) + assertions[DEFAULT_REF] = "second-assertion" + clock.advance(TOKEN_TTL - 29) + assert mint(exchange) == "new" + + assert endpoint.requests[1].json_body()["assertion"] == "second-assertion" + + +class TestMetricsEmission: + def test_cold_mint_emits_success_with_duration(self): + sink = RecordingMetricsSink() + endpoint = ScriptedTokenEndpoint([token_response()], on_request=lambda _request: time.sleep(0.05)) + + mint(make_exchange(endpoint, metrics_sink=sink)) + + ((call_type, duration),) = sink.successes + assert call_type == "cold_mint" + assert duration >= 0.05 + assert sink.failures == [] + assert sink.cache_hits == 0 + + def test_cache_hit_emits_counter_not_a_mint(self): + sink = RecordingMetricsSink() + exchange = make_exchange(ScriptedTokenEndpoint([token_response()]), metrics_sink=sink) + + mint(exchange) + mint(exchange) + + assert sink.cache_hits == 1 + assert len(sink.successes) == 1 + + def test_every_failed_cold_mint_attempt_emits_a_failure(self): + sink = RecordingMetricsSink() + exchange = make_exchange( + ScriptedTokenEndpoint([httpx2.Response(503, json={"error": "unavailable"})]), metrics_sink=sink + ) + + first = exchange.get_token(make_params(), EXCHANGE_BASE) + second = exchange.get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(first, TokenEndpointError) + assert isinstance(second, TokenEndpointError) + assert [call_type for call_type, _duration, _error in sink.failures] == ["cold_mint", "cold_mint"] + assert all(isinstance(error, TokenEndpointError) and error.status_code == 503 for _, _, error in sink.failures) + assert sink.successes == [] + + def test_a_401_retry_reports_the_failed_attempt_and_the_cold_mint(self): + sink = RecordingMetricsSink() + endpoint = ScriptedTokenEndpoint([httpx2.Response(401, json={"error": "invalid_grant"}), token_response()]) + + mint(make_exchange(endpoint, metrics_sink=sink)) + + assert [call_type for call_type, _duration, _error in sink.failures] == ["cold_mint"] + assert [call_type for call_type, _duration in sink.successes] == ["cold_mint"] + + def test_failure_payload_carries_no_assertion_material(self): + sink = RecordingMetricsSink() + + result = make_exchange(EchoingUnauthorizedEndpoint(), metrics_sink=sink).get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, TokenEndpointError) + assert len(sink.failures) == 2 + for failure in sink.failures: + assert DEFAULT_ASSERTION not in repr(failure) + assert DEFAULT_ASSERTION not in _error_summary(failure[2]) + + def test_raising_sink_never_breaks_mint_serve_or_failure(self): + exchange = make_exchange(ScriptedTokenEndpoint([token_response()]), metrics_sink=RaisingMetricsSink()) + + assert mint(exchange) == mint(exchange) + + failing = make_exchange(RaisingTokenEndpoint(httpx2.ConnectError("boom")), metrics_sink=RaisingMetricsSink()) + assert isinstance(failing.get_token(make_params(), EXCHANGE_BASE), TokenTransportError) + + +class TestDefaultAssertionReader: + def test_reads_through_litellm_secret_resolution(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("WIF_ASSERTION_FOR_DEFAULT_READER", "header.payload.signature") + + assert _default_assertion_reader("os.environ/WIF_ASSERTION_FOR_DEFAULT_READER") == "header.payload.signature" + + def test_an_unset_reference_reads_as_none(self): + assert _default_assertion_reader("os.environ/DEFINITELY_NOT_SET_WIF_ASSERTION_REF") is None + + +class TestErrorSummary: + def test_every_error_variant_summarises_without_carrying_a_secret(self): + summaries: Final = { + _error_summary(AssertionSourceError(kind="unreadable", source_ref="oidc/file/x")), + _error_summary(InsecureTokenUrl(host="token.internal")), + _error_summary(TokenEndpointError(status_code=401, redacted_body="invalid_grant")), + _error_summary(TokenTransportError(detail="ConnectError: refused")), + _error_summary(MalformedTokenResponse(detail="empty access_token")), + } + + assert {s.split(":")[0] for s in summaries} == { + "AssertionSourceError", + "InsecureTokenUrl", + "TokenEndpointError", + "TokenTransportError", + "MalformedTokenResponse", + }, "each variant names itself so a log line says which stage failed" + + +class TestExchangeClient: + def test_redirects_are_not_followed_and_timeouts_are_bounded(self): + """Only the initial token URL is validated, so a 3xx must not be allowed to replay the + assertion to an origin that was never checked; and a stalled endpoint must not pin a + worker for longer than the exchange budget.""" + client = new_exchange_client() + + assert client.follow_redirects is False + assert client.timeout.connect == EXCHANGE_CONNECT_TIMEOUT_SECONDS + assert client.timeout.read == EXCHANGE_TIMEOUT_SECONDS + + +class RecordingServiceHooks: + def __init__(self) -> None: + self.successes: list[tuple[ServiceTypes, str, float]] = [] + self.failures: list[tuple[ServiceTypes, float, str | Exception, str]] = [] + + async def async_service_success_hook(self, service: ServiceTypes, call_type: str, duration: float) -> None: + self.successes.append((service, call_type, duration)) + + async def async_service_failure_hook( + self, service: ServiceTypes, duration: float, error: str | Exception, call_type: str + ) -> None: + self.failures.append((service, duration, error, call_type)) + + +class RaisingServiceHooks: + """Every hook raises, and each call is recorded first so a test can prove the sink kept + calling through rather than bailing after the first failure.""" + + def __init__(self) -> None: + self.attempts: list[str] = [] # mutable-ok: a test spy accumulating calls in order + + async def async_service_success_hook(self, service: ServiceTypes, call_type: str, duration: float) -> None: + self.attempts.append(f"success:{call_type}") + raise RuntimeError("hook down") + + async def async_service_failure_hook( + self, service: ServiceTypes, duration: float, error: str | Exception, call_type: str + ) -> None: + self.attempts.append(f"failure:{call_type}") + raise RuntimeError("hook down") + + +class InlineExecutor(concurrent.futures.Executor): + def submit(self, fn, /, *args, **kwargs): + future: concurrent.futures.Future = concurrent.futures.Future() + future.set_result(fn(*args, **kwargs)) + return future + + +class NeverRunsExecutor(concurrent.futures.Executor): + """Accepts work and never runs it, standing in for a telemetry backend that has stalled, so a + test can show the backlog stops growing instead of consuming memory for as long as traffic lasts.""" + + def __init__(self) -> None: + self.submitted = 0 # mutable-ok: a test spy counting accepted work + + def submit(self, fn, /, *args, **kwargs): + self.submitted += 1 + return concurrent.futures.Future() + + +class TestServiceLoggingMetricsSink: + def _sink(self, hooks) -> ServiceLoggingMetricsSink: + return ServiceLoggingMetricsSink(service_logging_factory=lambda: hooks, executor=InlineExecutor()) + + def test_success_maps_to_anthropic_wif_service(self): + hooks = RecordingServiceHooks() + + self._sink(hooks).exchange_success(call_type="cold_mint", duration_seconds=0.2) + + assert hooks.successes == [(ServiceTypes.ANTHROPIC_WIF, "cold_mint", 0.2)] + + def test_a_stalled_backend_stops_accepting_work_instead_of_queueing_without_bound(self): + stalled: Final = NeverRunsExecutor() + sink: Final = ServiceLoggingMetricsSink(service_logging_factory=RecordingServiceHooks, executor=stalled) + + for _ in range(_METRICS_QUEUE_LIMIT + 500): + sink.cache_hit() + + assert stalled.submitted == _METRICS_QUEUE_LIMIT, ( + "once the backlog is full further events are dropped, so request volume cannot grow it" + ) + + def test_a_drained_backlog_accepts_work_again(self): + hooks: Final = RecordingServiceHooks() + sink: Final = ServiceLoggingMetricsSink(service_logging_factory=lambda: hooks, executor=InlineExecutor()) + + for _ in range(_METRICS_QUEUE_LIMIT + 10): + sink.cache_hit() + + assert len(hooks.successes) == _METRICS_QUEUE_LIMIT + 10, ( + "an executor that actually runs releases each slot, so nothing is dropped" + ) + + def test_failure_maps_variant_and_redacted_summary(self): + hooks = RecordingServiceHooks() + error = TokenEndpointError(status_code=503, redacted_body="error: unavailable") + + self._sink(hooks).exchange_failure(call_type="refresh", duration_seconds=0.1, error=error) + + ((service, duration, emitted, call_type),) = hooks.failures + assert service is ServiceTypes.ANTHROPIC_WIF + assert duration == 0.1 + assert call_type == "refresh" + assert isinstance(emitted, TokenExchangeEndpointFailure) + assert str(emitted) == _error_summary(error) + + def test_transport_failure_gets_its_own_error_class(self): + hooks = RecordingServiceHooks() + + self._sink(hooks).exchange_failure( + call_type="refresh", duration_seconds=0.05, error=TokenTransportError(detail="ConnectError: boom") + ) + + ((_service, _duration, emitted, _call_type),) = hooks.failures + assert isinstance(emitted, TokenExchangeTransportFailure) + + def test_cache_hit_maps_to_cache_service_with_zero_duration(self): + hooks = RecordingServiceHooks() + + self._sink(hooks).cache_hit() + + assert hooks.successes == [(ServiceTypes.ANTHROPIC_WIF_CACHE, CALL_TYPE_CACHE_HIT, 0.0)] + + def test_end_to_end_reflected_assertion_never_reaches_the_hook(self): + hooks = RecordingServiceHooks() + exchange = make_exchange(EchoingUnauthorizedEndpoint(), metrics_sink=self._sink(hooks)) + + result = exchange.get_token(make_params(), EXCHANGE_BASE) + + assert isinstance(result, TokenEndpointError) + assert [call_type for _service, _duration, _emitted, call_type in hooks.failures] == ["cold_mint", "cold_mint"] + for _service, _duration, emitted, _call_type in hooks.failures: + assert DEFAULT_ASSERTION not in str(emitted) + assert DEFAULT_ASSERTION not in repr(emitted) + + def test_raising_hooks_are_swallowed(self): + hooks: Final = RaisingServiceHooks() + sink: Final = self._sink(hooks) + + sink.exchange_success(call_type="cold_mint", duration_seconds=0.2) + sink.cache_hit() + sink.exchange_failure(call_type="cold_mint", duration_seconds=0.1, error=TokenTransportError(detail="boom")) + + assert hooks.attempts == ["success:cold_mint", "success:cache_hit", "failure:cold_mint"], ( + "every event is still handed to the hooks, and one raising hook does not stop the next" + ) diff --git a/tests/test_litellm/llms/base_llm/auth/test_client_credentials.py b/tests/test_litellm/llms/base_llm/auth/test_client_credentials.py index 76c309cae37c..1951a9d430ae 100644 --- a/tests/test_litellm/llms/base_llm/auth/test_client_credentials.py +++ b/tests/test_litellm/llms/base_llm/auth/test_client_credentials.py @@ -15,7 +15,7 @@ keycloak_assertion_source, ) from litellm.llms.base_llm.auth.identity_source import KeycloakSource, identity_source_ref -from litellm.llms.base_llm.auth.token_exchange import MAX_RESPONSE_BYTES +from litellm.llms.base_llm.auth.oauth_endpoint import MAX_RESPONSE_BYTES TOKEN_URL: Final = "https://keycloak.example/realms/litellm/protocol/openid-connect/token" CLIENT_ID: Final = "litellm" @@ -458,7 +458,7 @@ class TestTokenUrlIsNotEchoedWholesale: stops an operator putting a credential in the URL, and these errors reach model callers.""" def test_query_string_is_dropped_from_a_status_error(self): - from litellm.llms.base_llm.auth.token_exchange import endpoint_url_for_error_message + from litellm.llms.base_llm.auth.oauth_endpoint import endpoint_url_for_error_message rendered = endpoint_url_for_error_message("https://idp.example/token?client_secret=supersecret") @@ -466,7 +466,7 @@ def test_query_string_is_dropped_from_a_status_error(self): assert rendered == "https://idp.example/token" def test_userinfo_is_dropped_too(self): - from litellm.llms.base_llm.auth.token_exchange import endpoint_url_for_error_message + from litellm.llms.base_llm.auth.oauth_endpoint import endpoint_url_for_error_message rendered = endpoint_url_for_error_message("https://user:pw@idp.example:8443/token") diff --git a/tests/test_litellm/llms/base_llm/auth/test_oauth_endpoint.py b/tests/test_litellm/llms/base_llm/auth/test_oauth_endpoint.py new file mode 100644 index 000000000000..ab1f56305379 --- /dev/null +++ b/tests/test_litellm/llms/base_llm/auth/test_oauth_endpoint.py @@ -0,0 +1,277 @@ +import base64 +import json +from typing import Final +from urllib.parse import quote, urlencode + +import pytest +from pydantic import SecretStr + +from litellm.llms.base_llm.auth.oauth_endpoint import ( + MAX_RESPONSE_BYTES, + drop_reflected_credential, + redact_oauth_error_body, + validate_token_endpoint_url, +) +from litellm.llms.base_llm.auth.types import InsecureTokenUrl, TokenEndpointError + +REFLECTED_MESSAGE: Final = "" + + +class TestRedactOauthErrorBody: + def test_rfc6749_fields_are_kept_and_each_capped(self): + body = { + "error": "invalid_grant", + "error_description": "d" * 300, + "error_uri": "https://errors.example/e1", + } + result = redact_oauth_error_body(400, json.dumps(body)) + + assert result == TokenEndpointError(status_code=400, redacted_body=result.redacted_body) + assert "invalid_grant" in result.redacted_body + assert "d" * 256 in result.redacted_body + assert "d" * 257 not in result.redacted_body + assert "https://errors.example/e1" in result.redacted_body + + def test_non_oauth_fields_are_never_rendered(self): + body = {"detail": "internal id 4711", "error": "server_error"} + result = redact_oauth_error_body(500, json.dumps(body)) + + assert result.redacted_body == "error: server_error" + + def test_an_object_without_oauth_fields_gets_a_constant_message(self): + result = redact_oauth_error_body(400, json.dumps({"detail": "internal id 4711"})) + assert result.redacted_body == "error response carried no RFC 6749 fields" + + def test_nested_error_envelope_renders_readable_text(self): + body = { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": "federation_rule_id is not a well-formed fdrl_ tagged ID", + }, + } + result = redact_oauth_error_body(400, json.dumps(body)) + + assert "invalid_request_error" in result.redacted_body + assert "federation_rule_id is not a well-formed fdrl_ tagged ID" in result.redacted_body + assert "{'" not in result.redacted_body + + def test_flat_rfc6749_shape_still_renders(self): + body = {"error": "invalid_grant", "error_description": "bad request"} + result = redact_oauth_error_body(400, json.dumps(body)) + + assert result.redacted_body == "error: invalid_grant; error_description: bad request" + + def test_nested_error_message_is_capped_at_256_chars(self): + body = {"error": {"type": "invalid_request_error", "message": "m" * 500}} + result = redact_oauth_error_body(400, json.dumps(body)) + + assert "m" * 256 in result.redacted_body + assert "m" * 257 not in result.redacted_body + + def test_json_string_body_is_not_echoed(self): + """A free-text body can carry back whatever was sent, so only structured OAuth fields are + ever rendered into an error an operator or caller will see.""" + result = redact_oauth_error_body(400, json.dumps("s" * 500)) + assert result.redacted_body == "non-object error response omitted" + assert "s" * 32 not in result.redacted_body + + def test_json_array_body_constant_message(self): + result = redact_oauth_error_body(400, json.dumps(["a", "b"])) + assert result.redacted_body == "non-object error response omitted" + + def test_plain_text_body_is_not_echoed(self): + result = redact_oauth_error_body(502, "t" * 500) + assert result.status_code == 502 + assert result.redacted_body == "non-JSON error response omitted" + assert "t" * 32 not in result.redacted_body + + def test_oversized_body_is_never_parsed(self): + body = '{"error": "' + "x" * MAX_RESPONSE_BYTES + '"}' + result = redact_oauth_error_body(400, body) + + assert result.redacted_body == "oversized error response omitted" + + def test_sentinel_messages_pass_through_unchanged(self): + result = redact_oauth_error_body(400, "oversized error response omitted") + assert result.redacted_body == "oversized error response omitted" + + def test_reflected_assertion_is_dropped(self): + """An endpoint that echoes the submitted assertion must not put it in the log or the error.""" + assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9.REFLECTEDPAYLOAD.signature") + body = {"error": "invalid_grant", "error_description": f"bad assertion {assertion.get_secret_value()}"} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert result.redacted_body == REFLECTED_MESSAGE + assert "REFLECTEDPAYLOAD" not in result.redacted_body + + def test_assertion_reflected_from_an_offset_is_dropped(self): + """Regression: the probe only looked at the assertion's first 24 characters, so an + endpoint echoing it from any later offset shared no prefix and slipped through.""" + assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "A" * 40 + "PAYLOADMIDDLE" + "B" * 40 + ".signature") + tail = assertion.get_secret_value()[24:] + body = {"error": "invalid_grant", "error_description": tail} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert "PAYLOADMIDDLE" not in result.redacted_body + assert tail[:40] not in result.redacted_body + + def test_a_secret_carrying_spaces_is_dropped_when_echoed_whole(self): + """Regression on the redactor itself: comparing a compacted response against an + uncompacted secret stopped matching hand-set passphrases, which are exactly the secrets + most likely to be echoed and the ones an earlier contiguous match had caught.""" + assertion = SecretStr("correct horse battery staple, 42!") + echoed = assertion.get_secret_value() + body = {"error": "invalid_client", "error_description": f"secret {echoed} rejected"} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert echoed not in result.redacted_body + + def test_a_percent_encoded_secret_is_dropped(self): + """A form-encoded grant puts the secret on the wire percent-escaped, so an echo of that + shape has to be recognised without every caller enumerating it.""" + assertion = SecretStr("sUp3r+S3cret/Value=123") + echoed = quote(assertion.get_secret_value(), safe="") + body = {"error": "invalid_client", "error_description": f"rejected {echoed}"} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert echoed not in result.redacted_body + + def test_a_space_encoded_as_plus_is_dropped(self): + """A form-encoded body writes a space as "+", not %20, so percent-decoding alone does not + recover the secret and a passphrase echoed in its wire shape would travel on.""" + assertion = SecretStr("correct horse battery staple") + echoed = urlencode({"client_secret": assertion.get_secret_value()}).split("=", 1)[1] + body = {"error": "invalid_client", "error_description": f"rejected {echoed}"} + + assert "+" in echoed + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert echoed not in result.redacted_body + + def test_several_wire_forms_are_all_compared(self): + """The caller declares each shape it sent, since an encoding the redactor cannot reverse + (base64 of id:secret) is only knowable there.""" + raw = SecretStr("sUp3rS3cretValue123") + blob = SecretStr(base64.b64encode(b"litellm:sUp3rS3cretValue123").decode()) + body = {"error": "invalid_client", "error_description": f"bad {blob.get_secret_value()}"} + + result = redact_oauth_error_body(400, json.dumps(body), (raw, blob)) + + assert blob.get_secret_value() not in result.redacted_body + + def test_a_fragment_shorter_than_a_long_run_is_dropped(self): + """A slice too short to share a long contiguous run with the assertion is still assertion + material, and repeated errors would hand it over piece by piece.""" + assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "A" * 60 + ".sigsigsig") + fragment = assertion.get_secret_value()[30:48] + body = {"error": "invalid_grant", "error_description": f"rejected near {fragment}"} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert fragment not in result.redacted_body + + def test_a_fragment_broken_up_by_delimiters_is_dropped(self): + """Splitting the echo defeats a contiguous match, so the comparison ignores whatever the + endpoint put between the pieces.""" + assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "A" * 60 + ".sigsigsig") + piece = assertion.get_secret_value()[20:44] + spaced = " ".join(piece[i : i + 6] for i in range(0, 24, 6)) + body = {"error": "invalid_grant", "error_description": spaced} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert spaced not in result.redacted_body + + def test_a_short_secret_is_still_matched_whole(self): + """A Keycloak client secret can be shorter than the probe length; the whole value is + compared in that case rather than a truncated prefix.""" + secret = SecretStr("short-secret") + body = {"error": "invalid_client", "error_description": "rejected short-secret"} + + result = redact_oauth_error_body(400, json.dumps(body), secret) + + assert "short-secret" not in result.redacted_body + + def test_a_short_secret_echoed_in_its_wire_shape_is_dropped(self): + """Regression: the run scan only ever compared eight-character windows, so a secret + with fewer credential characters than that could never match once it came back + percent-encoded rather than verbatim, and the whole-value check needs the raw form.""" + secret = SecretStr("p@ss w0rd!") + echoed = quote(secret.get_secret_value(), safe="") + body = {"error": "invalid_client", "error_description": f"rejected {echoed}"} + + assert secret.get_secret_value() not in echoed + result = redact_oauth_error_body(400, json.dumps(body), secret) + + assert echoed not in result.redacted_body + + def test_an_unrelated_body_is_not_falsely_redacted(self): + """The scan must not fire on a body that merely shares short runs with the assertion.""" + assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "Z" * 60 + ".signature") + body = {"error": "invalid_grant", "error_description": "the federation rule was not found"} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert "the federation rule was not found" in result.redacted_body + + +class TestDropReflectedCredential: + def test_no_credential_passes_the_text_through(self): + assert drop_reflected_credential("token_type was mac", None) == "token_type was mac" + + def test_an_empty_credential_never_matches(self): + assert drop_reflected_credential("anything at all", SecretStr("")) == "anything at all" + + def test_a_credential_made_only_of_punctuation_never_matches(self): + assert drop_reflected_credential("??? ###", SecretStr("!!!")) == "??? ###" + + def test_a_verbatim_echo_is_replaced_by_the_sentinel(self): + secret = SecretStr("eyJhbGciOiJSUzI1NiJ9.PAYLOAD.signature") + rendered = f"could not parse {secret.get_secret_value()}" + + assert drop_reflected_credential(rendered, secret) == REFLECTED_MESSAGE + + def test_a_percent_encoded_echo_is_replaced_by_the_sentinel(self): + secret = SecretStr("sUp3r+S3cret/Value=123") + rendered = f"rejected {quote(secret.get_secret_value(), safe='')}" + + assert drop_reflected_credential(rendered, secret) == REFLECTED_MESSAGE + + def test_unrelated_text_is_returned_unchanged(self): + secret = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "Z" * 60 + ".signature") + + assert drop_reflected_credential("token_type was mac", secret) == "token_type was mac" + + +class TestValidateTokenEndpointUrl: + def test_https_is_returned_unchanged(self): + assert validate_token_endpoint_url("https://token.example/v1") == "https://token.example/v1" + + def test_plain_http_is_rejected_naming_only_the_host(self): + result = validate_token_endpoint_url("http://token.example/v1/oauth/token?client_secret=x") + + assert result == InsecureTokenUrl(host="token.example") + assert "/v1/oauth/token" not in str(result) + assert "client_secret" not in str(result) + + def test_a_url_without_a_scheme_is_rejected(self): + assert validate_token_endpoint_url("token.example/v1") == InsecureTokenUrl(host="") + + @pytest.mark.parametrize( + "url", + [ + "http://localhost:8080/v1/oauth/token", + "http://127.0.0.1/v1/oauth/token", + "http://[::1]/v1/oauth/token", + ], + ) + def test_localhost_http_is_allowed(self, url: str): + assert validate_token_endpoint_url(url) == url + + def test_a_localhost_lookalike_is_still_rejected(self): + assert validate_token_endpoint_url("http://localhost.example/v1") == InsecureTokenUrl(host="localhost.example") diff --git a/tests/test_litellm/llms/base_llm/auth/test_token_exchange.py b/tests/test_litellm/llms/base_llm/auth/test_token_exchange.py deleted file mode 100644 index 129eec88c2bb..000000000000 --- a/tests/test_litellm/llms/base_llm/auth/test_token_exchange.py +++ /dev/null @@ -1,1752 +0,0 @@ -import asyncio -import concurrent.futures -import base64 -import json -from urllib.parse import quote, urlencode -import logging -import threading -import time -from collections.abc import Callable, Mapping -from types import MappingProxyType -from typing import Final -from urllib.parse import parse_qsl - -import httpx -import pytest -from pydantic import SecretStr - -from litellm.llms.base_llm.auth.token_exchange import ( - _METRICS_QUEUE_LIMIT, - _REDACTION_CAP, - ADVISORY_REFRESH_BACKOFF_SECONDS, - CALL_TYPE_CACHE_HIT, - FALLBACK_TOKEN_TTL_SECONDS, - MAX_ASSERTION_BYTES, - MAX_RESPONSE_BYTES, - JwtBearerTokenExchangeEngine, - ServiceLoggingMetricsSink, - TokenExchangeEndpointFailure, - TokenExchangeTransportFailure, - _default_assertion_reader, - _error_summary, - _HttpxSyncTokenPoster, - _new_exchange_handler, - redact_oauth_error_body, -) -from litellm.llms.base_llm.auth.types import ( - AssertionSource, - AssertionSourceError, - BodyEncoding, - ExchangeError, - ExchangeResult, - InsecureTokenUrl, - MalformedTokenResponse, - MintedToken, - TokenEndpointError, - TokenExchangeSpec, - TokenTransportError, -) -from litellm.secret_managers.main import OidcPathNotAllowedError, _resolve_oidc_file_path -from litellm.types.services import ServiceTypes - -DEFAULT_REF: Final = "oidc/env/TEST_ASSERTION" -DEFAULT_ASSERTION: Final = "test-jwt-assertion" -EXCHANGE_URL: Final = "https://token.example/v1/oauth/token" - - -class FakeClock: - def __init__(self, start: float = 1_000.0) -> None: - self.now = start - - def __call__(self) -> float: - return self.now - - def advance(self, seconds: float) -> None: - self.now += seconds - - -class RecordedRequest: - def __init__(self, url: str, content: bytes, headers: Mapping[str, str], timeout: float) -> None: - self.url = url - self.content = content - self.headers = dict(headers) - self.timeout = timeout - - def json_body(self) -> dict: - return json.loads(self.content) - - -class ScriptedPoster: - """Returns scripted responses in order (repeating the last one); records requests.""" - - def __init__( - self, - responses: list[httpx.Response], - on_request: Callable[[RecordedRequest], None] | None = None, - ) -> None: - self.requests: list[RecordedRequest] = [] - self._responses = list(responses) - self._on_request = on_request - - def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: - recorded = RecordedRequest(url, content, headers, timeout) - self.requests.append(recorded) - if self._on_request is not None: - self._on_request(recorded) - if len(self._responses) > 1: - return self._responses.pop(0) - return self._responses[0] - - -class RaisingPoster: - def __init__(self, error: Exception) -> None: - self.calls = 0 - self._error = error - - def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: - self.calls += 1 - raise self._error - - -class ManualExecutor(concurrent.futures.Executor): - """Records submissions; runs them only when the test says so.""" - - def __init__(self) -> None: - self.pending: list[Callable[[], None]] = [] - - def submit(self, fn, /, *args, **kwargs): - future: concurrent.futures.Future = concurrent.futures.Future() - self.pending.append(lambda: fn(*args, **kwargs)) - return future - - def run_all(self) -> None: - drained = list(self.pending) - self.pending.clear() - for job in drained: - job() - - -class InlineExecutor(concurrent.futures.Executor): - def submit(self, fn, /, *args, **kwargs): - future: concurrent.futures.Future = concurrent.futures.Future() - future.set_result(fn(*args, **kwargs)) - return future - - -class NeverRunsExecutor(concurrent.futures.Executor): - """Accepts work and never runs it, standing in for a telemetry backend that has stalled, so a - test can show the backlog stops growing instead of consuming memory for as long as traffic lasts.""" - - def __init__(self) -> None: - self.submitted = 0 # mutable-ok: a test spy counting accepted work - - def submit(self, fn, /, *args, **kwargs): - self.submitted += 1 - return concurrent.futures.Future() - - -def token_response(token: str = "sk-ant-oat01-minted", expires_in: int | None = 3600) -> httpx.Response: - body: Final[dict[str, str | int]] = { - "access_token": token, - "token_type": "Bearer", - **({} if expires_in is None else {"expires_in": expires_in}), - } - return httpx.Response(200, json=body) - - -def make_spec( - *, - token_url: str = "https://token.example/v1/oauth/token", - assertion_ref: str = DEFAULT_REF, - assertion_field: str = "assertion", - static_body: Mapping[str, str] = MappingProxyType( - { - "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", - "federation_rule_id": "fdrl_1", - "organization_id": "org-1", - } - ), - body_encoding: BodyEncoding = "json", - request_headers: Mapping[str, str] = MappingProxyType( - {"anthropic-beta": "oauth-2025-04-20,oidc-federation-2026-04-01"} - ), - cache_key_identity: tuple[str, ...] = ("fdrl_1", "org-1", "", ""), - timeout_seconds: float = 2.0, - assertion_source: AssertionSource | None = None, -) -> TokenExchangeSpec: - return TokenExchangeSpec( - token_url=token_url, - assertion_ref=assertion_ref, - assertion_field=assertion_field, - static_body=static_body, - body_encoding=body_encoding, - request_headers=request_headers, - cache_key_identity=cache_key_identity, - timeout_seconds=timeout_seconds, - assertion_source=assertion_source, - ) - - -class RecordingMetricsSink: - def __init__(self) -> None: - self.successes: list[tuple[str, float]] = [] - self.failures: list[tuple[str, float, ExchangeError]] = [] - self.cache_hits = 0 - - def exchange_success(self, *, call_type: str, duration_seconds: float) -> None: - self.successes.append((call_type, duration_seconds)) - - def exchange_failure(self, *, call_type: str, duration_seconds: float, error: ExchangeError) -> None: - self.failures.append((call_type, duration_seconds, error)) - - def cache_hit(self) -> None: - self.cache_hits += 1 - - -def make_engine( - poster, - reader: Mapping[str, str] | Callable[[str], str | None] | None = None, - clock: FakeClock | None = None, - executor: concurrent.futures.Executor | None = None, - max_entries: int = 64, - metrics_sink=None, -) -> JwtBearerTokenExchangeEngine: - resolved_reader = reader if callable(reader) else (reader or {DEFAULT_REF: DEFAULT_ASSERTION}).get - return JwtBearerTokenExchangeEngine( - poster=poster, - assertion_reader=resolved_reader, - clock=clock if clock is not None else FakeClock(), - refresh_executor=executor if executor is not None else ManualExecutor(), - max_entries=max_entries, - metrics_sink=metrics_sink if metrics_sink is not None else RecordingMetricsSink(), - ) - - -def mint(engine: JwtBearerTokenExchangeEngine, spec: TokenExchangeSpec) -> MintedToken: - result = engine.get_token(spec) - assert isinstance(result, MintedToken) - return result - - -class TestFreshMintWireExact: - def test_json_body_and_headers(self): - poster = ScriptedPoster([token_response(expires_in=3600)]) - clock = FakeClock(start=1_000.0) - engine = make_engine(poster, clock=clock) - spec = make_spec() - - result = mint(engine, spec) - - assert result.access_token.get_secret_value() == "sk-ant-oat01-minted" - assert result.expires_at == 1_000.0 + 3600 - assert len(poster.requests) == 1 - request = poster.requests[0] - assert request.url == "https://token.example/v1/oauth/token" - assert request.timeout == 2.0 - assert request.headers == { - "content-type": "application/json", - "anthropic-beta": "oauth-2025-04-20,oidc-federation-2026-04-01", - } - assert request.json_body() == { - "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", - "federation_rule_id": "fdrl_1", - "organization_id": "org-1", - "assertion": DEFAULT_ASSERTION, - } - - def test_form_body_and_content_type(self): - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster) - spec = make_spec(body_encoding="form") - - mint(engine, spec) - - request = poster.requests[0] - assert request.headers["content-type"] == "application/x-www-form-urlencoded" - assert dict(parse_qsl(request.content.decode())) == { - "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", - "federation_rule_id": "fdrl_1", - "organization_id": "org-1", - "assertion": DEFAULT_ASSERTION, - } - - -def test_cache_hit_zero_posts(): - poster = ScriptedPoster([token_response(expires_in=3600)]) - clock = FakeClock() - engine = make_engine(poster, clock=clock) - spec = make_spec() - - first = mint(engine, spec) - clock.advance(100.0) - second = mint(engine, spec) - - assert len(poster.requests) == 1 - assert second.access_token.get_secret_value() == first.access_token.get_secret_value() - - -@pytest.mark.parametrize( - "remaining,expect_advisory_submit,expect_new_token", - [ - (121.0, False, False), - (120.0, True, False), - (119.0, True, False), - (31.0, True, False), - (30.0, False, True), - (29.0, False, True), - ], -) -def test_window_boundaries(remaining: float, expect_advisory_submit: bool, expect_new_token: bool): - poster = ScriptedPoster([token_response("old-token", expires_in=3600), token_response("new-token")]) - clock = FakeClock(start=1_000.0) - executor = ManualExecutor() - engine = make_engine(poster, clock=clock, executor=executor) - spec = make_spec() - - mint(engine, spec) - expires_at = 1_000.0 + 3600 - clock.now = expires_at - remaining - result = mint(engine, spec) - - assert len(executor.pending) == (1 if expect_advisory_submit else 0) - expected_token = "new-token" if expect_new_token else "old-token" - assert result.access_token.get_secret_value() == expected_token - assert len(poster.requests) == (2 if expect_new_token else 1) - - -def test_advisory_serve_stale_single_flight_backoff(caplog: pytest.LogCaptureFixture): - poster = ScriptedPoster( - [ - token_response("stale-token", expires_in=3600), - httpx.Response(500, json={"error": "server_error"}), - httpx.Response(500, json={"error": "server_error"}), - ] - ) - clock = FakeClock(start=1_000.0) - executor = ManualExecutor() - engine = make_engine(poster, clock=clock, executor=executor) - spec = make_spec() - - mint(engine, spec) - clock.now = 1_000.0 + 3600 - 100.0 - - first = mint(engine, spec) - second = mint(engine, spec) - assert first.access_token.get_secret_value() == "stale-token" - assert second.access_token.get_secret_value() == "stale-token" - assert len(executor.pending) == 1 - - with caplog.at_level(logging.WARNING, logger="LiteLLM"): - executor.run_all() - assert len(poster.requests) == 2 - warning_records = [r for r in caplog.records if r.levelno == logging.WARNING] - assert any("Advisory token refresh" in r.getMessage() for r in warning_records) - assert "server_error" in caplog.text - assert DEFAULT_ASSERTION not in caplog.text - assert "stale-token" not in caplog.text - - within_backoff = mint(engine, spec) - assert within_backoff.access_token.get_secret_value() == "stale-token" - assert len(executor.pending) == 0 - - clock.advance(ADVISORY_REFRESH_BACKOFF_SECONDS) - after_backoff = mint(engine, spec) - assert after_backoff.access_token.get_secret_value() == "stale-token" - assert len(executor.pending) == 1 - executor.run_all() - assert len(poster.requests) == 3 - - -class GatedPoster: - """Blocks the leader inside post() until the test releases it.""" - - def __init__(self, response: httpx.Response) -> None: - self.entered = threading.Event() - self.release = threading.Event() - self.calls = 0 - self._calls_lock = threading.Lock() - self._response = response - - def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: - with self._calls_lock: - self.calls += 1 - self.entered.set() - assert self.release.wait(timeout=10) - return self._response - - -def _run_concurrent_get_token( - engine: JwtBearerTokenExchangeEngine, spec: TokenExchangeSpec, poster: GatedPoster, thread_count: int -) -> list[ExchangeResult]: - results: list[ExchangeResult] = [] - results_lock = threading.Lock() - start_barrier = threading.Barrier(thread_count) - - def worker() -> None: - start_barrier.wait() - result = engine.get_token(spec) - with results_lock: - results.append(result) - - threads = [threading.Thread(target=worker, daemon=True) for _ in range(thread_count)] - for thread in threads: - thread.start() - assert poster.entered.wait(timeout=10) - time.sleep(0.3) - poster.release.set() - for thread in threads: - thread.join(timeout=10) - assert not thread.is_alive() - return results - - -def test_mandatory_single_leader(): - poster = GatedPoster(token_response("leader-token")) - engine = make_engine(poster) - spec = make_spec() - - results = _run_concurrent_get_token(engine, spec, poster, thread_count=5) - - assert poster.calls == 1 - assert len(results) == 5 - for result in results: - assert isinstance(result, MintedToken) - assert result.access_token.get_secret_value() == "leader-token" - - -def test_mandatory_failure_is_value(): - poster = GatedPoster(httpx.Response(500, json={"error": "server_error"})) - engine = make_engine(poster) - spec = make_spec() - - results = _run_concurrent_get_token(engine, spec, poster, thread_count=3) - - assert len(results) == 3 - for result in results: - assert isinstance(result, TokenEndpointError) - assert result.status_code == 500 - assert "server_error" in result.redacted_body - - -def test_lock_released_around_io(): - inner_spec = make_spec( - token_url="https://inner.example/v1/oauth/token", - cache_key_identity=("fdrl_inner", "org-1", "", ""), - ) - engine_holder: dict[str, JwtBearerTokenExchangeEngine] = {} - inner_results: list[ExchangeResult] = [] - - class ReentrantPoster: - def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: - if url == "https://token.example/v1/oauth/token": - inner_results.append(engine_holder["engine"].get_token(inner_spec)) - return token_response() - - engine = make_engine(ReentrantPoster()) - engine_holder["engine"] = engine - - outcome: list[ExchangeResult] = [] - thread = threading.Thread(target=lambda: outcome.append(engine.get_token(make_spec())), daemon=True) - thread.start() - thread.join(timeout=10) - - assert not thread.is_alive(), "engine held its lock across poster I/O and deadlocked" - assert len(outcome) == 1 - assert isinstance(outcome[0], MintedToken) - assert len(inner_results) == 1 - assert isinstance(inner_results[0], MintedToken) - - -def test_401_retry_once_with_reread(): - assertions = {DEFAULT_REF: "assertion-v1"} - - def rotate_on_first_request(request: RecordedRequest) -> None: - assertions[DEFAULT_REF] = "assertion-v2" - - poster = ScriptedPoster( - [httpx.Response(401, json={"error": "invalid_grant"}), token_response()], - on_request=rotate_on_first_request, - ) - engine = make_engine(poster, reader=assertions.get) - - result = mint(engine, make_spec()) - - assert result.access_token.get_secret_value() == "sk-ant-oat01-minted" - assert len(poster.requests) == 2 - assert poster.requests[0].json_body()["assertion"] == "assertion-v1" - assert poster.requests[1].json_body()["assertion"] == "assertion-v2" - - -class RotatingAssertionSource: - """A per-call assertion source that mints a fresh value on every read -- the shape - internal_issuer/keycloak identity sources take (a fresh JWT/token minted per call).""" - - def __init__(self, values: list[str]) -> None: - self._values = iter(values) - self.calls = 0 - - def __call__(self) -> str: - self.calls += 1 - return next(self._values) - - -class EchoingUnauthorizedPoster: - """401s every attempt, echoing the submitted assertion back into the error body -- a - token endpoint that reflects the request.""" - - def __init__(self) -> None: - self.requests: list[RecordedRequest] = [] - - def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: - recorded = RecordedRequest(url, content, headers, timeout) - self.requests.append(recorded) - submitted = recorded.json_body()["assertion"] - return httpx.Response(401, json={"error": "invalid_grant", "error_description": f"bad assertion {submitted}"}) - - -def test_401_retry_redacts_the_assertion_actually_sent_not_a_fresh_reread(): - """Regression: with a rotating identity source, the reflection-drop check must match the - assertion the failing (second) attempt actually sent. Re-reading for the check would mint a - THIRD value that was never sent, so the reflection probe would miss and the actually-sent, - actually-reflected second assertion would leak into the error.""" - poster = EchoingUnauthorizedPoster() - source = RotatingAssertionSource(["assertion-v1", "assertion-v2", "assertion-v3"]) - engine = make_engine(poster) - spec = make_spec(assertion_source=source) - - result = engine.get_token(spec) - - assert isinstance(result, TokenEndpointError) - assert len(poster.requests) == 2 - assert poster.requests[0].json_body()["assertion"] == "assertion-v1" - assert poster.requests[1].json_body()["assertion"] == "assertion-v2" - assert source.calls == 2, "the failing attempt's own assertion must be reused, never re-read a third time" - assert "assertion-v1" not in result.redacted_body - assert "assertion-v2" not in result.redacted_body - assert "assertion-v3" not in result.redacted_body - - -def test_401_twice_is_endpoint_error(): - poster = ScriptedPoster([httpx.Response(401, json={"error": "invalid_grant"})]) - engine = make_engine(poster) - - result = engine.get_token(make_spec()) - - assert isinstance(result, TokenEndpointError) - assert result.status_code == 401 - assert "invalid_grant" in result.redacted_body - assert len(poster.requests) == 2 - - -class TestRedactionAndCaps: - def test_object_body_reduced_to_rfc6749_fields(self): - poster = ScriptedPoster( - [ - httpx.Response( - 400, - json={ - "error": "invalid_grant", - "error_description": "d" * 500, - "error_uri": "https://errors.example/e1", - "assertion_echo": "LEAKED-ASSERTION", - }, - ) - ] - ) - result = make_engine(poster).get_token(make_spec()) - - assert isinstance(result, TokenEndpointError) - assert result.status_code == 400 - assert "invalid_grant" in result.redacted_body - assert "d" * 256 in result.redacted_body - assert "d" * 257 not in result.redacted_body - assert "https://errors.example/e1" in result.redacted_body - assert "LEAKED-ASSERTION" not in result.redacted_body - - def test_nested_error_envelope_renders_readable_text(self): - body = { - "type": "error", - "error": { - "type": "invalid_request_error", - "message": "federation_rule_id is not a well-formed fdrl_ tagged ID", - }, - } - result = redact_oauth_error_body(400, json.dumps(body)) - - assert "invalid_request_error" in result.redacted_body - assert "federation_rule_id is not a well-formed fdrl_ tagged ID" in result.redacted_body - assert "{'" not in result.redacted_body - - def test_flat_rfc6749_shape_still_renders(self): - body = {"error": "invalid_grant", "error_description": "bad request"} - result = redact_oauth_error_body(400, json.dumps(body)) - - assert result.redacted_body == "error: invalid_grant; error_description: bad request" - - def test_nested_error_message_is_capped_at_256_chars(self): - body = {"error": {"type": "invalid_request_error", "message": "m" * 500}} - result = redact_oauth_error_body(400, json.dumps(body)) - - assert "m" * 256 in result.redacted_body - assert "m" * 257 not in result.redacted_body - - def test_json_string_body_is_not_echoed(self): - """A free-text body can carry back whatever was sent, so only structured OAuth fields are - ever rendered into an error an operator or caller will see.""" - result = redact_oauth_error_body(400, json.dumps("s" * 500)) - assert result.redacted_body == "non-object error response omitted" - assert "s" * 32 not in result.redacted_body - - def test_plain_text_body_is_not_echoed(self): - result = redact_oauth_error_body(502, "t" * 500) - assert result.redacted_body == "non-JSON error response omitted" - assert "t" * 32 not in result.redacted_body - - def test_reflected_assertion_is_dropped(self): - """An endpoint that echoes the submitted assertion must not put it in the log or the error.""" - assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9.REFLECTEDPAYLOAD.signature") - body = {"error": "invalid_grant", "error_description": f"bad assertion {assertion.get_secret_value()}"} - - result = redact_oauth_error_body(400, json.dumps(body), assertion) - - assert assertion.get_secret_value() not in result.redacted_body - assert "REFLECTEDPAYLOAD" not in result.redacted_body - - def test_assertion_reflected_from_an_offset_is_dropped(self): - """Regression: the probe only looked at the assertion's first 24 characters, so an - endpoint echoing it from any later offset shared no prefix and slipped through.""" - assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "A" * 40 + "PAYLOADMIDDLE" + "B" * 40 + ".signature") - tail = assertion.get_secret_value()[24:] - body = {"error": "invalid_grant", "error_description": tail} - - result = redact_oauth_error_body(400, json.dumps(body), assertion) - - assert "PAYLOADMIDDLE" not in result.redacted_body - assert tail[:40] not in result.redacted_body - - def test_a_secret_carrying_spaces_is_dropped_when_echoed_whole(self): - """Regression on the redactor itself: comparing a compacted response against an - uncompacted secret stopped matching hand-set passphrases, which are exactly the secrets - most likely to be echoed and the ones an earlier contiguous match had caught.""" - assertion = SecretStr("correct horse battery staple, 42!") - echoed = assertion.get_secret_value() - body = {"error": "invalid_client", "error_description": f"secret {echoed} rejected"} - - result = redact_oauth_error_body(400, json.dumps(body), assertion) - - assert echoed not in result.redacted_body - - def test_a_percent_encoded_secret_is_dropped(self): - """A form-encoded grant puts the secret on the wire percent-escaped, so an echo of that - shape has to be recognised without every caller enumerating it.""" - assertion = SecretStr("sUp3r+S3cret/Value=123") - echoed = quote(assertion.get_secret_value(), safe="") - body = {"error": "invalid_client", "error_description": f"rejected {echoed}"} - - result = redact_oauth_error_body(400, json.dumps(body), assertion) - - assert echoed not in result.redacted_body - - def test_a_space_encoded_as_plus_is_dropped(self): - """A form-encoded body writes a space as "+", not %20, so percent-decoding alone does not - recover the secret and a passphrase echoed in its wire shape would travel on.""" - assertion = SecretStr("correct horse battery staple") - echoed = urlencode({"client_secret": assertion.get_secret_value()}).split("=", 1)[1] - body = {"error": "invalid_client", "error_description": f"rejected {echoed}"} - - assert "+" in echoed - result = redact_oauth_error_body(400, json.dumps(body), assertion) - - assert echoed not in result.redacted_body - - def test_several_wire_forms_are_all_compared(self): - """The caller declares each shape it sent, since an encoding the redactor cannot reverse - (base64 of id:secret) is only knowable there.""" - raw = SecretStr("sUp3rS3cretValue123") - blob = SecretStr(base64.b64encode(b"litellm:sUp3rS3cretValue123").decode()) - body = {"error": "invalid_client", "error_description": f"bad {blob.get_secret_value()}"} - - result = redact_oauth_error_body(400, json.dumps(body), (raw, blob)) - - assert blob.get_secret_value() not in result.redacted_body - - def test_a_fragment_shorter_than_a_long_run_is_dropped(self): - """A slice too short to share a long contiguous run with the assertion is still assertion - material, and repeated errors would hand it over piece by piece.""" - assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "A" * 60 + ".sigsigsig") - fragment = assertion.get_secret_value()[30:48] - body = {"error": "invalid_grant", "error_description": f"rejected near {fragment}"} - - result = redact_oauth_error_body(400, json.dumps(body), assertion) - - assert fragment not in result.redacted_body - - def test_a_fragment_broken_up_by_delimiters_is_dropped(self): - """Splitting the echo defeats a contiguous match, so the comparison ignores whatever the - endpoint put between the pieces.""" - assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "A" * 60 + ".sigsigsig") - piece = assertion.get_secret_value()[20:44] - spaced = " ".join(piece[i : i + 6] for i in range(0, 24, 6)) - body = {"error": "invalid_grant", "error_description": spaced} - - result = redact_oauth_error_body(400, json.dumps(body), assertion) - - assert spaced not in result.redacted_body - - def test_a_short_secret_is_still_matched_whole(self): - """A Keycloak client secret can be shorter than the probe length; the whole value is - compared in that case rather than a truncated prefix.""" - secret = SecretStr("short-secret") - body = {"error": "invalid_client", "error_description": "rejected short-secret"} - - result = redact_oauth_error_body(400, json.dumps(body), secret) - - assert "short-secret" not in result.redacted_body - - def test_a_short_secret_echoed_in_its_wire_shape_is_dropped(self): - """Regression: the run scan only ever compared eight-character windows, so a secret - with fewer credential characters than that could never match once it came back - percent-encoded rather than verbatim, and the whole-value check needs the raw form.""" - secret = SecretStr("p@ss w0rd!") - echoed = quote(secret.get_secret_value(), safe="") - body = {"error": "invalid_client", "error_description": f"rejected {echoed}"} - - assert secret.get_secret_value() not in echoed - result = redact_oauth_error_body(400, json.dumps(body), secret) - - assert echoed not in result.redacted_body - - def test_an_unrelated_body_is_not_falsely_redacted(self): - """The scan must not fire on a body that merely shares short runs with the assertion.""" - assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "Z" * 60 + ".signature") - body = {"error": "invalid_grant", "error_description": "the federation rule was not found"} - - result = redact_oauth_error_body(400, json.dumps(body), assertion) - - assert "the federation rule was not found" in result.redacted_body - - def test_json_array_body_constant_message(self): - result = redact_oauth_error_body(400, json.dumps(["a", "b"])) - assert result.redacted_body == "non-object error response omitted" - - def test_oversized_body_never_parsed(self): - poster = ScriptedPoster([httpx.Response(400, content=b'{"error": "' + b"x" * MAX_RESPONSE_BYTES + b'"}')]) - result = make_engine(poster).get_token(make_spec()) - - assert isinstance(result, TokenEndpointError) - assert result.redacted_body == "oversized error response omitted" - - def test_oversized_success_body_is_malformed(self): - poster = ScriptedPoster( - [httpx.Response(200, content=b'{"access_token": "' + b"x" * MAX_RESPONSE_BYTES + b'"}')] - ) - result = make_engine(poster).get_token(make_spec()) - - assert not isinstance(result, MintedToken) - assert b"x" * 10 not in str(result).encode() - - -@pytest.mark.parametrize("access_token", ["", " "]) -def test_empty_access_token_is_malformed(access_token: str): - poster = ScriptedPoster( - [httpx.Response(200, json={"access_token": access_token, "token_type": "Bearer", "expires_in": 3600})] - ) - result = make_engine(poster).get_token(make_spec()) - - assert isinstance(result, MalformedTokenResponse) - assert "empty access_token" in result.detail - - -def test_sentinel_leak_audit(caplog: pytest.LogCaptureFixture): - jwt_sentinel = "JWT-SENTINEL-2c9f1e7ab4" - token_sentinel = "sk-ant-oat01-TOKEN-SENTINEL-90d4c3aa17" - ref = "oidc/env/SENTINEL_ASSERTION" - - with caplog.at_level(logging.DEBUG): - success_poster = ScriptedPoster([token_response(token_sentinel, expires_in=3600)]) - success_clock = FakeClock() - success_executor = ManualExecutor() - engine = make_engine(success_poster, reader={ref: jwt_sentinel}, clock=success_clock, executor=success_executor) - spec = make_spec(assertion_ref=ref) - minted = mint(engine, spec) - - endpoint_error = make_engine( - ScriptedPoster([httpx.Response(400, json={"error": "invalid_grant"})]), reader={ref: jwt_sentinel} - ).get_token(spec) - transport_error = make_engine(RaisingPoster(RuntimeError("boom")), reader={ref: jwt_sentinel}).get_token(spec) - malformed_error = make_engine( - ScriptedPoster([httpx.Response(200, json={"unexpected": "shape"})]), reader={ref: jwt_sentinel} - ).get_token(spec) - oversized_error = make_engine( - ScriptedPoster([token_response()]), reader={ref: jwt_sentinel + "x" * MAX_ASSERTION_BYTES} - ).get_token(spec) - insecure_error = make_engine(ScriptedPoster([token_response()]), reader={ref: jwt_sentinel}).get_token( - make_spec(assertion_ref=ref, token_url="http://token.example/v1/oauth/token") - ) - - success_poster._responses = [httpx.Response(500, json={"error": "server_error"})] - success_clock.now = success_clock.now + 3600 - 100.0 - stale = engine.get_token(spec) - success_executor.run_all() - - audited_values = [ - str(minted), - repr(minted), - str(minted.access_token), - repr(minted.access_token), - str(endpoint_error), - repr(endpoint_error), - str(transport_error), - repr(transport_error), - str(malformed_error), - repr(malformed_error), - str(oversized_error), - repr(oversized_error), - str(insecure_error), - repr(insecure_error), - str(stale), - repr(stale), - caplog.text, - ] - assert isinstance(oversized_error, AssertionSourceError) - assert oversized_error.kind == "oversized" - for value in audited_values: - assert jwt_sentinel not in value - assert token_sentinel not in value - - -class TestAssertionGuards: - @pytest.mark.parametrize( - "assertion_value,expected_kind", - [ - ("x" * (MAX_ASSERTION_BYTES + 1), "oversized"), - (" \n\t ", "empty"), - (None, "missing"), - ], - ) - def test_bad_assertion_values(self, assertion_value: str | None, expected_kind: str): - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster, reader=lambda ref: assertion_value) - - result = engine.get_token(make_spec()) - - assert isinstance(result, AssertionSourceError) - assert result.kind == expected_kind - assert result.source_ref == DEFAULT_REF - assert len(poster.requests) == 0 - - @pytest.mark.parametrize( - "raised,expected_kind", - [ - (OidcPathNotAllowedError("path outside allowed credential directories"), "disallowed_path"), - (ValueError("Environment variable ANTHROPIC_IDENTITY_TOKEN not found"), "unreadable"), - (ImportError("needs PyJWT and cryptography: pip install 'litellm[proxy]'"), "unreadable"), - (OSError("permission denied"), "unreadable"), - ], - ) - def test_raising_reader(self, raised: Exception, expected_kind: str): - poster = ScriptedPoster([token_response()]) - - def reader(ref: str) -> str | None: - raise raised - - result = make_engine(poster, reader=reader).get_token(make_spec()) - - assert isinstance(result, AssertionSourceError) - assert result.kind == expected_kind - assert len(poster.requests) == 0 - - def test_value_error_message_is_captured_as_detail(self): - poster = ScriptedPoster([token_response()]) - - def reader(ref: str) -> str | None: - raise ValueError("Keycloak token endpoint returned invalid_client") - - result = make_engine(poster, reader=reader).get_token(make_spec()) - - assert isinstance(result, AssertionSourceError) - assert result.detail == "Keycloak token endpoint returned invalid_client" - - def test_import_error_message_is_captured_as_detail(self): - poster = ScriptedPoster([token_response()]) - - def reader(ref: str) -> str | None: - raise ImportError("the internal_issuer identity source needs PyJWT and cryptography: pip install 'litellm[proxy]'") - - result = make_engine(poster, reader=reader).get_token(make_spec()) - - assert isinstance(result, AssertionSourceError) - assert result.detail is not None - assert "litellm[proxy]" in result.detail - - @pytest.mark.parametrize( - "raised", - [OidcPathNotAllowedError("path outside allowed credential directories"), OSError("permission denied")], - ) - def test_non_value_error_never_populates_detail(self, raised: Exception): - """Only the ValueError branch carries operator-diagnosable text; every other reader failure - stays detail=None, matching today's file/env behavior byte-for-byte.""" - poster = ScriptedPoster([token_response()]) - - def reader(ref: str) -> str | None: - raise raised - - result = make_engine(poster, reader=reader).get_token(make_spec()) - - assert isinstance(result, AssertionSourceError) - assert result.detail is None - - def test_value_error_detail_is_capped(self): - poster = ScriptedPoster([token_response()]) - overlong_message = "x" * (_REDACTION_CAP + 100) - - def reader(ref: str) -> str | None: - raise ValueError(overlong_message) - - result = make_engine(poster, reader=reader).get_token(make_spec()) - - assert isinstance(result, AssertionSourceError) - assert result.detail == overlong_message[:_REDACTION_CAP] - - -class TestAssertionSourceOverridesEngineReader: - """``TokenExchangeSpec.assertion_source`` is the dispatch mechanism a per-config identity - source (internal_issuer, keycloak) plugs into the shared engine with -- it must win over the - engine-level reader, and failures must still be reported against ``assertion_ref``.""" - - def test_assertion_source_is_used_instead_of_the_reader(self): - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster, reader=lambda ref: "from-engine-reader") - spec = make_spec(assertion_source=lambda: "from-assertion-source") - - result = mint(engine, spec) - - assert result.access_token.get_secret_value() == "sk-ant-oat01-minted" - assert poster.requests[0].json_body()["assertion"] == "from-assertion-source" - - def test_reader_is_never_called_when_assertion_source_is_set(self): - poster = ScriptedPoster([token_response()]) - calls: list[str] = [] - - def reader(ref: str) -> str | None: - calls.append(ref) - return "from-engine-reader" - - engine = make_engine(poster, reader=reader) - spec = make_spec(assertion_source=lambda: "from-assertion-source") - - mint(engine, spec) - - assert calls == [] - - def test_assertion_source_failure_is_reported_against_assertion_ref(self): - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster, reader=lambda ref: "from-engine-reader") - - def raising_source() -> str | None: - raise ValueError("keycloak token endpoint returned invalid_client") - - spec = make_spec(assertion_source=raising_source, assertion_ref="oidc/keycloak/abc123") - - result = engine.get_token(spec) - - assert isinstance(result, AssertionSourceError) - assert result.source_ref == "oidc/keycloak/abc123" - assert result.detail == "keycloak token endpoint returned invalid_client" - assert len(poster.requests) == 0 - - def test_assertion_source_is_re_invoked_on_401_retry(self): - """The retry's second attempt must also prefer ``assertion_source`` for the assertion it - sends, not silently fall back to the engine reader.""" - values = iter(["assertion-v1", "assertion-v2"]) - poster = ScriptedPoster([httpx.Response(401, json={"error": "invalid_grant"}), token_response()]) - engine = make_engine(poster, reader=lambda ref: "from-engine-reader") - spec = make_spec(assertion_source=lambda: next(values)) - - result = mint(engine, spec) - - assert result.access_token.get_secret_value() == "sk-ant-oat01-minted" - assert poster.requests[0].json_body()["assertion"] == "assertion-v1" - assert poster.requests[1].json_body()["assertion"] == "assertion-v2" - - -class TestOidcFilePathAllowlistRaisesTypedError: - """The engine classifies assertion-source failures by exception type (see - TestAssertionGuards.test_raising_reader); that classification only works if the real - oidc/file allowlist actually raises OidcPathNotAllowedError rather than a bare ValueError.""" - - def test_out_of_allowlist_absolute_path(self, monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", raising=False) - - with pytest.raises(OidcPathNotAllowedError): - _resolve_oidc_file_path("/etc/not-a-credential-dir/token") - - def test_relative_path(self, monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", raising=False) - - with pytest.raises(OidcPathNotAllowedError): - _resolve_oidc_file_path("relative/token/path") - - -class TestHttpsEnforcement: - def test_plain_http_rejected_host_only_zero_posts(self): - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster) - - result = engine.get_token(make_spec(token_url="http://token.example/v1/oauth/token")) - - assert result == InsecureTokenUrl(host="token.example") - assert "/v1/oauth/token" not in str(result) - assert len(poster.requests) == 0 - - @pytest.mark.parametrize( - "url", - [ - "http://localhost:8080/v1/oauth/token", - "http://127.0.0.1/v1/oauth/token", - "http://[::1]/v1/oauth/token", - ], - ) - def test_localhost_http_allowed(self, url: str): - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster) - - result = engine.get_token(make_spec(token_url=url)) - - assert isinstance(result, MintedToken) - assert len(poster.requests) == 1 - - -def test_cache_key_semantics(): - poster = ScriptedPoster([token_response()]) - assertions = {DEFAULT_REF: DEFAULT_ASSERTION, "oidc/env/OTHER": "other-assertion"} - engine = make_engine(poster, reader=assertions.get) - base_spec = make_spec() - - mint(engine, base_spec) - mint(engine, make_spec(cache_key_identity=("fdrl_1", "org-1", "svc-2", ""))) - mint(engine, make_spec(token_url="https://other.example/v1/oauth/token")) - mint(engine, make_spec(assertion_ref="oidc/env/OTHER")) - assert len(poster.requests) == 4 - - assertions[DEFAULT_REF] = "rotated-assertion" - cached = mint(engine, base_spec) - assert len(poster.requests) == 4 - assert cached.access_token.get_secret_value() == "sk-ant-oat01-minted" - - -def test_the_cache_returns_to_its_bound_after_an_all_in_flight_burst(): - """An entry a leader owns is never evictable, so a burst of distinct identities can push the map - past max_entries. It must come back down once those entries are idle, rather than holding the - high-water mark for the life of the process.""" - clock = FakeClock() - engine = make_engine(ScriptedPoster([token_response(expires_in=3600)]), clock=clock, max_entries=4) - - def spec_for(index: int) -> TokenExchangeSpec: - return make_spec(cache_key_identity=("fdrl_1", f"org-{index}", "", "")) - - for index in range(12): - mint(engine, spec_for(index)) - - assert len(engine._entries) <= 4, ( # noqa: SLF001 # the bound under test is internal state - f"the cap is enforced once entries are idle, saw {len(engine._entries)}" - ) - - -def test_bounded_eviction(): - clock = FakeClock() - - class PerCallPoster: - def __init__(self) -> None: - self.calls = 0 - - def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: - self.calls += 1 - body = json.loads(content) - expires_in = 3600 + int(body["organization_id"].split("-")[1]) - return token_response(f"token-{body['organization_id']}", expires_in=expires_in) - - poster = PerCallPoster() - engine = make_engine(poster, clock=clock, max_entries=64) - - def spec_for(index: int) -> TokenExchangeSpec: - return make_spec( - static_body={ - "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", - "federation_rule_id": "fdrl_1", - "organization_id": f"org-{index}", - }, - cache_key_identity=("fdrl_1", f"org-{index}", "", ""), - ) - - for index in range(65): - mint(engine, spec_for(index)) - assert poster.calls == 65 - - mint(engine, spec_for(0)) - assert poster.calls == 66, "the earliest-expiring entry (index 0) should have been evicted" - - mint(engine, spec_for(2)) - assert poster.calls == 66, "a later-expiring entry should still be cached" - - mint(engine, spec_for(1)) - assert poster.calls == 67, "re-inserting index 0 should have evicted the next earliest-expiring entry" - - -@pytest.mark.parametrize("expires_in", [None, 0, -5]) -def test_missing_or_nonsense_expires_in_gets_fallback_ttl(expires_in: int | None): - poster = ScriptedPoster( - [token_response("short-lived", expires_in=expires_in), token_response("reminted", expires_in=3600)] - ) - clock = FakeClock(start=1_000.0) - engine = make_engine(poster, clock=clock) - spec = make_spec() - - first = mint(engine, spec) - assert first.expires_at == 1_000.0 + FALLBACK_TOKEN_TTL_SECONDS - - clock.advance(FALLBACK_TOKEN_TTL_SECONDS + 1.0) - second = mint(engine, spec) - - assert second.access_token.get_secret_value() == "reminted" - assert len(poster.requests) == 2, "a token without a sane expires_in must never be cached forever" - - -async def test_aget_token_loop_responsive(): - class SleepingPoster: - def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: - time.sleep(0.3) - return token_response() - - engine = make_engine(SleepingPoster()) - spec = make_spec() - ticks = {"count": 0} - stop = asyncio.Event() - - async def ticker() -> None: - while not stop.is_set(): - ticks["count"] += 1 - await asyncio.sleep(0.01) - - ticker_task = asyncio.create_task(ticker()) - result = await engine.aget_token(spec) - stop.set() - await ticker_task - - assert isinstance(result, MintedToken) - assert result.access_token.get_secret_value() == "sk-ant-oat01-minted" - assert ticks["count"] >= 5, "the event loop was blocked during aget_token" - sync_result = engine.get_token(spec) - assert sync_result == result - - -def test_invalidate_forces_refresh(): - poster = ScriptedPoster([token_response("token-1", expires_in=3600), token_response("token-2", expires_in=3600)]) - engine = make_engine(poster) - spec = make_spec() - - first = mint(engine, spec) - assert first.access_token.get_secret_value() == "token-1" - - engine.invalidate(spec) - second = mint(engine, spec) - assert second.access_token.get_secret_value() == "token-2" - assert len(poster.requests) == 2 - - third = mint(engine, spec) - assert third.access_token.get_secret_value() == "token-2" - assert len(poster.requests) == 2, "force_refresh must be one-shot" - - -def test_invalidate_unknown_spec_is_noop(): - poster = ScriptedPoster([token_response()]) - engine = make_engine(poster) - - engine.invalidate(make_spec()) - - assert len(poster.requests) == 0 - - -def test_advisory_failure_wakes_expired_follower_to_re_lead(): - poster = ScriptedPoster( - [ - token_response("initial-token", expires_in=3600), - httpx.Response(500, json={"error": "server_error"}), - token_response("recovered-token", expires_in=3600), - ] - ) - clock = FakeClock(start=1_000.0) - executor = ManualExecutor() - engine = make_engine(poster, clock=clock, executor=executor) - spec = make_spec() - - mint(engine, spec) - clock.now = 1_000.0 + 3600 - 100.0 - mint(engine, spec) - assert len(executor.pending) == 1 - - clock.advance(200.0) - results: list[ExchangeResult] = [] - follower = threading.Thread(target=lambda: results.append(engine.get_token(spec)), daemon=True) - follower.start() - time.sleep(0.3) - executor.run_all() - follower.join(timeout=10) - - assert not follower.is_alive() - assert len(results) == 1 - result = results[0] - assert isinstance(result, MintedToken), f"follower was handed {result!r} instead of re-leading a fresh mint" - assert result.access_token.get_secret_value() == "recovered-token" - assert len(poster.requests) == 3 - - -class TwoAttemptGatedPoster: - """401 on the first attempt, then blocks the leader's retry until released.""" - - def __init__(self) -> None: - self.entered_second = threading.Event() - self.release = threading.Event() - self.calls = 0 - - def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: - self.calls += 1 - if self.calls == 1: - return httpx.Response(401, json={"error": "invalid_grant"}) - self.entered_second.set() - assert self.release.wait(timeout=30) - return token_response("slow-leader-token") - - -def test_follower_budget_outlasts_slow_two_attempt_leader(): - poster = TwoAttemptGatedPoster() - engine = make_engine(poster) - spec = make_spec(timeout_seconds=1.0) - - leader_results: list[ExchangeResult] = [] - leader = threading.Thread(target=lambda: leader_results.append(engine.get_token(spec)), daemon=True) - leader.start() - assert poster.entered_second.wait(timeout=10) - - follower_results: list[ExchangeResult] = [] - follower = threading.Thread(target=lambda: follower_results.append(engine.get_token(spec)), daemon=True) - follower.start() - time.sleep(6.5) - poster.release.set() - leader.join(timeout=10) - follower.join(timeout=10) - - assert leader_results and isinstance(leader_results[0], MintedToken) - assert follower_results, "follower never returned" - follower_result = follower_results[0] - assert isinstance(follower_result, MintedToken), ( - f"follower gave up before the leader's two-attempt worst case: {follower_result!r}" - ) - assert follower_result.access_token.get_secret_value() == "slow-leader-token" - - -class FailThenGatePoster: - """500 on the first call, then blocks until released before succeeding.""" - - def __init__(self) -> None: - self.entered_gate = threading.Event() - self.release = threading.Event() - self.calls = 0 - - def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: - self.calls += 1 - if self.calls == 1: - return httpx.Response(500, json={"error": "server_error"}) - self.entered_gate.set() - assert self.release.wait(timeout=30) - return token_response("round-two-token") - - -def test_new_round_timed_out_follower_never_returns_previous_rounds_error(): - poster = FailThenGatePoster() - clock = FakeClock() - engine = make_engine(poster, clock=clock) - spec = make_spec(timeout_seconds=0.05) - - first = engine.get_token(spec) - assert isinstance(first, TokenEndpointError) - - clock.advance(ADVISORY_REFRESH_BACKOFF_SECONDS + 1.0) - leader = threading.Thread(target=lambda: engine.get_token(spec), daemon=True) - leader.start() - assert poster.entered_gate.wait(timeout=10) - - follower_result = engine.get_token(spec) - - assert isinstance(follower_result, TokenTransportError), ( - f"timed-out follower returned the previous round's error: {follower_result!r}" - ) - assert "timed out" in follower_result.detail - poster.release.set() - leader.join(timeout=10) - - -def test_lead_backoff_fails_fast_within_window_and_expires_after(): - poster = ScriptedPoster([httpx.Response(500, json={"error": "server_error"})]) - clock = FakeClock() - engine = make_engine(poster, clock=clock) - spec = make_spec() - - first = engine.get_token(spec) - assert isinstance(first, TokenEndpointError) - assert len(poster.requests) == 1 - - clock.advance(ADVISORY_REFRESH_BACKOFF_SECONDS - 1.0) - second = engine.get_token(spec) - assert second == first - assert len(poster.requests) == 1, "a request inside the backoff window must make zero POSTs" - - clock.advance(1.0) - third = engine.get_token(spec) - assert isinstance(third, TokenEndpointError) - assert len(poster.requests) == 2 - - -def test_invalidate_bypasses_lead_backoff(): - poster = ScriptedPoster( - [httpx.Response(500, json={"error": "server_error"}), token_response("post-invalidate", expires_in=3600)] - ) - clock = FakeClock() - engine = make_engine(poster, clock=clock) - spec = make_spec() - - first = engine.get_token(spec) - assert isinstance(first, TokenEndpointError) - - engine.invalidate(spec) - second = engine.get_token(spec) - - assert isinstance(second, MintedToken) - assert second.access_token.get_secret_value() == "post-invalidate" - assert len(poster.requests) == 2 - - -class StubExchangeHandler: - """Stands in for the HTTPHandler the default poster builds, so the poster's own contract is - testable without a socket.""" - - def __init__(self, result: httpx.Response | Exception | None) -> None: - self.calls = 0 - self._result = result - - def post(self, url: str, *, content: bytes, headers: dict[str, str], timeout: float) -> httpx.Response | None: - self.calls += 1 - if isinstance(self._result, Exception): - raise self._result - return self._result - - -class TestDefaultTokenPoster: - def test_builds_its_handler_once_and_reuses_it(self): - built: list[StubExchangeHandler] = [] - - def factory() -> StubExchangeHandler: - handler = StubExchangeHandler(httpx.Response(200, json={"access_token": "t"})) - built.append(handler) - return handler - - poster: Final = _HttpxSyncTokenPoster(handler_factory=factory) # pyright: ignore[reportArgumentType] # StubExchangeHandler stands in for the legacy-untyped HTTPHandler - for _ in range(3): - poster.post(EXCHANGE_URL, content=b"", headers={}, timeout=1.0) - - assert len(built) == 1 - assert built[0].calls == 3 - - def test_the_real_handler_refuses_to_follow_redirects(self): - assert _new_exchange_handler().client.follow_redirects is False, ( - "a redirected exchange POST would replay the workload assertion to the redirect target" - ) - - def test_an_http_status_error_becomes_its_response(self): - response: Final = httpx.Response( - 401, json={"error": "invalid_grant"}, request=httpx.Request("POST", EXCHANGE_URL) - ) - poster: Final = _HttpxSyncTokenPoster( - handler_factory=lambda: StubExchangeHandler( # pyright: ignore[reportArgumentType] # StubExchangeHandler stands in for the legacy-untyped HTTPHandler - httpx.HTTPStatusError("boom", request=response.request, response=response) - ) - ) - - assert poster.post(EXCHANGE_URL, content=b"", headers={}, timeout=1.0).status_code == 401 - - def test_a_missing_response_is_a_transport_error(self): - poster: Final = _HttpxSyncTokenPoster(handler_factory=lambda: StubExchangeHandler(None)) # pyright: ignore[reportArgumentType] # StubExchangeHandler stands in for the legacy-untyped HTTPHandler - - with pytest.raises(httpx.TransportError): - poster.post(EXCHANGE_URL, content=b"", headers={}, timeout=1.0) - - -class TestDefaultAssertionReader: - def test_reads_through_litellm_secret_resolution(self, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("WIF_ASSERTION_FOR_DEFAULT_READER", "header.payload.signature") - - assert _default_assertion_reader("os.environ/WIF_ASSERTION_FOR_DEFAULT_READER") == "header.payload.signature" - - def test_an_unset_reference_reads_as_none(self): - assert _default_assertion_reader("os.environ/DEFINITELY_NOT_SET_WIF_ASSERTION_REF") is None - - -class TestErrorSummary: - def test_every_error_variant_summarises_without_carrying_a_secret(self): - summaries: Final = { - _error_summary(AssertionSourceError(kind="unreadable", source_ref="oidc/file/x")), - _error_summary(InsecureTokenUrl(host="token.internal")), - _error_summary(TokenEndpointError(status_code=401, redacted_body="invalid_grant")), - _error_summary(TokenTransportError(detail="ConnectError: refused")), - _error_summary(MalformedTokenResponse(detail="empty access_token")), - } - - assert {s.split(":")[0] for s in summaries} == { - "AssertionSourceError", - "InsecureTokenUrl", - "TokenEndpointError", - "TokenTransportError", - "MalformedTokenResponse", - }, "each variant names itself so a log line says which stage failed" - - -class TestNonBearerTokenType: - def test_a_non_bearer_token_type_is_refused(self): - poster: Final = ScriptedPoster( - [httpx.Response(200, json={"access_token": "tok", "token_type": "mac", "expires_in": 300})] - ) - engine: Final = JwtBearerTokenExchangeEngine(poster=poster, assertion_reader=lambda _ref: DEFAULT_ASSERTION) - - result: Final = engine.get_token(make_spec()) - - assert isinstance(result, MalformedTokenResponse) - assert "non-bearer" in result.detail - - -class TestShortLivedRefreshWindows: - """A token whose lifetime is at or below the flat 120s advisory window used to be inside that - window from birth, so every request armed another background exchange. The windows now scale - with the observed lifetime; long-lived tokens must keep the flat 120s/30s behaviour.""" - - @staticmethod - def _engine_with( - expires_in: int | None, - ) -> tuple[JwtBearerTokenExchangeEngine, ScriptedPoster, FakeClock, ManualExecutor, TokenExchangeSpec]: - poster = ScriptedPoster([token_response("short-lived", expires_in=expires_in), token_response("reminted")]) - clock = FakeClock(start=1_000.0) - executor = ManualExecutor() - engine = make_engine(poster, clock=clock, executor=executor) - return engine, poster, clock, executor, make_spec() - - def test_fallback_ttl_token_is_served_without_arming_a_refresh(self): - engine, poster, clock, executor, spec = self._engine_with(expires_in=None) - - first = mint(engine, spec) - assert first.expires_at == 1_000.0 + FALLBACK_TOKEN_TTL_SECONDS - - for _ in range(5): - clock.advance(1.0) - assert mint(engine, spec).access_token.get_secret_value() == "short-lived" - - assert executor.pending == [], "a freshly minted fallback-TTL token must not arm a refresh on every request" - assert len(poster.requests) == 1 - - @pytest.mark.parametrize( - "elapsed,expect_advisory_submit", - [(29.0, False), (30.0, True), (52.0, True)], - ) - def test_fallback_ttl_token_refreshes_around_its_half_life(self, elapsed: float, expect_advisory_submit: bool): - engine, poster, clock, executor, spec = self._engine_with(expires_in=None) - - mint(engine, spec) - clock.advance(elapsed) - served = mint(engine, spec) - - assert served.access_token.get_secret_value() == "short-lived" - assert len(executor.pending) == (1 if expect_advisory_submit else 0) - executor.run_all() - assert len(poster.requests) == (2 if expect_advisory_submit else 1) - - @pytest.mark.parametrize( - "elapsed,expect_new_token", - [(52.0, False), (53.0, True)], - ) - def test_fallback_ttl_mandatory_wall_scales_with_the_lifetime(self, elapsed: float, expect_new_token: bool): - engine, poster, clock, executor, spec = self._engine_with(expires_in=None) - - mint(engine, spec) - clock.advance(elapsed) - served = mint(engine, spec) - - assert served.access_token.get_secret_value() == ("reminted" if expect_new_token else "short-lived") - assert len(executor.pending) == (0 if expect_new_token else 1) - - @pytest.mark.parametrize( - "elapsed,expect_advisory_submit", - [(89.0, False), (100.0, True)], - ) - def test_a_200s_token_scales_its_advisory_window_too(self, elapsed: float, expect_advisory_submit: bool): - engine, poster, clock, executor, spec = self._engine_with(expires_in=200) - - mint(engine, spec) - clock.advance(elapsed) - served = mint(engine, spec) - - assert served.access_token.get_secret_value() == "short-lived" - assert len(executor.pending) == (1 if expect_advisory_submit else 0) - - @pytest.mark.parametrize("expires_in", [240, 3600]) - @pytest.mark.parametrize( - "remaining,expect_advisory_submit,expect_new_token", - [ - (121.0, False, False), - (120.0, True, False), - (31.0, True, False), - (30.0, False, True), - ], - ) - def test_long_lived_tokens_keep_the_flat_windows( - self, expires_in: int, remaining: float, expect_advisory_submit: bool, expect_new_token: bool - ): - engine, poster, clock, executor, spec = self._engine_with(expires_in=expires_in) - - mint(engine, spec) - clock.now = 1_000.0 + expires_in - remaining - served = mint(engine, spec) - - assert len(executor.pending) == (1 if expect_advisory_submit else 0) - assert served.access_token.get_secret_value() == ("reminted" if expect_new_token else "short-lived") - assert len(poster.requests) == (2 if expect_new_token else 1) - - -class RaisingMetricsSink: - def exchange_success(self, *, call_type: str, duration_seconds: float) -> None: - raise RuntimeError("metrics sink down") - - def exchange_failure(self, *, call_type: str, duration_seconds: float, error: ExchangeError) -> None: - raise RuntimeError("metrics sink down") - - def cache_hit(self) -> None: - raise RuntimeError("metrics sink down") - - -class TestMetricsEmission: - def test_cold_mint_emits_success_with_duration(self): - clock = FakeClock() - sink = RecordingMetricsSink() - poster = ScriptedPoster([token_response()], on_request=lambda _request: clock.advance(0.25)) - engine = make_engine(poster, clock=clock, metrics_sink=sink) - - mint(engine, make_spec()) - - assert sink.successes == [("cold_mint", 0.25)] - assert sink.failures == [] - assert sink.cache_hits == 0 - - def test_cache_hit_emits_counter_not_a_mint(self): - clock = FakeClock() - sink = RecordingMetricsSink() - engine = make_engine(ScriptedPoster([token_response()]), clock=clock, metrics_sink=sink) - spec = make_spec() - - mint(engine, spec) - clock.advance(100.0) - mint(engine, spec) - - assert sink.cache_hits == 1 - assert len(sink.successes) == 1 - - def test_advisory_refresh_call_type(self): - clock = FakeClock(start=1_000.0) - sink = RecordingMetricsSink() - executor = ManualExecutor() - poster = ScriptedPoster([token_response("old", expires_in=3600), token_response("new")]) - engine = make_engine(poster, clock=clock, executor=executor, metrics_sink=sink) - spec = make_spec() - - mint(engine, spec) - clock.now = 1_000.0 + 3600 - 119.0 - mint(engine, spec) - executor.run_all() - - assert [call_type for call_type, _ in sink.successes] == ["cold_mint", "advisory_refresh"] - assert sink.cache_hits == 1 - - def test_mandatory_refresh_call_type(self): - clock = FakeClock(start=1_000.0) - sink = RecordingMetricsSink() - poster = ScriptedPoster([token_response("old", expires_in=3600), token_response("new")]) - engine = make_engine(poster, clock=clock, metrics_sink=sink) - spec = make_spec() - - mint(engine, spec) - clock.now = 1_000.0 + 3600 - 29.0 - mint(engine, spec) - - assert [call_type for call_type, _ in sink.successes] == ["cold_mint", "mandatory_refresh"] - assert sink.cache_hits == 0 - - def test_failed_exchange_emits_failure_once_and_negative_cache_does_not_reemit(self): - sink = RecordingMetricsSink() - poster = ScriptedPoster([httpx.Response(503, json={"error": "unavailable"})]) - engine = make_engine(poster, metrics_sink=sink) - spec = make_spec() - - first = engine.get_token(spec) - second = engine.get_token(spec) - - assert isinstance(first, TokenEndpointError) - assert isinstance(second, TokenEndpointError) - assert len(sink.failures) == 1 - call_type, _duration, error = sink.failures[0] - assert call_type == "cold_mint" - assert isinstance(error, TokenEndpointError) - assert error.status_code == 503 - assert sink.successes == [] - - def test_failure_payload_carries_no_assertion_material(self): - sink = RecordingMetricsSink() - engine = make_engine(EchoingUnauthorizedPoster(), metrics_sink=sink) - - result = engine.get_token(make_spec()) - - assert isinstance(result, TokenEndpointError) - (failure,) = sink.failures - assert DEFAULT_ASSERTION not in repr(failure) - assert DEFAULT_ASSERTION not in _error_summary(failure[2]) - - def test_raising_sink_never_breaks_mint_serve_or_failure(self): - clock = FakeClock() - engine = make_engine(ScriptedPoster([token_response()]), clock=clock, metrics_sink=RaisingMetricsSink()) - spec = make_spec() - - minted = mint(engine, spec) - clock.advance(100.0) - served = mint(engine, spec) - - assert served.access_token.get_secret_value() == minted.access_token.get_secret_value() - - failing = make_engine(RaisingPoster(httpx.ConnectError("boom")), metrics_sink=RaisingMetricsSink()) - result = failing.get_token(make_spec()) - assert isinstance(result, TokenTransportError) - - -class RecordingServiceHooks: - def __init__(self) -> None: - self.successes: list[tuple[ServiceTypes, str, float]] = [] - self.failures: list[tuple[ServiceTypes, float, str | Exception, str]] = [] - - async def async_service_success_hook(self, service: ServiceTypes, call_type: str, duration: float) -> None: - self.successes.append((service, call_type, duration)) - - async def async_service_failure_hook( - self, service: ServiceTypes, duration: float, error: str | Exception, call_type: str - ) -> None: - self.failures.append((service, duration, error, call_type)) - - -class RaisingServiceHooks: - """Every hook raises, and each call is recorded first so a test can prove the sink kept - calling through rather than bailing after the first failure.""" - - def __init__(self) -> None: - self.attempts: list[str] = [] # mutable-ok: a test spy accumulating calls in order - - async def async_service_success_hook(self, service: ServiceTypes, call_type: str, duration: float) -> None: - self.attempts.append(f"success:{call_type}") - raise RuntimeError("hook down") - - async def async_service_failure_hook( - self, service: ServiceTypes, duration: float, error: str | Exception, call_type: str - ) -> None: - self.attempts.append(f"failure:{call_type}") - raise RuntimeError("hook down") - - -class TestServiceLoggingMetricsSink: - def _sink(self, hooks) -> ServiceLoggingMetricsSink: - return ServiceLoggingMetricsSink(service_logging_factory=lambda: hooks, executor=InlineExecutor()) - - def test_success_maps_to_anthropic_wif_service(self): - hooks = RecordingServiceHooks() - - self._sink(hooks).exchange_success(call_type="cold_mint", duration_seconds=0.2) - - assert hooks.successes == [(ServiceTypes.ANTHROPIC_WIF, "cold_mint", 0.2)] - - def test_a_stalled_backend_stops_accepting_work_instead_of_queueing_without_bound(self): - stalled: Final = NeverRunsExecutor() - sink: Final = ServiceLoggingMetricsSink(service_logging_factory=RecordingServiceHooks, executor=stalled) - - for _ in range(_METRICS_QUEUE_LIMIT + 500): - sink.cache_hit() - - assert stalled.submitted == _METRICS_QUEUE_LIMIT, ( - "once the backlog is full further events are dropped, so request volume cannot grow it" - ) - - def test_a_drained_backlog_accepts_work_again(self): - hooks: Final = RecordingServiceHooks() - sink: Final = ServiceLoggingMetricsSink(service_logging_factory=lambda: hooks, executor=InlineExecutor()) - - for _ in range(_METRICS_QUEUE_LIMIT + 10): - sink.cache_hit() - - assert len(hooks.successes) == _METRICS_QUEUE_LIMIT + 10, ( - "an executor that actually runs releases each slot, so nothing is dropped" - ) - - def test_failure_maps_variant_and_redacted_summary(self): - hooks = RecordingServiceHooks() - error = TokenEndpointError(status_code=503, redacted_body="error: unavailable") - - self._sink(hooks).exchange_failure(call_type="mandatory_refresh", duration_seconds=0.1, error=error) - - ((service, duration, emitted, call_type),) = hooks.failures - assert service is ServiceTypes.ANTHROPIC_WIF - assert duration == 0.1 - assert call_type == "mandatory_refresh" - assert isinstance(emitted, TokenExchangeEndpointFailure) - assert str(emitted) == _error_summary(error) - - def test_transport_failure_gets_its_own_error_class(self): - hooks = RecordingServiceHooks() - - self._sink(hooks).exchange_failure( - call_type="advisory_refresh", duration_seconds=0.05, error=TokenTransportError(detail="ConnectError: boom") - ) - - ((_service, _duration, emitted, _call_type),) = hooks.failures - assert isinstance(emitted, TokenExchangeTransportFailure) - - def test_cache_hit_maps_to_cache_service_with_zero_duration(self): - hooks = RecordingServiceHooks() - - self._sink(hooks).cache_hit() - - assert hooks.successes == [(ServiceTypes.ANTHROPIC_WIF_CACHE, CALL_TYPE_CACHE_HIT, 0.0)] - - def test_end_to_end_reflected_assertion_never_reaches_the_hook(self): - hooks = RecordingServiceHooks() - sink = self._sink(hooks) - engine = make_engine(EchoingUnauthorizedPoster(), metrics_sink=sink) - - result = engine.get_token(make_spec()) - - assert isinstance(result, TokenEndpointError) - ((_service, _duration, emitted, call_type),) = hooks.failures - assert call_type == "cold_mint" - assert DEFAULT_ASSERTION not in str(emitted) - assert DEFAULT_ASSERTION not in repr(emitted) - - def test_raising_hooks_are_swallowed(self): - hooks: Final = RaisingServiceHooks() - sink: Final = self._sink(hooks) - - sink.exchange_success(call_type="cold_mint", duration_seconds=0.2) - sink.cache_hit() - sink.exchange_failure(call_type="cold_mint", duration_seconds=0.1, error=TokenTransportError(detail="boom")) - - assert hooks.attempts == ["success:cold_mint", "success:cache_hit", "failure:cold_mint"], ( - "every event is still handed to the hooks, and one raising hook does not stop the next" - ) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 63df501fb3e5..e8c4f88e06db 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, Mock, patch import httpx +import httpx2 import pytest import litellm @@ -2127,18 +2128,18 @@ def test_sync_retrieve_file_content_raises_on_http_error(): } -class _BlockingWifPoster: - """A token-endpoint poster that blocks until released, so the test can prove +class _BlockingWifTokenEndpoint: + """A token endpoint that blocks until released, so the test can prove the exchange ran off the event loop's own thread instead of freezing it.""" def __init__(self): self.release = threading.Event() self.thread_ids = [] - def post(self, url, *, content, headers, timeout): + def __call__(self, request: httpx2.Request) -> httpx2.Response: self.thread_ids.append(threading.get_ident()) self.release.wait(timeout=5) - return httpx.Response( + return httpx2.Response( 200, json={ "access_token": "sk-ant-oat01-llm-http-handler-seam", @@ -2157,23 +2158,23 @@ async def test_async_retrieve_file_content_wif_exchange_does_not_block_event_loo from litellm.llms.anthropic import common_utils as anthropic_common_utils from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token - from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + from litellm.llms.anthropic.wif_exchange import AnthropicWifTokenExchange for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): monkeypatch.delenv(name, raising=False) for name, value in _FILE_CONTENT_WIF_ENV.items(): monkeypatch.setenv(name, value) - poster = _BlockingWifPoster() - engine = JwtBearerTokenExchangeEngine(poster=poster) + token_endpoint = _BlockingWifTokenEndpoint() + exchange = AnthropicWifTokenExchange(http_client=httpx2.Client(transport=httpx2.MockTransport(token_endpoint))) sync_calls = [] def sync_shim(litellm_params, api_base, model): sync_calls.append(model) - return get_anthropic_wif_token(litellm_params, api_base, model, engine) + return get_anthropic_wif_token(litellm_params, api_base, model, exchange) async def async_shim(litellm_params, api_base, model): - return await aget_anthropic_wif_token(litellm_params, api_base, model, engine) + return await aget_anthropic_wif_token(litellm_params, api_base, model, exchange) monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim) monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim) @@ -2204,17 +2205,17 @@ async def ticker(): ) await asyncio.sleep(0.05) # The ticker kept advancing while the token exchange was still blocked on - # poster.release, proving the exchange did not run inline on the event loop. + # token_endpoint.release, proving the exchange did not run inline on the event loop. assert len(ticks) > 0 assert not retrieve_task.done() - poster.release.set() + token_endpoint.release.set() await retrieve_task await ticker_task assert sync_calls == [] - assert poster.thread_ids - assert poster.thread_ids[0] != threading.get_ident() + assert token_endpoint.thread_ids + assert token_endpoint.thread_ids[0] != threading.get_ident() sent_headers = client.get.call_args.kwargs["headers"] assert sent_headers["authorization"] == "Bearer sk-ant-oat01-llm-http-handler-seam" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 3823a14b7c70..9bb1ff7555c7 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -10,6 +10,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx +import httpx2 import pytest from fastapi import HTTPException, Request, Response from fastapi.responses import StreamingResponse @@ -4804,7 +4805,7 @@ async def test_wif_mint_goes_through_async_facade(self, monkeypatch): from litellm.llms.anthropic import common_utils as anthropic_common_utils from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token - from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + from litellm.llms.anthropic.wif_exchange import AnthropicWifTokenExchange from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( anthropic_proxy_route, ) @@ -4817,23 +4818,25 @@ async def test_wif_mint_goes_through_async_facade(self, monkeypatch): minted: Final = "sk-ant-oat01-route-minted" thread_ids: Final = [] - class ThreadRecordingPoster: - def post(self, url, *, content, headers, timeout): + class ThreadRecordingTokenEndpoint: + def __call__(self, request: httpx2.Request) -> httpx2.Response: thread_ids.append(threading.get_ident()) - return httpx.Response( + return httpx2.Response( 200, json={"access_token": minted, "token_type": "Bearer", "expires_in": 3600}, ) - engine = JwtBearerTokenExchangeEngine(poster=ThreadRecordingPoster()) + exchange = AnthropicWifTokenExchange( + http_client=httpx2.Client(transport=httpx2.MockTransport(ThreadRecordingTokenEndpoint())) + ) sync_calls: Final = [] def sync_shim(litellm_params, api_base, model): sync_calls.append(model) - return get_anthropic_wif_token(litellm_params, api_base, model, engine) + return get_anthropic_wif_token(litellm_params, api_base, model, exchange) async def async_shim(litellm_params, api_base, model): - return await aget_anthropic_wif_token(litellm_params, api_base, model, engine) + return await aget_anthropic_wif_token(litellm_params, api_base, model, exchange) monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim) monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim) @@ -4887,7 +4890,7 @@ def _clear_anthropic_env(self, monkeypatch) -> None: def _enable_wif(self, monkeypatch) -> None: from litellm.llms.anthropic import common_utils as anthropic_common_utils from litellm.llms.anthropic.wif import aget_anthropic_wif_token - from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + from litellm.llms.anthropic.wif_exchange import AnthropicWifTokenExchange monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_plan") monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-plan") @@ -4895,17 +4898,19 @@ def _enable_wif(self, monkeypatch) -> None: minted: Final = self._MINTED - class StubPoster: - def post(self, url, *, content, headers, timeout): - return httpx.Response( + class StubTokenEndpoint: + def __call__(self, request: httpx2.Request) -> httpx2.Response: + return httpx2.Response( 200, json={"access_token": minted, "token_type": "Bearer", "expires_in": 3600}, ) - engine: Final = JwtBearerTokenExchangeEngine(poster=StubPoster()) + exchange: Final = AnthropicWifTokenExchange( + http_client=httpx2.Client(transport=httpx2.MockTransport(StubTokenEndpoint())) + ) async def async_shim(litellm_params, api_base, model): - return await aget_anthropic_wif_token(litellm_params, api_base, model, engine) + return await aget_anthropic_wif_token(litellm_params, api_base, model, exchange) monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 225134b4e2bc..bd7a679dd6a2 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22180 + "limit": 22178 }, "LIT002": { "limit": 26745 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16462 + "limit": 16456 }, "LIT011": { - "limit": 5506 + "limit": 5498 }, "LIT012": { "limit": 4486 diff --git a/uv.lock b/uv.lock index 89205cd95278..7bcbecacd76e 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-02T16:58:34.594994Z" +exclude-newer = "2026-09-02T20:48:47.1088Z" exclude-newer-span = "P3D" [manifest] @@ -291,21 +291,20 @@ wheels = [ [[package]] name = "anthropic" -version = "0.84.0" +version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, - { name = "distro" }, { name = "docstring-parser" }, - { name = "httpx" }, + { name = "httpx2" }, { name = "jiter" }, { name = "pydantic" }, { name = "sniffio" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/04/ea/0869d6df9ef83dcf393aeefc12dd81677d091c6ffc86f783e51cf44062f2/anthropic-0.84.0.tar.gz", hash = "sha256:72f5f90e5aebe62dca316cb013629cfa24996b0f5a4593b8c3d712bc03c43c37", size = 539457, upload-time = "2026-02-25T05:22:38.54Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b4/50/463166f02179ab279edb61de1589a6f69cb3838d6a2fb6f2c92a3f8042f1/anthropic-1.3.0.tar.gz", hash = "sha256:6873492a77ede8849a161ab1bc78bc9a1e492a006d0b5bb4c57ac77845df838a", size = 1148177, upload-time = "2026-09-01T17:37:10.392Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/64/ca/218fa25002a332c0aa149ba18ffc0543175998b1f65de63f6d106689a345/anthropic-0.84.0-py3-none-any.whl", hash = "sha256:861c4c50f91ca45f942e091d83b60530ad6d4f98733bfe648065364da05d29e7", size = 455156, upload-time = "2026-02-25T05:22:40.468Z" }, + { url = "https://files.pythonhosted.org/packages/2c/5d/7863a9961d320c23787c7b594956afe4e878f9c0ae2376b11a20e416791d/anthropic-1.3.0-py3-none-any.whl", hash = "sha256:e7e7dbebf9f3c84a23954ab989378af6ae10a4d1804c81e9fea4b5ced695ce75", size = 1296959, upload-time = "2026-09-01T17:37:08.525Z" }, ] [package.optional-dependencies] @@ -3268,6 +3267,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, ] +[[package]] +name = "httpcore2" +version = "2.12.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "h11" }, + { name = "truststore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/be/ad/f4f0e57345f1870f3e8cb624e058d7eca6e5a27d33bcc3311d9b618734cd/httpcore2-2.12.0.tar.gz", hash = "sha256:9293522bba0aa7c4c8e9e3f040c16575bd8868e155a77fa30c7a9085a5eae648", size = 67548, upload-time = "2026-08-18T13:22:08.211Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d2/74/d370e55600d9bcfa0d9794b0166126d49291a3d2b20c268fc98c453a4948/httpcore2-2.12.0-py3-none-any.whl", hash = "sha256:7e04258ce01013d7d615e5b910a3b27fac937d7a95038227e79652b4ba3b4ceb", size = 83074, upload-time = "2026-08-18T13:22:05.854Z" }, +] + [[package]] name = "httplib2" version = "0.32.0" @@ -3304,6 +3316,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc", size = 8960, upload-time = "2025-10-10T21:48:21.158Z" }, ] +[[package]] +name = "httpx2" +version = "2.12.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio", marker = "sys_platform != 'emscripten'" }, + { name = "httpcore2", marker = "sys_platform != 'emscripten'" }, + { name = "httpx2-jsfetch", marker = "python_full_version >= '3.12' and sys_platform == 'emscripten'" }, + { name = "idna" }, + { name = "truststore", marker = "sys_platform != 'emscripten'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7f/f8/579a8b51e42e38ee32647df9f08aa25643ae788e275cc625b199829c4671/httpx2-2.12.0.tar.gz", hash = "sha256:7631fe9887a8a2275f4a2540e053aa670fcc50742864a9ae7c66e609fdcf12cf", size = 100040, upload-time = "2026-08-18T13:22:09.086Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/95/411ba65569158e862368917aaf56597f3e5fa3b91b0502919638465a08f3/httpx2-2.12.0-py3-none-any.whl", hash = "sha256:cc8b6eecb8661c146b8f89a60e97456ee086e91a784ed31ac450c3a9e613dd36", size = 95427, upload-time = "2026-08-18T13:22:06.834Z" }, +] + +[[package]] +name = "httpx2-jsfetch" +version = "1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/c4/0e5636363151a2a1795e0a77617168b9ca438e1748ec05fc9b5687f93d64/httpx2_jsfetch-1.0.tar.gz", hash = "sha256:70a0e3eabfef7cce5ad9c629f7d01ca05e418f586646f4ddf14782e4c1454c60", size = 6872, upload-time = "2026-08-07T00:13:07.492Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/43/832f631d32e4f1211caa2ba368317739fe71f0b8530e4c9d15dc454bac2a/httpx2_jsfetch-1.0-py3-none-any.whl", hash = "sha256:cb916b707601e69a07721aabc8f3f6659be3a6893bc1ff5c6f9e02241df2da32", size = 6382, upload-time = "2026-08-07T00:13:06.567Z" }, +] + [[package]] name = "huey" version = "2.6.0" @@ -3467,11 +3505,11 @@ wheels = [ [[package]] name = "idna" -version = "3.15" +version = "3.19" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/82/77/7b3966d0b9d1d31a36ddf1746926a11dface89a83409bf1483f0237aa758/idna-3.15.tar.gz", hash = "sha256:ca962446ea538f7092a95e057da437618e886f4d349216d2b1e294abfdb65fdc", size = 199245, upload-time = "2026-05-12T22:45:57.011Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5f/f7/abb373e5757eaec4b922b92f97ec8d6d7e057cf06778247604fbc4e7c3f3/idna-3.19.tar.gz", hash = "sha256:5e0811a4383b21dc5838069f801c4fb62113b7447663d2530d2bd6e77b49bf15", size = 215237, upload-time = "2026-08-18T05:14:24.27Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d2/23/408243171aa9aaba178d3e2559159c24c1171a641aa83b67bdd3394ead8e/idna-3.15-py3-none-any.whl", hash = "sha256:048adeaf8c2d788c40fee287673ccaa74c24ffd8dcf09ffa555a2fbb59f10ac8", size = 72340, upload-time = "2026-05-12T22:45:55.733Z" }, + { url = "https://files.pythonhosted.org/packages/57/b0/0e52c878c53f245edd3a11020f20979b3f490f245af532c7cae3027754b5/idna-3.19-py3-none-any.whl", hash = "sha256:815e7be7a7806d54abb586dc943addc79e8b2ee16915059658cbeff4b1b43bf4", size = 68550, upload-time = "2026-08-18T05:14:22.343Z" }, ] [[package]] @@ -4570,6 +4608,7 @@ healthcheck = [ ] proxy-dev = [ { name = "a2a-sdk" }, + { name = "anthropic", extra = ["vertex"] }, { name = "azure-identity" }, { name = "hypercorn" }, { name = "opentelemetry-api" }, @@ -4584,7 +4623,7 @@ proxy-dev = [ requires-dist = [ { name = "a2a-sdk", marker = "extra == 'extra-proxy'", specifier = ">=1.1.0,<2.0" }, { name = "aiohttp", specifier = ">=3.14.2,<4.0" }, - { name = "anthropic", extras = ["vertex"], marker = "extra == 'proxy-runtime'", specifier = ">=0.84.0,<1.0" }, + { name = "anthropic", extras = ["vertex"], marker = "extra == 'proxy-runtime'", specifier = ">=1.3.0,<2" }, { name = "apscheduler", marker = "extra == 'proxy'", specifier = ">=3.11.2,<4.0" }, { name = "audioread", marker = "extra == 'stt-nvidia-riva'", specifier = ">=3.0.1" }, { name = "aurelio-sdk", marker = "python_full_version < '3.14' and extra == 'semantic-router'", specifier = ">=0.0.19,<1.0" }, @@ -4681,7 +4720,7 @@ provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "mo [package.metadata.requires-dev] ci = [ { name = "aiodynamo", specifier = "==24.7" }, - { name = "anthropic", specifier = "==0.84.0" }, + { name = "anthropic", specifier = "==1.3.0" }, { name = "argon2-cffi", specifier = "==25.1.0" }, { name = "assemblyai", specifier = "==0.52.4" }, { name = "beautifulsoup4", specifier = "==4.14.3" }, @@ -4759,6 +4798,7 @@ healthcheck = [ ] proxy-dev = [ { name = "a2a-sdk", specifier = "==1.1.0" }, + { name = "anthropic", extras = ["vertex"], specifier = "==1.3.0" }, { name = "azure-identity", specifier = "==1.25.2" }, { name = "hypercorn", specifier = "==0.17.3" }, { name = "opentelemetry-api", specifier = "==1.28.0" }, @@ -9707,6 +9747,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ce/13/53c2ab6ac27804769314554a062e0651a44db2360be47e21cf0a29d202ee/traceloop_sdk-0.33.12-py3-none-any.whl", hash = "sha256:d47a474afbf4a68ff38a702dbaca7b17d2d4f0b0e14dc2f1560b6bdd3859ac75", size = 25932, upload-time = "2024-11-13T20:29:25.174Z" }, ] +[[package]] +name = "truststore" +version = "0.10.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/53/a3/1585216310e344e8102c22482f6060c7a6ea0322b63e026372e6dcefcfd6/truststore-0.10.4.tar.gz", hash = "sha256:9d91bd436463ad5e4ee4aba766628dd6cd7010cf3e2461756b3303710eebc301", size = 26169, upload-time = "2025-08-12T18:49:02.73Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/19/97/56608b2249fe206a67cd573bc93cd9896e1efb9e98bce9c163bcdc704b88/truststore-0.10.4-py3-none-any.whl", hash = "sha256:adaeaecf1cbb5f4de3b1959b42d41f6fab57b2b1666adb59e89cb0b53361d981", size = 18660, upload-time = "2025-08-12T18:49:01.46Z" }, +] + [[package]] name = "typer" version = "0.25.1"