feat(anthropic): workload identity federation and pluggable identity sources

Backend half of #38818 (internal copy of the fork PR #38013), rebuilt as one
commit on top of litellm_internal_staging without the dashboard changes.

Deployments on anthropic/ without a static api_key can exchange an OIDC
workload assertion for a short-lived sk-ant-oat01 token through a shared
RFC 7523 JWT-bearer engine. The assertion comes from a mounted token file,
an env token, a LiteLLM-signed issuer, or Keycloak, chosen per deployment,
per named credential, or through ANTHROPIC_IDENTITY_SOURCE. The federation
fields are server-owned: refused inline in request bodies and on
POST /model/new, proxy-admin only on credentials, and the token exchange
is pinned to api.anthropic.com unless LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS
adds a host. GET /credentials/{name}/jwks exports the public key set of a
LiteLLM-signed credential for the Claude Console.

The OpenAI federation trio from #39613 rides along on the backend side with
the same server-owned handling.

Fixes #28607
Resolves LIT-6107

Co-authored-by: derhornspieler <15236687+derhornspieler@users.noreply.github.com>
This commit is contained in:
mateo-berri 2026-09-05 12:05:17 -07:00
parent a0058ed157
commit 6e44b5b626
113 changed files with 14209 additions and 2683 deletions

View file

@ -7,7 +7,10 @@ from typing import Any, Final, Literal
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
from litellm.litellm_core_utils.get_litellm_params import (
ANTHROPIC_WIF_KWARGS_KEYS,
AWS_CREDENTIAL_KWARGS_KEYS,
)
from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
from litellm.types.llms.openai import Batch
from litellm.types.utils import CallTypes, ModelInfo, Usage
@ -507,6 +510,9 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict:
"max_retries",
"_litellm_internal_model_credentials",
*AWS_CREDENTIAL_KWARGS_KEYS,
# A federated deployment holds no api_key, so without these the fetch that reads a
# finished batch's output has nothing to authenticate with and its cost is never billed.
*sorted(ANTHROPIC_WIF_KWARGS_KEYS),
)
for key in credential_keys:
if key in litellm_params:

View file

@ -483,6 +483,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
)
api_key = optional_params.api_key or litellm.api_key or litellm.azure_key or get_secret_str("ANTHROPIC_API_KEY")
batch_params: Final = dict(litellm_params) # mutable-ok: handler contract, copied not shared
response = anthropic_batches_instance.retrieve_batch(
_is_async=_is_async,
batch_id=batch_id,
@ -490,6 +491,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
api_key=api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
litellm_params=batch_params,
)
else:
raise litellm.exceptions.BadRequestError(

View file

@ -24,10 +24,54 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset(
# The per-deployment Rust opt-in.
RUST_KWARG_KEY: Final = "rust"
# Anthropic workload identity federation config, read from litellm_params by the
# Anthropic auth tier. Registered like `rust`: here so the kwargs funnel carries
# them, and in `all_litellm_params` so they never leak into the provider body.
ANTHROPIC_WIF_KWARGS_KEYS: Final = frozenset(
{
"anthropic_federation_rule_id",
"anthropic_organization_id",
"anthropic_service_account_id",
"anthropic_workspace_id",
"anthropic_identity_token_file",
"anthropic_identity_token",
# Identity-source selection (Phase 1): absent means the legacy
# token_file/env resolver above, byte-identical to today.
"anthropic_identity_source",
# internal_issuer: litellm self-signs the workload assertion.
"anthropic_issuer_url",
"anthropic_issuer_subject",
"anthropic_issuer_audience",
"anthropic_issuer_ttl_seconds",
"anthropic_issuer_signing_key_ref",
# keycloak: litellm fetches the assertion via client_credentials.
"anthropic_keycloak_token_url",
"anthropic_keycloak_client_id",
"anthropic_keycloak_auth_method",
"anthropic_keycloak_client_secret_ref",
"anthropic_keycloak_scope",
# Set server-side when a client redirects api_base, to stop a federated deployment minting
# for a base the caller chose. It has to ride this funnel or it is dropped on the way and
# the deployment federates anyway; being carried here also request-bans it, which is right,
# since a caller must not be able to set it in either direction.
"anthropic_disable_workload_identity_federation",
}
)
OPENAI_WIF_KWARGS_KEYS: Final = frozenset(
{
"openai_identity_provider_id",
"openai_service_account_id",
"openai_identity_token_file",
}
)
# Keys `completion()` forwards from its own kwargs into `get_litellm_params`,
# which are otherwise invisible to it because that call site passes explicit
# named arguments rather than `**kwargs`.
FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({RUST_KWARG_KEY})
FORWARDED_KWARGS_KEYS: Final = (
AWS_CREDENTIAL_KWARGS_KEYS | ANTHROPIC_WIF_KWARGS_KEYS | OPENAI_WIF_KWARGS_KEYS | frozenset({RUST_KWARG_KEY})
)
# Pre-define optional kwargs keys as frozenset for O(1) lookups
# These are extracted from kwargs only if present, avoiding unnecessary .get() calls
@ -65,6 +109,8 @@ OPTIONAL_KWARGS_KEYS: Final = (
}
)
| AWS_CREDENTIAL_KWARGS_KEYS
| ANTHROPIC_WIF_KWARGS_KEYS
| OPENAI_WIF_KWARGS_KEYS
)
# Backward-compatible alias for existing imports/tests.

View file

@ -42,6 +42,7 @@ class AnthropicBatchesHandler:
timeout: float | httpx.Timeout,
max_retries: int | None,
logging_obj: LiteLLMLoggingObj | None = None,
litellm_params: dict | None = None, # mutable-ok: handed straight to validate_environment
) -> LiteLLMBatch:
"""
Async: Retrieve a batch from Anthropic.
@ -60,9 +61,7 @@ class AnthropicBatchesHandler:
# Resolve API credentials
api_base = api_base or self.anthropic_model_info.get_api_base(api_base)
api_key = api_key or self.anthropic_model_info.get_api_key()
if not api_key:
raise ValueError("Missing Anthropic API Key")
resolved_litellm_params: Final = litellm_params if litellm_params is not None else {}
# Create a minimal logging object if not provided
if logging_obj is None:
@ -85,16 +84,18 @@ class AnthropicBatchesHandler:
api_base=api_base,
batch_id=batch_id,
optional_params={},
litellm_params={},
litellm_params=resolved_litellm_params,
)
# Validate environment and get headers
headers: Final = self.provider_config.validate_environment(
# Validate environment and get headers. Offloaded to a worker thread: a WIF token
# exchange here would otherwise block the event loop.
headers: Final = await asyncio.to_thread(
self.provider_config.validate_environment,
headers={},
model="",
messages=[],
optional_params={},
litellm_params={},
litellm_params=resolved_litellm_params,
api_key=api_key,
api_base=api_base,
)
@ -130,6 +131,7 @@ class AnthropicBatchesHandler:
timeout: float | httpx.Timeout,
max_retries: int | None,
logging_obj: LiteLLMLoggingObj | None = None,
litellm_params: dict | None = None, # mutable-ok: handed straight to validate_environment
) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
"""
Retrieve a batch from Anthropic.
@ -154,6 +156,7 @@ class AnthropicBatchesHandler:
timeout=timeout,
max_retries=max_retries,
logging_obj=logging_obj,
litellm_params=litellm_params,
)
else:
return asyncio.run(
@ -164,5 +167,6 @@ class AnthropicBatchesHandler:
timeout=timeout,
max_retries=max_retries,
logging_obj=logging_obj,
litellm_params=litellm_params,
)
)

View file

@ -13,6 +13,8 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest
from litellm.types.utils import LiteLLMBatch, LlmProviders, ModelResponse
from ..common_utils import merge_anthropic_beta_headers, without_caller_credential_headers
if TYPE_CHECKING:
import tiktoken
@ -70,24 +72,30 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
api_base: str | None = None,
) -> dict:
"""Validate and prepare environment-specific headers and parameters."""
if api_base is None and isinstance(litellm_params, dict):
api_base = litellm_params.get("api_base")
auth_header: Final = self.anthropic_model_info.get_auth_header(api_key, api_base)
params_mapping: Final = litellm_params if isinstance(litellm_params, dict) else None
if api_base is None and params_mapping is not None:
api_base = params_mapping.get("api_base")
auth_header: Final = self.anthropic_model_info.get_auth_header(
api_key, api_base, litellm_params=params_mapping, allow_workload_identity=True
)
if auth_header is None:
raise ValueError(
"Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params"
)
_headers: Final = {
merged_beta: Final = merge_anthropic_beta_headers(
merge_anthropic_beta_headers(headers.get("anthropic-beta"), auth_header.get("anthropic-beta")),
"message-batches-2024-09-24",
)
# The deployment's own credential is applied below, so a caller-supplied one must not
# ride along: without this a minted federation Bearer travels beside the caller's x-api-key.
return {
**without_caller_credential_headers(headers),
"accept": "application/json",
"anthropic-version": "2023-06-01",
"content-type": "application/json",
**auth_header,
"anthropic-beta": merged_beta,
}
_headers.update(auth_header)
# Add beta header for message batches
if "anthropic-beta" not in headers:
headers["anthropic-beta"] = "message-batches-2024-09-24"
headers.update(_headers)
return headers
def get_complete_batch_url(
self,

View file

@ -3,7 +3,7 @@ import re
import time
from collections.abc import Callable, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, NoReturn, cast
from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, cast
import httpx
from pydantic import ValidationError
@ -284,6 +284,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
to pass metadata to anthropic, it's {"user_id": "any-relevant-information"}
"""
_workload_identity_eligible: ClassVar[bool] = True
max_tokens: int | None = None
stop_sequences: list | None = None
temperature: int | None = None

View file

@ -7,10 +7,11 @@ import re
from collections.abc import Mapping, MutableMapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Any, Final, Literal
from typing import Any, ClassVar, Final, Literal
from urllib.parse import quote
import httpx
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
import litellm
from litellm.constants import (
@ -25,8 +26,14 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
from litellm.litellm_core_utils.prompt_templates.factory import (
THOUGHT_SIGNATURE_SEPARATOR,
)
from litellm.llms.anthropic.wif import (
aget_anthropic_wif_token,
anthropic_base_without_chat_suffix,
get_anthropic_wif_token,
)
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.proxy._types import SpecialHeaders
from litellm.types.llms.anthropic import (
ANTHROPIC_HOSTED_TOOLS,
ANTHROPIC_OAUTH_BETA_HEADER,
@ -81,6 +88,29 @@ def _strip_bedrock_id_suffixes(model: str) -> str:
)
_SERVER_OWNED_AUTH_HEADERS: Final = SpecialHeaders.litellm_credential_header_names()
_WIF_ELIGIBILITY_ATTR: Final = "_workload_identity_eligible"
def without_caller_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
"""``headers`` minus every header that authenticates the caller to litellm.
The deployment's own credential is applied on top of the result, so a caller-supplied
credential must not survive into the upstream request: without this a minted federation
Bearer travels beside the caller's own ``x-api-key``, and Anthropic sees two credentials.
"""
return MappingProxyType(
{name: value for name, value in headers.items() if name.lower() not in _SERVER_OWNED_AUTH_HEADERS}
)
def config_allows_workload_identity(config: object) -> bool:
"""A federation token is an Anthropic-org credential and its exchange POSTs the workload's OIDC
assertion to the deployment's own host, so eligibility is declared per class and read from that
class's own ``__dict__``: a subclass written for another provider inherits nothing."""
return type(config).__dict__.get(_WIF_ELIGIBILITY_ATTR, False) is True
def is_anthropic_oauth_key(value: str | None) -> bool:
"""Check if a value contains an Anthropic OAuth token (sk-ant-oat*)."""
if value is None:
@ -90,12 +120,20 @@ def is_anthropic_oauth_key(value: str | None) -> bool:
return value.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX)
def _merge_beta_headers(existing: str | None, new_beta: str) -> str:
"""Merge a new beta value into an existing comma-separated anthropic-beta header."""
if not existing:
return new_beta
betas: Final = {b.strip() for b in existing.split(",") if b.strip()}
betas.add(new_beta)
def merge_anthropic_beta_headers(existing: str | Sequence[str] | None, new_beta: str | Sequence[str] | None) -> str:
"""Merge anthropic-beta header values, deduplicated and sorted.
Either side may arrive as a list rather than a comma-separated string: the Skills surface
accepted a list-valued header before it shared this helper, and callers still send one.
"""
values: Final = (
entry
for side in (existing, new_beta)
if side
for entry in ((side,) if isinstance(side, str) else side)
if isinstance(entry, str)
)
betas: Final = frozenset(b.strip() for value in values for b in value.split(",") if b.strip())
return ",".join(sorted(betas))
@ -122,7 +160,9 @@ def optionally_handle_anthropic_oauth(headers: dict, api_key: str | None) -> tup
):
headers.pop(name)
headers["authorization"] = auth_header
headers["anthropic-beta"] = _merge_beta_headers(headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER)
headers["anthropic-beta"] = merge_anthropic_beta_headers(
headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER
)
headers["anthropic-dangerous-direct-browser-access"] = "true"
return headers, api_key
# Check api_key directly (standard chat/completion flow)
@ -130,7 +170,9 @@ def optionally_handle_anthropic_oauth(headers: dict, api_key: str | None) -> tup
for name in tuple(header_name for header_name in headers if header_name.lower() == "x-api-key"):
headers.pop(name)
headers["authorization"] = f"Bearer {api_key}"
headers["anthropic-beta"] = _merge_beta_headers(headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER)
headers["anthropic-beta"] = merge_anthropic_beta_headers(
headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER
)
headers["anthropic-dangerous-direct-browser-access"] = "true"
return headers, api_key
@ -145,7 +187,79 @@ class AnthropicError(BaseLLMException):
super().__init__(status_code=status_code, message=message, headers=headers)
_MODEL_LIST_PAGE_CAP: Final = 20
def _litellm_params_str(litellm_params: Mapping[str, object] | None, key: str) -> str | None:
value: Final = litellm_params.get(key) if litellm_params is not None else None
return value if isinstance(value, str) else None
class _AnthropicModelListEntry(BaseModel):
id: str
class _AnthropicModelsPage(BaseModel):
data: Sequence[_AnthropicModelListEntry] = Field(default_factory=tuple)
has_more: bool = False
last_id: str | None = None
def _sanitized_anthropic_error(response: httpx.Response, detail: str | None = None) -> str:
"""A provider error detail built only from structured fields, never ``response.text``
verbatim: the raw body is untrusted content the caller of ``/v1/models`` did not ask for
and should not have echoed back to it wholesale."""
if detail is not None:
return f"HTTP {response.status_code}: {detail}"
try:
body: Final = response.json()
except ValueError:
return f"HTTP {response.status_code}"
error: Final = body.get("error") if isinstance(body, dict) else None
message: Final = error.get("message") if isinstance(error, dict) else None
return f"HTTP {response.status_code}: {message}" if isinstance(message, str) else f"HTTP {response.status_code}"
def _fetch_anthropic_models_page(
api_base: str, headers: Mapping[str, str], after_id: str | None
) -> _AnthropicModelsPage:
# after_id rides the URL because the client mutates the params mapping it is handed,
# which a read-only one cannot support
query: Final = f"?after_id={quote(after_id)}" if after_id else ""
response: Final = litellm.module_level_client.get(
url=f"{api_base}/v1/models{query}",
headers=headers,
follow_redirects=False,
)
try:
response.raise_for_status()
except httpx.HTTPStatusError:
raise Exception(f"Failed to fetch models from Anthropic. {_sanitized_anthropic_error(response)}") from None
try:
return _AnthropicModelsPage.model_validate(response.json())
except ValueError as e:
raise Exception(
f"Failed to fetch models from Anthropic. {_sanitized_anthropic_error(response, detail=str(e))}"
) from None
def _fetch_anthropic_model_ids(
api_base: str, headers: Mapping[str, str], after_id: str | None, pages_left: int
) -> tuple[str, ...]:
collected: tuple[str, ...] = () # rebind-ok: accumulates one page of ids per iteration
cursor: str | None = after_id # rebind-ok: advances to each page's last_id
for _ in range(max(pages_left, 0)):
page = _fetch_anthropic_models_page(api_base, headers, cursor) # rebind-ok: one page per iteration
collected += tuple(entry.id for entry in page.data)
if not page.has_more or page.last_id is None:
return collected
cursor = page.last_id
raise Exception(f"Anthropic /v1/models did not terminate within {_MODEL_LIST_PAGE_CAP} pages.")
class AnthropicModelInfo(BaseLLMModelInfo):
_workload_identity_eligible: ClassVar[bool] = True
def is_cache_control_set(self, messages: list[AllMessageValues]) -> bool:
"""
Return if {"cache_control": ..} in message content block
@ -720,7 +834,9 @@ class AnthropicModelInfo(BaseLLMModelInfo):
return list(set(betas))
@staticmethod
def _make_api_key_auth_header(api_key: str, api_base: str | None, use_bearer_for_custom_base: bool = False) -> dict:
def _make_api_key_auth_header(
api_key: str, api_base: str | None, use_bearer_for_custom_base: bool = False
) -> Mapping[str, str]:
if use_bearer_for_custom_base and (
api_base and "api.anthropic.com" not in api_base and not api_key.startswith("sk-ant-")
):
@ -728,6 +844,33 @@ class AnthropicModelInfo(BaseLLMModelInfo):
return {"authorization": value}
return {"x-api-key": api_key}
def _credential_headers(
self,
*,
api_key: str | None,
auth_token: str | None,
api_base: str | None,
use_bearer_for_custom_base: bool,
wif_minted: bool,
betas: set[str], # mutable-ok: the caller's beta accumulator, appended to by the oauth tier
) -> Mapping[str, str]:
"""The credential tier walk: a consumer OAuth token, then ANTHROPIC_AUTH_TOKEN, then an api key.
A server-minted federation token takes the same Bearer shape as a consumer OAuth token but is
not browser-forwarded, so it does not get the direct-browser-access header.
"""
if api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX):
betas.add(ANTHROPIC_OAUTH_BETA_HEADER)
oauth_headers: Final = {"authorization": f"Bearer {api_key}"}
if wif_minted:
return oauth_headers
return {**oauth_headers, "anthropic-dangerous-direct-browser-access": "true"}
if auth_token and not api_key:
return {"authorization": f"Bearer {auth_token}"}
if api_key:
return self._make_api_key_auth_header(api_key, api_base, use_bearer_for_custom_base)
return {}
def get_anthropic_headers(
self,
api_key: str | None = None,
@ -749,6 +892,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
container_with_skills_used: bool = False,
api_base: str | None = None,
use_bearer_for_custom_base: bool = False,
wif_minted: bool = False,
) -> dict:
betas: Final = set()
# Anthropic no longer requires the prompt-caching beta header
@ -784,20 +928,21 @@ class AnthropicModelInfo(BaseLLMModelInfo):
if container_with_skills_used:
betas.add("skills-2025-10-02")
_is_oauth: Final = api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX)
headers: Final = {
"anthropic-version": anthropic_version or "2023-06-01",
"accept": "application/json",
"content-type": "application/json",
}
if _is_oauth:
headers["authorization"] = f"Bearer {api_key}"
headers["anthropic-dangerous-direct-browser-access"] = "true"
betas.add(ANTHROPIC_OAUTH_BETA_HEADER)
elif auth_token and not api_key:
headers["authorization"] = f"Bearer {auth_token}"
elif api_key:
headers.update(self._make_api_key_auth_header(api_key, api_base, use_bearer_for_custom_base))
headers.update(
self._credential_headers(
api_key=api_key,
auth_token=auth_token,
api_base=api_base,
use_bearer_for_custom_base=use_bearer_for_custom_base,
wif_minted=wif_minted,
betas=betas,
)
)
if user_anthropic_beta_headers is not None:
betas.update(user_anthropic_beta_headers)
@ -824,10 +969,11 @@ class AnthropicModelInfo(BaseLLMModelInfo):
api_key: str | None = None,
api_base: str | None = None,
) -> dict:
if api_base is None and isinstance(litellm_params, dict):
api_base = litellm_params.get("api_base")
params_mapping: Final = litellm_params if isinstance(litellm_params, dict) else None
if api_base is None and params_mapping is not None:
api_base = params_mapping.get("api_base")
use_bearer_for_custom_base: Final[bool] = bool(
isinstance(litellm_params, dict) and litellm_params.get("use_bearer_for_custom_base", False)
params_mapping is not None and params_mapping.get("use_bearer_for_custom_base", False)
)
# Check for Anthropic OAuth token in headers
headers, api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key)
@ -836,9 +982,23 @@ class AnthropicModelInfo(BaseLLMModelInfo):
auth_token: str | None = None
if api_key is None:
auth_token = AnthropicModelInfo.get_auth_token()
if api_key is None and auth_token is None:
wif_token: Final = (
get_anthropic_wif_token(params_mapping, api_base, model)
if api_key is None and auth_token is None and config_allows_workload_identity(self)
else None
)
wif_minted: Final = wif_token is not None
resolved_api_key: Final = wif_token if wif_token is not None else api_key
if resolved_api_key is None and auth_token is None:
raise litellm.AuthenticationError(
message="Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` in your environment vars",
message=(
"Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the "
"environment variables or via params. Please set `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` "
"in your environment vars, or configure workload identity federation via "
"`ANTHROPIC_FEDERATION_RULE_ID`, `ANTHROPIC_ORGANIZATION_ID`, "
"`ANTHROPIC_SERVICE_ACCOUNT_ID` and "
"`ANTHROPIC_IDENTITY_TOKEN_FILE` (or `ANTHROPIC_IDENTITY_TOKEN`)"
),
llm_provider="anthropic",
model=model,
)
@ -863,7 +1023,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
computer_tool_used=computer_tool_used,
prompt_caching_set=prompt_caching_set,
pdf_used=pdf_used,
api_key=api_key,
api_key=resolved_api_key,
auth_token=auth_token,
file_id_used=file_id_used,
web_search_tool_used=web_search_tool_used,
@ -878,11 +1038,12 @@ class AnthropicModelInfo(BaseLLMModelInfo):
container_with_skills_used=container_with_skills_used,
api_base=api_base,
use_bearer_for_custom_base=use_bearer_for_custom_base,
wif_minted=wif_minted,
)
headers = {**headers, **anthropic_headers}
caller_headers: Final = without_caller_credential_headers(headers) if wif_minted else headers
return headers
return {**caller_headers, **anthropic_headers}
@staticmethod
def get_api_base(api_base: str | None = None) -> str | None:
@ -917,52 +1078,121 @@ class AnthropicModelInfo(BaseLLMModelInfo):
api_key: str | None = None,
api_base: str | None = None,
use_bearer_for_custom_base: bool = False,
) -> dict | None:
litellm_params: Mapping[str, object] | None = None,
allow_workload_identity: bool = False,
) -> Mapping[str, str] | None:
"""Resolve Anthropic credentials and return the appropriate auth header dict.
Checks ANTHROPIC_API_KEY first (-> x-api-key or Bearer depending on
use_bearer_for_custom_base), then ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer).
Returns None if neither is available.
use_bearer_for_custom_base), then ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer),
then workload identity federation (-> Authorization: Bearer with a minted
sk-ant-oat01 token, honoring anthropic_* litellm_params when provided). Every
Bearer built from an sk-ant-oat token carries the mandatory oauth anthropic-beta.
Returns None if no credential source is available.
"""
static_header: Final = AnthropicModelInfo._static_auth_header(api_key, api_base, use_bearer_for_custom_base)
if static_header is not None:
return static_header
if not allow_workload_identity:
return None
wif_token: Final = get_anthropic_wif_token(litellm_params, api_base, "")
if wif_token is not None:
return AnthropicModelInfo._oauth_bearer_header(wif_token)
return None
@staticmethod
async def aget_auth_header(
api_key: str | None = None,
api_base: str | None = None,
use_bearer_for_custom_base: bool = False,
litellm_params: Mapping[str, object] | None = None,
allow_workload_identity: bool = False,
) -> Mapping[str, str] | None:
"""Async counterpart of get_auth_header: the WIF tier can block on a token
exchange POST, so async callers await it off the event loop."""
static_header: Final = AnthropicModelInfo._static_auth_header(api_key, api_base, use_bearer_for_custom_base)
if static_header is not None:
return static_header
if not allow_workload_identity:
return None
wif_token: Final = await aget_anthropic_wif_token(litellm_params, api_base, "")
if wif_token is not None:
return AnthropicModelInfo._oauth_bearer_header(wif_token)
return None
@staticmethod
def _static_auth_header(
api_key: str | None,
api_base: str | None,
use_bearer_for_custom_base: bool,
) -> Mapping[str, str] | None:
resolved_key: Final = AnthropicModelInfo.get_api_key(api_key)
if resolved_key is not None:
if is_anthropic_oauth_key(resolved_key):
return {"authorization": f"Bearer {resolved_key}"}
return AnthropicModelInfo._oauth_bearer_header(resolved_key)
return AnthropicModelInfo._make_api_key_auth_header(resolved_key, api_base, use_bearer_for_custom_base)
auth_token: Final = AnthropicModelInfo.get_auth_token()
if auth_token is not None:
return {"authorization": f"Bearer {auth_token}"}
return None
@staticmethod
def _oauth_bearer_header(token: str) -> Mapping[str, str]:
return {"authorization": f"Bearer {token}", "anthropic-beta": ANTHROPIC_OAUTH_BETA_HEADER}
@staticmethod
def get_base_model(model: str | None = None) -> str | None:
return model.replace("anthropic/", "") if model else None
def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]:
api_base = AnthropicModelInfo.get_api_base(api_base)
auth_header: Final = AnthropicModelInfo.get_auth_header(api_key, api_base)
if api_base is None or auth_header is None:
raise ValueError(
"ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN is not set. Please set the environment variable, to query Anthropic's `/models` endpoint."
)
headers: Final = {"anthropic-version": "2023-06-01"}
headers.update(auth_header)
response: Final = litellm.module_level_client.get(
url=f"{api_base}/v1/models",
headers=headers,
return self._list_models(api_key=api_key, api_base=api_base, litellm_params=None)
def discover_models(
self, litellm_params: Mapping[str, object] | None = None
) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override
"""Live discovery for a configured deployment: unlike ``get_models``, this threads the
full ``litellm_params`` into ``get_auth_header`` so a workload-identity-federation source
configured on the deployment (rather than the environment) is honored, gated the same way
every other Anthropic auth surface is via ``config_allows_workload_identity``."""
return self._list_models(
api_key=_litellm_params_str(litellm_params, "api_key"),
api_base=_litellm_params_str(litellm_params, "api_base"),
litellm_params=litellm_params,
)
try:
response.raise_for_status()
except httpx.HTTPStatusError:
raise Exception(
f"Failed to fetch models from Anthropic. Status code: {response.status_code}, Response: {response.text}"
def _list_models(
self,
*,
api_key: str | None,
api_base: str | None,
litellm_params: Mapping[str, object] | None,
) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override
resolved_api_base: Final = AnthropicModelInfo.get_api_base(api_base)
auth_header: Final = AnthropicModelInfo.get_auth_header(
api_key,
resolved_api_base,
litellm_params=litellm_params,
allow_workload_identity=config_allows_workload_identity(self),
)
if resolved_api_base is None or auth_header is None:
raise ValueError(
"ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN (or workload "
"identity federation via ANTHROPIC_FEDERATION_RULE_ID/ANTHROPIC_ORGANIZATION_ID/"
"ANTHROPIC_IDENTITY_TOKEN_FILE) is not set. Please set the environment variable, to query "
"Anthropic's `/models` endpoint."
)
models: Final[Sequence[Mapping[str, str]]] = response.json()["data"]
litellm_model_names: Final = ["anthropic/" + model["id"] for model in models]
return litellm_model_names
headers: Final = MappingProxyType({"anthropic-version": "2023-06-01", **auth_header})
# /v1/models is appended below, so a base the operator already wrote as .../v1 or
# .../v1/messages would otherwise be asked for /v1/v1/models.
model_ids: Final = _fetch_anthropic_model_ids(
anthropic_base_without_chat_suffix(resolved_api_base),
headers,
after_id=None,
pages_left=_MODEL_LIST_PAGE_CAP,
)
return [ # mutable-ok: matches get_models' list[str] contract shared by every provider override
"anthropic/" + model_id for model_id in model_ids
]
def get_token_counter(self) -> BaseTokenCounter | None:
"""

View file

@ -41,7 +41,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
model: The model identifier (e.g., "claude-3-5-sonnet-20241022")
messages: The messages to count tokens for
api_key: The Anthropic API key
api_base: Optional custom API base URL
api_base: Optional deployment api_base the count-tokens path is appended to
timeout: Optional timeout for the request (defaults to litellm.request_timeout)
Returns:
@ -67,7 +67,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
verbose_logger.debug("Transformed request: %s", request_body)
# Get endpoint URL
endpoint_url: Final = api_base or self.get_anthropic_count_tokens_endpoint()
endpoint_url: Final = self.get_anthropic_count_tokens_endpoint(api_base)
verbose_logger.debug("Making request to: %s", endpoint_url)

View file

@ -2,10 +2,10 @@
Anthropic Token Counter implementation using the CountTokens API.
"""
import os
from typing import Any, Final
from litellm._logging import verbose_logger
from litellm.exceptions import AuthenticationError
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.types.utils import LlmProviders, TokenCountResponse
@ -46,28 +46,33 @@ class AnthropicTokenCounter(BaseTokenCounter):
Returns:
TokenCountResponse with token count, or None if counting fails
"""
from litellm.llms.anthropic.common_utils import AnthropicError
from litellm.llms.anthropic.common_utils import AnthropicError, AnthropicModelInfo
from litellm.llms.anthropic.wif import aget_anthropic_wif_token
if not messages:
return None
deployment = deployment or {}
litellm_params: Final = deployment.get("litellm_params", {})
# Get Anthropic API key from deployment config or environment
api_key = litellm_params.get("api_key")
if not api_key:
api_key = os.getenv("ANTHROPIC_API_KEY")
if not api_key:
verbose_logger.warning("No Anthropic API key found for token counting")
return None
api_base: Final = litellm_params.get("api_base")
static_key: Final = AnthropicModelInfo.get_api_key(litellm_params.get("api_key"))
auth_token_configured: Final = AnthropicModelInfo.get_auth_token() is not None
try:
api_key: Final = (
static_key
if static_key or auth_token_configured
else await aget_anthropic_wif_token(litellm_params, api_base, model_to_use)
)
if not api_key:
verbose_logger.warning("No Anthropic credential found for token counting")
return None
result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request(
model=model_to_use,
messages=messages,
api_key=api_key,
api_base=api_base,
tools=tools,
system=system,
)
@ -80,8 +85,8 @@ class AnthropicTokenCounter(BaseTokenCounter):
tokenizer_type="anthropic_api",
original_response=result,
)
except AnthropicError as e:
verbose_logger.warning("Anthropic CountTokens API error: status=%s, message=%s", e.status_code, e.message)
except (AnthropicError, AuthenticationError) as e:
verbose_logger.warning("Anthropic CountTokens error: status=%s, message=%s", e.status_code, e.message)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,

View file

@ -7,6 +7,7 @@ This module handles the transformation of requests to Anthropic's CountTokens AP
from typing import Any, Final
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
from litellm.llms.anthropic.wif import resolve_anthropic_base
class AnthropicCountTokensConfig:
@ -19,14 +20,21 @@ class AnthropicCountTokensConfig:
- Response: {"input_tokens": <number>}
"""
def get_anthropic_count_tokens_endpoint(self) -> str:
def get_anthropic_count_tokens_endpoint(self, api_base: str | None = None) -> str:
"""
Get the Anthropic CountTokens API endpoint.
Args:
api_base: The deployment's api_base, which names the chat surface (a host, or a
base already carrying ``/v1`` or ``/v1/messages``); the count-tokens path is
appended to it, so it is never the full count-tokens URL. Unset or empty falls
back to ``ANTHROPIC_API_BASE`` / ``ANTHROPIC_BASE_URL`` and then Anthropic's
host, the same resolution chat and the federated exchange use
Returns:
The endpoint URL for the CountTokens API
"""
return "https://api.anthropic.com/v1/messages/count_tokens"
return resolve_anthropic_base(api_base) + "/v1/messages/count_tokens"
def transform_request_to_count_tokens(
self,

View file

@ -1,5 +1,5 @@
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import Any, Final
from typing import Any, ClassVar, Final
import httpx
@ -23,12 +23,24 @@ from litellm.types.router import GenericLiteLLMParams
from ...common_utils import (
AnthropicError,
AnthropicModelInfo,
merge_anthropic_beta_headers,
optionally_handle_anthropic_oauth,
strip_advisor_blocks_from_messages,
)
DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01"
_CALLER_CREDENTIAL_HEADERS: Final = frozenset({"x-api-key", "authorization"})
def _carries_caller_credential(headers: Mapping[str, str]) -> bool:
"""Whether the caller sent their own Anthropic credential, in which case this passthrough
honors it and never mints. Matched case-insensitively: an SDK caller passing ``X-Api-Key``
through extra_headers would otherwise slip the check and end up sending their key beside a
minted federation Bearer."""
return any(name.lower() in _CALLER_CREDENTIAL_HEADERS for name in headers)
DROP_UNSUPPORTED_ADAPTIVE_EFFORT_WARNING: Final = (
"Dropping adaptive `thinking`/`output_config.effort` for model=%s: the model "
"does not support extended thinking, or max_tokens is too small to fit the "
@ -42,6 +54,8 @@ DROP_UNFITTING_REASONING_EFFORT_WARNING: Final = (
class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
_workload_identity_eligible: ClassVar[bool] = True
@property
def custom_llm_provider(self) -> str | None:
return "anthropic"
@ -308,32 +322,103 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
# Check for Anthropic OAuth token in Authorization header
headers, api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key)
header_names: Final = frozenset(name.lower() for name in headers)
if "x-api-key" not in header_names and "authorization" not in header_names:
auth_header: Final = AnthropicModelInfo.get_auth_header(api_key)
if auth_header is None:
raise AuthenticationError(
message=(
"Missing Anthropic API Key - A call is being made to anthropic but no key is set "
"either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` "
"or `ANTHROPIC_AUTH_TOKEN` in your environment vars"
if not _carries_caller_credential(headers):
self._apply_env_auth_header(
headers,
self._require_auth_header(
AnthropicModelInfo.get_auth_header(
api_key,
api_base=api_base,
litellm_params=litellm_params,
allow_workload_identity=self._allows_workload_identity,
),
llm_provider=self._resolved_provider,
model=model,
)
headers.update(auth_header)
),
)
return self._finalize_messages_headers(headers, optional_params), api_base
async def avalidate_anthropic_messages_environment(
self,
headers: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
model: str,
messages: list[Any], # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
optional_params: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
litellm_params: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict, str | None]: # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
if type(self).validate_anthropic_messages_environment is not (
AnthropicMessagesConfig.validate_anthropic_messages_environment
):
# a subclass sync override must keep winning on the async path
return self.validate_anthropic_messages_environment(
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
api_key=api_key,
api_base=api_base,
)
oauth_headers, oauth_api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key)
if not _carries_caller_credential(oauth_headers):
self._apply_env_auth_header(
oauth_headers,
self._require_auth_header(
await AnthropicModelInfo.aget_auth_header(
oauth_api_key,
api_base=api_base,
litellm_params=litellm_params,
allow_workload_identity=self._allows_workload_identity,
),
model=model,
),
)
return self._finalize_messages_headers(oauth_headers, optional_params), api_base
def _require_auth_header(self, auth_header: Mapping[str, str] | None, model: str) -> Mapping[str, str]:
if auth_header is None:
raise AuthenticationError(
message=(
"Missing Anthropic API Key - A call is being made to anthropic but no key is set "
"either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` "
"or `ANTHROPIC_AUTH_TOKEN` in your environment vars"
),
llm_provider=self._resolved_provider,
model=model,
)
return auth_header
@staticmethod
def _apply_env_auth_header(headers: dict, auth_header: Mapping[str, str] | None) -> None: # mutable-ok: out-param
if auth_header is None:
return
merged_beta: Final = merge_anthropic_beta_headers(
headers.get("anthropic-beta"), auth_header.get("anthropic-beta")
)
headers.update(auth_header)
if merged_beta:
headers["anthropic-beta"] = merged_beta
@property
def _allows_workload_identity(self) -> bool:
"""Subclasses reuse this validate step for their own /v1/messages-compatible providers, so
eligibility is declared per class and never inherited."""
from litellm.llms.anthropic.common_utils import config_allows_workload_identity
return config_allows_workload_identity(self)
def _finalize_messages_headers(self, headers: dict, optional_params: dict) -> dict: # mutable-ok: out-param
if "anthropic-version" not in headers:
headers["anthropic-version"] = DEFAULT_ANTHROPIC_API_VERSION
if "content-type" not in headers:
headers["content-type"] = "application/json"
headers = self._update_headers_with_anthropic_beta(
return self._update_headers_with_anthropic_beta(
headers=headers,
optional_params=optional_params,
)
return headers, api_base
@staticmethod
def _translate_reasoning_effort_to_anthropic(
model: str, optional_params: dict, max_tokens: int | None, custom_llm_provider: str

View file

@ -73,7 +73,9 @@ class AnthropicFilesHandler:
# Get Anthropic API credentials
api_base = self.anthropic_model_info.get_api_base(api_base)
auth_header: Final = self.anthropic_model_info.get_auth_header(api_key, api_base)
auth_header: Final = await self.anthropic_model_info.aget_auth_header(
api_key, api_base, allow_workload_identity=True
)
if auth_header is None:
raise ValueError("Missing Anthropic API Key")

View file

@ -14,6 +14,7 @@ Anthropic Files API endpoints:
import calendar
import time
from collections.abc import Mapping
from typing import Final, cast
import httpx
@ -35,7 +36,12 @@ from litellm.types.llms.openai import (
)
from litellm.types.utils import LlmProviders
from ..common_utils import AnthropicError, AnthropicModelInfo
from ..common_utils import (
AnthropicError,
AnthropicModelInfo,
merge_anthropic_beta_headers,
without_caller_credential_headers,
)
ANTHROPIC_FILES_API_BASE: Final = "https://api.anthropic.com"
ANTHROPIC_FILES_BETA_HEADER: Final = "files-api-2025-04-14"
@ -94,21 +100,55 @@ class AnthropicFilesConfig(BaseFilesConfig):
api_key: str | None = None,
api_base: str | None = None,
) -> dict:
if api_base is None and isinstance(litellm_params, dict):
api_base = litellm_params.get("api_base")
auth_header: Final = AnthropicModelInfo.get_auth_header(api_key, api_base)
params_mapping, resolved_api_base = self._resolve_params(litellm_params, api_base)
auth_header: Final = AnthropicModelInfo.get_auth_header(
api_key, resolved_api_base, litellm_params=params_mapping, allow_workload_identity=True
)
return self._finalize_headers(headers, auth_header)
async def avalidate_environment(
self,
headers: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
model: str,
messages: list, # mutable-ok: mirrors the sync validate_environment contract this overrides
optional_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
litellm_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
api_key: str | None = None,
api_base: str | None = None,
) -> dict: # mutable-ok: mirrors the sync validate_environment contract this overrides
"""Async counterpart of validate_environment: the WIF tier can block on a token
exchange POST, so async callers await it off the event loop."""
params_mapping, resolved_api_base = self._resolve_params(litellm_params, api_base)
auth_header: Final = await AnthropicModelInfo.aget_auth_header(
api_key, resolved_api_base, litellm_params=params_mapping, allow_workload_identity=True
)
return self._finalize_headers(headers, auth_header)
@staticmethod
def _resolve_params(
litellm_params: dict, api_base: str | None
) -> tuple[dict | None, str | None]: # mutable-ok: mirrors the sync validate_environment contract this overrides
params_mapping: Final = litellm_params if isinstance(litellm_params, dict) else None
if api_base is None and params_mapping is not None:
api_base = params_mapping.get("api_base")
return params_mapping, api_base
@staticmethod
def _finalize_headers(headers: dict, auth_header: Mapping[str, str] | None) -> dict: # mutable-ok: out-param
if auth_header is None:
raise ValueError(
"Anthropic API key is required. Set ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN environment variable or pass api_key parameter."
)
headers.update(
{
**auth_header,
"anthropic-version": "2023-06-01",
"anthropic-beta": ANTHROPIC_FILES_BETA_HEADER,
}
merged_beta: Final = merge_anthropic_beta_headers(
merge_anthropic_beta_headers(headers.get("anthropic-beta"), auth_header.get("anthropic-beta")),
ANTHROPIC_FILES_BETA_HEADER,
)
return headers
return {
**without_caller_credential_headers(headers),
**auth_header,
"anthropic-version": "2023-06-01",
"anthropic-beta": merged_beta,
}
def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]:
return ["purpose"]

View file

@ -2,6 +2,7 @@
Anthropic Skills API configuration and transformations
"""
from types import MappingProxyType
from typing import Final
import httpx
@ -35,40 +36,35 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict:
"""Add Anthropic-specific headers"""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
from litellm.constants import ANTHROPIC_SKILLS_API_BETA_VERSION
from litellm.llms.anthropic.common_utils import (
AnthropicModelInfo,
merge_anthropic_beta_headers,
without_caller_credential_headers,
)
# Get API key from litellm_params if available
api_key = None
api_base = None
if litellm_params is not None:
api_key = litellm_params.api_key
api_base = litellm_params.api_base
auth_header: Final = AnthropicModelInfo.get_auth_header(api_key, api_base)
auth_header: Final = AnthropicModelInfo.get_auth_header(
api_key=litellm_params.api_key if litellm_params is not None else None,
api_base=litellm_params.api_base if litellm_params is not None else None,
litellm_params=MappingProxyType(dict(litellm_params)) if litellm_params is not None else None,
allow_workload_identity=True,
)
if auth_header is None:
raise ValueError("ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN is required for Skills API")
headers.update(auth_header)
headers["anthropic-version"] = "2023-06-01"
# Add beta header for skills API
from litellm.constants import ANTHROPIC_SKILLS_API_BETA_VERSION
if "anthropic-beta" not in headers:
headers["anthropic-beta"] = ANTHROPIC_SKILLS_API_BETA_VERSION
elif isinstance(headers["anthropic-beta"], list):
if ANTHROPIC_SKILLS_API_BETA_VERSION not in headers["anthropic-beta"]:
headers["anthropic-beta"].append(ANTHROPIC_SKILLS_API_BETA_VERSION)
elif isinstance(headers["anthropic-beta"], str):
if ANTHROPIC_SKILLS_API_BETA_VERSION not in headers["anthropic-beta"]:
headers["anthropic-beta"] = [
headers["anthropic-beta"],
ANTHROPIC_SKILLS_API_BETA_VERSION,
]
headers["content-type"] = "application/json"
return headers
merged_beta: Final = merge_anthropic_beta_headers(
merge_anthropic_beta_headers(headers.get("anthropic-beta"), auth_header.get("anthropic-beta")),
ANTHROPIC_SKILLS_API_BETA_VERSION,
)
# The deployment's own credential is applied here, so a caller-supplied one must not ride
# along upstream beside a minted federation Bearer.
return { # mutable-ok: validate_environment's contract returns a real dict, which httpx then consumes
**without_caller_credential_headers(headers),
**auth_header,
"anthropic-version": "2023-06-01",
"anthropic-beta": merged_beta,
"content-type": "application/json",
}
def get_complete_url(
self,

View file

@ -0,0 +1,501 @@
"""Anthropic workload identity federation: exchanges an external OIDC identity
token for a short-lived ``sk-ant-oat01`` token via the shared RFC 7523 engine."""
import os
from collections.abc import Callable, Mapping
from itertools import chain
from types import MappingProxyType
from typing import Final, NoReturn, TypeVar
from urllib.parse import urlsplit, urlunsplit
from pydantic import BaseModel, ConfigDict, ValidationError
from typing_extensions import assert_never
import litellm
from litellm.llms.base_llm.auth.client_credentials import keycloak_assertion_source
from litellm.llms.base_llm.auth.identity_source import (
AnthropicIdentitySourceKind,
InternalIssuerSource,
KeycloakSource,
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"
_DEFAULT_API_BASE: Final = "https://api.anthropic.com"
_INLINE_ENV_VAR: Final = "ANTHROPIC_IDENTITY_TOKEN"
_DISABLE_WIF_PARAM: Final = "anthropic_disable_workload_identity_federation"
_ACCEPTED_REF_PREFIX: Final = "oidc/"
_CHAT_BASE_SUFFIXES: Final = ("/v1/messages", "/v1")
# Hosts a federated exchange may talk to. api_base decides where the workload's assertion is sent
# AND where the minted org-scoped token is presented, so anyone able to write api_base on a
# federated deployment could otherwise redirect both. Gating each write path does not terminate:
# a deployment, a referenced credential and a future endpoint all reach the same value. This is the
# one place a federated exchange is built, so the trust decision is enforced here instead, and the
# allowlist is server-owned -- read from the environment, never from a model or credential API.
_TRUSTED_EXCHANGE_HOSTS_ENV: Final = "LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS"
_DEFAULT_TRUSTED_EXCHANGE_HOST: Final = "api.anthropic.com"
_REJECTED_REF_PREFIX: Final = "oidc/env_path/"
_IDENTITY_SOURCE_PARAM: Final = "anthropic_identity_source"
_IDENTITY_SOURCE_ENV: Final = "ANTHROPIC_IDENTITY_SOURCE"
_IDENTITY_TOKEN_FILE_PARAM: Final = "anthropic_identity_token_file"
_IDENTITY_TOKEN_PARAM: Final = "anthropic_identity_token"
# litellm_params key -> InternalIssuerSource/KeycloakSource field name. Every key here must
# also be listed in ANTHROPIC_WIF_KWARGS_KEYS (get_litellm_params.py), which is what makes it
# request-banned and cleared on a client-redirected api_base -- see types/utils.py's
# anthropic_wif_litellm_params, derived from that same set.
_INTERNAL_ISSUER_FIELD_MAP: Final[Mapping[str, str]] = MappingProxyType(
{
"anthropic_issuer_url": "issuer_url",
"anthropic_issuer_subject": "subject",
"anthropic_issuer_audience": "audience",
"anthropic_issuer_ttl_seconds": "ttl_seconds",
"anthropic_issuer_signing_key_ref": "signing_key_ref",
}
)
_KEYCLOAK_FIELD_MAP: Final[Mapping[str, str]] = MappingProxyType(
{
"anthropic_keycloak_token_url": "token_url",
"anthropic_keycloak_client_id": "client_id",
"anthropic_keycloak_auth_method": "auth_method",
"anthropic_keycloak_client_secret_ref": "client_secret_ref",
"anthropic_keycloak_scope": "scope",
}
)
_DENIAL_HINT: Final = (
" Anthropic answers every denied exchange with the same 401; the reason (for example"
" workspace_id_required or jti_reused) is only shown in the Claude Console under"
" Settings > Workload identity, in the rule's authentication history."
)
_WORKSPACE_HINT: Final = (
" If the federation rule is enabled in more than one workspace, set anthropic_workspace_id"
" (or ANTHROPIC_WORKSPACE_ID) to the wrkspc_ id of the workspace to mint tokens for, or to 'default'."
)
_SERVICE_ACCOUNT_HINT: Final = (
" Anthropic's reference lists service_account_id as required: set anthropic_service_account_id"
" (or ANTHROPIC_SERVICE_ACCOUNT_ID) to the svac_ id the federation rule targets."
)
_MISSING_IDS_HINT: Final = (
" Copy them from the federation rule's detail page under Settings > Workload identity in the"
" Claude Console, or set ANTHROPIC_FEDERATION_RULE_ID and ANTHROPIC_ORGANIZATION_ID."
)
_ALLOWLIST_HINT: Final = (
" Identity token files must sit under an allowed credential directory"
" (/var/run/secrets or /run/secrets by default);"
" set LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS to extend the allowlist."
)
_EMPTY_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
_IdentitySourceVariant = TypeVar("_IdentitySourceVariant", bound="InternalIssuerSource | KeycloakSource")
class AnthropicWifParams(BaseModel):
model_config = ConfigDict(frozen=True)
federation_rule_id: str
organization_id: str
service_account_id: str | None = None
workspace_id: str | None = None
assertion_ref: str
assertion_source: Callable[[], str | None] | None = None
def resolve_anthropic_wif_params(litellm_params: Mapping[str, object] | None) -> AnthropicWifParams | None:
if litellm_params is not None and litellm_params.get(_DISABLE_WIF_PARAM) is True:
return None
federation_rule_id: Final = _config_value(
litellm_params, "anthropic_federation_rule_id", "ANTHROPIC_FEDERATION_RULE_ID"
)
organization_id: Final = _config_value(litellm_params, "anthropic_organization_id", "ANTHROPIC_ORGANIZATION_ID")
if federation_rule_id is None or organization_id is None:
_raise_if_identity_source_configured(litellm_params, federation_rule_id, organization_id)
return None
identity_source: Final = _resolve_identity_source(litellm_params)
if identity_source is None:
return None
assertion_ref, assertion_source = identity_source
return AnthropicWifParams(
federation_rule_id=federation_rule_id,
organization_id=organization_id,
service_account_id=_config_value(
litellm_params, "anthropic_service_account_id", "ANTHROPIC_SERVICE_ACCOUNT_ID"
),
workspace_id=_config_value(litellm_params, "anthropic_workspace_id", "ANTHROPIC_WORKSPACE_ID"),
assertion_ref=assertion_ref,
assertion_source=assertion_source,
)
def _resolve_identity_source(
litellm_params: Mapping[str, object] | None,
) -> tuple[str, Callable[[], str] | None] | None:
"""Dispatches on ``anthropic_identity_source``. Absent (the default) keeps today's
token_file/env resolution byte-identical, with no ``assertion_source`` closure -- the engine
falls back to its own reader exactly as it does today. A recognized kind builds the matching
frozen config, hashes it into the ``oidc/<kind>/<hash>`` cache-key ref (``identity_source_ref``),
and closes the source's fetch/mint function over it. An unset-but-invalid config (unknown
kind, a missing required field, or a field from the other variant) fails closed here rather
than silently falling back to token_file. A deployment whose params carry a legacy token or
token_file ref stays on legacy resolution even when ``ANTHROPIC_IDENTITY_SOURCE`` names a
fleet-wide kind: the env kind only governs deployments that set no identity params of their own."""
source_kind: Final = _resolve_source_kind(litellm_params)
if source_kind is None:
legacy_ref: Final = _resolve_assertion_ref(litellm_params)
return (legacy_ref, None) if legacy_ref is not None else None
params: Final[Mapping[str, object]] = MappingProxyType(
{key: value for key, value in (litellm_params or _EMPTY_PARAMS).items() if value is not None}
)
match source_kind:
case AnthropicIdentitySourceKind.internal_issuer.value:
_reject_foreign_variant_fields(params, foreign_field_map=_KEYCLOAK_FIELD_MAP, chosen_kind=source_kind)
issuer_config: Final = _build_variant(InternalIssuerSource, params, _INTERNAL_ISSUER_FIELD_MAP)
return identity_source_ref(issuer_config), internal_issuer_assertion_source(issuer_config)
case AnthropicIdentitySourceKind.keycloak.value:
_reject_foreign_variant_fields(
params, foreign_field_map=_INTERNAL_ISSUER_FIELD_MAP, chosen_kind=source_kind
)
keycloak_config: Final = _build_variant(KeycloakSource, params, _KEYCLOAK_FIELD_MAP)
return identity_source_ref(keycloak_config), keycloak_assertion_source(keycloak_config)
case _:
_raise_unknown_source_kind(source_kind)
def _raise_unknown_source_kind(source_kind: str) -> NoReturn:
raise litellm.AuthenticationError(
message=(
f"{_IDENTITY_SOURCE_PARAM} must be one of "
f"{', '.join(kind.value for kind in AnthropicIdentitySourceKind)}; got {source_kind!r}."
),
llm_provider="anthropic",
model="",
)
def _raise_if_identity_source_configured(
litellm_params: Mapping[str, object] | None, federation_rule_id: str | None, organization_id: str | None
) -> None:
"""A configured identity source is an explicit request to federate, so a missing rule or
organization id fails closed with the ids named, rather than silently skipping federation
and surfacing later as a missing API key."""
source_kind: Final = _resolve_source_kind(litellm_params)
if source_kind is None:
return
if source_kind not in {kind.value for kind in AnthropicIdentitySourceKind}:
_raise_unknown_source_kind(source_kind)
missing: Final = tuple(
param
for param, value in (
("anthropic_federation_rule_id", federation_rule_id),
("anthropic_organization_id", organization_id),
)
if value is None
)
raise litellm.AuthenticationError(
message=(
f"{_IDENTITY_SOURCE_PARAM} is {source_kind!r}, but {' and '.join(missing)} "
f"{'is' if len(missing) == 1 else 'are'} not set.{_MISSING_IDS_HINT}"
),
llm_provider="anthropic",
model="",
)
def _resolve_source_kind(litellm_params: Mapping[str, object] | None) -> str | None:
param_kind: Final = _param_str(litellm_params, _IDENTITY_SOURCE_PARAM)
if param_kind is not None:
return param_kind
has_param_legacy_ref: Final = any(
_param_str(litellm_params, key) is not None for key in (_IDENTITY_TOKEN_FILE_PARAM, _IDENTITY_TOKEN_PARAM)
)
return None if has_param_legacy_ref else _env_str(_IDENTITY_SOURCE_ENV)
def _reject_foreign_variant_fields(
litellm_params: Mapping[str, object], foreign_field_map: Mapping[str, str], chosen_kind: str
) -> None:
foreign_keys_present: Final = tuple(param for param in foreign_field_map if param in litellm_params)
if foreign_keys_present:
raise litellm.AuthenticationError(
message=(
f"{_IDENTITY_SOURCE_PARAM} is {chosen_kind!r}, but {', '.join(sorted(foreign_keys_present))} "
"belongs to a different identity source and cannot be set alongside it."
),
llm_provider="anthropic",
model="",
)
def _build_variant(
model: type[_IdentitySourceVariant],
litellm_params: Mapping[str, object],
field_map: Mapping[str, str],
) -> _IdentitySourceVariant:
fields: Final = MappingProxyType(
{field_map[key]: value for key, value in litellm_params.items() if key in field_map}
)
try:
return model.model_validate(fields)
except ValidationError as e:
# hide_input_in_errors=True on both variant models keeps a secret pasted into the
# wrong field (e.g. a client_secret typed as signing_key_ref) out of str(e).
raise litellm.AuthenticationError(
message=f"Invalid {_IDENTITY_SOURCE_PARAM} configuration: {e}",
llm_provider="anthropic",
model="",
) 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,
) -> 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))
return _token_from_result(result, model, params)
async def aget_anthropic_wif_token(
litellm_params: Mapping[str, object] | None,
api_base: str | None,
model: str,
engine: JwtBearerTokenExchangeEngine = default_token_exchange_engine,
) -> 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))
return _token_from_result(result, model, params)
def _token_from_result(result: ExchangeResult, model: str, params: AnthropicWifParams) -> str:
match result:
case MintedToken():
return result.access_token.get_secret_value()
case _:
_raise_anthropic_wif_error(
result,
model=model,
workspace_id_set=params.workspace_id is not None,
service_account_id_set=params.service_account_id is not None,
)
def resolve_anthropic_base(api_base: str | None) -> str:
"""The base every Anthropic tier derives its URLs from: the deployment api_base when set,
else ``ANTHROPIC_API_BASE`` / ``ANTHROPIC_BASE_URL``, else Anthropic's host, with trailing
slashes and chat-appended ``/v1/messages`` suffixes stripped, so the token URL, the cache key
and the count-tokens URL all agree for the same deployment."""
return anthropic_base_without_chat_suffix(api_base or _resolve_default_api_base())
def _trusted_exchange_hosts() -> frozenset[str]:
"""Hostnames a federated exchange may reach: Anthropic's own, plus whatever the operator put in
the environment. Comma separated, case folded, entries given as a URL reduced to their host."""
configured: Final = os.getenv(_TRUSTED_EXCHANGE_HOSTS_ENV) or ""
extra: Final = (entry.strip() for entry in configured.split(",") if entry.strip())
return frozenset(
chain(
(_DEFAULT_TRUSTED_EXCHANGE_HOST,),
((urlsplit(entry).hostname or entry.split("/")[0]).lower() for entry in extra),
)
)
def _raise_if_exchange_host_untrusted(exchange_base: str, model: str) -> None:
"""The federated exchange refuses any host the operator has not vouched for, whatever wrote the
deployment's api_base. Exact hostname match, never a substring: ``api.anthropic.com.evil.test``
contains the real host and must not pass."""
host: Final = (urlsplit(exchange_base).hostname or "").lower()
if host and host in _trusted_exchange_hosts():
return
raise litellm.AuthenticationError(
message=(
f"Anthropic workload identity federation refused to use host {host or exchange_base!r}. "
f"A federated exchange sends the workload's identity token to this host and presents the "
f"minted token to it, so only {_DEFAULT_TRUSTED_EXCHANGE_HOST} is trusted by default. To "
f"use a private Anthropic-compatible gateway, add its hostname to the "
f"{_TRUSTED_EXCHANGE_HOSTS_ENV} environment variable (comma separated); that is a "
f"decision to trust it with org-scoped credentials, so it is deliberately server-owned "
f"and cannot be set through the model or credential APIs."
),
llm_provider="anthropic",
model=model,
)
def _resolve_default_api_base() -> str:
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
return AnthropicModelInfo.get_api_base(None) or _DEFAULT_API_BASE
def anthropic_base_without_chat_suffix(base: str) -> str:
"""A deployment base with its chat-surface suffix removed, so the token URL and model
discovery both derive from the same value whatever form the operator configured."""
parts: Final = urlsplit(base)
if not parts.scheme or not parts.netloc:
return base.rstrip("/")
return urlunsplit((parts.scheme, parts.netloc, _strip_path_suffixes(parts.path), "", ""))
def _strip_path_suffixes(path: str) -> str:
"""Drop the chat-surface suffixes a deployment base may carry, so every tier derives the same
token URL. Each pass removes at most one suffix, so the loop is bounded by the segment count."""
trimmed = path.rstrip("/") # rebind-ok: fixed-point strip, one suffix per pass
while True:
shortened = next( # rebind-ok: one suffix removed per iteration
(trimmed.removesuffix(suffix) for suffix in _CHAT_BASE_SUFFIXES if trimmed.endswith(suffix)),
trimmed,
)
if shortened == trimmed:
return trimmed
# Re-strip: a doubled suffix leaves a trailing slash that would stop the next match.
trimmed = shortened.rstrip("/")
def _config_value(litellm_params: Mapping[str, object] | None, param_key: str, env_name: str) -> str | None:
return _param_str(litellm_params, param_key) or _env_str(env_name)
def _param_str(litellm_params: Mapping[str, object] | None, key: str) -> str | None:
if litellm_params is None:
return None
value: Final = litellm_params.get(key)
return value if isinstance(value, str) and value else None
def _env_str(name: str) -> str | None:
from litellm.secret_managers.main import get_secret_str
value: Final = get_secret_str(name)
return value if isinstance(value, str) and value else None
def _resolve_assertion_ref(litellm_params: Mapping[str, object] | None) -> str | None:
file_param: Final = _param_str(litellm_params, _IDENTITY_TOKEN_FILE_PARAM)
if file_param is not None:
return f"oidc/file/{file_param}"
inline_param: Final = _param_str(litellm_params, _IDENTITY_TOKEN_PARAM)
if inline_param is not None:
return _validated_inline_ref(inline_param)
file_env: Final = _env_str("ANTHROPIC_IDENTITY_TOKEN_FILE")
if file_env is not None:
return f"oidc/file/{file_env}"
if _env_str(_INLINE_ENV_VAR) is not None:
return f"oidc/env/{_INLINE_ENV_VAR}"
return None
def _validated_inline_ref(value: str) -> str:
if value.startswith(_ACCEPTED_REF_PREFIX) and not value.startswith(_REJECTED_REF_PREFIX):
return value
raise litellm.AuthenticationError(
message=(
"anthropic_identity_token must be an oidc/ secret reference such as oidc/env/VAR_NAME,"
" oidc/file//absolute/path, oidc/github/<audience>, or oidc/google/<audience>."
" Raw identity tokens and oidc/env_path/ references are not accepted;"
" to pass a token directly, export it and reference it as oidc/env/VAR_NAME"
),
llm_provider="anthropic",
model="",
)
def _raise_anthropic_wif_error(
error: ExchangeError, model: str, workspace_id_set: bool, service_account_id_set: bool
) -> NoReturn:
detail: Final = _error_detail(
error, workspace_id_set=workspace_id_set, service_account_id_set=service_account_id_set
)
raise litellm.AuthenticationError(
message=f"Anthropic workload identity federation failed. {detail}",
llm_provider="anthropic",
model=model,
)
def _denial_hints(workspace_id_set: bool, service_account_id_set: bool) -> str:
return "".join(
(
_DENIAL_HINT,
"" if workspace_id_set else _WORKSPACE_HINT,
"" if service_account_id_set else _SERVICE_ACCOUNT_HINT,
)
)
def _error_detail(error: ExchangeError, workspace_id_set: bool, service_account_id_set: bool) -> str:
match error:
case AssertionSourceError() if error.kind == "disallowed_path":
return f"Could not read the OIDC identity token from {error.source_ref}.{_ALLOWLIST_HINT}"
case AssertionSourceError():
base: Final = f"Could not obtain the OIDC identity token ({error.kind}) from {error.source_ref}."
return f"{base} {error.detail}" if error.detail else base
case InsecureTokenUrl():
return f"The token endpoint must use https; refusing to send the identity token to host {error.host!r}."
case TokenEndpointError() if error.status_code == 401:
hints: Final = _denial_hints(workspace_id_set, service_account_id_set)
return f"The token endpoint returned HTTP 401: {error.redacted_body}{hints}"
case TokenEndpointError():
return f"The token endpoint returned HTTP {error.status_code}: {error.redacted_body}"
case TokenTransportError():
return f"Could not reach the token endpoint: {error.detail}."
case MalformedTokenResponse():
return f"The token endpoint returned an unusable response: {error.detail}."
case _:
assert_never(error)

View file

@ -1,3 +1,4 @@
from collections.abc import Mapping
from typing import Final
from urllib.parse import urlsplit, urlunsplit
@ -218,6 +219,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
aembedding=None,
max_retries: int | None = None,
shared_session=None,
litellm_params: Mapping[str, object] | None = None,
) -> EmbeddingResponse:
"""
- Separate image url from text

View file

@ -41,6 +41,29 @@ class BaseAnthropicMessagesConfig(ABC):
"""
return headers, api_base
async def avalidate_anthropic_messages_environment(
self,
headers: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
model: str,
messages: list[Any], # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
optional_params: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
litellm_params: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict, str | None]: # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
"""Async counterpart used by the async handler. The default delegates to the
sync implementation; providers whose sync path can block the event loop
(e.g. a WIF token exchange) override this."""
return self.validate_anthropic_messages_environment(
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
api_key=api_key,
api_base=api_base,
)
@abstractmethod
def get_complete_url(
self,

View file

@ -0,0 +1,99 @@
from litellm.llms.base_llm.auth.client_credentials import (
SecretReader,
fetch_keycloak_assertion,
keycloak_assertion_source,
)
from litellm.llms.base_llm.auth.identity_source import (
AnthropicIdentitySourceConfig,
AnthropicIdentitySourceKind,
InternalIssuerSource,
KeycloakSource,
identity_source_config_adapter,
identity_source_ref,
)
from litellm.llms.base_llm.auth.internal_issuer import (
SigningKeyReader,
internal_issuer_assertion_source,
internal_issuer_jwks_document,
mint_internal_issuer_assertion,
)
from litellm.llms.base_llm.auth.jwt_signing import (
ALG,
build_jwk,
build_jwks,
jwks_document_json,
load_es256_private_key,
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,
MAX_RESPONSE_BYTES,
JwtBearerTokenExchangeEngine,
default_token_exchange_engine,
redact_oauth_error_body,
validate_token_endpoint_url,
)
from litellm.llms.base_llm.auth.types import (
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",
"internal_issuer_assertion_source",
"internal_issuer_jwks_document",
"jwks_document_json",
"keycloak_assertion_source",
"load_es256_private_key",
"mint_internal_issuer_assertion",
"redact_oauth_error_body",
"rfc7638_thumbprint",
"sign_es256_jwt",
"validate_token_endpoint_url",
)

View file

@ -0,0 +1,227 @@
"""Fetches a fresh RFC 6749 client_credentials assertion for Anthropic's ``keycloak`` identity
source: LiteLLM authenticates to Keycloak as its own confidential client and presents the
resulting ``access_token`` as the workload assertion (Phase 1 decision 2).
The client secret is the operator-supplied pointer at ``KeycloakSource.client_secret_ref``,
resolved the same way every other WIF secret pointer already is (env, a Credential, or whatever
secret manager ``litellm.secret_manager_client`` is globally configured to, Vault included).
Every fetch is a fresh HTTP POST; nothing here caches a fetched token, since the outer
token-exchange engine already caches the Anthropic token it buys with one -- see decision 2's
"no Keycloak-side cache" ruling.
"""
import base64
import threading
from collections.abc import Callable, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeAlias
from urllib.parse import quote, quote_plus, urlencode
import httpx
from pydantic import BaseModel, SecretStr, ValidationError
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 (
MAX_RESPONSE_BYTES,
endpoint_url_for_error_message,
redact_oauth_error_body,
require_posted_response,
validate_token_endpoint_url,
)
from litellm.llms.base_llm.auth.types import InsecureTokenUrl, SyncTokenPoster
if TYPE_CHECKING:
from litellm.llms.custom_httpx.http_handler import HTTPHandler
SecretReader: TypeAlias = Callable[[str], str | None] # mutable-ok: Callable param-list syntax, not a list
_GRANT_TYPE: Final = "client_credentials"
_TIMEOUT_SECONDS: Final = 30.0
_FORM_CONTENT_TYPE: Final = "application/x-www-form-urlencoded"
class _ClientCredentialsResponse(BaseModel):
access_token: str
def _default_secret_reader(ref: str) -> str | None:
from litellm.secret_managers.main import get_secret_str
return get_secret_str(ref)
def _new_keycloak_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
class _HttpxSyncKeycloakPoster:
"""Dedicated HTTPHandler for the Keycloak token POST: no ``logging_obj`` (so litellm's
request/response logging never sees the client secret or the fetched token), redirects
disabled. A separate instance from the outer engine's own poster, since this is a genuinely
new HTTP call site whose no-logging guarantee must be built here, not assumed inherited."""
def __init__(self, handler_factory: Callable[[], "HTTPHandler"] = _new_keycloak_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, "keycloak token endpoint")
_DEFAULT_POSTER: Final[SyncTokenPoster] = _HttpxSyncKeycloakPoster()
def _form_encode(value: str) -> str:
"""RFC 6749 Appendix B before RFC 6749 2.3.1's base64: application/x-www-form-urlencoded
with spaces as ``%20`` rather than ``+``, else a reserved character (":", "+", "%", " ") in
the id or secret corrupts the credential the far side decodes back out of Basic auth."""
return quote(value, safe="")
def _basic_auth_header(client_id: str, client_secret: str) -> str:
encoded_pair: Final = f"{_form_encode(client_id)}:{_form_encode(client_secret)}"
return "Basic " + base64.b64encode(encoded_pair.encode()).decode("ascii")
def _prepared_request(config: KeycloakSource, client_secret: str) -> tuple[bytes, Mapping[str, str]]:
scope_field: Final[Mapping[str, str]] = (
MappingProxyType({"scope": config.scope}) if config.scope else MappingProxyType({})
)
match config.auth_method:
case "client_secret_basic":
return (
urlencode(MappingProxyType({"grant_type": _GRANT_TYPE, **scope_field})).encode(),
MappingProxyType(
{
"content-type": _FORM_CONTENT_TYPE,
"authorization": _basic_auth_header(config.client_id, client_secret),
}
),
)
case "client_secret_post":
return (
urlencode(
MappingProxyType(
{
"grant_type": _GRANT_TYPE,
"client_id": config.client_id,
"client_secret": client_secret,
**scope_field,
}
)
).encode(),
MappingProxyType({"content-type": _FORM_CONTENT_TYPE}),
)
case _:
assert_never(config.auth_method)
def _resolve_client_secret(config: KeycloakSource, secret_reader: SecretReader) -> str:
secret: Final = secret_reader(config.client_secret_ref)
if not secret:
raise ValueError(f"keycloak client secret {ref_for_error_message(config.client_secret_ref)} could not be read")
return secret
def _wire_forms_of_secret(config: KeycloakSource, client_secret: str) -> tuple[SecretStr, ...]:
"""Every shape the secret leaves this process in, so an echo of any of them is caught.
Neither grant sends the secret verbatim. client_secret_basic base64s ``id:secret``, which
decodes straight back to it, and client_secret_post percent-escapes it. An endpoint echoing
either shape hands over reversible material a raw comparison would miss.
"""
raw: Final = SecretStr(client_secret)
match config.auth_method:
case "client_secret_basic":
encoded_pair: Final = f"{_form_encode(config.client_id)}:{_form_encode(client_secret)}"
return (raw, SecretStr(base64.b64encode(encoded_pair.encode()).decode("ascii")))
case "client_secret_post":
# urlencode escapes reserved characters and writes a space as "+", so a secret
# containing either leaves in a shape the raw comparison would not recognise coming
# back. quote_plus is what urlencode itself applies.
return (raw, SecretStr(quote_plus(client_secret)))
case _:
assert_never(config.auth_method)
def _endpoint_error_message(config: KeycloakSource, response: httpx.Response, client_secret: str) -> str:
endpoint_error: Final = redact_oauth_error_body(
response.status_code, response.text, _wire_forms_of_secret(config, client_secret)
)
return (
f"keycloak token endpoint {endpoint_url_for_error_message(config.token_url)} "
f"returned HTTP {endpoint_error.status_code}: {endpoint_error.redacted_body}"
)
def _parse_success_body(response: httpx.Response) -> str:
if len(response.content) > MAX_RESPONSE_BYTES:
raise ValueError("keycloak token response exceeded the size cap")
try:
parsed: Final = _ClientCredentialsResponse.model_validate_json(response.content)
except ValidationError as e:
raise ValueError("keycloak token response failed schema validation") from e
token: Final = parsed.access_token.strip()
if not token:
raise ValueError("keycloak token response carried an empty access_token")
return token
def fetch_keycloak_assertion(
config: KeycloakSource,
*,
poster: SyncTokenPoster = _DEFAULT_POSTER,
secret_reader: SecretReader = _default_secret_reader,
) -> str:
"""POSTs one fresh client_credentials grant and returns the resulting ``access_token`` as the
workload assertion; the caller must not cache the result -- see the module docstring."""
match validate_token_endpoint_url(config.token_url):
case InsecureTokenUrl(host=host):
raise ValueError(f"keycloak token_url must use https; refusing to send the client secret to host {host!r}")
case _:
pass
client_secret: Final = _resolve_client_secret(config, secret_reader)
content, headers = _prepared_request(config, client_secret)
try:
response: Final = poster.post(config.token_url, content=content, headers=headers, timeout=_TIMEOUT_SECONDS)
except Exception as e: # noqa: BLE001 # injected posters may raise beyond httpx; every failure becomes a ValueError
raise ValueError(
f"could not reach the keycloak token endpoint {endpoint_url_for_error_message(config.token_url)}: "
f"{type(e).__name__}"
) from e
if not 200 <= response.status_code < 300:
raise ValueError(_endpoint_error_message(config, response, client_secret))
return _parse_success_body(response)
def keycloak_assertion_source(
config: KeycloakSource,
*,
poster: SyncTokenPoster = _DEFAULT_POSTER,
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)
-- 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)

View file

@ -0,0 +1,76 @@
"""Tagged-union identity-source configs for Anthropic workload identity federation, beyond the
existing token_file/env resolver in ``litellm/llms/anthropic/wif.py``.
Each variant only ever carries secret *pointer names* (``signing_key_ref``, ``client_secret_ref``),
never a resolved secret value, so ``identity_source_ref`` can safely hash a variant into the short,
content-derived ``oidc/<kind>/<hash>`` string used elsewhere as a get_secret ref, a token-exchange
cache-key discriminator, and an operator-facing error pointer.
"""
import hashlib
from enum import Enum
from typing import Annotated, Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
_REF_HASH_HEX_LENGTH: Final = 16
_MAX_TTL_SECONDS: Final = 3600
_DEFAULT_TTL_SECONDS: Final = 300
class AnthropicIdentitySourceKind(str, Enum):
internal_issuer = "internal_issuer"
keycloak = "keycloak"
class InternalIssuerSource(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", hide_input_in_errors=True)
kind: Literal[AnthropicIdentitySourceKind.internal_issuer] = AnthropicIdentitySourceKind.internal_issuer
issuer_url: str
subject: str
audience: str | None = None
ttl_seconds: Annotated[int, Field(gt=0, le=_MAX_TTL_SECONDS)] = _DEFAULT_TTL_SECONDS
signing_key_ref: str
class KeycloakSource(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", hide_input_in_errors=True)
kind: Literal[AnthropicIdentitySourceKind.keycloak] = AnthropicIdentitySourceKind.keycloak
token_url: str
client_id: str
auth_method: Literal["client_secret_basic", "client_secret_post"] = "client_secret_basic"
client_secret_ref: str
scope: str | None = None
AnthropicIdentitySourceConfig: TypeAlias = Annotated[InternalIssuerSource | KeycloakSource, Field(discriminator="kind")]
identity_source_config_adapter: Final = TypeAdapter[AnthropicIdentitySourceConfig](AnthropicIdentitySourceConfig)
def identity_source_ref(config: AnthropicIdentitySourceConfig) -> str:
"""``oidc/<kind>/<hash>``: a short, secret-free pointer, stable for identical config and rolling
whenever any field does, including a ``*_ref`` pointer NAME (never the secret it points to)."""
digest: Final = hashlib.sha256(config.model_dump_json().encode()).hexdigest()[:_REF_HASH_HEX_LENGTH]
return f"oidc/{config.kind.value}/{digest}"
_POINTER_REF_PREFIXES: Final = (
"oidc/",
"os.environ/",
"hashicorp_vault/",
"aws_secret_manager/",
"google_secret_manager/",
)
def ref_for_error_message(ref: str) -> str:
"""A ``*_ref`` rendered for an operator-facing error.
Naming the pointer is deliberate: it is what tells an operator which setting failed to
resolve. But these fields only ever fail to resolve when what was written is not a pointer,
and an operator who pasted the secret itself has made the field's value the secret. So the
value is echoed only when it is recognizably a pointer, and withheld otherwise.
"""
return ref if ref.startswith(_POINTER_REF_PREFIXES) else "<withheld: not a secret reference>"

View file

@ -0,0 +1,86 @@
"""Mints a self-issued workload assertion for Anthropic's ``internal_issuer`` identity source:
LiteLLM signs its own short-lived ES256 JWT instead of reading one from a mounted OIDC file.
Signing custody is the operator-supplied PEM at ``InternalIssuerSource.signing_key_ref``,
resolved the same way every other WIF secret pointer already is (env, a Credential, or
whatever secret manager ``litellm.secret_manager_client`` is globally configured to, Vault
included) -- see Phase 1 decision 1. Every mint is fresh; nothing here caches a minted JWT,
since the outer token-exchange engine already caches the Anthropic token it buys with one.
"""
import time
import uuid
from collections.abc import Callable, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias
from litellm.llms.base_llm.auth.identity_source import InternalIssuerSource, ref_for_error_message
from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json, sign_es256_jwt
SigningKeyReader: TypeAlias = Callable[[str], str | None] # mutable-ok: Callable param-list syntax, not a list
def _default_signing_key_reader(ref: str) -> str | None:
from litellm.secret_managers.main import get_secret_str
return get_secret_str(ref)
def _claims(config: InternalIssuerSource, issued_at: int) -> Mapping[str, object]:
return MappingProxyType(
{
key: value
for key, value in (
("sub", config.subject),
("iss", config.issuer_url),
("aud", config.audience),
("iat", issued_at),
("exp", issued_at + config.ttl_seconds),
("jti", str(uuid.uuid4())),
)
if value is not None
}
)
def _resolve_signing_key(config: InternalIssuerSource, key_reader: SigningKeyReader) -> str:
pem: Final = key_reader(config.signing_key_ref)
if not pem:
raise ValueError(
f"internal_issuer signing key {ref_for_error_message(config.signing_key_ref)} could not be read"
)
return pem
def mint_internal_issuer_assertion(
config: InternalIssuerSource,
*,
key_reader: SigningKeyReader = _default_signing_key_reader,
clock: Callable[[], float] = time.time,
) -> str:
"""Signs one fresh, short-lived assertion; the caller must not cache the result, since a
cached copy would defeat the point of re-minting on every exchange."""
pem: Final = _resolve_signing_key(config, key_reader)
return sign_es256_jwt(pem, _claims(config, issued_at=int(clock())))
def internal_issuer_assertion_source(
config: InternalIssuerSource,
*,
key_reader: SigningKeyReader = _default_signing_key_reader,
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)
-- 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)
def internal_issuer_jwks_document(
config: InternalIssuerSource,
*,
key_reader: SigningKeyReader = _default_signing_key_reader,
) -> str:
"""The operator-facing JWKS export, resolved from a configured identity source rather than
a raw PEM in hand -- the JSON document to register as Anthropic's inline federation issuer."""
return jwks_document_json(_resolve_signing_key(config, key_reader))

View file

@ -0,0 +1,115 @@
"""ES256 JWT signing primitives for Anthropic workload identity federation's
``internal_issuer`` identity source (see ``identity_source.InternalIssuerSource``).
Pure functions over an already-resolved PEM string: no I/O, no secret-manager awareness, no
caching. Given the signing key at, say, $ISSUER_SIGNING_KEY_PEM, an operator publishes the
JWKS document Anthropic's inline federation issuer needs with one line:
python -c "from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json; \\
import os; print(jwks_document_json(os.environ['ISSUER_SIGNING_KEY_PEM']))"
"""
import base64
import hashlib
import json
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeAlias
if TYPE_CHECKING:
from cryptography.hazmat.primitives.asymmetric import ec
ALG: Final = "ES256"
MISSING_SIGNING_DEPENDENCIES_MESSAGE: Final = (
"the internal_issuer identity source needs PyJWT and cryptography, which a base litellm install "
"does not include: pip install 'litellm[proxy]'"
)
_JWK_CURVE_NAME: Final = "P-256"
_JWK_KEY_TYPE: Final = "EC"
_COORDINATE_BYTE_LENGTH: Final = 32 # P-256 field element width, RFC 7518 6.2.1.2/6.2.1.3
Jwk: TypeAlias = Mapping[str, str]
Jwks: TypeAlias = Mapping[str, tuple[Jwk, ...]]
def load_es256_private_key(pem: str) -> "ec.EllipticCurvePrivateKey":
"""Parses an unencrypted PEM EC private key. Never echoes the key material in an error."""
try:
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.serialization import load_pem_private_key
except ImportError as e:
raise ImportError(MISSING_SIGNING_DEPENDENCIES_MESSAGE) from e
try:
key: Final = load_pem_private_key(pem.encode(), password=None)
except (ValueError, TypeError) as e:
raise ValueError("internal_issuer signing key is not a valid unencrypted PEM private key") from e
if not isinstance(key, ec.EllipticCurvePrivateKey) or not isinstance(key.curve, ec.SECP256R1):
raise ValueError( # noqa: TRY004 # the reader classifies ValueError into a readable config error; TypeError would not
"internal_issuer signing key must be an EC P-256 (secp256r1) private key for ES256"
)
return key
def _b64url_coordinate(value: int) -> str:
return base64.urlsafe_b64encode(value.to_bytes(_COORDINATE_BYTE_LENGTH, "big")).rstrip(b"=").decode("ascii")
def _jwk_thumbprint_members(public_key: "ec.EllipticCurvePublicKey") -> Jwk:
"""RFC 7638 3.2's exact EC member set (crv, kty, x, y) and nothing else: an extra member
here would change the thumbprint and desync it from the ``kid`` published in the JWKS."""
numbers: Final = public_key.public_numbers()
return MappingProxyType(
{
"crv": _JWK_CURVE_NAME,
"kty": _JWK_KEY_TYPE,
"x": _b64url_coordinate(numbers.x),
"y": _b64url_coordinate(numbers.y),
}
)
def rfc7638_thumbprint(public_key: "ec.EllipticCurvePublicKey") -> str:
"""RFC 7638: SHA-256 over the lexicographically member-ordered, whitespace-free JSON
rendering of the thumbprint members, base64url-encoded without padding."""
canonical: Final = json.dumps(
dict(sorted(_jwk_thumbprint_members(public_key).items())), # mutable-ok: json.dumps needs a real dict
separators=(",", ":"),
)
return base64.urlsafe_b64encode(hashlib.sha256(canonical.encode()).digest()).rstrip(b"=").decode("ascii")
def build_jwk(public_key: "ec.EllipticCurvePublicKey", kid: str) -> Jwk:
return MappingProxyType({**_jwk_thumbprint_members(public_key), "use": "sig", "alg": ALG, "kid": kid})
def build_jwks(public_key: "ec.EllipticCurvePublicKey") -> Jwks:
kid: Final = rfc7638_thumbprint(public_key)
return MappingProxyType({"keys": (build_jwk(public_key, kid),)})
def jwks_document_json(pem: str) -> str:
"""The operator-facing export: the JSON document to register as Anthropic's inline JWKS.
``build_jwks`` returns ``MappingProxyType``/tuple values per this repo's no-mutation
convention; the ``json`` module only knows plain ``dict``/``list``, so those are converted
at this one serialization boundary rather than giving up immutability throughout the module.
"""
key: Final = load_es256_private_key(pem)
jwks: Final = build_jwks(key.public_key())
return json.dumps(
{"keys": [dict(jwk) for jwk in jwks["keys"]]}, # mutable-ok: json.dumps needs real dicts/lists
indent=2,
)
def sign_es256_jwt(pem: str, claims: Mapping[str, object]) -> str:
"""Signs ``claims`` with the PEM key, stamping ``kid`` as its RFC 7638 thumbprint so a
verifier can look the signing key up in the published JWKS by ``kid`` alone."""
try:
import jwt
except ImportError as e:
raise ImportError(MISSING_SIGNING_DEPENDENCIES_MESSAGE) from e
key: Final = load_es256_private_key(pem)
kid: Final = rfc7638_thumbprint(key.public_key())
headers: Final = {"kid": kid} # mutable-ok: PyJWT requires a real dict, not a Mapping
return jwt.encode(dict(claims), key, algorithm=ALG, headers=headers) # mutable-ok: PyJWT requires a real dict

View file

@ -0,0 +1,861 @@
"""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 = "<redacted: response echoed the request>"
# 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()

View file

@ -0,0 +1,99 @@
"""Provider-agnostic types for the RFC 7523 JWT-bearer token exchange engine."""
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"]
source_ref: str
detail: str | None = None
@dataclass(frozen=True, slots=True)
class InsecureTokenUrl:
host: str
@dataclass(frozen=True, slots=True)
class TokenEndpointError:
status_code: int
redacted_body: str
@dataclass(frozen=True, slots=True)
class TokenTransportError:
detail: str
@dataclass(frozen=True, slots=True)
class MalformedTokenResponse:
detail: str
ExchangeError: TypeAlias = (
AssertionSourceError | InsecureTokenUrl | TokenEndpointError | TokenTransportError | MalformedTokenResponse
)
ExchangeResult: TypeAlias = MintedToken | ExchangeError
ExchangeCallType: TypeAlias = Literal["cold_mint", "mandatory_refresh", "advisory_refresh"]
class TokenExchangeMetricsSink(Protocol):
"""Observability seam for the exchange engine. 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."""
def exchange_success(self, *, call_type: ExchangeCallType, duration_seconds: float) -> None: ...
def exchange_failure(
self, *, call_type: ExchangeCallType, duration_seconds: float, error: ExchangeError
) -> None: ...
def cache_hit(self) -> None: ...
class SyncTokenPoster(Protocol):
"""Returns the response for ANY status; never raises for status."""
def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: ...

View file

@ -5,6 +5,7 @@ Utility functions for base LLM classes.
import copy
import json
from abc import ABC, abstractmethod
from collections.abc import Mapping
from typing import Any, Final
from openai.lib import _parsing, _pydantic
@ -57,6 +58,22 @@ class BaseLLMModelInfo(ABC):
"""
return []
def discover_models(
self, litellm_params: Mapping[str, object] | None = None
) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override
"""
Live model discovery for a configured deployment. Defaults to the api_key/api_base
facade every provider already implements via ``get_models``; a provider whose
discovery needs more of ``litellm_params`` (e.g. Anthropic's workload identity
federation) overrides this instead of widening ``get_models`` for every provider.
"""
api_key: Final = litellm_params.get("api_key") if litellm_params is not None else None
api_base: Final = litellm_params.get("api_base") if litellm_params is not None else None
return self.get_models(
api_key=api_key if isinstance(api_key, str) else None,
api_base=api_base if isinstance(api_base, str) else None,
)
@staticmethod
@abstractmethod
def get_api_key(api_key: str | None = None) -> str | None:

View file

@ -1289,7 +1289,7 @@ class HTTPHandler:
self,
url: str,
params: dict | None = None,
headers: dict | None = None,
headers: Mapping[str, Any] | None = None,
follow_redirects: bool | None = None,
timeout: float | httpx.Timeout | None = None,
):

View file

@ -1,11 +1,25 @@
import asyncio
import inspect
import json
import ssl
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager
from functools import lru_cache
from types import MappingProxyType, ModuleType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, Union, cast, get_type_hints
from typing import (
TYPE_CHECKING,
Any,
Final,
Literal,
Optional,
Protocol,
TypedDict,
TypeVar,
Union,
cast,
get_type_hints,
runtime_checkable,
)
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
import httpx
@ -219,6 +233,55 @@ class _MediaUploadKwargs(TypedDict, total=False):
timeout: float | httpx.Timeout
@runtime_checkable
class _AsyncFilesEnvironmentValidator(Protocol):
async def avalidate_environment(
self,
headers: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
model: str,
messages: list, # mutable-ok: mirrors the sync validate_environment contract this overrides
optional_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
litellm_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
api_key: str | None = None,
api_base: str | None = None,
) -> dict: ... # mutable-ok: mirrors the sync validate_environment contract this overrides
async def _avalidate_files_environment(
provider_config: BaseFilesConfig | BaseBatchesConfig,
*,
headers: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
model: str,
messages: list, # mutable-ok: mirrors the sync validate_environment contract this overrides
optional_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
litellm_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
api_key: str | None,
) -> dict: # mutable-ok: mirrors the sync validate_environment contract this overrides
"""Await the provider's async credential hook when it has one (e.g. Anthropic's workload
identity token exchange); otherwise offload the sync hook to a worker thread. Either way
the caller, an async file handler, never blocks the event loop on it."""
if isinstance(provider_config, _AsyncFilesEnvironmentValidator) and inspect.iscoroutinefunction(
provider_config.avalidate_environment
):
return await provider_config.avalidate_environment(
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
api_key=api_key,
)
return await asyncio.to_thread(
provider_config.validate_environment,
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
api_key=api_key,
)
def _google_genai_streaming_hidden_params(
*,
api_base: str,
@ -2130,7 +2193,7 @@ class BaseLLMHTTPHandler:
(
headers,
api_base,
) = anthropic_messages_provider_config.validate_anthropic_messages_environment(
) = await anthropic_messages_provider_config.avalidate_anthropic_messages_environment(
headers=merged_headers or {},
model=model,
messages=messages,
@ -3520,6 +3583,19 @@ class BaseLLMHTTPHandler:
"""
Creates a file using Gemini's two-step upload process
"""
if _is_async:
return self._avalidate_and_create_file(
create_file_data=create_file_data,
litellm_params=litellm_params,
provider_config=provider_config,
headers=headers,
api_base=api_base,
api_key=api_key,
logging_obj=logging_obj,
client=client,
timeout=timeout,
)
# get config from model, custom llm provider
headers = provider_config.validate_environment(
api_key=api_key,
@ -3549,18 +3625,6 @@ class BaseLLMHTTPHandler:
optional_params={},
)
if _is_async:
return self.async_create_file(
transformed_request=transformed_request,
litellm_params=litellm_params,
provider_config=provider_config,
headers=headers,
api_base=api_base,
logging_obj=logging_obj,
client=client,
timeout=timeout,
)
if client is None or not isinstance(client, HTTPHandler):
sync_httpx_client = _get_httpx_client()
else:
@ -3688,6 +3752,52 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params_with_url,
)
async def _avalidate_and_create_file(
self,
*,
create_file_data: CreateFileRequest,
litellm_params: dict, # mutable-ok: mirrors the create_file contract this dispatches for
provider_config: BaseFilesConfig,
headers: dict, # mutable-ok: mirrors the create_file contract this dispatches for
api_base: str | None,
api_key: str | None,
logging_obj: LiteLLMLoggingObj,
client: HTTPHandler | AsyncHTTPHandler | None,
timeout: float | httpx.Timeout | None,
) -> OpenAIFileObject:
validated_headers: Final = await _avalidate_files_environment(
provider_config,
headers=headers,
model="",
messages=[], # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
litellm_params=litellm_params,
api_key=api_key,
)
complete_api_base: Final = provider_config.get_complete_file_url(
api_base=api_base,
api_key=api_key,
model="",
optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
litellm_params=litellm_params,
data=create_file_data,
)
return await self.async_create_file(
transformed_request=provider_config.transform_create_file_request(
model="",
create_file_data=create_file_data,
litellm_params=litellm_params,
optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
),
litellm_params=litellm_params,
provider_config=provider_config,
headers=validated_headers,
api_base=complete_api_base,
logging_obj=logging_obj,
client=client,
timeout=timeout,
)
async def async_create_file(
self,
transformed_request: Union[bytes, str, dict, "TwoStepFileUploadConfig"],
@ -3938,6 +4048,20 @@ class BaseLLMHTTPHandler:
if model is None:
raise ValueError("model is required for create_batch")
if _is_async:
return self._avalidate_and_create_batch(
create_batch_data=create_batch_data,
litellm_params=litellm_params,
provider_config=provider_config,
headers=headers,
api_base=api_base,
api_key=api_key,
logging_obj=logging_obj,
client=client,
timeout=timeout,
model=model,
)
headers = provider_config.validate_environment(
api_key=api_key,
headers=headers,
@ -3966,19 +4090,6 @@ class BaseLLMHTTPHandler:
optional_params={},
)
if _is_async:
return self.async_create_batch(
transformed_request=transformed_request,
litellm_params=litellm_params,
provider_config=provider_config,
headers=headers,
api_base=api_base,
logging_obj=logging_obj,
client=client,
timeout=timeout,
create_batch_data=create_batch_data,
)
if client is None or not isinstance(client, HTTPHandler):
sync_httpx_client = _get_httpx_client()
else:
@ -4115,6 +4226,54 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
)
async def _avalidate_and_create_batch(
self,
*,
create_batch_data: "CreateBatchRequest",
litellm_params: dict, # mutable-ok: mirrors the create_batch contract this dispatches for
provider_config: "BaseBatchesConfig",
headers: dict, # mutable-ok: mirrors the create_batch contract this dispatches for
api_base: str | None,
api_key: str | None,
logging_obj: "LiteLLMLoggingObj",
client: Union["HTTPHandler", "AsyncHTTPHandler"] | None,
timeout: float | httpx.Timeout | None,
model: str,
) -> "LiteLLMBatch":
validated_headers: Final = await _avalidate_files_environment(
provider_config,
headers=headers,
model=model,
messages=[], # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
litellm_params=litellm_params,
api_key=api_key,
)
complete_api_base: Final = provider_config.get_complete_batch_url(
api_base=api_base,
api_key=api_key,
model=model,
optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
litellm_params=litellm_params,
data=create_batch_data,
)
return await self.async_create_batch(
transformed_request=provider_config.transform_create_batch_request(
model=model,
create_batch_data=create_batch_data,
litellm_params=litellm_params,
optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
),
litellm_params=litellm_params,
provider_config=provider_config,
headers=validated_headers,
api_base=complete_api_base,
logging_obj=logging_obj,
client=client,
timeout=timeout,
create_batch_data=create_batch_data,
)
async def async_create_batch(
self,
transformed_request: bytes | str | dict,
@ -4712,7 +4871,8 @@ class BaseLLMHTTPHandler:
)
# Validate environment and get headers
headers = provider_config.validate_environment(
headers = await _avalidate_files_environment(
provider_config,
api_key=litellm_params.get("api_key"),
headers=headers,
model="",
@ -4836,7 +4996,8 @@ class BaseLLMHTTPHandler:
)
# Validate environment and get headers
headers = provider_config.validate_environment(
headers = await _avalidate_files_environment(
provider_config,
api_key=litellm_params.get("api_key"),
headers=headers,
model="",
@ -4960,7 +5121,8 @@ class BaseLLMHTTPHandler:
)
# Validate environment and get headers
headers = provider_config.validate_environment(
headers = await _avalidate_files_environment(
provider_config,
api_key=litellm_params.get("api_key"),
headers=headers,
model="",
@ -5091,7 +5253,8 @@ class BaseLLMHTTPHandler:
)
# Validate environment and get headers
headers = provider_config.validate_environment(
headers = await _avalidate_files_environment(
provider_config,
api_key=litellm_params.get("api_key"),
headers=headers,
model="",

View file

@ -54,6 +54,7 @@ from litellm.types.utils import (
from litellm.utils import convert_to_model_response_object
from ..common_utils import OpenAIError
from ..workload_identity import get_workload_identity_bearer_token, resolve_openai_workload_identity_config
if TYPE_CHECKING:
import tiktoken
@ -70,6 +71,11 @@ else:
_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
def _litellm_params_str(litellm_params: Mapping[str, object] | None, key: str) -> str | None:
value: Final = litellm_params.get(key) if litellm_params is not None else None
return value if isinstance(value, str) else None
class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
"""
Reference: https://platform.openai.com/docs/api-reference/chat/create
@ -747,28 +753,39 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
"""
Calls OpenAI's `/v1/models` endpoint and returns the list of models.
"""
if api_base is None:
api_base = "https://api.openai.com"
if api_key is None:
api_key = get_secret_str("OPENAI_API_KEY")
# Strip api_base to just the base URL (scheme + host + port)
parsed_url: Final = httpx.URL(api_base)
base_url = f"{parsed_url.scheme}://{parsed_url.host}"
if parsed_url.port:
base_url += f":{parsed_url.port}"
response: Final = litellm.module_level_client.get(
url=f"{base_url}/v1/models",
headers={"Authorization": f"Bearer {api_key}"},
return self._fetch_model_ids(
api_base=api_base, bearer_token=get_secret_str("OPENAI_API_KEY") if api_key is None else api_key
)
def discover_models(
self, litellm_params: Mapping[str, object] | None = None
) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override
if type(self) is not OpenAIGPTConfig:
return super().discover_models(litellm_params)
api_key: Final = _litellm_params_str(litellm_params, "api_key")
api_base: Final = _litellm_params_str(litellm_params, "api_base")
workload_identity_config: Final = resolve_openai_workload_identity_config(
api_key=api_key, api_base=api_base, litellm_params=litellm_params
)
if workload_identity_config is None:
return self.get_models(api_key=api_key, api_base=api_base)
return self._fetch_model_ids(
api_base=api_base, bearer_token=get_workload_identity_bearer_token(workload_identity_config)
)
@staticmethod
def _fetch_model_ids(
api_base: str | None, bearer_token: str | None
) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override
parsed_url: Final = httpx.URL(api_base or "https://api.openai.com")
port_suffix: Final = f":{parsed_url.port}" if parsed_url.port else ""
response: Final = litellm.module_level_client.get(
url=f"{parsed_url.scheme}://{parsed_url.host}{port_suffix}/v1/models",
headers={"Authorization": f"Bearer {bearer_token}"},
)
if response.status_code != 200:
raise Exception(f"Failed to get models: {response.text}")
models: Final = response.json()["data"]
return [model["id"] for model in models]
return [model["id"] for model in response.json()["data"]]
@staticmethod
def get_api_key(api_key: str | None = None) -> str | None:

View file

@ -382,8 +382,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
organization: str | None = None,
client: OpenAI | AsyncOpenAI | None = None,
shared_session: Optional["ClientSession"] = None,
litellm_params: Mapping[str, object] | None = None,
) -> OpenAI | AsyncOpenAI | None:
workload_identity_config: Final = resolve_openai_workload_identity_config(api_key=api_key, api_base=api_base)
workload_identity_config: Final = resolve_openai_workload_identity_config(
api_key=api_key, api_base=api_base, litellm_params=litellm_params
)
client_initialization_params: Final[dict] = locals()
if client is None:
if not isinstance(max_retries, int):
@ -773,6 +776,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
max_retries=max_retries,
organization=organization,
stream_options=stream_options,
litellm_params=litellm_params,
)
else:
if not isinstance(max_retries, int):
@ -786,6 +790,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
max_retries=max_retries,
organization=organization,
client=client,
litellm_params=litellm_params,
)
## LOGGING
@ -927,6 +932,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
organization=organization,
client=client,
shared_session=shared_session,
litellm_params=litellm_params,
)
## LOGGING
@ -1022,6 +1028,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
max_retries=None,
headers=None,
stream_options: dict | None = None,
litellm_params: Mapping[str, object] | None = None,
):
data["stream"] = True
data.update(self.get_stream_options(stream_options=stream_options, api_base=api_base))
@ -1035,6 +1042,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
max_retries=max_retries,
organization=organization,
client=client,
litellm_params=litellm_params,
)
## LOGGING
logging_obj.pre_call(
@ -1107,6 +1115,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
organization=organization,
client=client,
shared_session=shared_session,
litellm_params=litellm_params,
)
## LOGGING
logging_obj.pre_call(
@ -1241,6 +1250,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
client: AsyncOpenAI | None = None,
max_retries=None,
shared_session: Optional["ClientSession"] = None,
litellm_params: Mapping[str, object] | None = None,
):
try:
openai_aclient: Final[AsyncOpenAI] = self._get_openai_client(
@ -1251,6 +1261,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
max_retries=max_retries,
client=client,
shared_session=shared_session,
litellm_params=litellm_params,
)
raw_response: Final = await self.make_openai_embedding_request(
openai_aclient=openai_aclient,
@ -1314,6 +1325,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
aembedding=None,
max_retries: int | None = None,
shared_session: Optional["ClientSession"] = None,
litellm_params: Mapping[str, object] | None = None,
) -> EmbeddingResponse:
super().embedding()
try:
@ -1340,6 +1352,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
client=client,
max_retries=max_retries,
shared_session=shared_session,
litellm_params=litellm_params,
)
openai_client: Final[OpenAI] = self._get_openai_client(
@ -1349,6 +1362,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
timeout=timeout,
max_retries=max_retries,
client=client,
litellm_params=litellm_params,
)
## embedding CALL

View file

@ -10,6 +10,7 @@ from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.get_litellm_params import OPENAI_WIF_KWARGS_KEYS
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
_safe_convert_created_field,
)
@ -494,7 +495,11 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
headers.setdefault("Content-Type", "application/json")
workload_identity_config: Final = (
resolve_openai_workload_identity_config(api_key=api_key, api_base=litellm_params.api_base)
resolve_openai_workload_identity_config(
api_key=api_key,
api_base=litellm_params.api_base,
litellm_params=litellm_params.model_dump(include=set(OPENAI_WIF_KWARGS_KEYS)),
)
if self.custom_llm_provider is LlmProviders.OPENAI
else None
)

View file

@ -1,5 +1,6 @@
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from functools import lru_cache
from typing import TYPE_CHECKING, Final
@ -43,6 +44,7 @@ class OpenAIWorkloadIdentityConfig:
def resolve_openai_workload_identity_config(
api_key: str | None,
api_base: str | None,
litellm_params: Mapping[str, object] | None = None,
) -> OpenAIWorkloadIdentityConfig | None:
static_api_key: Final = normalize_nonempty_secret_str(api_key) or normalize_nonempty_secret_str(
get_secret_str("OPENAI_API_KEY")
@ -54,10 +56,12 @@ def resolve_openai_workload_identity_config(
)
if not _targets_openai_api(effective_api_base):
return None
identity_provider_id: Final = get_secret_str("OPENAI_IDENTITY_PROVIDER_ID")
service_account_id: Final = get_secret_str("OPENAI_SERVICE_ACCOUNT_ID")
token_file: Final = get_secret_str("OPENAI_IDENTITY_TOKEN_FILE")
if not identity_provider_id or not service_account_id or not token_file:
identity_provider_id: Final = _config_value(
litellm_params, "openai_identity_provider_id", "OPENAI_IDENTITY_PROVIDER_ID"
)
service_account_id: Final = _config_value(litellm_params, "openai_service_account_id", "OPENAI_SERVICE_ACCOUNT_ID")
token_file: Final = _config_value(litellm_params, "openai_identity_token_file", "OPENAI_IDENTITY_TOKEN_FILE")
if identity_provider_id is None or service_account_id is None or token_file is None:
return None
return OpenAIWorkloadIdentityConfig(
identity_provider_id=identity_provider_id,
@ -70,6 +74,13 @@ def get_workload_identity_bearer_token(config: OpenAIWorkloadIdentityConfig) ->
return _workload_identity_auth(config).get_token()
def _config_value(litellm_params: Mapping[str, object] | None, param_key: str, env_name: str) -> str | None:
param_value: Final = litellm_params.get(param_key) if litellm_params is not None else None
if isinstance(param_value, str) and param_value:
return param_value
return normalize_nonempty_secret_str(get_secret_str(env_name))
def _targets_openai_api(api_base: str | None) -> bool:
if api_base is None:
return True

View file

@ -6326,6 +6326,7 @@ def embedding(
aembedding=aembedding,
max_retries=max_retries,
shared_session=shared_session,
litellm_params=litellm_params_dict,
)
elif custom_llm_provider == "databricks":
api_base = api_base or litellm.api_base or get_secret("DATABRICKS_API_BASE")

View file

@ -5,7 +5,7 @@ These are the canonical credential types for the proxy. They live in the model
layer; ``litellm.types.utils`` re-exports them for backwards compatibility.
"""
from pydantic import BaseModel, model_validator
from pydantic import BaseModel, Field, model_validator
class CredentialBase(BaseModel):
@ -15,6 +15,10 @@ class CredentialBase(BaseModel):
class CredentialItem(CredentialBase):
credential_values: dict
# PATCH-only instruction naming keys to drop from the stored credential_values. It describes an
# edit rather than the credential, so it stays out of dumps: those feed config loading, the DB
# write, and the in-memory list, none of which have a place for it.
credential_values_to_delete: tuple[str, ...] | None = Field(default=None, exclude=True)
class CreateCredentialItem(CredentialBase):

View file

@ -33,7 +33,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
)
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS
from litellm.types.utils import CustomPricingLiteLLMParams
from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params
from litellm.types.utils import CustomPricingLiteLLMParams, server_owned_wif_litellm_params
def is_invalid_virtual_key_error(exception: BaseException | None) -> bool:
@ -226,6 +227,16 @@ def _allow_model_level_clientside_configurable_parameters(
# ``extra_body.aws_web_identity_token``) without re-validating, so the
# banned-key check has to descend into it the same way it descends into
# ``litellm_embedding_config``.
_SERVER_OWNED_WIF_UNCONDITIONAL_BANNED: Final[tuple[str, ...]] = server_owned_wif_litellm_params
# The Bedrock Claude Platform route reads a workspace from workspace_id or aws_workspace_id as
# well, and neither is a federation parameter, so say so rather than leaving that caller stuck.
# Re-exported from litellm.types.router, where it lives so the router can call it on a
# post-authentication merge without core importing from the proxy package.
reject_server_owned_wif_params = _reject_server_owned_wif_params
_NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body")
# Metadata containers that carry per-request configuration consumed by the
@ -379,6 +390,7 @@ def _check_banned_params(
Shared between the root-level check and the nested-config check so a
new banned param only needs to be added in one place.
"""
reject_server_owned_wif_params(body)
for param in _BANNED_REQUEST_BODY_PARAMS:
if param not in body:
continue
@ -519,6 +531,7 @@ def is_request_body_safe(request_body: dict, general_settings: dict, llm_router:
_reject_url_valued_fallback_target(target)
litellm_params: Final = _coerce_metadata_to_dict(request_body.get("litellm_params"))
if litellm_params is not None:
reject_server_owned_wif_params(litellm_params)
litellm_params_metadata: Final = _coerce_metadata_to_dict(litellm_params.get("metadata"))
if litellm_params_metadata is not None:
_check_banned_params(

View file

@ -0,0 +1,158 @@
"""Shared helper for resolving a named Credential's values server-side.
Memory first (``litellm.credential_list``, already decrypted -- matching
``CredentialAccessor.get_credential_values``), then a DB decrypt fallback for a pod whose
in-memory list has not yet picked up a credential another pod just wrote or updated.
"""
import asyncio
from collections.abc import Mapping
from itertools import chain
from types import MappingProxyType
from typing import Final
import litellm
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
from litellm.proxy.utils import PrismaClient
from litellm.repositories.credentials_repository import CredentialsRepository
from litellm.types.router import (
GenericLiteLLMParams,
server_owned_wif_fields_named,
server_owned_wif_fields_present,
)
from litellm.types.utils import CredentialItem, LlmProviders
_LITELLM_PROVIDER_IDS: Final = frozenset(provider.value for provider in LlmProviders)
def stored_credential_provider(credential_provider: object) -> str | None:
"""The dashboard stores its display casing (``Anthropic``) on credentials it creates, so the
provider a credential names is the lowercased value when that is a litellm provider id."""
if not isinstance(credential_provider, str):
return None
lowered: Final = credential_provider.lower()
return lowered if lowered in _LITELLM_PROVIDER_IDS else None
def decrypted_or_stored(key: str, value: str) -> str:
"""The stored value decrypted, or as stored when it was never encrypted (a config.yaml value)."""
decrypted: Final = decrypt_value_helper(value=value, key=key)
return value if decrypted is None else decrypted
def _decrypted(db_credential: CredentialItem) -> CredentialItem:
"""The stored credential with every value decrypted, leaving already-plaintext values alone."""
decrypted_values: Final = MappingProxyType(
{key: decrypted_or_stored(key, value) for key, value in db_credential.credential_values.items()}
)
return CredentialItem(
credential_name=db_credential.credential_name,
credential_values=decrypted_values, # pyright: ignore[reportArgumentType] # declared dict[str, str], and pydantic copies this mapping into one on validation; LIT002 rules out building that dict here
credential_info=db_credential.credential_info,
)
async def hydrate_named_credential_authoritative(
credential_name: str,
prisma_client: PrismaClient | None,
) -> CredentialItem | None:
"""The stored credential, preferring the row over this pod's in-memory copy.
``hydrate_named_credential`` reads memory first, which is right when serving a request. A
management operation cannot: on a pod whose in-memory copy predates another pod's update, it
would export the superseded JWKS, or discover models against superseded values. Same reason
``named_credential_wif_fields`` reads both.
"""
if prisma_client is None:
return await hydrate_named_credential(credential_name, prisma_client)
db_credential: Final = await CredentialsRepository(prisma_client).find_by_name(credential_name)
if db_credential is None:
return await hydrate_named_credential(credential_name, prisma_client)
return _decrypted(db_credential)
async def hydrate_named_credential(
credential_name: str,
prisma_client: PrismaClient | None,
) -> CredentialItem | None:
for credential in litellm.credential_list:
if credential.credential_name == credential_name:
return credential
if prisma_client is None:
return None
db_credential: Final = await CredentialsRepository(prisma_client).find_by_name(credential_name)
if db_credential is None:
return None
return _decrypted(db_credential)
async def named_credential_wif_fields(
credential_name: str,
prisma_client: PrismaClient | None,
) -> tuple[str, ...]:
"""Federation field names a write to ``credential_name`` would touch, from memory AND the row.
Resolution reads memory first and stops there, which is right when serving a request. An
authorization decision cannot: a pod whose in-memory copy predates an admin adding federation
fields would see none and allow the write. This reads both and returns the union, so the gate
refuses whenever either side says the credential is server-owned.
"""
in_memory: Final = tuple(
name
for credential in litellm.credential_list
if credential.credential_name == credential_name
for name in server_owned_wif_fields_named(credential.credential_values)
)
if prisma_client is None:
return in_memory
db_credential: Final = await CredentialsRepository(prisma_client).find_by_name(credential_name)
stored: Final = () if db_credential is None else server_owned_wif_fields_named(db_credential.credential_values)
return tuple(dict.fromkeys(in_memory + stored))
async def effective_server_owned_wif_fields(
stored: Mapping[str, object] | None,
incoming: GenericLiteLLMParams | None,
prisma_client: PrismaClient | None,
) -> tuple[str, ...]:
"""Federation field names the deployment would carry AFTER this write.
Authorization has to read the resulting deployment, not the submitted payload. A patch that
names no federation field still lands on a deployment that has them, and a patch that only
attaches ``litellm_credential_name`` inherits whatever that credential holds.
The two sides are matched differently on purpose. ``stored`` is matched by VALUE, because
``GenericLiteLLMParams`` declares every federation field, so matching it by key would
report every deployment on the proxy as federated. ``incoming`` is matched by the keys the
write actually set, so an explicit null still counts as touching the field.
"""
from_stored: Final = () if stored is None else server_owned_wif_fields_present(stored)
from_incoming: Final = () if incoming is None else server_owned_wif_fields_named(incoming.model_fields_set)
from_credential: Final = tuple(
chain.from_iterable(
await asyncio.gather(
*(
named_credential_wif_fields(credential_name, prisma_client)
for credential_name in _effective_credential_names(stored, incoming)
)
)
)
)
return tuple(dict.fromkeys(from_stored + from_incoming + from_credential))
def _effective_credential_names(
stored: Mapping[str, object] | None,
incoming: GenericLiteLLMParams | None,
) -> tuple[str, ...]:
"""Both the credential the deployment already carries and the one this write names.
Taking only the incoming name would let a write clear its way out: detaching a federated
credential, by sending ``litellm_credential_name: null`` alongside an api_key or api_base of
the caller's choosing, would leave nothing federated to find and the write would be allowed.
Detaching an administrator's federated credential is itself an administrator's action, so the
stored name counts whatever the write says.
"""
from_stored: Final = None if stored is None else stored.get("litellm_credential_name")
from_incoming: Final = None if incoming is None else incoming.litellm_credential_name
return tuple(dict.fromkeys(name for name in (from_stored, from_incoming) if isinstance(name, str)))

View file

@ -13,16 +13,111 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.llms.anthropic.wif import (
_IDENTITY_SOURCE_PARAM, # pyright: ignore[reportPrivateUsage] # one canonical param name, shared with the litellm_params identity-source resolver
_INTERNAL_ISSUER_FIELD_MAP, # pyright: ignore[reportPrivateUsage] # one canonical field map, shared with the litellm_params identity-source resolver
_build_variant, # pyright: ignore[reportPrivateUsage] # one canonical builder, shared with the litellm_params identity-source resolver
)
from litellm.llms.base_llm.auth.identity_source import (
AnthropicIdentitySourceKind,
InternalIssuerSource,
)
from litellm.llms.base_llm.auth.internal_issuer import internal_issuer_jwks_document
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.credential_hydration import (
hydrate_named_credential,
hydrate_named_credential_authoritative,
named_credential_wif_fields,
stored_credential_provider,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.proxy.utils import handle_exception_on_proxy, jsonify_object
from litellm.repositories.credentials_repository import CredentialsRepository
from litellm.types.router import server_owned_wif_fields_named
from litellm.types.utils import CreateCredentialItem, CredentialItem
router: Final = APIRouter()
def _reject_non_admin_wif_fields(
wif_fields: tuple[str, ...],
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""A credential referenced by ``litellm_credential_name`` feeds its values into the same
workload identity federation resolution as a deployment's own ``litellm_params``. Only proxy
admins may touch a server-owned WIF field, whether they write it, drop it, or edit a stored
credential that already carries one.
"""
if not wif_fields or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return
raise HTTPException(
status_code=403,
detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict
"error": (
f"Only proxy admins can change {wif_fields[0]!r}, a server-owned workload identity federation "
"parameter."
)
},
)
def _incoming_wif_fields(credential: CredentialItem) -> tuple[str, ...]:
"""WIF fields the request payload itself touches: the ones it sets (to any value, ``None``
included, since the key alone is what the federation resolver reacts to), plus the ones it
names in ``credential_values_to_delete``, since dropping a federation field off the stored
credential breaks every deployment referencing it just as installing one would redirect them.
"""
return server_owned_wif_fields_named(credential.credential_values) + server_owned_wif_fields_named(
credential.credential_values_to_delete or ()
)
def _stored_wif_fields(stored_credential: CredentialItem) -> tuple[str, ...]:
return server_owned_wif_fields_named(stored_credential.credential_values)
def _reject_overlapping_credential_values(credential: CredentialItem) -> None:
overlap: Final = frozenset(credential.credential_values) & frozenset(credential.credential_values_to_delete or ())
if overlap:
raise HTTPException(
status_code=400,
detail=f"credential_values_to_delete overlaps credential_values for key(s): {sorted(overlap)}",
)
def _sync_in_memory_credential(credential: CredentialItem, credential_name: str, new_name: str) -> None:
"""Mirror a DB credential update into the in-memory ``credential_list`` used by request-time
resolution; a no-op if the credential isn't loaded in memory (e.g. proxy restarted since boot).
"""
existing_in_memory: CredentialItem | None = None
for cred in litellm.credential_list:
if cred.credential_name == credential_name:
existing_in_memory = cred
break
if existing_in_memory is None:
return
in_memory_values: Final = dict(existing_in_memory.credential_values or {})
if credential.credential_values:
in_memory_values.update(credential.credential_values)
for key in credential.credential_values_to_delete or ():
in_memory_values.pop(key, None)
in_memory_info: Final = dict(existing_in_memory.credential_info or {})
if credential.credential_info:
in_memory_info.update(credential.credential_info)
updated_in_memory: Final = CredentialItem(
credential_name=new_name,
credential_values=in_memory_values,
credential_info=in_memory_info,
)
# Remove old entry if renamed, then use upsert_credentials to handle duplicates
if new_name != credential_name:
litellm.credential_list = [c for c in litellm.credential_list if c.credential_name != credential_name]
CredentialAccessor.upsert_credentials([updated_in_memory])
class CredentialHelperUtils:
@staticmethod
def encrypt_credential_values(credential: CredentialItem, new_encryption_key: str | None = None) -> CredentialItem:
@ -84,13 +179,19 @@ async def create_credential(
status_code=400,
detail="Credential values are required. Unable to infer credential values from model ID.",
)
_reject_non_admin_wif_fields(server_owned_wif_fields_named(credential.credential_values), user_api_key_dict)
_reject_non_admin_wif_fields(
await named_credential_wif_fields(credential.credential_name, prisma_client), user_api_key_dict
)
processed_credential: Final = CredentialItem(
credential_name=credential.credential_name,
credential_values=credential.credential_values,
credential_info=credential.credential_info,
)
encrypted_credential: Final = CredentialHelperUtils.encrypt_credential_values(processed_credential)
credentials_dict: Final = encrypted_credential.model_dump()
# exclude_none: wif.py rejects foreign-variant fields by presence, so persisting a null
# for every unset variant field would fail the next request against this credential
credentials_dict: Final = encrypted_credential.model_dump(exclude_none=True)
credentials_dict_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str
"dict[str, object]", jsonify_object(credentials_dict)
)
@ -175,6 +276,75 @@ async def get_credential_by_name(
raise handle_exception_on_proxy(e)
@router.get(
"/credentials/{credential_name:path}/jwks",
dependencies=(Depends(user_api_key_auth),),
tags=["credential management"], # mutable-ok: FastAPI's include_router does self.tags.copy(), needs a real list
)
async def get_credential_internal_issuer_jwks(
credential_name: str = Path(..., description="The credential name, percent-decoded; may contain slashes"),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI resolves the dependency from the default
):
"""
Export the public JWKS for an anthropic ``internal_issuer`` credential, so the operator can
register it on the Anthropic federation issuer from the UI. Never touches the private signing
key: only its derived public JWKS leaves this process. 404s for any other credential shape.
"""
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict
"error": "Only proxy admins can export a credential's JWKS."
},
)
try:
credential: Final = await hydrate_named_credential_authoritative(credential_name, prisma_client)
credential_provider: Final = (
None
if credential is None
else stored_credential_provider(credential.credential_info.get("custom_llm_provider"))
)
if credential is None or credential_provider != "anthropic":
raise HTTPException(
status_code=404,
detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict
"error": f"No anthropic credential named {credential_name!r}."
},
)
configured_source: Final = credential.credential_values.get(_IDENTITY_SOURCE_PARAM)
if configured_source != AnthropicIdentitySourceKind.internal_issuer.value:
raise HTTPException(
status_code=404,
detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict
"error": (
f"Credential {credential_name!r} is not configured with "
f"{_IDENTITY_SOURCE_PARAM}={AnthropicIdentitySourceKind.internal_issuer.value!r}."
)
},
)
try:
issuer_source: Final = _build_variant(
InternalIssuerSource, credential.credential_values, _INTERNAL_ISSUER_FIELD_MAP
)
jwks_document: Final = internal_issuer_jwks_document(issuer_source)
except (litellm.AuthenticationError, ValueError) as e:
raise HTTPException(
status_code=400,
detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict
"error": str(e)
},
) from e
return Response(content=jwks_document, media_type="application/json")
except HTTPException:
raise
except Exception as e: # noqa: BLE001 # endpoint boundary: every failure becomes the proxy's error contract
verbose_proxy_logger.exception(e)
raise handle_exception_on_proxy(e)
@router.get(
"/credentials/by_model/{model_id}",
dependencies=[Depends(user_api_key_auth)],
@ -239,6 +409,9 @@ async def delete_credential(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
_reject_non_admin_wif_fields(
await named_credential_wif_fields(credential_name, prisma_client), user_api_key_dict
)
deleted: Final = await CredentialsRepository(prisma_client).delete_by_name(credential_name)
if deleted is None:
raise HTTPException(
@ -249,6 +422,8 @@ async def delete_credential(
## DELETE FROM LITELLM ##
litellm.credential_list = [cred for cred in litellm.credential_list if cred.credential_name != credential_name]
return {"success": True, "message": "Credential deleted successfully"}
except HTTPException:
raise
except Exception as e:
raise handle_exception_on_proxy(e)
@ -282,6 +457,9 @@ def update_db_credential(
merged_credential.credential_values.update(encrypted_params)
for key in updated_patch.credential_values_to_delete or ():
merged_credential.credential_values.pop(key, None)
# update model info
if encrypted_credential.credential_info:
"""Update credential info"""
@ -310,6 +488,8 @@ async def update_credential(
from litellm.proxy.proxy_server import prisma_client
try:
_reject_overlapping_credential_values(credential)
_reject_non_admin_wif_fields(_incoming_wif_fields(credential), user_api_key_dict)
if prisma_client is None:
raise HTTPException(
status_code=500,
@ -319,9 +499,14 @@ async def update_credential(
db_credential: Final = await credentials_repository.find_by_name(credential_name)
if db_credential is None:
raise HTTPException(status_code=404, detail="Credential not found in DB.")
_reject_non_admin_wif_fields(_stored_wif_fields(db_credential), user_api_key_dict)
if credential.credential_name != credential_name:
shadowed_credential: Final = await hydrate_named_credential(credential.credential_name, prisma_client)
if shadowed_credential is not None:
_reject_non_admin_wif_fields(_stored_wif_fields(shadowed_credential), user_api_key_dict)
merged_credential: Final = update_db_credential(db_credential, credential)
credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str
"dict[str, object]", jsonify_object(merged_credential.model_dump())
"dict[str, object]", jsonify_object(merged_credential.model_dump(exclude_none=True))
)
await credentials_repository.update_by_name(
credential_name,
@ -332,29 +517,7 @@ async def update_credential(
)
# Sync in-memory credential_list (skip if not in memory - e.g., proxy restarted)
new_name: Final = merged_credential.credential_name
existing_in_memory: CredentialItem | None = None
for cred in litellm.credential_list:
if cred.credential_name == credential_name:
existing_in_memory = cred
break
if existing_in_memory is not None:
in_memory_values: Final = dict(existing_in_memory.credential_values or {})
if credential.credential_values:
in_memory_values.update(credential.credential_values)
in_memory_info: Final = dict(existing_in_memory.credential_info or {})
if credential.credential_info:
in_memory_info.update(credential.credential_info)
updated_in_memory: Final = CredentialItem(
credential_name=new_name,
credential_values=in_memory_values,
credential_info=in_memory_info,
)
# Remove old entry if renamed, then use upsert_credentials to handle duplicates
if new_name != credential_name:
litellm.credential_list = [c for c in litellm.credential_list if c.credential_name != credential_name]
CredentialAccessor.upsert_credentials([updated_in_memory])
_sync_in_memory_credential(credential, credential_name, merged_credential.credential_name)
return {"success": True, "message": "Credential updated successfully"}
except Exception as e:

View file

@ -38,6 +38,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.auth_utils import (
_BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check
reject_server_owned_wif_params,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
@ -2065,6 +2066,7 @@ async def test_model_connection(
"Could not find model %s in router: %s. Proceeding with request params only.", model_name, e
)
reject_server_owned_wif_params(request_litellm_params)
# Merge: config params (from proxy config) as base, request params override
litellm_params = {
**_config_base_for_health_check(
@ -2091,6 +2093,9 @@ async def test_model_connection(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
premium_user=premium_user,
# The Deployment above already carries the caller's merged params, so the effective
# state is model_params itself; there is no separate incoming patch here.
incoming_params=None,
)
mode = mode or litellm_params.pop("mode", None)

View file

@ -56,6 +56,10 @@ from litellm.proxy.common_utils.config_sync_pubsub import (
coordination_redis_cache,
publish_config_change,
)
from litellm.proxy.common_utils.credential_hydration import (
effective_server_owned_wif_fields,
hydrate_named_credential,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
@ -289,6 +293,23 @@ def _raise_on_strategy_router_write_violation(
)
def _reject_non_admin_blocked_flag_on_create(
blocked: bool | None,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""Same proxy-admin-only rule patch_model applies to the blocked flag: a team admin passed
the team-scoped auth check above, but must not be able to create a model already paused
(or explicitly unpaused) out from under the proxy admin.
"""
if blocked is not None and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise ProxyException(
message="Only proxy admins can set a model's blocked flag.",
type=ProxyErrorTypes.auth_error.value,
code=status.HTTP_403_FORBIDDEN,
param="blocked",
)
AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301
_CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)"
_STORED_LITELLM_PARAMS_SQL: Final = (
@ -831,6 +852,7 @@ async def patch_model(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
premium_user=premium_user,
incoming_params=patch_data.litellm_params,
)
# Pause/resume (`blocked`) is a proxy-admin-only privilege. Team admins
@ -1141,6 +1163,8 @@ async def _add_model_to_db(
}
if model_params.model_info.id is not None:
_data["model_id"] = model_params.model_info.id
if model_params.blocked is not None:
_data["blocked"] = model_params.blocked
_create_data: Final = cast("Mapping[str, object]", _data) # cast-ok: str-keyed json payload built just above
if not should_create_model_in_db:
return LiteLLM_ProxyModelTable(**_data)
@ -1691,14 +1715,65 @@ class ModelManagementAuthChecks:
)
return True
@staticmethod
async def _reject_non_admin_wif_write(
*,
model_params: Deployment,
incoming_params: GenericLiteLLMParams | None,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
) -> None:
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return
stored: Final = model_params.litellm_params.model_dump(exclude_none=True)
wif_fields: Final = await effective_server_owned_wif_fields(stored, incoming_params, prisma_client)
if wif_fields:
# ProxyException rather than HTTPException so the offending field stays a structured
# `param`, which is the contract the narrower gate this replaced already published.
raise ProxyException(
message=(
f"Only proxy admins can modify a deployment configured for workload identity "
f"federation ({wif_fields[0]!r})."
),
type=ProxyErrorTypes.auth_error.value,
code=status.HTTP_403_FORBIDDEN,
param=wif_fields[0],
)
# A name the caller expects an admin to create later would resolve to nothing today and
# start federating the moment it exists, so a non-admin may only attach one that is already there.
if incoming_params is not None and "litellm_credential_name" in incoming_params.model_fields_set:
named: Final = incoming_params.litellm_credential_name
if isinstance(named, str) and await hydrate_named_credential(named, prisma_client) is None:
raise ProxyException(
message=f"No credential named {named!r} exists.",
type=ProxyErrorTypes.bad_request_error.value,
code=status.HTTP_400_BAD_REQUEST,
param="litellm_credential_name",
)
@staticmethod
async def can_user_make_model_call(
model_params: Deployment,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
premium_user: bool,
*,
incoming_params: GenericLiteLLMParams | None,
allow_missing_team: bool = False,
) -> Literal[True]:
# Federation fields choose which server-side secret is read and where the org-scoped token
# it buys is sent, so only a proxy admin may touch a deployment that has them. Evaluated on
# the RESULTING deployment: a patch naming no federation field still lands on one that has
# them, and a patch attaching a credential by name inherits whatever that credential holds.
# `incoming_params` is keyword-only with no default so a new write path cannot typecheck
# without deciding what it writes.
await ModelManagementAuthChecks._reject_non_admin_wif_write(
model_params=model_params,
incoming_params=incoming_params,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
## Check team model auth
if model_params.model_info is not None and model_params.model_info.team_id is not None:
team_obj_row: Final = await _repo_team_table(prisma_client).find_unique(
@ -1793,6 +1868,7 @@ async def delete_model(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
premium_user=premium_user,
incoming_params=None,
allow_missing_team=True,
)
@ -1914,7 +1990,6 @@ async def delete_team_model_alias(
return removed_model_aliases
#### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964
@router.post(
"/model/new",
description="Allows adding new models to the model list in the config.yaml",
@ -1983,8 +2058,11 @@ async def add_new_model(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
premium_user=premium_user,
incoming_params=model_params.litellm_params,
)
_reject_non_admin_blocked_flag_on_create(model_params.blocked, user_api_key_dict)
ModelManagementAuthChecks.can_user_attach_credential(
litellm_params=model_params.litellm_params,
user_api_key_dict=user_api_key_dict,
@ -2166,6 +2244,7 @@ async def update_model(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
premium_user=premium_user,
incoming_params=model_params.litellm_params,
)
ModelManagementAuthChecks.can_user_attach_credential(

View file

@ -32,7 +32,11 @@ from litellm.constants import (
BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES,
)
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
from litellm.llms.anthropic.common_utils import (
_SERVER_OWNED_AUTH_HEADERS, # pyright: ignore[reportPrivateUsage] # canonical set, must not be duplicated here
AnthropicModelInfo,
merge_anthropic_beta_headers,
)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
@ -644,6 +648,57 @@ async def is_streaming_request_fn(request: Request) -> bool:
return False
def _anthropic_passthrough_headers(auth_header: Mapping[str, str] | None, client_beta: str | None) -> Mapping[str, str]:
"""Custom headers take priority over forwarded client headers, so merge the
client's anthropic-beta into the auth header's instead of clobbering it."""
if auth_header is None:
return MappingProxyType({})
auth_beta: Final = auth_header.get("anthropic-beta")
if auth_beta is None or client_beta is None:
return auth_header
return MappingProxyType({**auth_header, "anthropic-beta": merge_anthropic_beta_headers(client_beta, auth_beta)})
def _configured_litellm_key_header_name() -> str | None:
from litellm.proxy.proxy_server import (
general_settings, # pyright: ignore[reportUnknownVariableType] # proxy_server declares it as a bare dict
)
configured: Final = general_settings.get( # pyright: ignore[reportUnknownMemberType] # proxy_server general_settings is a bare dict
"litellm_key_header_name"
)
return configured if isinstance(configured, str) else None
def _anthropic_passthrough_header_plan(
request: Request, auth_header: Mapping[str, str] | None, litellm_key_header_name: str | None
) -> tuple[Mapping[str, str], bool]:
"""Returns the headers to send upstream plus whether the relay should still forward the
caller's own headers. Once the server owns the Anthropic credential, none of the headers
the proxy accepts a LiteLLM key in (``SpecialHeaders`` plus the configured custom name)
may ride upstream beside it, so the forward merge runs here with those stripped and the
relay is told not to merge again. With no server credential the caller's key is the only
one there is, so forwarding stays on (BYOK)."""
server_headers: Final = _anthropic_passthrough_headers(auth_header, request.headers.get("anthropic-beta"))
if auth_header is None:
return server_headers, True
caller_owned: Final = _SERVER_OWNED_AUTH_HEADERS | frozenset(
(litellm_key_header_name.lower(),) if litellm_key_header_name else ()
)
caller_headers: Final[dict[str, str]] = { # mutable-ok: forward_headers_from_request takes a concrete dict
name: value for name, value in request.headers.items() if name.lower() not in caller_owned
}
merged: Final = cast( # cast-ok: forward_headers_from_request is untyped upstream, its result is a header dict
"dict[str, str]",
HttpPassThroughEndpointHelpers.forward_headers_from_request( # pyright: ignore[reportUnknownMemberType] # untyped upstream
request_headers=caller_headers,
headers=dict(server_headers), # mutable-ok: forward_headers_from_request takes a concrete dict
forward_headers=True,
),
)
return MappingProxyType(merged), False
@router.api_route(
"/anthropic/{endpoint:path}",
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
@ -682,12 +737,17 @@ async def anthropic_proxy_route(
is_streaming_request: Final = await is_streaming_request_fn(request)
## CREATE PASS-THROUGH
auth_header: Final = AnthropicModelInfo.get_auth_header(anthropic_api_key or None)
auth_header: Final = await AnthropicModelInfo.aget_auth_header(
anthropic_api_key or None, allow_workload_identity=True
)
upstream_headers, forward_caller_headers = _anthropic_passthrough_header_plan(
request, auth_header, _configured_litellm_key_header_name()
)
endpoint_func: Final = create_pass_through_route(
endpoint=endpoint,
target=str(updated_url),
custom_headers=auth_header if auth_header is not None else {},
_forward_headers=True,
custom_headers=upstream_headers,
_forward_headers=forward_caller_headers,
is_streaming_request=is_streaming_request,
) # dynamically construct pass-through endpoint based on incoming path
received_value: Final = await endpoint_func(

View file

@ -336,6 +336,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
)
from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber
from litellm.proxy.common_utils.credential_hydration import decrypted_or_stored
from litellm.proxy.common_utils.debug_utils import init_verbose_loggers
from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
@ -719,6 +720,7 @@ from litellm.types.router import (
RouterGeneralSettings,
RoutingPlugin,
SearchToolTypedDict,
holds_secret_pointer,
updateDeployment,
)
from litellm.types.router import ModelInfo as RouterModelInfo
@ -4679,7 +4681,7 @@ class ProxyConfig:
if isinstance(item, dict):
item = self._check_for_os_environ_vars(config=item, depth=depth + 1, max_depth=max_depth)
# if the value is a string and starts with "os.environ/" - then it's an environment variable
elif isinstance(value, str) and value.startswith("os.environ/"):
elif isinstance(value, str) and value.startswith("os.environ/") and not holds_secret_pointer(key):
resolved = get_secret(value)
if resolved is None and secret_manager_would_be_consulted(value):
verbose_proxy_logger.warning("%s is absent from the configured secret manager", value)
@ -5777,7 +5779,7 @@ class ProxyConfig:
for model in model_list:
### LOAD FROM os.environ/ ###
for k, v in model["litellm_params"].items():
if isinstance(v, str) and v.startswith("os.environ/"):
if isinstance(v, str) and v.startswith("os.environ/") and not holds_secret_pointer(k):
model["litellm_params"][k] = get_secret(v)
validate_deployment_max_agentic_loops(model)
validate_deployment_complexity_router_placement(model)
@ -6170,7 +6172,7 @@ class ProxyConfig:
for model in model_list:
### LOAD FROM os.environ/ ###
for k, v in model["litellm_params"].items():
if isinstance(v, str) and v.startswith("os.environ/"):
if isinstance(v, str) and v.startswith("os.environ/") and not holds_secret_pointer(k):
model["litellm_params"][k] = get_secret(v)
## check if they have model-id's ##
@ -6198,7 +6200,11 @@ class ProxyConfig:
return value
decrypted_value: Final = decrypt_value_helper(value=value, key=key, return_original_value=True)
if isinstance(decrypted_value, str) and decrypted_value.startswith("os.environ/"):
if (
isinstance(decrypted_value, str)
and decrypted_value.startswith("os.environ/")
and not holds_secret_pointer(key)
):
return get_secret(decrypted_value)
return decrypted_value
@ -7967,7 +7973,7 @@ class ProxyConfig:
decrypted_credential_values: Final = {}
for k, v in credential_object.credential_values.items():
decrypted_credential_values[k] = decrypt_value_helper(value=v, key=k) or v
decrypted_credential_values[k] = decrypted_or_stored(k, v)
credential_object.credential_values = decrypted_credential_values
return credential_object

View file

@ -235,6 +235,7 @@ from litellm.types.router import (
RoutingStrategy,
SearchToolTypedDict,
TaggedPreRoutingStrategy,
holds_secret_pointer,
)
from litellm.types.services import ServiceTypes
from litellm.types.utils import (
@ -8650,13 +8651,10 @@ class Router:
if ptu_error is not None and is_ptu_cost_attribution_enabled():
raise ValueError(ptu_error)
zeroed_pricing: Final = zeroed_ptu_pricing(_model_info, _litellm_params) if config_sourced else None
litellm_params: Final[LiteLLM_Params] = LiteLLM_Params(
**(
_litellm_params
if zeroed_pricing is None
else MappingProxyType({**_litellm_params, **zeroed_pricing})
)
merged_params: Final[Mapping[str, Any]] = (
_litellm_params if zeroed_pricing is None else MappingProxyType({**_litellm_params, **zeroed_pricing})
)
litellm_params: Final[LiteLLM_Params] = LiteLLM_Params(**merged_params)
warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=_litellm_params)
deployment = Deployment(
**deployment_info,
@ -9226,7 +9224,7 @@ class Router:
## check if litellm params in os.environ
if isinstance(_litellm_params, dict):
for k, v in _litellm_params.items():
if isinstance(v, str) and v.startswith("os.environ/"):
if isinstance(v, str) and v.startswith("os.environ/") and not holds_secret_pointer(k):
_litellm_params[k] = get_secret(v)
_model_info: dict = model.pop("model_info", {})

View file

@ -13,8 +13,16 @@ Ensures cooldowns are applied correctly.
from typing import Final
from litellm.types.utils import server_owned_wif_litellm_params
clientside_credential_keys: Final = ["api_key", "api_base", "base_url"]
# Set on a deployment whose api_base was client-redirected, so the Anthropic auth path refuses to
# mint a federation token there even when WIF is configured only through ANTHROPIC_* env vars (which
# cannot be cleared from litellm_params).
DISABLE_WORKLOAD_IDENTITY_PARAM: Final = "anthropic_disable_workload_identity_federation"
_WIF_CLEAR_ON_BASE_OVERRIDE: Final = tuple(sorted(server_owned_wif_litellm_params))
def _admin_config_fields_to_clear_on_base_override() -> list[str]:
"""
@ -59,6 +67,14 @@ def _admin_config_fields_to_clear_on_base_override() -> list[str]:
# ``api_base`` for the same reason as the OCI entries above.
"nvcf_function_id",
"use_ssl",
# Workload-identity federation minting fields, restated here from
# server_owned_wif_litellm_params the same way azure_ad_token above is restated
# despite also being declared on CredentialLiteLLMParams (hence covered by
# typed_fields too): a federation token minted for a client-redirected api_base
# would send the workload's OIDC assertion, and then the minted bearer, to the
# caller-chosen host, so this list must stay correct even if a field is ever
# dropped from the typed model.
*_WIF_CLEAR_ON_BASE_OVERRIDE,
]
return typed_fields + kwargs_only_fields
@ -101,5 +117,6 @@ def get_dynamic_litellm_params(litellm_params: dict, request_kwargs: dict) -> di
litellm_params.pop(field, None)
if field in request_kwargs:
litellm_params[field] = request_kwargs[field]
litellm_params[DISABLE_WORKLOAD_IDENTITY_PARAM] = True
return litellm_params

View file

@ -24,7 +24,7 @@ from litellm.router_utils.cooldown_handlers import (
from litellm.router_utils.router_callbacks.track_deployment_metrics import (
increment_deployment_failures_for_current_minute,
)
from litellm.types.router import LiteLLMParamsTypedDict
from litellm.types.router import LiteLLMParamsTypedDict, reject_server_owned_wif_params
if TYPE_CHECKING:
from litellm.router import Router as _Router
@ -507,6 +507,14 @@ async def run_async_fallback(
failed_model_group: Final = get_pre_routing_selection(kwargs) or original_model_group
attempted.record(failed_model_group)
# A dict target is merged straight into kwargs below, and kwargs win over the deployment's own
# params, so a stored key/team/global fallback could otherwise set a federation field that the
# request itself is forbidden to carry. Checked here rather than at the merge: inside the loop
# the refusal would be caught as a per-target failure and quietly skipped to the next one.
for target in fallback_model_group:
if isinstance(target, dict):
reject_server_owned_wif_params(target)
for mg in fallback_model_group:
if mg == failed_model_group:
continue

View file

@ -49,6 +49,10 @@ def _oidc_token_cache_ttl(oidc_token: str, max_ttl: int) -> int:
_DEFAULT_OIDC_ALLOWED_CREDENTIAL_DIRS: Final = ("/var/run/secrets", "/run/secrets")
class OidcPathNotAllowedError(ValueError):
"""An ``oidc/file/`` path was rejected by the credential-directory allowlist."""
def _get_oidc_allowed_credential_dirs() -> list[str]:
"""
Return the absolute, normalized list of directories from which
@ -73,7 +77,7 @@ def _resolve_oidc_file_path(requested_path: str) -> str:
credential directories. Raises ``ValueError`` otherwise.
"""
if not os.path.isabs(requested_path):
raise ValueError(
raise OidcPathNotAllowedError(
"oidc/file path must be absolute. Use the format "
"'oidc/file//var/run/secrets/<name>' (note the leading slash "
"after 'oidc/file/')."
@ -87,7 +91,7 @@ def _resolve_oidc_file_path(requested_path: str) -> str:
# commonpath raises when paths are on different drives (Windows);
# treat as not-matching and continue.
continue
raise ValueError(
raise OidcPathNotAllowedError(
"oidc/file path is outside the allowed credential directories. "
"Set LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS to extend the allowlist."
)

View file

@ -750,5 +750,6 @@ ANTHROPIC_EFFORT_BETA_HEADER: Final = "effort-2025-11-24"
# 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"

View file

@ -4,7 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc
import datetime
import enum
from collections.abc import Mapping
from collections.abc import Container, Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints
@ -27,6 +27,10 @@ from .utils import (
ModelResponse,
StandardLoggingRoutingDecision,
)
from .utils import (
# private alias: `from .types.router import *` would rebind a public Final in litellm/__init__.py
server_owned_wif_litellm_params as _server_owned_wif_litellm_params,
)
class ConfigurableClientsideParamsCustomAuth(TypedDict):
@ -297,6 +301,69 @@ class CredentialLiteLLMParams(BaseModel):
## IBM WATSONX ##
watsonx_region_name: str | None = None
## ANTHROPIC WORKLOAD IDENTITY FEDERATION ##
# Without these, get_deployment_credentials_with_provider silently drops a
# litellm_params-configured WIF setup before files/batches/passthrough callers see
# it, the same #30235-shaped gap azure_ad_token above was added to close.
anthropic_federation_rule_id: str | None = None
anthropic_organization_id: str | None = None
anthropic_service_account_id: str | None = None
anthropic_workspace_id: str | None = None
anthropic_identity_token_file: str | None = None
anthropic_identity_token: str | None = None
anthropic_identity_source: str | None = None
anthropic_issuer_url: str | None = None
anthropic_issuer_subject: str | None = None
anthropic_issuer_audience: str | None = None
anthropic_issuer_ttl_seconds: int | None = None
anthropic_issuer_signing_key_ref: str | None = None
anthropic_keycloak_token_url: str | None = None
anthropic_keycloak_client_id: str | None = None
anthropic_keycloak_auth_method: str | None = None
anthropic_keycloak_client_secret_ref: str | None = None
anthropic_keycloak_scope: str | None = None
# Server-set when a client redirects api_base. Declared so it survives the strict dump the
# other federation fields above are declared for, rather than being rebuilt away in transit.
anthropic_disable_workload_identity_federation: bool | None = None
## OPENAI WORKLOAD IDENTITY FEDERATION ##
openai_identity_provider_id: str | None = None
openai_service_account_id: str | None = None
openai_identity_token_file: str | None = None
def server_owned_wif_fields_present(fields: Mapping[str, object]) -> tuple[str, ...]:
"""Server-owned workload identity federation field names set in ``fields``.
``fields`` is a ``litellm_params`` dict (or a credential's ``credential_values`` mapping,
which feeds the same resolution when referenced by name). Derived from
``server_owned_wif_litellm_params`` rather than hand-copied, so a persistence gate built on
this stays correct when a new WIF field is added there.
"""
return tuple(name for name in _server_owned_wif_litellm_params if fields.get(name) is not None)
def server_owned_wif_fields_named(keys: Container[str]) -> tuple[str, ...]:
"""Server-owned workload identity federation field names that appear in ``keys``, whatever
value they carry.
The write gates on credentials need this key-based sibling of ``server_owned_wif_fields_present``:
``get_litellm_params`` forwards a WIF kwarg on key presence and the federation resolver rejects
a foreign variant's field by key, so a persisted ``{"anthropic_issuer_url": None}`` wedges every
deployment that references the credential even though no value is set. Pass a mapping (its keys
are tested) or a plain collection of key names.
"""
return tuple(name for name in _server_owned_wif_litellm_params if name in keys)
_WIF_POINTER_FIELDS: Final = frozenset(name for name in _server_owned_wif_litellm_params if name.endswith("_ref"))
def holds_secret_pointer(param_name: str) -> bool:
"""A ``*_ref`` federation field is a secret POINTER the identity source dereferences at use
time, so a loader expanding ``os.environ/`` values must leave it as written."""
return param_name in _WIF_POINTER_FIELDS
_RESERVED_INIT_KEYS: Final = frozenset({"self", "params", "__class__"})
@ -542,6 +609,9 @@ class Deployment(BaseModel):
model_name: str
litellm_params: LiteLLM_Params
model_info: ModelInfo
# admin-toggled pause flag; mirrors LiteLLM_ProxyModelTable.blocked. None means "don't set it
# on create" -- the Prisma column defaults to False -- rather than "explicitly unblocked".
blocked: bool | None = None
model_config = ConfigDict(extra="allow", protected_namespaces=())
@ -1064,3 +1134,23 @@ class AdaptiveRouterPreferences(BaseModel):
quality_tier: int = Field(ge=1, le=3)
strengths: list[RequestType] = Field(default_factory=list)
_BEDROCK_WORKSPACE_HINT: Final = " On the Bedrock Claude Platform route, pass workspace_id or aws_workspace_id instead."
def reject_server_owned_wif_params(body: Mapping[str, object]) -> None:
"""Raise ``ValueError`` if a mapping that did not come from deployment config carries a
server-owned workload identity federation field.
These are never client-settable on any surface, with or without a client-side credential
opt-in. This lives here rather than under ``litellm.proxy`` so the router can call it on a
post-authentication merge without core importing from the proxy package.
"""
for param in _server_owned_wif_litellm_params:
if param in body:
raise ValueError(
f"Rejected Request: {param} is a server-owned workload identity federation parameter "
"and cannot be set in a request body; configure it on the deployment instead."
+ (_BEDROCK_WORKSPACE_HINT if param == "anthropic_workspace_id" else "")
)

View file

@ -25,6 +25,8 @@ class ServiceTypes(str, enum.Enum):
AUTH = "auth"
PROXY_PRE_CALL = "proxy_pre_call"
POD_LOCK_MANAGER = "pod_lock_manager"
ANTHROPIC_WIF = "anthropic_wif"
ANTHROPIC_WIF_CACHE = "anthropic_wif_cache"
"""
Operational metrics for DB Transaction Queues
@ -67,6 +69,13 @@ DEFAULT_SERVICE_CONFIGS: Final = {
ServiceTypes.ROUTER.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]},
ServiceTypes.AUTH.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]},
ServiceTypes.PROXY_PRE_CALL.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]},
ServiceTypes.ANTHROPIC_WIF.value: { # mutable-ok: ServiceConfig mandates the dict-of-list shape
"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] # mutable-ok: ServiceConfig mandates a list
},
# cache hits are counter-only: no HTTP call happens, so observing a latency would be a lie
ServiceTypes.ANTHROPIC_WIF_CACHE.value: { # mutable-ok: ServiceConfig mandates the dict-of-list shape
"metrics": [ServiceMetrics.COUNTER] # mutable-ok: ServiceConfig mandates a list
},
# Operational metrics for DB Transaction Queues
ServiceTypes.POD_LOCK_MANAGER.value: {"metrics": [ServiceMetrics.GAUGE]},
ServiceTypes.IN_MEMORY_DAILY_SPEND_UPDATE_QUEUE.value: {"metrics": [ServiceMetrics.GAUGE]},

View file

@ -3643,9 +3643,27 @@ bedrock_batch_litellm_params: Final = (
"bedrock_tags",
)
# Anthropic workload identity federation config, read from litellm_params by the
# Anthropic auth tier. Listed for the same reason as the fields above: an
# unrecognized top-level key is swept into extra_body and sent to /v1/messages.
# Derived from get_litellm_params.ANTHROPIC_WIF_KWARGS_KEYS (not hand-typed) so the
# request-body ban list and the clear-on-api_base-override list can never drift from
# the set the kwargs funnel actually forwards. Imported here rather than at module top:
# get_litellm_params.py's own import chain (llms/openai/data_residency -> llms/__init__)
# reaches back into this module for CallTypes, which by this point in the file is
# already bound on the partially-initialized module.
from ..litellm_core_utils.get_litellm_params import ( # noqa: E402 # deferred past CallTypes to break the import cycle
ANTHROPIC_WIF_KWARGS_KEYS,
OPENAI_WIF_KWARGS_KEYS,
)
anthropic_wif_litellm_params: Final = tuple(sorted(ANTHROPIC_WIF_KWARGS_KEYS))
openai_wif_litellm_params: Final = tuple(sorted(OPENAI_WIF_KWARGS_KEYS))
server_owned_wif_litellm_params: Final = anthropic_wif_litellm_params + openai_wif_litellm_params
all_litellm_params = (
agentic_loop_internal_litellm_params
+ [TRUSTED_CALLBACK_VARS_FIELD, *bedrock_batch_litellm_params]
+ [TRUSTED_CALLBACK_VARS_FIELD, *bedrock_batch_litellm_params, *server_owned_wif_litellm_params]
+ [
"metadata",
"litellm_metadata",

View file

@ -7361,9 +7361,8 @@ def _get_valid_models_from_provider_api(
if cached_result is not None:
return cached_result
models: Final = provider_config.get_models(
api_key=litellm_params.api_key if litellm_params is not None else None,
api_base=litellm_params.api_base if litellm_params is not None else None,
models: Final = provider_config.discover_models(
litellm_params=litellm_params.model_dump(exclude_none=True) if litellm_params is not None else None
)
_model_cache.set_cached_model_info(custom_llm_provider, litellm_params, models)

View file

@ -36,6 +36,7 @@ GET /user/spend/report
# Admin UI helper endpoints; serve UI forms and caller-scoped views, not desired state
GET /budget/settings
GET /credentials/{credential_name}/jwks
GET /router/fields
GET /guardrails/ui/add_guardrail_settings
GET /guardrails/ui/category_yaml/{category_name}

View file

@ -1235,6 +1235,7 @@ async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch):
result set - zero cost, zero usage, no models - instead of letting the file
fetch raise "Output file id is None" on every aretrieve_batch logging poll.
"""
# The output-file fetch must not even be attempted when there is no output file.
async def _must_not_fetch(*args, **kwargs):
pytest.fail("_fetch_batch_output_file_content should not be called")
@ -1361,7 +1362,10 @@ def test_anthropic_response_body_is_result_message():
def test_anthropic_usage_conversion_includes_cache_tokens():
body = {"model": "claude-sonnet-4-5-20250929", "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)}
body = {
"model": "claude-sonnet-4-5-20250929",
"usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000),
}
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="anthropic")
assert usage.prompt_tokens == 11000
assert usage.completion_tokens == 200
@ -1376,7 +1380,9 @@ def test_bedrock_model_output_line_success_check():
"modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}},
}
assert bu._batch_response_was_successful(row, custom_llm_provider="bedrock") is True
assert bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6"
assert (
bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6"
)
def test_bedrock_cost_uses_deployment_model_name():
@ -1430,7 +1436,13 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch):
rows = [
{
"custom_id": "req-1",
"response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}},
"response": {
"status_code": 200,
"body": {
"model": "gpt-5.2",
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
},
},
}
]
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
@ -1674,7 +1686,10 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) ->
def test_bedrock_converse_shaped_batch_usage_is_parsed():
body = {"model": "us.amazon.nova-lite-v1:0", "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}}
body = {
"model": "us.amazon.nova-lite-v1:0",
"usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742},
}
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (2202, 540, 2742)
@ -1718,3 +1733,37 @@ def test_unparsable_bedrock_batch_usage_warns(caplog):
assert usage.total_tokens == 0
assert "does not understand" in caplog.text
assert "inputTextTokenCount" in caplog.text
class TestFileAccessCredentialsCarryFederation:
"""A federated deployment holds no api_key, so the fetch that reads a finished batch's output
has to inherit the federation fields or it cannot authenticate and the batch is never billed."""
def test_federation_fields_survive_extraction(self):
from litellm.batches.batch_utils import _extract_file_access_credentials
credentials = _extract_file_access_credentials(
{
"model": "anthropic/claude-sonnet-4-5",
"anthropic_federation_rule_id": "fdrl_x",
"anthropic_organization_id": "org-x",
"anthropic_identity_token_file": "/var/run/secrets/anthropic.com/token",
"something_unrelated": "dropped",
}
)
assert credentials["anthropic_federation_rule_id"] == "fdrl_x"
assert credentials["anthropic_organization_id"] == "org-x"
assert credentials["anthropic_identity_token_file"] == "/var/run/secrets/anthropic.com/token"
assert "something_unrelated" not in credentials
def test_every_federation_field_is_carried(self):
"""Derived from the kwargs set, so a new federation field is carried without an edit here."""
from litellm.batches.batch_utils import _extract_file_access_credentials
from litellm.litellm_core_utils.get_litellm_params import ANTHROPIC_WIF_KWARGS_KEYS
params = {name: f"value-{name}" for name in ANTHROPIC_WIF_KWARGS_KEYS}
credentials = _extract_file_access_credentials(params)
assert set(credentials) == set(ANTHROPIC_WIF_KWARGS_KEYS)

View file

@ -135,3 +135,28 @@ def test_services_logger_custom_latency_buckets():
REGISTRY.unregister(collector)
except Exception:
pass
def test_anthropic_wif_services_are_wired_into_the_registry():
"""Reverting the ANTHROPIC_WIF/ANTHROPIC_WIF_CACHE ServiceTypes members or their
DEFAULT_SERVICE_CONFIGS entries must fail here: the exchange service gets counters plus a
latency histogram, while the cache-hit service is counter-only so a hit can never fake a latency."""
from litellm.types.services import DEFAULT_SERVICE_CONFIGS
assert ServiceTypes.ANTHROPIC_WIF.value == "anthropic_wif"
assert ServiceTypes.ANTHROPIC_WIF_CACHE.value == "anthropic_wif_cache"
assert DEFAULT_SERVICE_CONFIGS["anthropic_wif"]["metrics"] == [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]
assert DEFAULT_SERVICE_CONFIGS["anthropic_wif_cache"]["metrics"] == [ServiceMetrics.COUNTER]
pl = PrometheusServicesLogger()
wif_names = {obj._name for obj in pl.payload_to_prometheus_map["anthropic_wif"]}
assert wif_names == {
"litellm_anthropic_wif_latency",
"litellm_anthropic_wif_failed_requests",
"litellm_anthropic_wif_total_requests",
}
cache_names = {obj._name for obj in pl.payload_to_prometheus_map["anthropic_wif_cache"]}
assert cache_names == {
"litellm_anthropic_wif_cache_failed_requests",
"litellm_anthropic_wif_cache_total_requests",
}

View file

@ -244,3 +244,129 @@ class TestRustOptIn:
from litellm.types.utils import all_litellm_params
assert "rust" in all_litellm_params
class TestAnthropicWifKeys:
"""The six anthropic_* WIF keys need the same dual registration as `rust`:
carried by the kwargs funnel into litellm_params (where the Anthropic auth
tier reads them) AND listed in all_litellm_params (so the extra_body sweep
never sends them to /v1/messages)."""
SIX_KEYS = {
"anthropic_federation_rule_id": "fdrl_1",
"anthropic_organization_id": "org-1",
"anthropic_service_account_id": "svcacct_1",
"anthropic_workspace_id": "wrkspc_1",
"anthropic_identity_token_file": "/var/run/secrets/tok",
"anthropic_identity_token": "oidc/env/TOK",
}
def test_keys_survive_into_litellm_params(self):
params = get_litellm_params(**self.SIX_KEYS)
for key, value in self.SIX_KEYS.items():
assert params[key] == value
def test_keys_are_forwarded_from_completion_kwargs(self):
from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS
assert set(self.SIX_KEYS) <= FORWARDED_KWARGS_KEYS
def test_keys_stay_out_of_the_provider_body(self):
from litellm.types.utils import all_litellm_params
for key in self.SIX_KEYS:
assert key in all_litellm_params
def test_keys_absent_when_not_configured(self):
params = get_litellm_params()
for key in self.SIX_KEYS:
assert key not in params
class TestAnthropicWifIdentitySourceKeys:
"""Phase 1 adds 11 more anthropic_* WIF keys (the anthropic_identity_source discriminator
plus the internal_issuer/keycloak identity-source fields) that need the same dual
registration as the original six tested above."""
NEW_KEYS = {
"anthropic_identity_source": "keycloak",
"anthropic_issuer_url": "https://issuer.example",
"anthropic_issuer_subject": "svc-account",
"anthropic_issuer_audience": "https://api.anthropic.com",
"anthropic_issuer_ttl_seconds": "300",
"anthropic_issuer_signing_key_ref": "oidc/env/ISSUER_KEY",
"anthropic_keycloak_token_url": "https://kc.example/realms/r/protocol/openid-connect/token",
"anthropic_keycloak_client_id": "litellm",
"anthropic_keycloak_auth_method": "client_secret_basic",
"anthropic_keycloak_client_secret_ref": "oidc/env/KC_SECRET",
"anthropic_keycloak_scope": "anthropic-wif",
# Server-set when a client redirects api_base; carried here so it is not dropped in transit
"anthropic_disable_workload_identity_federation": True,
}
def test_new_keys_are_exactly_the_non_legacy_registered_set(self):
"""Fails the moment a key is added to ANTHROPIC_WIF_KWARGS_KEYS without a matching entry
here (or vice versa), catching drift between what wif.py dispatches on and what this
test (and the funnel/provider-body tests below) actually exercises."""
from litellm.litellm_core_utils.get_litellm_params import ANTHROPIC_WIF_KWARGS_KEYS
assert set(self.NEW_KEYS) == ANTHROPIC_WIF_KWARGS_KEYS - set(TestAnthropicWifKeys.SIX_KEYS)
def test_keys_survive_into_litellm_params(self):
params = get_litellm_params(**self.NEW_KEYS)
for key, value in self.NEW_KEYS.items():
assert params[key] == value
def test_keys_are_forwarded_from_completion_kwargs(self):
from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS
assert set(self.NEW_KEYS) <= FORWARDED_KWARGS_KEYS
def test_keys_stay_out_of_the_provider_body(self):
from litellm.types.utils import all_litellm_params
for key in self.NEW_KEYS:
assert key in all_litellm_params
def test_keys_absent_when_not_configured(self):
params = get_litellm_params()
for key in self.NEW_KEYS:
assert key not in params
class TestOpenAIWifKeys:
"""The three openai_* WIF keys carry a deployment's federation identity through the kwargs
funnel into litellm_params (where the OpenAI client factory reads them) and stay out of the
provider body, exactly like the anthropic_* keys above."""
THREE_KEYS = {
"openai_identity_provider_id": "idp_1",
"openai_service_account_id": "user-1",
"openai_identity_token_file": "/var/run/secrets/tokens/openai",
}
def test_keys_are_exactly_the_registered_set(self):
from litellm.litellm_core_utils.get_litellm_params import OPENAI_WIF_KWARGS_KEYS
assert set(self.THREE_KEYS) == OPENAI_WIF_KWARGS_KEYS
def test_keys_survive_into_litellm_params(self):
params = get_litellm_params(**self.THREE_KEYS)
for key, value in self.THREE_KEYS.items():
assert params[key] == value
def test_keys_are_forwarded_from_completion_kwargs(self):
from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS
assert set(self.THREE_KEYS) <= FORWARDED_KWARGS_KEYS
def test_keys_stay_out_of_the_provider_body(self):
from litellm.types.utils import all_litellm_params
for key in self.THREE_KEYS:
assert key in all_litellm_params
def test_keys_absent_when_not_configured(self):
params = get_litellm_params()
for key in self.THREE_KEYS:
assert key not in params

View file

@ -14,6 +14,8 @@ asyncio.run) is exercised directly, mirroring the dispatch-contract discipline i
tests/test_litellm/batches/test_main.py.
"""
import asyncio
import threading
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -35,9 +37,7 @@ def _ok_batch_response():
"ended_at": "2024-09-24T11:00:00Z",
"request_counts": {"succeeded": 2, "errored": 0},
},
request=httpx.Request(
"GET", "https://api.anthropic.com/v1/messages/batches/msgbatch_abc"
),
request=httpx.Request("GET", "https://api.anthropic.com/v1/messages/batches/msgbatch_abc"),
)
@ -59,9 +59,7 @@ def patched_client():
@pytest.mark.asyncio
async def test_aretrieve_batch_fires_get_with_correct_url_and_headers(
handler, patched_client
):
async def test_aretrieve_batch_fires_get_with_correct_url_and_headers(handler, patched_client):
fake_client, factory = patched_client
batch = await handler.aretrieve_batch(
@ -76,9 +74,7 @@ async def test_aretrieve_batch_fires_get_with_correct_url_and_headers(
fake_client.get.assert_awaited_once()
_, call_kwargs = fake_client.get.call_args
# Exact URL built by get_retrieve_batch_url.
assert call_kwargs["url"] == (
"https://api.anthropic.com/v1/messages/batches/msgbatch_abc"
)
assert call_kwargs["url"] == ("https://api.anthropic.com/v1/messages/batches/msgbatch_abc")
# Auth + version + beta headers built by validate_environment.
headers = call_kwargs["headers"]
assert headers["x-api-key"] == "sk-ant-test"
@ -93,9 +89,7 @@ async def test_aretrieve_batch_fires_get_with_correct_url_and_headers(
@pytest.mark.asyncio
async def test_aretrieve_batch_uses_anthropic_provider_for_client(
handler, patched_client
):
async def test_aretrieve_batch_uses_anthropic_provider_for_client(handler, patched_client):
from litellm.types.utils import LlmProviders
_, factory = patched_client
@ -111,14 +105,10 @@ async def test_aretrieve_batch_uses_anthropic_provider_for_client(
@pytest.mark.asyncio
async def test_aretrieve_batch_resolves_api_key_from_model_info(
handler, patched_client
):
async def test_aretrieve_batch_resolves_api_key_from_model_info(handler, patched_client):
fake_client, _ = patched_client
# api_key=None -> handler falls back to AnthropicModelInfo.get_api_key().
with patch.object(
handler.anthropic_model_info, "get_api_key", return_value="sk-from-env"
):
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="sk-from-env"):
await handler.aretrieve_batch(
batch_id="msgbatch_abc",
api_base="https://api.anthropic.com",
@ -134,9 +124,7 @@ async def test_aretrieve_batch_resolves_api_key_from_model_info(
async def test_aretrieve_batch_missing_api_key_raises(handler, patched_client):
fake_client, _ = patched_client
# No api_key and resolver yields None -> hard error before any network call.
with patch.object(
handler.anthropic_model_info, "get_api_key", return_value=None
):
with patch.object(handler.anthropic_model_info, "get_api_key", return_value=None):
with pytest.raises(ValueError, match="Missing Anthropic API Key"):
await handler.aretrieve_batch(
batch_id="msgbatch_abc",
@ -165,9 +153,7 @@ async def test_aretrieve_batch_resolves_default_api_base(handler, patched_client
max_retries=0,
)
_, call_kwargs = fake_client.get.call_args
assert call_kwargs["url"] == (
"https://api.anthropic.com/v1/messages/batches/msgbatch_abc"
)
assert call_kwargs["url"] == ("https://api.anthropic.com/v1/messages/batches/msgbatch_abc")
@pytest.mark.asyncio
@ -176,9 +162,7 @@ async def test_aretrieve_batch_raises_for_status(handler):
error_response = httpx.Response(
status_code=404,
json={"error": "not found"},
request=httpx.Request(
"GET", "https://api.anthropic.com/v1/messages/batches/missing"
),
request=httpx.Request("GET", "https://api.anthropic.com/v1/messages/batches/missing"),
)
fake_client = MagicMock()
fake_client.get = AsyncMock(return_value=error_response)
@ -213,21 +197,15 @@ async def test_aretrieve_batch_invokes_pre_call_logging(handler, patched_client)
assert pre_kwargs["input"] == "msgbatch_abc"
assert pre_kwargs["api_key"] == "sk-ant-test"
# The logged api_base is the full retrieve URL, not the bare base.
assert pre_kwargs["additional_args"]["api_base"] == (
"https://api.anthropic.com/v1/messages/batches/msgbatch_abc"
)
assert pre_kwargs["additional_args"]["api_base"] == ("https://api.anthropic.com/v1/messages/batches/msgbatch_abc")
@pytest.mark.asyncio
async def test_aretrieve_batch_builds_default_logging_obj_when_absent(
handler, patched_client
):
async def test_aretrieve_batch_builds_default_logging_obj_when_absent(handler, patched_client):
# logging_obj=None -> handler constructs a real Logging object; the call
# must still complete (no AttributeError on a missing logger).
_, _ = patched_client
with patch(
"litellm.litellm_core_utils.litellm_logging.Logging"
) as logging_cls:
with patch("litellm.litellm_core_utils.litellm_logging.Logging") as logging_cls:
logging_cls.return_value = MagicMock()
batch = await handler.aretrieve_batch(
batch_id="msgbatch_abc",
@ -281,3 +259,99 @@ def test_retrieve_batch_sync_runs_to_result(handler, patched_client):
assert isinstance(batch, LiteLLMBatch)
assert batch.id == "msgbatch_abc"
assert batch.status == "completed"
# =========================================================================== #
# aretrieve_batch must not block the event loop on a WIF token exchange
# =========================================================================== #
_WIF_ENV = {
"ANTHROPIC_FEDERATION_RULE_ID": "fdrl_batches_seam",
"ANTHROPIC_ORGANIZATION_ID": "org-batches-seam",
"ANTHROPIC_IDENTITY_TOKEN": "batches-seam-inline-jwt",
}
class _BlockingPoster:
"""A token-endpoint poster 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):
self.thread_ids.append(threading.get_ident())
self.release.wait(timeout=5)
return httpx.Response(
200,
json={
"access_token": "sk-ant-oat01-batches-seam",
"token_type": "Bearer",
"expires_in": 3600,
},
)
@pytest.mark.asyncio
async def test_aretrieve_batch_wif_exchange_does_not_block_event_loop(handler, patched_client, monkeypatch):
"""Regression: aretrieve_batch called the synchronous validate_environment
directly, so a cold WIF mint ran inline on the event loop and froze every
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
fake_client, _ = patched_client
for name in (
"ANTHROPIC_API_KEY",
"ANTHROPIC_AUTH_TOKEN",
"ANTHROPIC_API_BASE",
"ANTHROPIC_BASE_URL",
):
monkeypatch.delenv(name, raising=False)
for name, value in _WIF_ENV.items():
monkeypatch.setenv(name, value)
poster = _BlockingPoster()
engine = JwtBearerTokenExchangeEngine(poster=poster)
def routed_through_injected_engine(litellm_params, api_base, model):
return get_anthropic_wif_token(litellm_params, api_base, model, engine)
monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", routed_through_injected_engine)
ticks = []
async def ticker():
for i in range(20):
await asyncio.sleep(0.005)
ticks.append(i)
ticker_task = asyncio.create_task(ticker())
await asyncio.sleep(0.02)
retrieve_task = asyncio.create_task(
handler.aretrieve_batch(
batch_id="msgbatch_abc",
api_base="https://api.anthropic.com",
api_key=None,
timeout=60.0,
max_retries=0,
)
)
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.
assert len(ticks) > 0
assert not retrieve_task.done()
poster.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()
sent_headers = fake_client.get.call_args.kwargs["headers"]
assert sent_headers["authorization"] == "Bearer sk-ant-oat01-batches-seam"

View file

@ -80,8 +80,11 @@ def test_validate_environment_preserves_existing_beta_header(config):
litellm_params={},
api_key="sk-ant-test",
)
# Existing beta header must NOT be overwritten.
assert headers["anthropic-beta"] == "custom-beta-value"
# Existing beta values are preserved and the batches beta is merged in.
assert set(headers["anthropic-beta"].split(",")) == {
"custom-beta-value",
"message-batches-2024-09-24",
}
def test_validate_environment_oauth_key_uses_bearer(config):
@ -100,9 +103,7 @@ def test_validate_environment_oauth_key_uses_bearer(config):
def test_validate_environment_missing_key_raises(config):
# No api_key passed and no env credentials -> get_auth_header returns None.
with patch.object(
config.anthropic_model_info, "get_auth_header", return_value=None
):
with patch.object(config.anthropic_model_info, "get_auth_header", return_value=None):
with pytest.raises(ValueError, match="Missing Anthropic API Key"):
config.validate_environment(
headers={},
@ -241,12 +242,7 @@ def test_get_retrieve_batch_url_uses_default_api_base(config):
def test_transform_retrieve_batch_request_returns_empty_dict(config):
assert (
config.transform_retrieve_batch_request(
batch_id="msgbatch_123", optional_params={}, litellm_params={}
)
== {}
)
assert config.transform_retrieve_batch_request(batch_id="msgbatch_123", optional_params={}, litellm_params={}) == {}
# =========================================================================== #
@ -455,9 +451,7 @@ def test_transform_retrieve_response_unparseable_json_raises(config):
def test_get_error_class_with_dict_headers(config):
err = config.get_error_class(
error_message="rate limited", status_code=429, headers={"x-ratelimit": "0"}
)
err = config.get_error_class(error_message="rate limited", status_code=429, headers={"x-ratelimit": "0"})
from litellm.llms.anthropic.common_utils import AnthropicError
assert isinstance(err, AnthropicError)
@ -467,9 +461,7 @@ def test_get_error_class_with_dict_headers(config):
def test_get_error_class_with_httpx_headers(config):
hdrs = httpx.Headers({"retry-after": "5"})
err = config.get_error_class(
error_message="server error", status_code=500, headers=hdrs
)
err = config.get_error_class(error_message="server error", status_code=500, headers=hdrs)
assert err.status_code == 500
assert err.message == "server error"
@ -543,9 +535,7 @@ def test_transform_response_skips_malformed_lines(config):
def fake_transform_parsed(*, completion_response, raw_response, model_response):
mr = ModelResponse()
setattr(
mr, "usage", Usage(prompt_tokens=7, completion_tokens=3, total_tokens=10)
)
setattr(mr, "usage", Usage(prompt_tokens=7, completion_tokens=3, total_tokens=10))
return mr
with patch.object(
@ -588,13 +578,16 @@ def test_transform_response_reraises_unexpected_error(config):
# A non-JSONDecodeError raised during usage aggregation must propagate
# (the outer `except Exception: raise e`), not be swallowed.
with patch.object(
config.anthropic_chat_config,
"transform_parsed_response",
side_effect=fake_transform_parsed,
), patch(
"litellm.cost_calculator.BaseTokenUsageProcessor.combine_usage_objects",
side_effect=RuntimeError("boom"),
with (
patch.object(
config.anthropic_chat_config,
"transform_parsed_response",
side_effect=fake_transform_parsed,
),
patch(
"litellm.cost_calculator.BaseTokenUsageProcessor.combine_usage_objects",
side_effect=RuntimeError("boom"),
),
):
with pytest.raises(RuntimeError, match="boom"):
config.transform_response(

View file

@ -48,9 +48,7 @@ class MockDynamicGuardrail(CustomGuardrail):
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
self.dynamic_params = self.get_guardrail_dynamic_request_body_params(
request_data
)
self.dynamic_params = self.get_guardrail_dynamic_request_body_params(request_data)
return inputs
@ -197,9 +195,7 @@ class TestAnthropicMessagesHandlerStreamingRequestData:
assert guardrail.request_data is not None
assert guardrail.request_data["response"] is mock_response
assert (
guardrail.request_data["litellm_metadata"]["user_api_key_user_id"] == "u-1"
)
assert guardrail.request_data["litellm_metadata"]["user_api_key_user_id"] == "u-1"
@pytest.mark.asyncio
async def test_mid_stream_chunk_passes_responses_so_far_and_metadata(self):
@ -211,9 +207,7 @@ class TestAnthropicMessagesHandlerStreamingRequestData:
with (
patch.object(handler, "_check_streaming_has_ended", return_value=False),
patch.object(
handler, "get_streaming_string_so_far", return_value="partial text"
),
patch.object(handler, "get_streaming_string_so_far", return_value="partial text"),
):
await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
@ -225,9 +219,7 @@ class TestAnthropicMessagesHandlerStreamingRequestData:
assert guardrail.request_data is not None
assert guardrail.request_data["responses"] is responses_so_far
assert (
guardrail.request_data["litellm_metadata"]["user_api_key_user_id"] == "u-1"
)
assert guardrail.request_data["litellm_metadata"]["user_api_key_user_id"] == "u-1"
class TestAnthropicMessagesHandlerStreamingOutputProcessing:
@ -276,17 +268,11 @@ class TestAnthropicMessagesHandlerInputProcessing:
data = {
"model": "claude-3-5-sonnet-20241022",
"messages": [{"role": "user", "content": "hello"}],
"litellm_metadata": {
"guardrails": [
{"cygnal-monitor": {"extra_body": {"policy_id": "policy-123"}}}
]
},
"litellm_metadata": {"guardrails": [{"cygnal-monitor": {"extra_body": {"policy_id": "policy-123"}}}]},
}
with patch("litellm.proxy.proxy_server.premium_user", True):
await handler.process_input_messages(
data=data, guardrail_to_apply=guardrail
)
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert data.get("litellm_metadata", {}).get("guardrails")
assert guardrail.dynamic_params == {"policy_id": "policy-123"}
@ -1183,9 +1169,7 @@ class TestAnthropicMessagesHandlerInputProcessing:
# Mock _check_streaming_has_ended to return False (stream not ended)
with (
patch.object(handler, "_check_streaming_has_ended", return_value=False),
patch.object(
handler, "get_streaming_string_so_far", return_value="partial text"
),
patch.object(handler, "get_streaming_string_so_far", return_value="partial text"),
):
responses_so_far = [b"data: some chunk"]
@ -1216,9 +1200,7 @@ class TestAnthropicMessagesHandlerInputProcessing:
data = {
"model": "claude-opus-4-6",
"messages": [
{"role": "user", "content": "What is the weather in San Francisco?"}
],
"messages": [{"role": "user", "content": "What is the weather in San Francisco?"}],
"tools": [
{
"type": "tool_search_tool_regex_20251119",
@ -1381,17 +1363,11 @@ class TestAnthropicMessagesIncrementalScan:
]
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
await handler.process_input_messages(
data=self._data(turn1, sid), guardrail_to_apply=guardrail
)
await handler.process_input_messages(data=self._data(turn1, sid), guardrail_to_apply=guardrail)
assert mock_api.call_count == 1
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [
"What is the capital of France?"
]
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == ["What is the capital of France?"]
mock_api.reset_mock()
await handler.process_input_messages(
data=self._data(turn2, sid), guardrail_to_apply=guardrail
)
await handler.process_input_messages(data=self._data(turn2, sid), guardrail_to_apply=guardrail)
assert mock_api.call_count == 1
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [
"Paris.",
@ -1830,9 +1806,7 @@ class TestAnthropicMessagesScanOnlyToolResults:
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.seen_texts == ["fetched POISON page"], (
"only the tool_result payload may reach the guardrail"
)
assert guardrail.seen_texts == ["fetched POISON page"], "only the tool_result payload may reach the guardrail"
assert guardrail.captured_inputs is not None
assert guardrail.captured_inputs.get("tools") is None
assert [m["role"] for m in guardrail.captured_inputs["structured_messages"]] == ["tool"]

View file

@ -23,9 +23,7 @@ async def test_make_call_passes_logging_obj_to_client_post():
mock_client = AsyncMock()
mock_response = MagicMock()
mock_response.aiter_lines = MagicMock(
return_value=iter(
[b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n']
)
return_value=iter([b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n'])
)
mock_client.post.return_value = mock_response
@ -94,9 +92,7 @@ def test_redacted_thinking_content_block_delta():
"data": "EuoBCoYBGAIiQJ/SxkPAgqxhKok29YrpJHRUJ0OT8ahCHKAwyhmRuUhtdmDX9+mn4gDzKNv3fVpQdB01zEPMzNY3QuTCd+1bdtEqQK6JuKHqdndbwpr81oVWb4wxd1GqF/7Jkw74IlQa27oobX+KuRkopr9Dllt/RDe7Se0sI1IkU7tJIAQCoP46OAwSDF51P09q67xhHlQ3ihoM2aOVlkghq/X0w8NlIjBMNvXYNbjhyrOcIg6kPFn2ed/KK7Cm5prYAtXCwkb4Wr5tUSoSHu9T5hKdJRbr6WsqEc7Lle7FULqMLZGkhqXyc3BA",
},
}
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=False, json_mode=False
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
model_response = model_response_iterator.chunk_parser(chunk=chunk)
print(f"\n\nmodel_response: {model_response}\n\n")
assert model_response.choices[0].delta.thinking_blocks is not None
@ -104,19 +100,14 @@ def test_redacted_thinking_content_block_delta():
print(
f"\n\nmodel_response.choices[0].delta.thinking_blocks[0]: {model_response.choices[0].delta.thinking_blocks[0]}\n\n"
)
assert (
model_response.choices[0].delta.thinking_blocks[0]["type"]
== "redacted_thinking"
)
assert model_response.choices[0].delta.thinking_blocks[0]["type"] == "redacted_thinking"
assert model_response.choices[0].delta.provider_specific_fields is not None
assert "thinking_blocks" in model_response.choices[0].delta.provider_specific_fields
def test_streaming_thinking_blocks_are_replayable_after_signature_delta():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
chunks = [
{
"type": "content_block_start",
@ -140,17 +131,12 @@ def test_streaming_thinking_blocks_are_replayable_after_signature_delta():
},
]
parsed_chunks = [
model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
]
parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks]
reasoning_content = "".join(
getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
for chunk in parsed_chunks
getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks
)
thinking_blocks = tuple(
block
for chunk in parsed_chunks
for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
)
expected_delta_blocks = (
{"type": "thinking", "thinking": "Step 1. "},
@ -164,18 +150,12 @@ def test_streaming_thinking_blocks_are_replayable_after_signature_delta():
assert reasoning_content == "Step 1. Step 2."
assert thinking_blocks == (*expected_delta_blocks, expected_thinking_block)
assert parsed_chunks[1].choices[0].delta.provider_specific_fields == {
"thinking_blocks": [expected_delta_blocks[0]]
}
assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == {
"thinking_blocks": [expected_thinking_block]
}
assert parsed_chunks[1].choices[0].delta.provider_specific_fields == {"thinking_blocks": [expected_delta_blocks[0]]}
assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == {"thinking_blocks": [expected_thinking_block]}
def test_streaming_unsigned_thinking_deltas_keep_reasoning_content():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
chunks = [
{
"type": "content_block_start",
@ -195,17 +175,12 @@ def test_streaming_unsigned_thinking_deltas_keep_reasoning_content():
{"type": "content_block_stop", "index": 0},
]
parsed_chunks = [
model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
]
parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks]
reasoning_content = "".join(
getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
for chunk in parsed_chunks
getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks
)
thinking_blocks = tuple(
block
for chunk in parsed_chunks
for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
)
assert reasoning_content == "Step 1. Step 2."
@ -216,9 +191,7 @@ def test_streaming_unsigned_thinking_deltas_keep_reasoning_content():
def test_streaming_truncated_thinking_deltas_keep_reasoning_content():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
chunks = [
{
"type": "content_block_start",
@ -237,17 +210,12 @@ def test_streaming_truncated_thinking_deltas_keep_reasoning_content():
},
]
parsed_chunks = [
model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
]
parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks]
reasoning_content = "".join(
getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
for chunk in parsed_chunks
getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks
)
thinking_blocks = tuple(
block
for chunk in parsed_chunks
for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
)
assert reasoning_content == "Step 1. Step 2."
@ -258,9 +226,7 @@ def test_streaming_truncated_thinking_deltas_keep_reasoning_content():
def test_handle_json_mode_chunk_response_format_tool():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
response_format_tool = ChatCompletionToolCallChunk(
id="tool_123",
type="function",
@ -271,9 +237,7 @@ def test_handle_json_mode_chunk_response_format_tool():
index=0,
)
text, tool_use = model_response_iterator._handle_json_mode_chunk(
"", response_format_tool
)
text, tool_use = model_response_iterator._handle_json_mode_chunk("", response_format_tool)
print(f"\n\nresponse_format_tool text: {text}\n\n")
print(f"\n\nresponse_format_tool tool_use: {tool_use}\n\n")
@ -282,15 +246,11 @@ def test_handle_json_mode_chunk_response_format_tool():
def test_handle_json_mode_chunk_regular_tool():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
regular_tool = ChatCompletionToolCallChunk(
id="tool_456",
type="function",
function=ChatCompletionToolCallFunctionChunk(
name="get_weather", arguments='{"location": "San Francisco, CA"}'
),
function=ChatCompletionToolCallFunctionChunk(name="get_weather", arguments='{"location": "San Francisco, CA"}'),
index=0,
)
@ -304,17 +264,13 @@ def test_handle_json_mode_chunk_regular_tool():
def test_handle_json_mode_chunk_streaming_response_format_tool():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
# First chunk: response_format tool with id and name, but no arguments
first_chunk = ChatCompletionToolCallChunk(
id="tool_123",
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=RESPONSE_FORMAT_TOOL_NAME, arguments=""
),
function=ChatCompletionToolCallFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME, arguments=""),
index=0,
)
@ -322,9 +278,7 @@ def test_handle_json_mode_chunk_streaming_response_format_tool():
second_chunk = ChatCompletionToolCallChunk(
id=None,
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=None, arguments='{"question": "What is the weather?"'
),
function=ChatCompletionToolCallFunctionChunk(name=None, arguments='{"question": "What is the weather?"'),
index=0,
)
@ -332,9 +286,7 @@ def test_handle_json_mode_chunk_streaming_response_format_tool():
third_chunk = ChatCompletionToolCallChunk(
id=None,
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=None, arguments=', "answer": "It is sunny"}'
),
function=ChatCompletionToolCallFunctionChunk(name=None, arguments=', "answer": "It is sunny"}'),
index=0,
)
@ -365,9 +317,7 @@ def test_handle_json_mode_chunk_streaming_response_format_tool():
def test_handle_json_mode_chunk_streaming_regular_tool():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
# First chunk: regular tool with id and name, but no arguments
first_chunk = ChatCompletionToolCallChunk(
@ -381,9 +331,7 @@ def test_handle_json_mode_chunk_streaming_regular_tool():
second_chunk = ChatCompletionToolCallChunk(
id=None,
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=None, arguments='{"location": "San Francisco, CA"}'
),
function=ChatCompletionToolCallFunctionChunk(name=None, arguments='{"location": "San Francisco, CA"}'),
index=0,
)
@ -408,27 +356,19 @@ def test_handle_json_mode_chunk_streaming_regular_tool():
def test_response_format_tool_finish_reason():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
# First chunk: response_format tool
response_format_tool = ChatCompletionToolCallChunk(
id="tool_123",
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=RESPONSE_FORMAT_TOOL_NAME, arguments='{"answer": "test"}'
),
function=ChatCompletionToolCallFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME, arguments='{"answer": "test"}'),
index=0,
)
# Process the tool call (should set converted_response_format_tool flag)
text, tool_use = model_response_iterator._handle_json_mode_chunk(
"", response_format_tool
)
print(
f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n"
)
text, tool_use = model_response_iterator._handle_json_mode_chunk("", response_format_tool)
print(f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n")
# Simulate message_delta chunk with tool_use stop_reason
message_delta_chunk = {
@ -447,25 +387,19 @@ def test_response_format_tool_finish_reason():
def test_regular_tool_finish_reason():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
# First chunk: regular tool (not response_format)
regular_tool = ChatCompletionToolCallChunk(
id="tool_456",
type="function",
function=ChatCompletionToolCallFunctionChunk(
name="get_weather", arguments='{"location": "San Francisco, CA"}'
),
function=ChatCompletionToolCallFunctionChunk(name="get_weather", arguments='{"location": "San Francisco, CA"}'),
index=0,
)
# Process the tool call (should NOT set converted_response_format_tool flag)
text, tool_use = model_response_iterator._handle_json_mode_chunk("", regular_tool)
print(
f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n"
)
print(f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n")
# Simulate message_delta chunk with tool_use stop_reason
message_delta_chunk = {
@ -525,9 +459,7 @@ def test_text_only_streaming_has_index_zero():
for chunk in chunks:
parsed = iterator.chunk_parser(chunk)
if parsed.choices:
assert (
parsed.choices[0].index == 0
), f"Expected index=0, got {parsed.choices[0].index}"
assert parsed.choices[0].index == 0, f"Expected index=0, got {parsed.choices[0].index}"
def test_streaming_thinking_deltas_count_reasoning_tokens_in_usage():
@ -704,9 +636,7 @@ def test_anthropic_completion_streaming_usage_matches_non_streaming_with_thinkin
]
self._write_response(
content_type="text/event-stream",
body="".join(
f"data: {json.dumps(event)}\n\n" for event in events
).encode("utf-8"),
body="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode("utf-8"),
)
return
@ -787,13 +717,9 @@ def test_anthropic_completion_streaming_usage_matches_non_streaming_with_thinkin
assert content_chunks == [answer_text]
assert stream_usage is not None
stream_completion_details = stream_usage["completion_tokens_details"]
assert (
stream_completion_details["reasoning_tokens"]
== non_stream_details.reasoning_tokens
)
assert stream_completion_details["reasoning_tokens"] == non_stream_details.reasoning_tokens
assert stream_completion_details["text_tokens"] == (
stream_usage["completion_tokens"]
- stream_completion_details["reasoning_tokens"]
stream_usage["completion_tokens"] - stream_completion_details["reasoning_tokens"]
)
assert requests_seen == [
{
@ -885,9 +811,9 @@ def test_text_and_tool_streaming_has_index_zero():
for chunk in chunks:
parsed = iterator.chunk_parser(chunk)
if parsed.choices:
assert (
parsed.choices[0].index == 0
), f"Expected index=0 for chunk type {chunk.get('type')}, got {parsed.choices[0].index}"
assert parsed.choices[0].index == 0, (
f"Expected index=0 for chunk type {chunk.get('type')}, got {parsed.choices[0].index}"
)
def test_multiple_tools_streaming_has_index_zero():
@ -940,15 +866,11 @@ def test_multiple_tools_streaming_has_index_zero():
for chunk in chunks:
parsed = iterator.chunk_parser(chunk)
if parsed.choices:
assert (
parsed.choices[0].index == 0
), f"Expected index=0, got {parsed.choices[0].index}"
assert parsed.choices[0].index == 0, f"Expected index=0, got {parsed.choices[0].index}"
def test_streaming_chunks_have_stable_ids():
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=False, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
first_chunk = {
"type": "content_block_delta",
"index": 0,
@ -973,9 +895,7 @@ def test_partial_json_chunk_accumulation():
This tests the fix for https://github.com/BerriAI/litellm/issues/17473
where network fragmentation can cause SSE data to arrive in partial chunks.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
partial_chunk_1 = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hel'
partial_chunk_2 = 'lo"}}'
@ -983,31 +903,21 @@ def test_partial_json_chunk_accumulation():
# First partial chunk should return None (still accumulating)
result1 = iterator._parse_sse_data(f"data:{partial_chunk_1}")
assert result1 is None, "First partial chunk should return None while accumulating"
assert (
iterator.chunk_type == "accumulated_json"
), "Should switch to accumulated_json mode"
assert (
iterator.accumulated_json == partial_chunk_1
), "Should have accumulated first part"
assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode"
assert iterator.accumulated_json == partial_chunk_1, "Should have accumulated first part"
# Second partial chunk should complete the JSON and return a parsed result
result2 = iterator._parse_sse_data(f"data:{partial_chunk_2}")
assert result2 is not None, "Second chunk should return parsed result"
assert (
iterator.accumulated_json == ""
), "Buffer should be cleared after successful parse"
assert (
result2.choices[0].delta.content == "Hello"
), f"Expected 'Hello', got '{result2.choices[0].delta.content}'"
assert iterator.accumulated_json == "", "Buffer should be cleared after successful parse"
assert result2.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result2.choices[0].delta.content}'"
def test_complete_json_chunk_no_accumulation():
"""
Test that complete JSON chunks are parsed immediately without accumulation.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
complete_chunk = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}'
@ -1015,18 +925,14 @@ def test_complete_json_chunk_no_accumulation():
assert result is not None, "Complete chunk should return parsed result immediately"
assert iterator.chunk_type == "valid_json", "Should remain in valid_json mode"
assert iterator.accumulated_json == "", "Buffer should remain empty"
assert (
result.choices[0].delta.content == "Hello"
), f"Expected 'Hello', got '{result.choices[0].delta.content}'"
assert result.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result.choices[0].delta.content}'"
def test_multiple_partial_chunks_accumulation():
"""
Test that multiple partial chunks can be accumulated across several iterations.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# Split a JSON chunk into three parts
part1 = '{"type":"content_block_del'
@ -1194,9 +1100,7 @@ def test_web_search_tool_result_no_extra_tool_calls():
The issue was that web_search_tool_result blocks have input_json_delta events with {}
that were incorrectly being converted to tool calls.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# Simulate the streaming sequence:
# 1. server_tool_use block starts (web_search)
@ -1271,9 +1175,7 @@ def test_web_search_tool_result_no_extra_tool_calls():
# Should have exactly 2 tool calls:
# 1. From content_block_start (server_tool_use) with id and name
# 2. From content_block_delta with the actual query
assert (
len(tool_calls_emitted) == 2
), f"Expected 2 tool calls, got {len(tool_calls_emitted)}"
assert len(tool_calls_emitted) == 2, f"Expected 2 tool calls, got {len(tool_calls_emitted)}"
# First tool call should have the id and name
assert tool_calls_emitted[0]["id"] == "srvtoolu_01ABC123"
@ -1289,9 +1191,7 @@ def test_current_content_block_type_tracking():
"""
Test that current_content_block_type is properly tracked and reset.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# Initially should be None
assert iterator.current_content_block_type is None
@ -1344,9 +1244,7 @@ def test_web_search_tool_result_captured_in_provider_specific_fields():
The web_search_tool_result content comes ALL AT ONCE in content_block_start,
not in deltas, so we need to capture it there.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# Simulate the streaming sequence with web_search_tool_result
chunks = [
@ -1417,23 +1315,15 @@ def test_web_search_tool_result_captured_in_provider_specific_fields():
and parsed.choices[0].delta.provider_specific_fields
and "web_search_results" in parsed.choices[0].delta.provider_specific_fields
):
web_search_results = parsed.choices[0].delta.provider_specific_fields[
"web_search_results"
]
web_search_results = parsed.choices[0].delta.provider_specific_fields["web_search_results"]
# Verify web_search_results was captured
assert web_search_results is not None, "web_search_results should be captured"
assert len(web_search_results) == 1, "Should have 1 web_search_tool_result block"
assert (
web_search_results[0]["type"] == "web_search_tool_result"
), "Block type should be web_search_tool_result"
assert (
web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123"
), "tool_use_id should match"
assert web_search_results[0]["type"] == "web_search_tool_result", "Block type should be web_search_tool_result"
assert web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123", "tool_use_id should match"
assert len(web_search_results[0]["content"]) == 2, "Should have 2 search results"
assert (
web_search_results[0]["content"][0]["title"] == "Fun Otter Facts"
), "First result title should match"
assert web_search_results[0]["content"][0]["title"] == "Fun Otter Facts", "First result title should match"
def test_web_fetch_tool_result_captured_in_provider_specific_fields():
@ -1447,9 +1337,7 @@ def test_web_fetch_tool_result_captured_in_provider_specific_fields():
The web_fetch_tool_result content comes ALL AT ONCE in content_block_start,
not in deltas, so we need to capture it there.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# Simulate the streaming sequence with web_fetch_tool_result
chunks = [
@ -1520,25 +1408,15 @@ def test_web_fetch_tool_result_captured_in_provider_specific_fields():
and parsed.choices[0].delta.provider_specific_fields
and "web_search_results" in parsed.choices[0].delta.provider_specific_fields
):
web_search_results = parsed.choices[0].delta.provider_specific_fields[
"web_search_results"
]
web_search_results = parsed.choices[0].delta.provider_specific_fields["web_search_results"]
# Verify web_fetch_tool_result was captured (stored in web_search_results list)
assert web_search_results is not None, "web_search_results should be captured"
assert len(web_search_results) == 1, "Should have 1 web_fetch_tool_result block"
assert (
web_search_results[0]["type"] == "web_fetch_tool_result"
), "Block type should be web_fetch_tool_result"
assert (
web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123"
), "tool_use_id should match"
assert (
web_search_results[0]["content"]["url"] == "https://example.com"
), "URL should match"
assert (
web_search_results[0]["content"]["content"]["title"] == "Example Page"
), "Title should match"
assert web_search_results[0]["type"] == "web_fetch_tool_result", "Block type should be web_fetch_tool_result"
assert web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123", "tool_use_id should match"
assert web_search_results[0]["content"]["url"] == "https://example.com", "URL should match"
assert web_search_results[0]["content"]["content"]["title"] == "Example Page", "Title should match"
def test_web_fetch_tool_result_no_extra_tool_calls():
@ -1551,9 +1429,7 @@ def test_web_fetch_tool_result_no_extra_tool_calls():
The issue was that web_fetch_tool_result blocks have input_json_delta events with {}
that were incorrectly being converted to tool calls.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# to verify it doesn't emit tool calls
chunks = [
@ -1597,9 +1473,9 @@ def test_web_fetch_tool_result_no_extra_tool_calls():
tool_call_count += 1
# Should have 0 tool calls - web_fetch_tool_result should not emit tool calls
assert (
tool_call_count == 0
), f"Expected 0 tool calls, got {tool_call_count}. web_fetch_tool_result should not emit tool calls"
assert tool_call_count == 0, (
f"Expected 0 tool calls, got {tool_call_count}. web_fetch_tool_result should not emit tool calls"
)
def test_container_in_provider_specific_fields_streaming():
@ -1609,9 +1485,7 @@ def test_container_in_provider_specific_fields_streaming():
When container with skills is used, the container field should be present in
the provider_specific_fields of the message_delta chunk.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# Simulate streaming chunks
chunks = [
@ -1679,20 +1553,12 @@ def test_container_in_provider_specific_fields_streaming():
and parsed.choices[0].delta.provider_specific_fields
and "container" in parsed.choices[0].delta.provider_specific_fields
):
container_field = parsed.choices[0].delta.provider_specific_fields[
"container"
]
container_field = parsed.choices[0].delta.provider_specific_fields["container"]
# Verify container was captured
assert (
container_field is not None
), "container should be captured in provider_specific_fields"
assert (
container_field["id"] == "container_011CW9hA9zpZ8xD3bjjShy4p"
), "container id should match"
assert (
container_field["expires_at"] == "2025-12-16T04:57:16.913181Z"
), "expires_at should match"
assert container_field is not None, "container should be captured in provider_specific_fields"
assert container_field["id"] == "container_011CW9hA9zpZ8xD3bjjShy4p", "container id should match"
assert container_field["expires_at"] == "2025-12-16T04:57:16.913181Z", "expires_at should match"
assert len(container_field["skills"]) == 1, "Should have 1 skill"
assert container_field["skills"][0]["skill_id"] == "pptx", "skill_id should be pptx"
assert container_field["skills"][0]["version"] == "20251013", "version should match"
@ -1705,9 +1571,7 @@ def test_container_in_provider_specific_fields_non_streaming():
When container with skills is used in non-streaming, the container field should be
present in the provider_specific_fields of the response.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=False, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
# Simulate a message_delta chunk with container (as it would appear in non-streaming)
message_delta_chunk = {
@ -1743,21 +1607,13 @@ def test_container_in_provider_specific_fields_non_streaming():
# Verify container is in provider_specific_fields
assert model_response.choices[0].delta.provider_specific_fields is not None
assert "container" in model_response.choices[0].delta.provider_specific_fields
container_field = model_response.choices[0].delta.provider_specific_fields[
"container"
]
container_field = model_response.choices[0].delta.provider_specific_fields["container"]
assert container_field["id"] == "container_abc123xyz", "container id should match"
assert (
container_field["expires_at"] == "2025-12-20T10:30:00.000000Z"
), "expires_at should match"
assert container_field["expires_at"] == "2025-12-20T10:30:00.000000Z", "expires_at should match"
assert len(container_field["skills"]) == 2, "Should have 2 skills"
assert (
container_field["skills"][0]["skill_id"] == "code_execution"
), "First skill_id should be code_execution"
assert (
container_field["skills"][1]["skill_id"] == "pptx"
), "Second skill_id should be pptx"
assert container_field["skills"][0]["skill_id"] == "code_execution", "First skill_id should be code_execution"
assert container_field["skills"][1]["skill_id"] == "pptx", "Second skill_id should be pptx"
def test_container_absent_when_not_provided():
@ -1766,9 +1622,7 @@ def test_container_absent_when_not_provided():
This ensures we don't add empty or None container fields.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=False, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
# message_delta without container
message_delta_chunk = {
@ -1787,9 +1641,9 @@ def test_container_absent_when_not_provided():
# Verify container is NOT in provider_specific_fields when not provided
if model_response.choices[0].delta.provider_specific_fields:
assert (
"container" not in model_response.choices[0].delta.provider_specific_fields
), "container should not be present when not provided in delta"
assert "container" not in model_response.choices[0].delta.provider_specific_fields, (
"container should not be present when not provided in delta"
)
def test_streaming_code_execution_produces_code_interpreter_results():
@ -1985,8 +1839,7 @@ def test_streaming_multiple_code_executions_no_duplicates():
# Second (final) emission: cumulative list with BOTH results
# This is what stream_chunk_builder will pick as "last value wins"
assert len(emissions[1]) == 2, (
f"Expected final emission to have 2 results, got {len(emissions[1])}. "
f"IDs: {[r.id for r in emissions[1]]}"
f"Expected final emission to have 2 results, got {len(emissions[1])}. IDs: {[r.id for r in emissions[1]]}"
)
assert emissions[1][0].id == "srvtoolu_01AAA"
assert emissions[1][0].code == "echo first"
@ -2150,9 +2003,7 @@ def test_empty_output_produces_null_outputs():
assert code_results is not None, "No code_interpreter_results emitted"
assert len(code_results) == 1
assert code_results[0].id == "srvtoolu_01AAA"
assert (
code_results[0].outputs is None
), f"Expected outputs=None for empty execution, got {code_results[0].outputs}"
assert code_results[0].outputs is None, f"Expected outputs=None for empty execution, got {code_results[0].outputs}"
def test_non_bash_tool_result_skipped():
@ -2215,12 +2066,10 @@ def test_non_bash_tool_result_skipped():
code_results = psf["code_interpreter_results"]
# code_interpreter_results should be emitted but empty (no bash results)
assert (
code_results is not None
), "Expected code_interpreter_results key to be emitted"
assert (
len(code_results) == 0
), f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}"
assert code_results is not None, "Expected code_interpreter_results key to be emitted"
assert len(code_results) == 0, (
f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}"
)
class TestRustChatCompletionsHook:
@ -2257,13 +2106,9 @@ class TestRustChatCompletionsHook:
from litellm.rust_bridge import chat_completions as bridge
monkeypatch.delenv("LITELLM_RUST", raising=False)
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
yield
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
@staticmethod
def _completion_kwargs(**overrides):
@ -2361,9 +2206,7 @@ class TestRustChatCompletionsHook:
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
seen = self._inject()
AnthropicChatCompletion().completion(
**self._completion_kwargs(optional_params={"max_tokens": 7})
)
AnthropicChatCompletion().completion(**self._completion_kwargs(optional_params={"max_tokens": 7}))
assert seen["call"][0]["optional_params"]["max_tokens"] == 7
def test_without_the_opt_in_the_core_is_never_consulted(self):
@ -2371,15 +2214,14 @@ class TestRustChatCompletionsHook:
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
seen = self._inject()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
) as transform, patch.object(
AnthropicChatCompletion, "acompletion_function"
with (
patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
) as transform,
patch.object(AnthropicChatCompletion, "acompletion_function"),
):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(litellm_params={})
)
AnthropicChatCompletion().completion(**self._completion_kwargs(litellm_params={}))
except Exception:
# The Python path goes on to make an HTTP call; reaching it is
# the assertion, so the network failure below is expected.
@ -2393,9 +2235,7 @@ class TestRustChatCompletionsHook:
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
seen = self._inject(decline_reason="unrecognized request parameter")
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
):
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
try:
AnthropicChatCompletion().completion(**self._completion_kwargs())
except Exception:
@ -2408,9 +2248,7 @@ class TestRustChatCompletionsHook:
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
seen = self._inject()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
):
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(optional_params={"max_tokens": 16, "stream": True})
@ -2424,9 +2262,7 @@ class TestRustChatCompletionsHook:
seen = self._inject()
logging_obj = MagicMock()
AnthropicChatCompletion().completion(
**self._completion_kwargs(logging_obj=logging_obj)
)
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
assert logging_obj.pre_call.call_count == 1
assert len(seen["call"]) == 1
@ -2440,9 +2276,7 @@ class TestRustChatCompletionsHook:
self._inject()
logging_obj = MagicMock()
AnthropicChatCompletion().completion(
**self._completion_kwargs(logging_obj=logging_obj)
)
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
assert logging_obj.post_call.call_count == 1
logged = logging_obj.post_call.call_args.kwargs["original_response"]
@ -2467,18 +2301,12 @@ class TestRustChatCompletionsHook:
raise _Declined("blank message text")
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, chat_completions=declining_native
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native)
logging_obj, calls = self._recording_logging_obj()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
):
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(logging_obj=logging_obj)
)
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
except Exception:
# The Python path goes on to make an HTTP call; the log count is
# the assertion, so a failure past this point is expected.
@ -2503,21 +2331,15 @@ class TestRustChatCompletionsHook:
async def declining_native(**_kwargs):
raise _Declined("blank message text")
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, achat_completions=declining_native
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=declining_native)
sentinel = object()
async def python_path(**_kwargs):
return sentinel
with patch.object(
AnthropicChatCompletion, "acompletion_function", side_effect=python_path
) as python_call:
result = await AnthropicChatCompletion().completion(
**self._completion_kwargs(acompletion=True)
)
with patch.object(AnthropicChatCompletion, "acompletion_function", side_effect=python_path) as python_call:
result = await AnthropicChatCompletion().completion(**self._completion_kwargs(acompletion=True))
assert result is sentinel
assert python_call.called, "a failing rust call must re-enter the python path"
@ -2530,20 +2352,15 @@ class TestRustChatCompletionsHook:
async def native(**_kwargs):
return dict(self.RUST_RESPONSE)
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, achat_completions=native
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=native)
with patch.object(AnthropicChatCompletion, "acompletion_function") as python_call:
result = await AnthropicChatCompletion().completion(
**self._completion_kwargs(acompletion=True)
)
result = await AnthropicChatCompletion().completion(**self._completion_kwargs(acompletion=True))
assert result.choices[0].message.content == "hello from rust"
assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
assert not python_call.called
def test_pre_call_logging_fires_once_when_the_sync_rust_call_declines(self, monkeypatch):
"""One request, one pre_call, on the synchronous path too. Without the
suppression the Python path logs a second time for the same attempt."""
@ -2563,27 +2380,19 @@ class TestRustChatCompletionsHook:
def declining_native(**_kwargs):
raise _Declined("blank message text")
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, chat_completions=declining_native
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native)
logging_obj, calls = self._recording_logging_obj()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
):
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(logging_obj=logging_obj)
)
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
except Exception:
# The Python path goes on to make an HTTP call; the log count is
# the assertion, so a failure past this point is expected.
pass
assert len(calls["pre_call"]) == 1
assert calls["pre_call"][0]["additional_args"]["complete_input_dict"]["model"] == (
"claude-sonnet-4-5"
)
assert calls["pre_call"][0]["additional_args"]["complete_input_dict"]["model"] == ("claude-sonnet-4-5")
def test_pre_call_logging_still_fires_when_rust_is_not_involved(self, monkeypatch):
"""The suppression must not swallow the log on the ordinary path."""
@ -2592,9 +2401,7 @@ class TestRustChatCompletionsHook:
self._inject()
logging_obj, calls = self._recording_logging_obj()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
):
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(litellm_params={}, logging_obj=logging_obj)

View file

@ -1,4 +1,3 @@
import pytest
from unittest.mock import MagicMock, patch
@ -33,13 +32,9 @@ def test_response_format_transformation_unit_test():
"additionalProperties": False,
}
result = config._create_json_tool_call_for_response_format(
json_schema=response_format_json_schema
)
result = config._create_json_tool_call_for_response_format(json_schema=response_format_json_schema)
assert result["input_schema"]["properties"] == {
"agent_doing": {"title": "Agent Doing", "type": "string"}
}
assert result["input_schema"]["properties"] == {"agent_doing": {"title": "Agent Doing", "type": "string"}}
print(result)
@ -550,9 +545,7 @@ def test_extract_response_content_with_citations():
},
}
_, citations, _, _, _, _, _, _ = config.extract_response_content(
completion_response
)
_, citations, _, _, _, _, _, _ = config.extract_response_content(completion_response)
assert citations == [
[
{
@ -625,12 +618,8 @@ def test_web_search_tool_transformation():
assert anthropic_web_search_tool["user_location"]["city"] == "San Francisco"
@pytest.mark.parametrize(
"search_context_size, expected_max_uses", [("low", 1), ("medium", 5), ("high", 10)]
)
def test_web_search_tool_transformation_with_search_context_size(
search_context_size, expected_max_uses
):
@pytest.mark.parametrize("search_context_size, expected_max_uses", [("low", 1), ("medium", 5), ("high", 10)])
def test_web_search_tool_transformation_with_search_context_size(search_context_size, expected_max_uses):
from litellm.types.llms.openai import OpenAIWebSearchOptions
config = AnthropicConfig()
@ -805,10 +794,7 @@ def test_web_search_tool_result_in_provider_specific_fields():
assert "web_search_results" in provider_fields
assert len(provider_fields["web_search_results"]) == 1
assert provider_fields["web_search_results"][0]["type"] == "web_search_tool_result"
assert (
provider_fields["web_search_results"][0]["tool_use_id"]
== "srvtoolu_provider_test"
)
assert provider_fields["web_search_results"][0]["tool_use_id"] == "srvtoolu_provider_test"
def test_multiple_web_search_tool_results():
@ -1032,10 +1018,7 @@ def test_transform_response_with_prefix_prompt():
)
assert result is not None
assert (
result.choices[0].message.content
== "You are a helpful assistant. The grass is green."
)
assert result.choices[0].message.content == "You are a helpful assistant. The grass is green."
def test_get_supported_params_thinking():
@ -1150,18 +1133,12 @@ def test_anthropic_beta_header_merging_with_output_format():
}
}
result_headers = config.update_headers_with_optional_anthropic_beta(
headers, optional_params
)
result_headers = config.update_headers_with_optional_anthropic_beta(headers, optional_params)
# Both beta headers should be present
beta_value = result_headers["anthropic-beta"]
assert (
"context-1m-2025-08-07" in beta_value
), f"User's context-1m beta header missing from: {beta_value}"
assert (
"structured-outputs-2025-11-13" in beta_value
), f"Structured output beta header missing from: {beta_value}"
assert "context-1m-2025-08-07" in beta_value, f"User's context-1m beta header missing from: {beta_value}"
assert "structured-outputs-2025-11-13" in beta_value, f"Structured output beta header missing from: {beta_value}"
def test_anthropic_beta_header_merging_with_multiple_features():
@ -1183,9 +1160,7 @@ def test_anthropic_beta_header_merging_with_multiple_features():
"tools": [{"type": "web_fetch_20250910", "name": "web_fetch"}],
}
result_headers = config.update_headers_with_optional_anthropic_beta(
headers, optional_params
)
result_headers = config.update_headers_with_optional_anthropic_beta(headers, optional_params)
beta_value = result_headers["anthropic-beta"]
@ -1228,9 +1203,7 @@ def test_anthropic_structured_output_beta_header():
"strict": True,
"schema": {
"description": 'Progress report for the thinking process\n\nThis model represents a snapshot of the agent\'s current progress during\nthe thinking process, providing a brief description of the current activity.\n\nAttributes:\n agent_doing: Brief description of what the agent is currently doing.\n Should be kept under 10 words. Example: "Learning about home automation"',
"properties": {
"agent_doing": {"title": "Agent Doing", "type": "string"}
},
"properties": {"agent_doing": {"title": "Agent Doing", "type": "string"}},
"required": ["agent_doing"],
"title": "ThinkingStep",
"type": "object",
@ -1244,10 +1217,7 @@ def test_anthropic_structured_output_beta_header():
assert response is not None
print(f"response: {response}")
print(f"raw_request_headers: {response['raw_request_headers']}")
assert (
"structured-outputs-2025-11-13"
in response["raw_request_headers"]["anthropic-beta"]
)
assert "structured-outputs-2025-11-13" in response["raw_request_headers"]["anthropic-beta"]
@pytest.mark.parametrize(
@ -1383,9 +1353,7 @@ def test_tool_search_regex_detection():
config = AnthropicModelInfo()
# Test with tool search regex tool
tools = [
{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}
]
tools = [{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}]
assert config.is_tool_search_used(tools) is True
# Test without tool search
@ -1400,9 +1368,7 @@ def test_tool_search_bm25_detection():
config = AnthropicModelInfo()
# Test with tool search BM25 tool
tools = [
{"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"}
]
tools = [{"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"}]
assert config.is_tool_search_used(tools) is True
@ -1594,9 +1560,7 @@ def test_tool_search_complete_response_parsing():
"tool_use_id": "srvtoolu_015i6aVA2niwzv4RG4DtnxDJ",
"content": {
"type": "tool_search_tool_search_result",
"tool_references": [
{"type": "tool_reference", "tool_name": "get_weather"}
],
"tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}],
},
},
{"type": "text", "text": "Great! I found a weather tool."},
@ -1647,9 +1611,7 @@ def test_tool_search_complete_response_parsing():
assert usage.server_tool_use is not None
assert usage.server_tool_use.web_search_requests == 0
assert (
usage.server_tool_use.tool_search_requests == 1
) # Counted from server_tool_use blocks
assert usage.server_tool_use.tool_search_requests == 1 # Counted from server_tool_use blocks
def test_allowed_callers_field_preservation():
@ -1701,9 +1663,7 @@ def test_programmatic_tool_calling_beta_header():
assert is_programmatic is True
# Test header generation
headers = model_info.get_anthropic_headers(
api_key="test-key", programmatic_tool_calling_used=True
)
headers = model_info.get_anthropic_headers(api_key="test-key", programmatic_tool_calling_used=True)
assert "anthropic-beta" in headers
assert "advanced-tool-use-2025-11-20" in headers["anthropic-beta"]
@ -1847,9 +1807,7 @@ def test_input_examples_beta_header():
assert is_examples_used is True
# Test header generation
headers = model_info.get_anthropic_headers(
api_key="test-key", input_examples_used=True
)
headers = model_info.get_anthropic_headers(api_key="test-key", input_examples_used=True)
assert "anthropic-beta" in headers
assert "advanced-tool-use-2025-11-20" in headers["anthropic-beta"]
@ -1935,10 +1893,7 @@ def test_input_examples_empty_list_not_added():
transformed_tool, _ = config._map_tool_helper(tool)
assert transformed_tool is not None
# Empty list should not be added
assert (
"input_examples" not in transformed_tool
or len(transformed_tool.get("input_examples", [])) == 0
)
assert "input_examples" not in transformed_tool or len(transformed_tool.get("input_examples", [])) == 0
# ============ Effort Parameter Tests ============
@ -1998,9 +1953,7 @@ def test_effort_beta_header_injection():
effort_used = model_info.is_effort_used(optional_params=optional_params, custom_llm_provider="anthropic")
assert effort_used is True
headers = model_info.get_anthropic_headers(
api_key="test-key", effort_used=effort_used
)
headers = model_info.get_anthropic_headers(api_key="test-key", effort_used=effort_used)
assert "anthropic-beta" in headers
assert "effort-2025-11-24" in headers["anthropic-beta"]
@ -2026,9 +1979,7 @@ def test_effort_validation():
optional_params = {"output_config": {"effort": "invalid"}}
with pytest.raises(
litellm.exceptions.BadRequestError, match="Invalid effort value"
):
with pytest.raises(litellm.exceptions.BadRequestError, match="Invalid effort value"):
config.transform_request(
model="claude-opus-4-5-20251101",
messages=messages,
@ -2264,16 +2215,8 @@ def test_anthropic_model_supports_speed_param_rejects_non_anthropic_providers(
):
"""Fast mode is direct-Anthropic-only. Vertex/Azure/Bedrock strip their prefix
before the shared transform runs, so the bare Opus id must still be rejected."""
assert (
AnthropicConfig._model_supports_speed_param(
"claude-opus-4-8", custom_llm_provider
)
is False
)
assert (
AnthropicConfig._model_supports_speed_param("claude-opus-4-8", "anthropic")
is True
)
assert AnthropicConfig._model_supports_speed_param("claude-opus-4-8", custom_llm_provider) is False
assert AnthropicConfig._model_supports_speed_param("claude-opus-4-8", "anthropic") is True
def test_vertex_anthropic_drops_speed_for_opus_with_drop_params(monkeypatch):
@ -2572,9 +2515,7 @@ def test_supports_effort_level_handles_provider_prefixes(model, level, expected)
("claude-opus-4-5-20251101", None, False),
],
)
def test_validate_effort_for_model_centralises_per_model_gating(
model, effort, expect_error
):
def test_validate_effort_for_model_centralises_per_model_gating(model, effort, expect_error):
err = AnthropicConfig._validate_effort_for_model(model, effort, "anthropic")
if expect_error:
assert err is not None
@ -2623,11 +2564,7 @@ def test_transform_request_injects_dummy_tool_without_tools_param():
litellm.modify_params = prev_modify_params
assert "tools" in result
names = [
t.get("name")
for t in result["tools"]
if isinstance(t, dict) and t.get("name") is not None
]
names = [t.get("name") for t in result["tools"] if isinstance(t, dict) and t.get("name") is not None]
assert "dummy_tool" in names
@ -2716,13 +2653,9 @@ def test_calculate_usage_completion_tokens_details_with_reasoning():
"output_tokens": 500,
}
# Simulating reasoning content that would count as ~50 tokens
reasoning_content = (
"Let me think about this step by step. " * 10
) # Roughly 50 tokens
reasoning_content = "Let me think about this step by step. " * 10 # Roughly 50 tokens
usage = config.calculate_usage(
usage_object=usage_object, reasoning_content=reasoning_content
)
usage = config.calculate_usage(usage_object=usage_object, reasoning_content=reasoning_content)
# completion_tokens_details should be populated with both reasoning and text tokens
assert usage.completion_tokens_details is not None
@ -2773,9 +2706,7 @@ def test_reasoning_effort_maps_to_adaptive_thinking_for_claude_4_6_models():
# reasoning_effort should not be in the result (it's transformed to thinking)
assert "reasoning_effort" not in result
# Should set output_config with the mapped effort value
assert (
"output_config" in result
), f"output_config missing for {model} with effort={effort}"
assert "output_config" in result, f"output_config missing for {model} with effort={effort}"
assert result["output_config"]["effort"] == effort_map[effort]
@ -2851,7 +2782,6 @@ def test_raw_adaptive_thinking_untouched_for_46_plus_model():
assert result["thinking"] == {"type": "adaptive"}
@pytest.mark.parametrize(
"model, expected",
[
@ -2877,9 +2807,7 @@ def test_raw_adaptive_thinking_untouched_for_46_plus_model():
("gpt-4o", False),
],
)
def test_is_adaptive_thinking_model_is_sourced_from_cost_map(
local_model_cost_map, model, expected
):
def test_is_adaptive_thinking_model_is_sourced_from_cost_map(local_model_cost_map, model, expected):
"""Adaptive thinking resolves from the cost map first (an explicit
supports_adaptive_thinking entry, or the anthropic-claude fallback rule for unmapped
future Claudes), then from a date-safe opus/sonnet/haiku >= 4.6 name version as a
@ -2995,9 +2923,7 @@ def test_reasoning_effort_sets_output_config_for_46_models():
drop_params=False,
)
assert (
"output_config" in result
), f"output_config missing for {model} with effort={effort}"
assert "output_config" in result, f"output_config missing for {model} with effort={effort}"
assert result["output_config"]["effort"] == effort
@ -3036,9 +2962,7 @@ def test_reasoning_effort_does_not_set_output_config_for_older_models():
drop_params=False,
)
assert (
"output_config" not in result
), f"output_config should not be set for {model}"
assert "output_config" not in result, f"output_config should not be set for {model}"
@pytest.mark.parametrize(
@ -3078,14 +3002,10 @@ def test_reasoning_effort_accepts_dict_shape_for_adaptive_model(reasoning_effort
)
# thinking must be set (adaptive for 4.6+)
assert (
"thinking" in result
), f"thinking missing for reasoning_effort={reasoning_effort_value!r}"
assert "thinking" in result, f"thinking missing for reasoning_effort={reasoning_effort_value!r}"
assert result["thinking"]["type"] == "adaptive"
# output_config must carry the mapped effort
assert (
"output_config" in result
), f"output_config missing for reasoning_effort={reasoning_effort_value!r}"
assert "output_config" in result, f"output_config missing for reasoning_effort={reasoning_effort_value!r}"
assert result["output_config"]["effort"] == "low"
@ -3114,16 +3034,13 @@ def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model(
drop_params=False,
)
assert (
"thinking" in result
), f"thinking missing for reasoning_effort={reasoning_effort_value!r}"
assert "thinking" in result, f"thinking missing for reasoning_effort={reasoning_effort_value!r}"
assert result["thinking"]["type"] == "enabled"
assert "budget_tokens" in result["thinking"]
assert result["thinking"]["budget_tokens"] > 0
# Older models must not get adaptive-thinking output_config
assert "output_config" not in result, (
f"output_config should not be set for non-adaptive model "
f"(reasoning_effort={reasoning_effort_value!r})"
f"output_config should not be set for non-adaptive model (reasoning_effort={reasoning_effort_value!r})"
)
@ -3174,12 +3091,8 @@ def test_reasoning_effort_unparseable_dict_is_dropped(bad_value):
model="claude-sonnet-4-6-20260219",
drop_params=False,
)
assert (
"thinking" not in result
), f"thinking should not be set for bad value {bad_value!r}"
assert (
"output_config" not in result
), f"output_config should not be set for bad value {bad_value!r}"
assert "thinking" not in result, f"thinking should not be set for bad value {bad_value!r}"
assert "output_config" not in result, f"output_config should not be set for bad value {bad_value!r}"
@pytest.mark.parametrize(
@ -3310,9 +3223,7 @@ def test_reasoning_effort_garbage_raises_bad_request(effort):
("max", DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET),
],
)
def test_reasoning_effort_xhigh_max_maps_to_budget_on_budget_model(
effort, expected_budget
):
def test_reasoning_effort_xhigh_max_maps_to_budget_on_budget_model(effort, expected_budget):
"""``xhigh`` / ``max`` extend the budget_tokens progression on budget-mode models."""
config = AnthropicConfig()
@ -3459,17 +3370,11 @@ def test_code_execution_tool_results_extraction():
# Verify first tool call
assert transformed_response.choices[0].message.tool_calls[0].id == "srvtoolu_01ABC"
assert (
transformed_response.choices[0].message.tool_calls[0].function.name
== "bash_code_execution"
)
assert transformed_response.choices[0].message.tool_calls[0].function.name == "bash_code_execution"
# Verify second tool call
assert transformed_response.choices[0].message.tool_calls[1].id == "srvtoolu_01DEF"
assert (
transformed_response.choices[0].message.tool_calls[1].function.name
== "text_editor_code_execution"
)
assert transformed_response.choices[0].message.tool_calls[1].function.name == "text_editor_code_execution"
# Verify tool results are in provider_specific_fields
provider_fields = transformed_response.choices[0].message.provider_specific_fields
@ -3492,10 +3397,7 @@ def test_code_execution_tool_results_extraction():
assert editor_result["content"]["is_file_update"] is False
# Verify text content is properly concatenated
assert (
"I'll calculate that for you."
in transformed_response.choices[0].message.content
)
assert "I'll calculate that for you." in transformed_response.choices[0].message.content
assert "Done!" in transformed_response.choices[0].message.content
@ -3563,10 +3465,7 @@ def test_code_execution_tool_results_in_hidden_params():
assert "provider_specific_fields" in hidden
assert "tool_results" in hidden["provider_specific_fields"]
assert len(hidden["provider_specific_fields"]["tool_results"]) == 1
assert (
hidden["provider_specific_fields"]["tool_results"][0]["content"]["stdout"]
== "hello\n"
)
assert hidden["provider_specific_fields"]["tool_results"][0]["content"]["stdout"] == "hello\n"
def test_tool_search_tool_result_not_in_tool_results():
@ -3762,10 +3661,7 @@ def test_compaction_block_in_provider_specific_fields():
assert "compaction_blocks" in provider_fields
assert len(provider_fields["compaction_blocks"]) == 1
assert provider_fields["compaction_blocks"][0]["type"] == "compaction"
assert (
"Summary of the conversation"
in provider_fields["compaction_blocks"][0]["content"]
)
assert "Summary of the conversation" in provider_fields["compaction_blocks"][0]["content"]
def test_multiple_compaction_blocks():
@ -3813,9 +3709,7 @@ def test_compaction_block_request_transformation():
{"role": "user", "content": "What is the weather in San Francisco?"},
{
"role": "assistant",
"content": [
{"type": "text", "text": "I don't have access to real-time data."}
],
"content": [{"type": "text", "text": "I don't have access to real-time data."}],
"provider_specific_fields": {
"compaction_blocks": [
{
@ -3828,9 +3722,7 @@ def test_compaction_block_request_transformation():
{"role": "user", "content": "What about New York?"},
]
result = anthropic_messages_pt(
messages=messages, model="claude-opus-4-6", llm_provider="anthropic"
)
result = anthropic_messages_pt(messages=messages, model="claude-opus-4-6", llm_provider="anthropic")
# Find the assistant message
assistant_message = None
@ -3944,9 +3836,7 @@ def test_map_openai_context_management_to_anthropic():
"instructions": "Focus on preserving code snippets",
}
]
result = config.map_openai_context_management_to_anthropic(
openai_format_with_instructions
)
result = config.map_openai_context_management_to_anthropic(openai_format_with_instructions)
assert result is not None
assert result["edits"][0]["trigger"]["value"] == 150000
@ -3973,9 +3863,7 @@ def test_map_openai_params_with_context_management():
config = AnthropicConfig()
# Test with OpenAI list format
non_default_params = {
"context_management": [{"type": "compaction", "compact_threshold": 200000}]
}
non_default_params = {"context_management": [{"type": "compaction", "compact_threshold": 200000}]}
optional_params = {}
result = config.map_openai_params(
@ -4012,10 +3900,7 @@ def test_map_openai_params_with_context_management():
)
assert "context_management" in result
assert (
result["context_management"]
== non_default_params_anthropic["context_management"]
)
assert result["context_management"] == non_default_params_anthropic["context_management"]
def test_cache_control_in_supported_params():
@ -4126,10 +4011,7 @@ def test_compaction_block_empty_list_not_added():
# Verify compaction_blocks is not in provider_specific_fields when there are none
provider_fields = result.choices[0].message.provider_specific_fields
if provider_fields:
assert (
"compaction_blocks" not in provider_fields
or provider_fields.get("compaction_blocks") is None
)
assert "compaction_blocks" not in provider_fields or provider_fields.get("compaction_blocks") is None
def test_fast_mode_beta_header():
@ -4178,9 +4060,7 @@ def test_fast_mode_usage_calculation():
"output_tokens": 500,
}
usage = config.calculate_usage(
usage_object=usage_object, reasoning_content=None, speed="fast"
)
usage = config.calculate_usage(usage_object=usage_object, reasoning_content=None, speed="fast")
assert usage.prompt_tokens == 1000
assert usage.completion_tokens == 500
@ -4201,9 +4081,7 @@ def test_fast_mode_cost_calculation():
base_completion = 0.025
with (
patch(
"litellm.llms.anthropic.cost_calculation.generic_cost_per_token"
) as mock_cost,
patch("litellm.llms.anthropic.cost_calculation.generic_cost_per_token") as mock_cost,
patch("litellm.get_model_info") as mock_info,
):
mock_cost.return_value = (base_prompt, base_completion)
@ -4243,9 +4121,7 @@ def test_fast_mode_with_inference_geo():
base_completion = 0.025
with (
patch(
"litellm.llms.anthropic.cost_calculation.generic_cost_per_token"
) as mock_cost,
patch("litellm.llms.anthropic.cost_calculation.generic_cost_per_token") as mock_cost,
patch("litellm.get_model_info") as mock_info,
):
mock_cost.return_value = (base_prompt, base_completion)
@ -4436,9 +4312,7 @@ def test_map_tool_helper_enforces_object_type_when_missing():
"name": "search_code",
"description": "Search for code patterns",
"parameters": {
"properties": {
"query": {"type": "string", "description": "Search query"}
},
"properties": {"query": {"type": "string", "description": "Search query"}},
"required": ["query"],
},
},
@ -4451,9 +4325,9 @@ def test_map_tool_helper_enforces_object_type_when_missing():
assert "properties" in result["input_schema"]
assert "query" in result["input_schema"]["properties"]
# Original parameters dict must not be modified in place
assert (
tool["function"]["parameters"] == original_params
), "parameters dict was mutated; _map_tool_helper should not modify caller data"
assert tool["function"]["parameters"] == original_params, (
"parameters dict was mutated; _map_tool_helper should not modify caller data"
)
def test_map_tool_helper_enforces_object_type_when_wrong_type():
@ -4479,13 +4353,13 @@ def test_map_tool_helper_enforces_object_type_when_wrong_type():
result, _ = config._map_tool_helper(tool)
assert result is not None
assert result["input_schema"]["type"] == "object"
assert (
result["input_schema"].get("properties") == {}
), "properties should be injected as {} when schema has non-object type and no properties key"
assert result["input_schema"].get("properties") == {}, (
"properties should be injected as {} when schema has non-object type and no properties key"
)
# Original parameters dict must not be modified in place
assert (
tool["function"]["parameters"] == original_params
), "parameters dict was mutated; _map_tool_helper should not modify caller data"
assert tool["function"]["parameters"] == original_params, (
"parameters dict was mutated; _map_tool_helper should not modify caller data"
)
def test_map_tool_helper_preserves_valid_object_schema():
@ -4552,12 +4426,8 @@ def test_extract_response_content_thinking_block_null_thinking():
{"type": "text", "text": "Hello"},
]
}
text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(
completion_response_null
)
assert (
thinking_blocks is not None
), "thinking blocks should not be None when thinking=null"
text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(completion_response_null)
assert thinking_blocks is not None, "thinking blocks should not be None when thinking=null"
assert len(thinking_blocks) == 1
assert "Hello" in text
@ -4568,12 +4438,8 @@ def test_extract_response_content_thinking_block_null_thinking():
{"type": "text", "text": "World"},
]
}
text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(
completion_response_missing
)
assert (
thinking_blocks is not None
), "thinking blocks should not be None when thinking key is absent"
text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(completion_response_missing)
assert thinking_blocks is not None, "thinking blocks should not be None when thinking key is absent"
assert len(thinking_blocks) == 1
assert "World" in text
@ -4584,9 +4450,7 @@ def test_extract_response_content_thinking_block_null_thinking():
{"type": "text", "text": "Done"},
]
}
text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(
completion_response_text
)
text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(completion_response_text)
assert thinking_blocks is not None
assert len(thinking_blocks) == 1
assert thinking_blocks[0]["thinking"] == "Let me think..."
@ -4645,12 +4509,8 @@ def test_advisor_beta_header_injected():
}
]
}
result = config.update_headers_with_optional_anthropic_beta(
headers, optional_params
)
assert ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value in result.get(
"anthropic-beta", ""
)
result = config.update_headers_with_optional_anthropic_beta(headers, optional_params)
assert ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value in result.get("anthropic-beta", "")
def test_advisor_beta_header_not_injected_without_tool():
@ -4658,9 +4518,7 @@ def test_advisor_beta_header_not_injected_without_tool():
config = AnthropicConfig()
headers: dict = {}
optional_params: dict = {"tools": []}
result = config.update_headers_with_optional_anthropic_beta(
headers, optional_params
)
result = config.update_headers_with_optional_anthropic_beta(headers, optional_params)
assert "advisor-tool-2026-03-01" not in result.get("anthropic-beta", "")
@ -4687,9 +4545,7 @@ def test_advisor_tool_result_preserved_in_response():
{"type": "text", "text": "Here is the implementation."},
]
}
text, _, _, _, tool_calls, _, tool_results, _ = config.extract_response_content(
completion_response
)
text, _, _, _, tool_calls, _, tool_results, _ = config.extract_response_content(completion_response)
assert "Consulting advisor." in text
assert "Here is the implementation." in text
# server_tool_use (advisor) should be a tool_call
@ -4804,9 +4660,7 @@ def test_basic_sanitize_anthropic_tool_name_replaces_invalid_chars():
)
assert (
_basic_sanitize_anthropic_tool_name(
"github_openapi_mcp-actions/download-job-logs-for-workflow-run"
)
_basic_sanitize_anthropic_tool_name("github_openapi_mcp-actions/download-job-logs-for-workflow-run")
== "github_openapi_mcp-actions_download-job-logs-for-workflow-run"
)
# other punctuation
@ -4835,9 +4689,7 @@ def test_build_anthropic_tool_name_maps_no_collisions():
]
)
assert forward == {
"actions/download-job-logs-for-workflow-run": (
"actions_download-job-logs-for-workflow-run"
),
"actions/download-job-logs-for-workflow-run": ("actions_download-job-logs-for-workflow-run"),
"pulls/list-files": "pulls_list-files",
}
assert reverse == {v: k for k, v in forward.items()}
@ -4888,9 +4740,7 @@ def test_build_anthropic_tool_name_maps_three_way_collision():
_build_anthropic_tool_name_maps,
)
forward, reverse = _build_anthropic_tool_name_maps(
["foo_bar", "foo/bar", "foo.bar"]
)
forward, reverse = _build_anthropic_tool_name_maps(["foo_bar", "foo/bar", "foo.bar"])
assert "foo_bar" not in forward # untouched
assert forward["foo/bar"] == "foo_bar_2"
assert forward["foo.bar"] == "foo_bar_3"
@ -4963,16 +4813,13 @@ def test_map_openai_params_does_not_pollute_optional_params_with_internal_keys()
)
# No internal keys may appear in optional_params for ANY input.
for key in optional_params:
assert not key.startswith(
"_anthropic_tool_name"
), f"optional_params leaked internal key {key!r}: {optional_params}"
assert not key.startswith("_anthropic_tool_name"), (
f"optional_params leaked internal key {key!r}: {optional_params}"
)
# And no key starting with `_` either; optional_params should only
# contain documented Anthropic Messages API parameters.
for key in optional_params:
assert not key.startswith("_"), (
f"optional_params leaked underscore-prefixed key {key!r}: "
f"{optional_params}"
)
assert not key.startswith("_"), f"optional_params leaked underscore-prefixed key {key!r}: {optional_params}"
def test_map_openai_params_no_maps_when_all_names_already_valid():
@ -5001,11 +4848,7 @@ def test_map_openai_params_no_maps_when_all_names_already_valid():
def test_rewrite_tool_names_in_messages_uses_forward_map():
config = AnthropicConfig()
forward_map = {
"actions/download-job-logs-for-workflow-run": (
"actions_download-job-logs-for-workflow-run"
)
}
forward_map = {"actions/download-job-logs-for-workflow-run": ("actions_download-job-logs-for-workflow-run")}
messages = [
{"role": "user", "content": "go"},
{
@ -5028,15 +4871,9 @@ def test_rewrite_tool_names_in_messages_uses_forward_map():
out = config._rewrite_tool_names_in_messages(messages, forward_map)
# input list must not be mutated
assert (
messages[1]["tool_calls"][0]["function"]["name"]
== "actions/download-job-logs-for-workflow-run"
)
assert messages[1]["tool_calls"][0]["function"]["name"] == "actions/download-job-logs-for-workflow-run"
# output rewritten according to forward map
assert (
out[1]["tool_calls"][0]["function"]["name"]
== "actions_download-job-logs-for-workflow-run"
)
assert out[1]["tool_calls"][0]["function"]["name"] == "actions_download-job-logs-for-workflow-run"
# non-tool-call messages pass through unchanged (same object)
assert out[0] is messages[0]
assert out[2] is messages[2]
@ -5112,9 +4949,7 @@ def test_sanitize_tool_names_in_request_does_not_mutate_caller_tool_dicts():
caller_tools = [caller_tool]
optional_params: dict = {"tools": caller_tools}
forward, reverse = config._sanitize_tool_names_in_request(
optional_params=optional_params
)
forward, reverse = config._sanitize_tool_names_in_request(optional_params=optional_params)
assert forward.get(original_name)
sanitized = forward[original_name]
@ -5263,10 +5098,7 @@ def test_streaming_iterator_reverse_maps_tool_use_name():
parsed = iterator.chunk_parser(chunk=chunk)
tool_calls = parsed.choices[0].delta.tool_calls
assert tool_calls is not None and len(tool_calls) == 1
assert (
tool_calls[0]["function"]["name"]
== "actions/download-job-logs-for-workflow-run"
)
assert tool_calls[0]["function"]["name"] == "actions/download-job-logs-for-workflow-run"
def test_streaming_iterator_passthrough_when_name_not_in_map():
@ -5362,9 +5194,9 @@ def test_transform_request_does_not_leak_internal_keys_into_body():
for tool in data.get("tools", []):
name = tool.get("name")
assert isinstance(name, str)
assert _re.fullmatch(
r"[a-zA-Z0-9_-]{1,128}", name
), f"sanitized tool name {name!r} still violates Anthropic regex"
assert _re.fullmatch(r"[a-zA-Z0-9_-]{1,128}", name), (
f"sanitized tool name {name!r} still violates Anthropic regex"
)
# Sent name for the bad tool is the disambiguated form, valid name passes through.
sent_names = {t["name"] for t in data["tools"]}
@ -5500,9 +5332,7 @@ def test_transform_request_rewrites_tool_names_in_history():
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_use":
tool_use_names.append(block.get("name"))
assert (
tool_use_names
), "expected at least one tool_use block in transformed messages"
assert tool_use_names, "expected at least one tool_use block in transformed messages"
for name in tool_use_names:
assert name == "actions_download-job-logs-for-workflow-run", (
f"history tool_use.name {name!r} not rewritten -- Anthropic will "
@ -5526,19 +5356,12 @@ def test_sanitize_tool_names_in_request_skips_hosted_tools():
}
forward, reverse = AnthropicConfig._sanitize_tool_names_in_request(optional_params)
# Only the custom tool was rewritten.
assert forward == {
"actions/download-job-logs-for-workflow-run": "actions_download-job-logs-for-workflow-run"
}
assert reverse == {
"actions_download-job-logs-for-workflow-run": "actions/download-job-logs-for-workflow-run"
}
assert forward == {"actions/download-job-logs-for-workflow-run": "actions_download-job-logs-for-workflow-run"}
assert reverse == {"actions_download-job-logs-for-workflow-run": "actions/download-job-logs-for-workflow-run"}
# Hosted tool's name unchanged.
assert optional_params["tools"][0]["name"] == "web_search"
# Custom tool's name updated in place.
assert (
optional_params["tools"][1]["name"]
== "actions_download-job-logs-for-workflow-run"
)
assert optional_params["tools"][1]["name"] == "actions_download-job-logs-for-workflow-run"
def test_sanitize_tool_names_in_request_no_tools_is_noop():
@ -5772,9 +5595,7 @@ def test_translate_system_message_keeps_billing_header_for_first_party_anthropic
assert config.should_strip_billing_metadata() is False
result = config.translate_system_message(
messages=_system_with_billing_header(
"You are Claude Code, Anthropic's official CLI for Claude."
)
messages=_system_with_billing_header("You are Claude Code, Anthropic's official CLI for Claude.")
)
texts = [block["text"] for block in result]
@ -5790,9 +5611,7 @@ def test_translate_system_message_strips_billing_header_for_bedrock():
config = BedrockClaudePlatformConfig()
assert config.should_strip_billing_metadata() is True
result = config.translate_system_message(
messages=_system_with_billing_header("real system prompt")
)
result = config.translate_system_message(messages=_system_with_billing_header("real system prompt"))
texts = [block["text"] for block in result]
assert all(not t.startswith("x-anthropic-billing-header:") for t in texts)
@ -5858,9 +5677,7 @@ def test_translate_system_message_strips_billing_header_for_bedrock_invoke():
config = AmazonAnthropicClaudeConfig()
assert config.should_strip_billing_metadata() is True
result = config.translate_system_message(
messages=_system_with_billing_header("real system prompt")
)
result = config.translate_system_message(messages=_system_with_billing_header("real system prompt"))
texts = [block["text"] for block in result]
assert all(not t.startswith("x-anthropic-billing-header:") for t in texts)
@ -5914,9 +5731,7 @@ def test_translate_system_message_strips_billing_header_for_bedrock_invoke():
),
],
)
def test_should_strip_billing_metadata_by_provider(
module_path, class_name, expected_strip
):
def test_should_strip_billing_metadata_by_provider(module_path, class_name, expected_strip):
import importlib
config_cls = getattr(importlib.import_module(module_path), class_name)
@ -6088,12 +5903,8 @@ def test_sampling_param_gating_driven_by_model_map_flag(monkeypatch):
"""The drop/raise decision must come from ``supports_sampling_params`` in
the model map, not just name matching: a flagged entry gates a model whose
name says nothing, and an explicit ``true`` overrides the name fallback."""
monkeypatch.setitem(
litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False}
)
monkeypatch.setitem(
litellm.model_cost, "claude-fable-5-test", {"supports_sampling_params": True}
)
monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False})
monkeypatch.setitem(litellm.model_cost, "claude-fable-5-test", {"supports_sampling_params": True})
config = AnthropicConfig()
flagged_off = config.map_openai_params(
@ -6213,9 +6024,7 @@ def test_is_anthropic_usage_object_rejects_responses_api_usage():
("claude-sonnet-4-5-20250929", False),
],
)
def test_disabled_thinking_omitted_only_for_always_on_models(
local_model_cost_map, model, expected_dropped
):
def test_disabled_thinking_omitted_only_for_always_on_models(local_model_cost_map, model, expected_dropped):
"""``thinking={"type": "disabled"}`` is omitted for always-on-thinking models
(Fable/Mythos, which 400 on it: the API remedy is to omit the param) and is
forwarded verbatim for every model that accepts it."""

View file

@ -136,22 +136,14 @@ def test_in_place_substitution_preserves_ordering():
responses_output = [msg_item, fc_exec1, fc_regular, fc_exec2]
# Apply the same logic as _transform_chat_completion_choices_to_responses_output
tool_result_items = (
LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(resp)
)
tool_result_items = LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(resp)
if tool_result_items:
result_by_id = {
(item.get("id") if isinstance(item, dict) else item.id): item
for item in tool_result_items
}
result_by_id = {(item.get("id") if isinstance(item, dict) else item.id): item for item in tool_result_items}
replaced_ids = set(result_by_id.keys())
responses_output = [
(
result_by_id[getattr(item, "call_id", None)]
if (
getattr(item, "type", None) == "function_call"
and getattr(item, "call_id", None) in replaced_ids
)
if (getattr(item, "type", None) == "function_call" and getattr(item, "call_id", None) in replaced_ids)
else item
)
for item in responses_output
@ -255,9 +247,7 @@ def test_end_to_end_streaming_chunks_to_code_interpreter_output():
assert code_results[0]["code"] == "echo e2e_test"
# Step 3: Extract via _extract_tool_result_output_items (Responses API layer)
tool_result_items = (
LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(assembled)
)
tool_result_items = LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(assembled)
assert len(tool_result_items) == 1
item = tool_result_items[0]
# Items are reconstructed as Pydantic OutputCodeInterpreterCall objects

View file

@ -6,7 +6,6 @@ import pytest
import litellm
from litellm.litellm_core_utils.prompt_templates.common_utils import (
TOOL_RESULT_IMAGE_PLACEHOLDER,
)
@ -53,9 +52,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block():
tool_calls=[
ChatCompletionDeltaToolCall(
id="call_d581d130-e234-4315-94e8-27e7ff7c4e55",
function=Function(
arguments='{"location": "Boston"}', name="get_weather"
),
function=Function(arguments='{"location": "Boston"}', name="get_weather"),
type="function",
index=0,
)
@ -69,9 +66,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block():
(
block_type,
content_block_start,
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
choices=choices
)
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
print(content_block_start)
@ -101,9 +96,7 @@ def test_translate_streaming_openai_chunk_strips_gemini_thought_from_tool_call_i
tool_calls=[
ChatCompletionDeltaToolCall(
id=combined,
function=Function(
arguments='{"a": 17, "b": 25}', name="add_numbers"
),
function=Function(arguments='{"a": 17, "b": 25}', name="add_numbers"),
type="function",
index=0,
)
@ -117,9 +110,7 @@ def test_translate_streaming_openai_chunk_strips_gemini_thought_from_tool_call_i
(
block_type,
content_block_start,
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
choices=choices
)
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
assert block_type == "tool_use"
assert content_block_start["id"] == base
@ -164,9 +155,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_content_block():
(
block_type,
content_block_start,
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
choices=choices
)
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
assert block_type == "thinking"
assert content_block_start == {
@ -202,9 +191,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_reasoning_content_only_co
(
block_type,
content_block_start,
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
choices=choices
)
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
assert block_type == "thinking"
assert content_block_start == {
@ -250,9 +237,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_signature_block(
(
block_type,
content_block_start,
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
choices=choices
)
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
assert block_type == "thinking"
assert content_block_start == {
@ -305,9 +290,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block_thinking_an
(
block_type,
content_block_start,
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
choices=choices
)
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
assert block_type == "thinking"
@ -350,10 +333,7 @@ def test_translate_anthropic_messages_to_openai_thinking_blocks():
assert "thinking_blocks" in result[1]
assert len(result[1]["thinking_blocks"]) == 2
assert result[1]["thinking_blocks"][0]["type"] == "thinking"
assert (
result[1]["thinking_blocks"][0]["thinking"]
== "I will call the get_weather tool."
)
assert result[1]["thinking_blocks"][0]["thinking"] == "I will call the get_weather tool."
assert result[1]["thinking_blocks"][0]["signature"] == "sigsig"
assert result[1]["thinking_blocks"][1]["type"] == "redacted_thinking"
assert result[1]["thinking_blocks"][1]["data"] == "REDACTED"
@ -456,9 +436,7 @@ def test_translate_anthropic_messages_to_openai_tool_message_placement():
assert tool_message_idx is not None, "Tool message not found"
assert user_message_idx is not None, "User message not found"
assert (
tool_message_idx < user_message_idx
), "Tool message should be placed before user message"
assert tool_message_idx < user_message_idx, "Tool message should be placed before user message"
@pytest.mark.parametrize(
@ -733,9 +711,7 @@ def test_translate_anthropic_to_openai_skips_prompt_cache_key_when_provider_lack
def test_translate_anthropic_to_openai_skips_prompt_cache_key_for_chained_litellm_proxy():
assert "prompt_cache_key" in litellm.get_supported_openai_params(
model="xai", custom_llm_provider="litellm_proxy"
)
assert "prompt_cache_key" in litellm.get_supported_openai_params(model="xai", custom_llm_provider="litellm_proxy")
openai_request = _translate_with_metadata("litellm_proxy/xai", {"user_id": "session-abc"}, "litellm_proxy")
assert openai_request["user"] == "session-abc"
assert "prompt_cache_key" not in openai_request
@ -780,7 +756,8 @@ def test_translate_openai_content_to_anthropic_empty_function_arguments():
id="call_empty_args",
type="function",
function=Function(
name="test_function", arguments="" # empty arguments string
name="test_function",
arguments="", # empty arguments string
),
)
],
@ -795,9 +772,7 @@ def test_translate_openai_content_to_anthropic_empty_function_arguments():
assert result[0]["type"] == "tool_use"
assert result[0]["id"] == "call_empty_args"
assert result[0]["name"] == "test_function"
assert (
result[0]["input"] == {}
), "Empty function arguments should result in empty dict"
assert result[0]["input"] == {}, "Empty function arguments should result in empty dict"
def test_translate_openai_content_to_anthropic_text_and_tool_calls():
@ -917,9 +892,7 @@ def test_translate_openai_response_to_anthropic_text_and_tool_calls():
ChatCompletionAssistantToolCall(
id="call_tool_combo",
type="function",
function=Function(
name="get_weather", arguments='{"location": "Paris"}'
),
function=Function(name="get_weather", arguments='{"location": "Paris"}'),
)
],
),
@ -929,9 +902,7 @@ def test_translate_openai_response_to_anthropic_text_and_tool_calls():
)
adapter = LiteLLMAnthropicMessagesAdapter()
anthropic_response = adapter.translate_openai_response_to_anthropic(
response=openai_response
)
anthropic_response = adapter.translate_openai_response_to_anthropic(response=openai_response)
anthropic_content = anthropic_response.get("content")
assert anthropic_content is not None
@ -972,9 +943,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_with_partial_json():
(
type_of_content,
content_block_delta,
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(
choices=choices
)
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices)
print("Type of content:", type_of_content)
print("Content block delta:", content_block_delta)
@ -1083,9 +1052,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_delta():
(
type_of_content,
content_block_delta,
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(
choices=choices
)
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices)
assert type_of_content == "thinking_delta"
assert content_block_delta["type"] == "thinking_delta"
@ -1128,9 +1095,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_with_thinking():
(
type_of_content,
content_block_delta,
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(
choices=choices
)
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices)
assert type_of_content == "signature_delta"
assert content_block_delta["type"] == "signature_delta"
@ -1194,9 +1159,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_emits_signature_when_thin
(
block_type,
content_block_start,
) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(
choices=choices
)
) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
assert block_type == "thinking"
@ -1236,9 +1199,7 @@ def test_translate_anthropic_messages_to_openai_user_message_with_base64_image()
# Check image content
assert result[0]["content"][1]["type"] == "image_url"
assert "image_url" in result[0]["content"][1]
assert result[0]["content"][1]["image_url"]["url"].startswith(
"data:image/png;base64,"
)
assert result[0]["content"][1]["image_url"]["url"].startswith("data:image/png;base64,")
assert (
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
in result[0]["content"][1]["image_url"]["url"]
@ -1276,18 +1237,14 @@ def test_translate_anthropic_messages_to_openai_user_message_with_url_image():
# Check image content
assert result[0]["content"][1]["type"] == "image_url"
assert "image_url" in result[0]["content"][1]
assert (
result[0]["content"][1]["image_url"]["url"] == "https://example.com/forest.jpg"
)
assert result[0]["content"][1]["image_url"]["url"] == "https://example.com/forest.jpg"
def test_translate_anthropic_messages_to_openai_tool_result_with_base64_image():
"""Test that base64 images in tool results are correctly translated to OpenAI format."""
anthropic_messages = [
AnthropicMessagesUserMessageParam(
role="user", content=[{"type": "text", "text": "Take a screenshot"}]
),
AnthropicMessagesUserMessageParam(role="user", content=[{"type": "text", "text": "Take a screenshot"}]),
AnthopicMessagesAssistantMessageParam(
role="assistant",
content=[
@ -1439,9 +1396,7 @@ def test_translate_anthropic_messages_to_openai_mixed_content_with_image():
# Check first image (base64)
assert result[0]["content"][1]["type"] == "image_url"
assert result[0]["content"][1]["image_url"]["url"].startswith(
"data:image/png;base64,"
)
assert result[0]["content"][1]["image_url"]["url"].startswith("data:image/png;base64,")
# Check middle text
assert result[0]["content"][2]["type"] == "text"
@ -1449,9 +1404,7 @@ def test_translate_anthropic_messages_to_openai_mixed_content_with_image():
# Check second image (URL)
assert result[0]["content"][3]["type"] == "image_url"
assert (
result[0]["content"][3]["image_url"]["url"] == "https://example.com/image2.jpg"
)
assert result[0]["content"][3]["image_url"]["url"] == "https://example.com/image2.jpg"
# Check final text
assert result[0]["content"][4]["type"] == "text"
@ -1497,10 +1450,7 @@ def test_translate_anthropic_messages_to_openai_tool_use_with_signature():
assert tool_call["id"] == "call_386f67af31f9415781bc35071405"
assert "function" in tool_call
assert "provider_specific_fields" in tool_call["function"]
assert (
tool_call["function"]["provider_specific_fields"]["thought_signature"]
== test_signature
)
assert tool_call["function"]["provider_specific_fields"]["thought_signature"] == test_signature
def test_translate_anthropic_messages_to_openai_tool_result_with_multiple_content_items():
@ -1558,9 +1508,7 @@ def test_translate_anthropic_messages_to_openai_tool_result_with_multiple_conten
result = adapter.translate_anthropic_messages_to_openai(messages=anthropic_messages)
# Count how many tool messages have the same tool_call_id
tool_messages = [
msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"
]
tool_messages = [msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"]
tool_call_ids = [msg.get("tool_call_id") for msg in tool_messages]
# The critical assertion: each tool_call_id should appear only ONCE
@ -1576,12 +1524,8 @@ def test_translate_anthropic_messages_to_openai_tool_result_with_multiple_conten
# The content should be a list with all items combined
tool_message = tool_messages[0]
assert tool_message["tool_call_id"] == "toolu_016hYHBkTf4JDF3p22UoYk5C"
assert isinstance(
tool_message["content"], list
), "Multiple content items should be combined into a list"
assert (
len(tool_message["content"]) == 3
), f"Expected 3 content items, got {len(tool_message['content'])}"
assert isinstance(tool_message["content"], list), "Multiple content items should be combined into a list"
assert len(tool_message["content"]) == 3, f"Expected 3 content items, got {len(tool_message['content'])}"
# Verify content types
assert tool_message["content"][0]["type"] == "text"
@ -1630,17 +1574,14 @@ def test_translate_anthropic_messages_to_openai_tool_result_single_item_backward
adapter = LiteLLMAnthropicMessagesAdapter()
result = adapter.translate_anthropic_messages_to_openai(messages=anthropic_messages)
tool_messages = [
msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"
]
tool_messages = [msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"]
assert len(tool_messages) == 1
tool_message = tool_messages[0]
# Single item should be a string for backward compatibility
assert isinstance(tool_message["content"], str), (
f"Single content item should be a string for backward compatibility, "
f"got {type(tool_message['content'])}"
f"Single content item should be a string for backward compatibility, got {type(tool_message['content'])}"
)
assert tool_message["content"] == "72°F and sunny"
@ -1689,9 +1630,7 @@ def test_streaming_chunk_with_both_text_and_tool_calls_issue_18238():
(
block_type,
content_block_start,
) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(
choices=choices
)
) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
assert block_type == "tool_use"
assert content_block_start["name"] == "Bash"
@ -1735,9 +1674,7 @@ def test_streaming_chunk_with_text_and_empty_tool_calls_returns_text_delta():
(
block_type,
content_block_start,
) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(
choices=choices
)
) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
assert block_type == "text"
assert content_block_start == {"type": "text", "text": ""}
@ -1748,15 +1685,12 @@ def test_streaming_chunk_with_text_and_empty_tool_calls_returns_text_delta():
# ============================================================================
# Model constant for cache control tests
CACHE_CONTROL_BEDROCK_CONVERSE_MODEL = (
"bedrock/converse/global.anthropic.claude-opus-4-5-20251101-v1:0"
)
CACHE_CONTROL_BEDROCK_CONVERSE_MODEL = "bedrock/converse/global.anthropic.claude-opus-4-5-20251101-v1:0"
CACHE_CONTROL_NON_ANTHROPIC_MODEL = "gpt-4"
# Bedrock Application Inference Profile ARN: the string contains neither
# "anthropic" nor "claude", so the model can only be recognized via its ARN shape
CACHE_CONTROL_BEDROCK_ARN_MODEL = (
"bedrock/converse/arn:aws:bedrock:us-east-1:123456789012:"
"application-inference-profile/abcdef123456"
"bedrock/converse/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456"
)
@ -1772,9 +1706,7 @@ def test_should_add_cache_control_for_anthropic_model():
"vertex_ai/claude-3-sonnet@20240229",
]:
target = {}
adapter._add_cache_control_if_applicable(
{"cache_control": cache_control}, target, model
)
adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model)
assert "cache_control" in target
assert target["cache_control"] == cache_control
@ -1790,9 +1722,7 @@ def test_should_not_add_cache_control_for_non_anthropic_model():
"gemini-pro",
]:
target = {}
adapter._add_cache_control_if_applicable(
{"cache_control": cache_control}, target, model
)
adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model)
assert "cache_control" not in target
@ -1807,9 +1737,7 @@ def test_should_not_add_cache_control_when_none():
{},
]:
target = {}
adapter._add_cache_control_if_applicable(
source, target, CACHE_CONTROL_BEDROCK_CONVERSE_MODEL
)
adapter._add_cache_control_if_applicable(source, target, CACHE_CONTROL_BEDROCK_CONVERSE_MODEL)
assert "cache_control" not in target
@ -1820,9 +1748,7 @@ def test_should_not_add_cache_control_when_model_none():
for model in [None, ""]:
target = {}
adapter._add_cache_control_if_applicable(
{"cache_control": cache_control}, target, model
)
adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model)
assert "cache_control" not in target
@ -1928,12 +1854,7 @@ def test_cache_control_fix_does_not_broaden_claude_detection():
make is_anthropic_claude_model treat ARN profiles as Claude, which would route
thinking params through unmodified and break non-Claude Bedrock profiles.
"""
assert (
LiteLLMAnthropicMessagesAdapter.is_anthropic_claude_model(
CACHE_CONTROL_BEDROCK_ARN_MODEL
)
is False
)
assert LiteLLMAnthropicMessagesAdapter.is_anthropic_claude_model(CACHE_CONTROL_BEDROCK_ARN_MODEL) is False
def test_thinking_preserved_for_bedrock_arn_inference_profile():
@ -2531,9 +2452,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_reasoning_content_without
(
type_of_content,
content_block_delta,
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(
choices=choices
)
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices)
assert type_of_content == "thinking_delta"
assert content_block_delta["type"] == "thinking_delta"
@ -2565,9 +2484,7 @@ def test_translate_openai_response_to_anthropic_with_reasoning_content_only():
)
adapter = LiteLLMAnthropicMessagesAdapter()
anthropic_response = adapter.translate_openai_response_to_anthropic(
response=openai_response
)
anthropic_response = adapter.translate_openai_response_to_anthropic(response=openai_response)
anthropic_content = anthropic_response.get("content")
assert anthropic_content is not None
@ -2580,9 +2497,7 @@ def test_translate_openai_response_to_anthropic_with_reasoning_content_only():
# Second block should be text
assert anthropic_content[1]["type"] == "text"
assert (
anthropic_content[1]["text"] == 'There are **3** "r"s in the word strawberry.'
)
assert anthropic_content[1]["text"] == 'There are **3** "r"s in the word strawberry.'
assert anthropic_response.get("stop_reason") == "end_turn"
@ -2634,9 +2549,7 @@ def test_truncate_tool_name_deterministic():
def test_truncate_tool_name_avoids_collisions():
"""Similar long names should produce different truncated names."""
name1 = "process_user_data_with_validation_and_error_handling_for_production_environment"
name2 = (
"process_user_data_with_validation_and_error_handling_for_staging_environment"
)
name2 = "process_user_data_with_validation_and_error_handling_for_staging_environment"
result1 = truncate_tool_name(name1)
result2 = truncate_tool_name(name2)
@ -2656,9 +2569,7 @@ def test_create_tool_name_mapping_no_long_names():
def test_create_tool_name_mapping_with_long_names():
"""Mapping should contain entries for truncated names."""
long_name = (
"a_very_long_tool_name_that_exceeds_the_64_character_limit_imposed_by_openai"
)
long_name = "a_very_long_tool_name_that_exceeds_the_64_character_limit_imposed_by_openai"
tools = [
{"name": "short_name"},
{"name": long_name},
@ -2683,9 +2594,7 @@ def test_translate_anthropic_tools_with_long_names():
]
adapter = LiteLLMAnthropicMessagesAdapter()
result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(
tools=tools, model="gpt-4"
)
result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(tools=tools, model="gpt-4")
assert len(result) == 1
# The tool name should be truncated
@ -2707,9 +2616,7 @@ def test_translate_anthropic_tools_mixed_names():
]
adapter = LiteLLMAnthropicMessagesAdapter()
result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(
tools=tools, model="gpt-4"
)
result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(tools=tools, model="gpt-4")
assert len(result) == 2
# Short name unchanged
@ -2723,9 +2630,7 @@ def test_translate_anthropic_tools_mixed_names():
def test_translate_openai_response_restores_tool_names():
"""Tool names in responses should be restored to original."""
original_name = (
"a_very_long_tool_name_that_needs_truncation_for_openai_api_compatibility"
)
original_name = "a_very_long_tool_name_that_needs_truncation_for_openai_api_compatibility"
truncated_name = truncate_tool_name(original_name)
tool_name_mapping = {truncated_name: original_name}
@ -2757,9 +2662,7 @@ def test_translate_openai_response_restores_tool_names():
)
adapter = LiteLLMAnthropicMessagesAdapter()
result = adapter.translate_openai_response_to_anthropic(
response=response, tool_name_mapping=tool_name_mapping
)
result = adapter.translate_openai_response_to_anthropic(response=response, tool_name_mapping=tool_name_mapping)
# Find the tool_use block in the response
tool_use_blocks = [c for c in result["content"] if c.get("type") == "tool_use"]
@ -2925,9 +2828,7 @@ def test_translate_openai_usage_to_anthropic_cache_tokens_from_dict_details_with
"cache_write_tokens": 20.0,
}
anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
usage
)
anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(usage)
assert anthropic_usage["input_tokens"] == 70
assert anthropic_usage["output_tokens"] == 50
@ -2946,9 +2847,7 @@ def test_translate_openai_usage_to_anthropic_ignores_fractional_cache_tokens():
"cache_creation_tokens": 20.25,
}
anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
usage
)
anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(usage)
assert anthropic_usage["input_tokens"] == 120
assert anthropic_usage["output_tokens"] == 50
@ -2965,9 +2864,7 @@ def test_translate_openai_usage_to_anthropic_ignores_bool_cache_tokens():
usage.cache_read_input_tokens = True
usage.cache_creation_input_tokens = True
anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
usage
)
anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(usage)
assert anthropic_usage["input_tokens"] == 120
assert anthropic_usage["output_tokens"] == 50
@ -3186,9 +3083,7 @@ def test_translate_streaming_openai_response_to_anthropic_cache_tokens_with_appl
assert message_delta["usage"]["output_tokens"] == 50
assert message_delta["usage"]["cache_read_input_tokens"] == 30
assert message_delta["usage"]["cache_creation_input_tokens"] == 20
assert message_delta["context_management"]["applied_edits"][0]["type"] == (
"compact_20260112"
)
assert message_delta["context_management"]["applied_edits"][0]["type"] == ("compact_20260112")
# =====================================================================
@ -3363,15 +3258,8 @@ class TestTranslateAnthropicOutputFormatToOpenAI:
assert schema["required"] == ["user"]
assert schema["properties"]["user"]["additionalProperties"] is False
assert schema["properties"]["user"]["required"] == ["name", "address"]
assert (
schema["properties"]["user"]["properties"]["address"][
"additionalProperties"
]
is False
)
assert schema["properties"]["user"]["properties"]["address"]["required"] == [
"city"
]
assert schema["properties"]["user"]["properties"]["address"]["additionalProperties"] is False
assert schema["properties"]["user"]["properties"]["address"]["required"] == ["city"]
def test_array_items_object_adds_additional_properties_false(self):
output_format = {
@ -3446,19 +3334,9 @@ class TestTranslateAnthropicOutputFormatToOpenAI:
assert sorted(schema["required"]) == ["age", "email", "name"]
def test_invalid_output_format_returns_none(self):
assert (
self.adapter.translate_anthropic_output_format_to_openai("invalid") is None
)
assert (
self.adapter.translate_anthropic_output_format_to_openai({"type": "text"})
is None
)
assert (
self.adapter.translate_anthropic_output_format_to_openai(
{"type": "json_schema"}
)
is None
)
assert self.adapter.translate_anthropic_output_format_to_openai("invalid") is None
assert self.adapter.translate_anthropic_output_format_to_openai({"type": "text"}) is None
assert self.adapter.translate_anthropic_output_format_to_openai({"type": "json_schema"}) is None
class TestAnthropicStreamWrapperToolArgs:
@ -3662,9 +3540,7 @@ def test_translate_openai_response_to_anthropic_with_polyfill_compaction_block()
)
response = _make_simple_openai_response(text="Hello after compaction.")
adapter = LiteLLMAnthropicMessagesAdapter()
result = adapter.translate_openai_response_to_anthropic(
response=response, polyfill_result=polyfill
)
result = adapter.translate_openai_response_to_anthropic(response=response, polyfill_result=polyfill)
content = result.get("content")
assert content is not None
@ -3696,9 +3572,7 @@ def test_translate_openai_response_to_anthropic_with_polyfill_iterations_usage()
)
response = _make_simple_openai_response(prompt_tokens=100, completion_tokens=30)
adapter = LiteLLMAnthropicMessagesAdapter()
result = adapter.translate_openai_response_to_anthropic(
response=response, polyfill_result=polyfill
)
result = adapter.translate_openai_response_to_anthropic(response=response, polyfill_result=polyfill)
usage = result.get("usage")
assert usage is not None
@ -3753,13 +3627,9 @@ def test_translate_openai_response_to_anthropic_with_polyfill_both_compaction_an
{"type": "compaction", "input_tokens": 300, "output_tokens": 75},
],
)
response = _make_simple_openai_response(
text="After compaction.", prompt_tokens=120, completion_tokens=40
)
response = _make_simple_openai_response(text="After compaction.", prompt_tokens=120, completion_tokens=40)
adapter = LiteLLMAnthropicMessagesAdapter()
result = adapter.translate_openai_response_to_anthropic(
response=response, polyfill_result=polyfill
)
result = adapter.translate_openai_response_to_anthropic(response=response, polyfill_result=polyfill)
# compaction block must come first
content = result.get("content")
@ -3861,7 +3731,9 @@ def test_translate_anthropic_tools_to_openai_omits_unset_strict():
assert function["parameters"]["required"] == ["query"]
TOOL_RESULT_IMAGE_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
TOOL_RESULT_IMAGE_B64 = (
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
)
TOOL_RESULT_IMAGE_URL = "https://example.com/screenshot.png"
@ -3869,8 +3741,7 @@ def _anthropic_tool_use_turn(*tool_use_ids):
return AnthopicMessagesAssistantMessageParam(
role="assistant",
content=[
{"type": "tool_use", "id": tid, "name": "read_file", "input": {"path": "img.png"}}
for tid in tool_use_ids
{"type": "tool_use", "id": tid, "name": "read_file", "input": {"path": "img.png"}} for tid in tool_use_ids
],
)
@ -3990,9 +3861,7 @@ def test_tool_result_parallel_tool_calls_keep_tool_message_adjacency():
result = _run_chat_completions_pipeline(
[
_anthropic_tool_use_turn("toolu_01", "toolu_02"),
_anthropic_tool_result_turn(
{"toolu_01": [_base64_image_block()], "toolu_02": [_url_image_block()]}
),
_anthropic_tool_result_turn({"toolu_01": [_base64_image_block()], "toolu_02": [_url_image_block()]}),
]
)
@ -4156,7 +4025,9 @@ def test_translate_anthropic_to_openai_without_prompt_cache_breakpoint_adds_noth
def test_translate_anthropic_messages_to_openai_carries_midturn_system_prompt_cache_breakpoint():
explicit = {"mode": "explicit"}
result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(
messages=[{"role": "system", "content": [{"type": "text", "text": "fix", "prompt_cache_breakpoint": explicit}]}],
messages=[
{"role": "system", "content": [{"type": "text", "text": "fix", "prompt_cache_breakpoint": explicit}]}
],
model="gpt-5.6",
)
assert result == [

View file

@ -34,9 +34,7 @@ import pytest
# Anchor sys.path to this file's location — not the working-directory-relative
# pattern Greptile flagged on PR #23706. Resolves correctly regardless of
# where pytest is invoked from.
sys.path.insert(
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))
)
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../..")))
from litellm.llms.anthropic.experimental_pass_through.adapters.handler import (
ANTHROPIC_ONLY_REQUEST_KEYS,
@ -174,9 +172,7 @@ class TestOutputConfigStrippedFromCompletionKwargs:
result = _call_prepare(
extra_kwargs={
"custom_llm_provider": "azure",
"output_config": {
"format": {"type": "json_schema", "schema": losing_schema}
},
"output_config": {"format": {"type": "json_schema", "schema": losing_schema}},
},
output_format={"type": "json_schema", "schema": winning_schema},
)

View file

@ -31,9 +31,7 @@ from litellm.types.utils import (
)
def _build_fake_stream(
content: str, finish_reason: str = "stop"
) -> MockResponseIterator:
def _build_fake_stream(content: str, finish_reason: str = "stop") -> MockResponseIterator:
"""Mimic a Vertex Gemma `:predict` fake stream: one collapsed chunk."""
model_response = ModelResponse()
model_response.choices = [
@ -133,9 +131,7 @@ def test_delayed_usage_chunk_preserves_cache_tokens():
wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="gpt-4o")
events = list(wrapper)
message_delta = next(
event for event in events if event.get("type") == "message_delta"
)
message_delta = next(event for event in events if event.get("type") == "message_delta")
assert message_delta["usage"]["input_tokens"] == 70
assert message_delta["usage"]["output_tokens"] == 5
@ -145,13 +141,7 @@ def test_delayed_usage_chunk_preserves_cache_tokens():
def test_splitter_passes_through_non_combined_chunks():
"""A chunk with content but no finish_reason is not split."""
chunk = ModelResponseStream(
choices=[
StreamingChoices(
index=0, delta=Delta(content="partial"), finish_reason=None
)
]
)
chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="partial"), finish_reason=None)])
chunks = list(_CombinedChunkSplitter(iter([chunk])))
assert len(chunks) == 1
assert chunks[0].choices[0].delta.content == "partial"
@ -159,11 +149,7 @@ def test_splitter_passes_through_non_combined_chunks():
def test_splitter_splits_combined_chunk_into_content_then_finish():
"""A chunk with both content and finish_reason becomes two chunks."""
chunk = ModelResponseStream(
choices=[
StreamingChoices(index=0, delta=Delta(content="done"), finish_reason="stop")
]
)
chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="done"), finish_reason="stop")])
content_chunk, finish_chunk = list(_CombinedChunkSplitter(iter([chunk])))
assert content_chunk.choices[0].delta.content == "done"
@ -193,9 +179,7 @@ def test_split_clears_reasoning_and_thinking_on_finish_chunk():
reasoning_content="some reasoning",
thinking_blocks=[{"type": "thinking"}],
)
chunk = SimpleNamespace(
choices=[SimpleNamespace(finish_reason="stop", delta=delta)]
)
chunk = SimpleNamespace(choices=[SimpleNamespace(finish_reason="stop", delta=delta)])
content_chunk, finish_chunk = _CombinedChunkSplitter._split(chunk)

View file

@ -22,9 +22,7 @@ def _make_text_chunk(
StreamingChoices(
finish_reason=finish_reason,
index=0,
delta=Delta(
content=text, role="assistant" if text else None, tool_calls=None
),
delta=Delta(content=text, role="assistant" if text else None, tool_calls=None),
logprobs=None,
)
]
@ -73,34 +71,23 @@ async def test_stream_emits_compaction_block_before_text():
compaction_start = next(
e
for e in events
if e.get("type") == "content_block_start"
and e.get("content_block", {}).get("type") == "compaction"
if e.get("type") == "content_block_start" and e.get("content_block", {}).get("type") == "compaction"
)
assert compaction_start["index"] == 0
compaction_delta = next(
e
for e in events
if e.get("type") == "content_block_delta"
and e.get("delta", {}).get("type") == "compaction_delta"
if e.get("type") == "content_block_delta" and e.get("delta", {}).get("type") == "compaction_delta"
)
assert compaction_delta["index"] == 0
assert (
compaction_delta["delta"]["content"] == "Summary of prior conversation turns."
)
assert compaction_delta["delta"]["content"] == "Summary of prior conversation turns."
compaction_stop = next(
e
for e in events
if e.get("type") == "content_block_stop" and e.get("index") == 0
)
compaction_stop = next(e for e in events if e.get("type") == "content_block_stop" and e.get("index") == 0)
assert compaction_stop is not None
text_start = next(
e
for e in events
if e.get("type") == "content_block_start"
and e.get("content_block", {}).get("type") == "text"
e for e in events if e.get("type") == "content_block_start" and e.get("content_block", {}).get("type") == "text"
)
assert text_start["index"] == 1
@ -177,14 +164,9 @@ async def test_stream_without_compaction_block_unchanged():
events = await _collect_events_async(wrapper)
assert not any(
e.get("content_block", {}).get("type") == "compaction"
for e in events
if e.get("type") == "content_block_start"
e.get("content_block", {}).get("type") == "compaction" for e in events if e.get("type") == "content_block_start"
)
text_start = next(
e
for e in events
if e.get("type") == "content_block_start"
and e.get("content_block", {}).get("type") == "text"
e for e in events if e.get("type") == "content_block_start" and e.get("content_block", {}).get("type") == "text"
)
assert text_start["index"] == 0

View file

@ -19,9 +19,7 @@ from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Us
def _text_chunk(text: str) -> ModelResponseStream:
return ModelResponseStream(
choices=[StreamingChoices(index=0, delta=Delta(content=text), finish_reason=None)]
)
return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=text), finish_reason=None)])
def _finish_chunk() -> ModelResponseStream:
@ -61,9 +59,7 @@ def test_leading_metadata_chunk_without_choices_does_not_kill_stream():
wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="mock-model")
events = list(wrapper)
text = "".join(
event["delta"]["text"] for event in events if event.get("type") == "content_block_delta"
)
text = "".join(event["delta"]["text"] for event in events if event.get("type") == "content_block_delta")
assert text == "Hello there"
assert events[-1]["type"] == "message_stop"

View file

@ -507,11 +507,7 @@ def _thinking_first_chunks() -> List[MagicMock]:
def _assert_thinking_first_block_opens_at_index_zero(events: List[dict]) -> None:
starts = [
(e["index"], e["content_block"]["type"])
for e in events
if e.get("type") == "content_block_start"
]
starts = [(e["index"], e["content_block"]["type"]) for e in events if e.get("type") == "content_block_start"]
assert starts == [(0, "thinking"), (1, "text")], starts
assert "" not in _text_deltas(events)
assert _thinking_deltas(events) == ["Let me think", "about it."]
@ -980,9 +976,7 @@ def test_tool_block_start_emitted_without_awaiting_the_next_chunk_sync():
"name": "Write",
"input": {},
}
assert stream.pulled == 1, (
f"content_block_start was withheld until {stream.pulled} upstream chunks had arrived"
)
assert stream.pulled == 1, f"content_block_start was withheld until {stream.pulled} upstream chunks had arrived"
@pytest.mark.asyncio
@ -997,9 +991,7 @@ async def test_tool_block_start_emitted_without_awaiting_the_next_chunk_async():
start = await wrapper.__anext__()
assert start["type"] == "content_block_start"
assert start["content_block"]["name"] == "Write"
assert stream.pulled == 1, (
f"content_block_start was withheld until {stream.pulled} upstream chunks had arrived"
)
assert stream.pulled == 1, f"content_block_start was withheld until {stream.pulled} upstream chunks had arrived"
@pytest.mark.parametrize("is_async", [False, True])

View file

@ -127,21 +127,17 @@ async def test_async_stream_emits_input_json_delta_for_bundled_tool_args():
):
input_json_delta_idx = i
assert (
tool_start_idx is not None
), f"Expected content_block_start with type=tool_use; events: {event_types}"
assert (
input_json_delta_idx is not None
), f"Expected content_block_delta with input_json_delta; events: {event_types}"
assert (
input_json_delta_idx == tool_start_idx + 1
), "input_json_delta should immediately follow the tool_use content_block_start"
assert tool_start_idx is not None, f"Expected content_block_start with type=tool_use; events: {event_types}"
assert input_json_delta_idx is not None, (
f"Expected content_block_delta with input_json_delta; events: {event_types}"
)
assert input_json_delta_idx == tool_start_idx + 1, (
"input_json_delta should immediately follow the tool_use content_block_start"
)
# Verify the delta carries the tool arguments
delta_event = events[input_json_delta_idx]
assert delta_event["delta"][
"partial_json"
], "input_json_delta should have non-empty partial_json"
assert delta_event["delta"]["partial_json"], "input_json_delta should have non-empty partial_json"
@pytest.mark.asyncio
@ -230,8 +226,7 @@ async def test_async_stream_no_extra_delta_when_tool_args_empty():
and e["delta"].get("type") == "input_json_delta"
]
assert len(input_json_deltas) == 1, (
f"Expected exactly 1 input_json_delta (from the follow-up chunk), "
f"got {len(input_json_deltas)}"
f"Expected exactly 1 input_json_delta (from the follow-up chunk), got {len(input_json_deltas)}"
)
assert input_json_deltas[0]["delta"]["partial_json"] == '{"location": "NYC"}'
@ -291,15 +286,13 @@ def test_sync_stream_emits_input_json_delta_for_bundled_tool_args():
):
input_json_delta_idx = i
assert (
tool_start_idx is not None
), f"Expected content_block_start with type=tool_use; events: {event_types}"
assert (
input_json_delta_idx is not None
), f"Expected content_block_delta with input_json_delta; events: {event_types}"
assert (
input_json_delta_idx == tool_start_idx + 1
), "input_json_delta should immediately follow the tool_use content_block_start"
assert tool_start_idx is not None, f"Expected content_block_start with type=tool_use; events: {event_types}"
assert input_json_delta_idx is not None, (
f"Expected content_block_delta with input_json_delta; events: {event_types}"
)
assert input_json_delta_idx == tool_start_idx + 1, (
"input_json_delta should immediately follow the tool_use content_block_start"
)
assert events[input_json_delta_idx]["delta"]["partial_json"]
@ -343,9 +336,7 @@ def test_sync_stream_no_extra_delta_when_tool_args_empty():
)
wrapper = AnthropicStreamWrapper(
completion_stream=iter(
[text_chunk, tool_name_chunk, tool_args_chunk, finish_chunk]
),
completion_stream=iter([text_chunk, tool_name_chunk, tool_args_chunk, finish_chunk]),
model="test-model",
)
@ -374,7 +365,6 @@ def test_sync_stream_no_extra_delta_when_tool_args_empty():
and e["delta"].get("type") == "input_json_delta"
]
assert len(input_json_deltas) == 1, (
f"Expected exactly 1 input_json_delta (from the follow-up chunk), "
f"got {len(input_json_deltas)}"
f"Expected exactly 1 input_json_delta (from the follow-up chunk), got {len(input_json_deltas)}"
)
assert input_json_deltas[0]["delta"]["partial_json"] == '{"location": "NYC"}'

View file

@ -258,9 +258,7 @@ def test_tool_result_list_content_shape_preserved():
{"role": "user", "content": "Hi"},
{
"role": "assistant",
"content": [
{"type": "tool_use", "id": "toolu_a", "name": "f", "input": {}}
],
"content": [{"type": "tool_use", "id": "toolu_a", "name": "f", "input": {}}],
},
{
"role": "user",
@ -274,9 +272,7 @@ def test_tool_result_list_content_shape_preserved():
},
{
"role": "assistant",
"content": [
{"type": "tool_use", "id": "toolu_b", "name": "f", "input": {}}
],
"content": [{"type": "tool_use", "id": "toolu_b", "name": "f", "input": {}}],
},
{
"role": "user",

View file

@ -193,10 +193,7 @@ def test_select_last_user_question_strips_tool_result_from_mixed_turn():
content = selected[0]["content"]
assert isinstance(content, list)
assert all(b.get("type") != "tool_result" for b in content)
assert any(
b.get("type") == "text" and b.get("text") == "follow-up question"
for b in content
)
assert any(b.get("type") == "text" and b.get("text") == "follow-up question" for b in content)
def test_select_last_user_question_skips_pure_tool_result_turn():
@ -425,9 +422,7 @@ def test_client_compaction_block_history_without_context_management():
def test_client_compaction_block_history_no_compaction_returns_none():
result = apply_client_compaction_block_history(
messages=_simple_messages(), system="base"
)
result = apply_client_compaction_block_history(messages=_simple_messages(), system="base")
assert result is None
@ -512,9 +507,7 @@ async def test_slice_only_no_compaction_block_under_threshold():
async def test_full_summary_path():
"""Over threshold: summary call fires, compaction_block and iterations_usage returned."""
messages = _simple_messages()
mock_response = _make_mock_response(
"<summary>Condensed history</summary>", prompt_tokens=200, completion_tokens=50
)
mock_response = _make_mock_response("<summary>Condensed history</summary>", prompt_tokens=200, completion_tokens=50)
with (
patch(
@ -1068,13 +1061,11 @@ async def test_summary_call_does_not_emit_consecutive_user_turns():
)
summary_messages = captured_calls[0]["summary_messages"]
user_indices = [
idx for idx, msg in enumerate(summary_messages) if msg.get("role") == "user"
]
user_indices = [idx for idx, msg in enumerate(summary_messages) if msg.get("role") == "user"]
# No two adjacent indices.
assert all(
b - a > 1 for a, b in zip(user_indices, user_indices[1:])
), f"two consecutive user turns produced: {summary_messages}"
assert all(b - a > 1 for a, b in zip(user_indices, user_indices[1:])), (
f"two consecutive user turns produced: {summary_messages}"
)
async def test_summary_call_sends_default_max_tokens():
@ -1157,9 +1148,9 @@ def test_summary_max_tokens_setting_falls_back_for_invalid_values():
"litellm.proxy.proxy_server.general_settings",
{"context_management_summary_max_tokens": bad},
):
assert (
_read_summary_max_tokens_setting() == COMPACT_SUMMARY_MAX_TOKENS
), f"expected default for invalid override {bad!r}"
assert _read_summary_max_tokens_setting() == COMPACT_SUMMARY_MAX_TOKENS, (
f"expected default for invalid override {bad!r}"
)
async def test_summary_call_sends_default_timeout():
@ -1282,9 +1273,7 @@ async def test_summary_model_denied_when_team_not_in_allowlist():
tools=None,
system=None,
edit_spec=_EDIT_SPEC_DEFAULT,
user_api_key_auth=_fake_user_api_key_auth(
key_models=["all-proxy-models"], team_models=["gpt-4o"]
),
user_api_key_auth=_fake_user_api_key_auth(key_models=["all-proxy-models"], team_models=["gpt-4o"]),
)
mock_call.assert_not_awaited()
@ -1313,9 +1302,7 @@ async def test_summary_model_allowed_when_in_key_allowlist():
tools=None,
system=None,
edit_spec=_EDIT_SPEC_DEFAULT,
user_api_key_auth=_fake_user_api_key_auth(
key_models=["claude-haiku-4-5", "gpt-4o"]
),
user_api_key_auth=_fake_user_api_key_auth(key_models=["claude-haiku-4-5", "gpt-4o"]),
)
mock_call.assert_awaited_once()
@ -1521,9 +1508,7 @@ async def test_summary_model_denied_when_key_over_model_budget():
limiter = MagicMock()
limiter.is_key_within_model_budget = AsyncMock(
side_effect=litellm.BudgetExceededError(
message="over budget", current_cost=10, max_budget=5
)
side_effect=litellm.BudgetExceededError(message="over budget", current_cost=10, max_budget=5)
)
with (
@ -1574,9 +1559,7 @@ async def test_summary_model_denied_when_user_over_model_budget():
limiter = MagicMock()
limiter.is_user_within_model_budget = AsyncMock(
side_effect=litellm.BudgetExceededError(
message="over budget", current_cost=10, max_budget=5
)
side_effect=litellm.BudgetExceededError(message="over budget", current_cost=10, max_budget=5)
)
with (
@ -1617,9 +1600,7 @@ async def test_summary_model_denied_when_user_over_model_budget():
_PROXY_VirtualKeyModelMaxBudgetLimiter,
)
real_params = inspect.signature(
_PROXY_VirtualKeyModelMaxBudgetLimiter.is_user_within_model_budget
).parameters
real_params = inspect.signature(_PROXY_VirtualKeyModelMaxBudgetLimiter.is_user_within_model_budget).parameters
for kwarg in ("user_id", "user_model_max_budget", "model"):
assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter no longer accepts"
@ -1641,9 +1622,7 @@ async def test_summary_model_denied_when_end_user_over_model_budget():
limiter = MagicMock()
limiter.is_key_within_model_budget = AsyncMock(return_value=True)
limiter.is_end_user_within_model_budget = AsyncMock(
side_effect=litellm.BudgetExceededError(
message="over budget", current_cost=10, max_budget=5
)
side_effect=litellm.BudgetExceededError(message="over budget", current_cost=10, max_budget=5)
)
with (
@ -1956,9 +1935,7 @@ async def test_model_budget_metadata_propagated_to_summary_call():
parent_litellm_metadata = {
"user_api_key": "sk-test",
"user_api_key_model_max_budget": {"claude-haiku-4-5": {"budget_limit": 5}},
"user_api_key_end_user_model_max_budget": {
"claude-haiku-4-5": {"budget_limit": 2}
},
"user_api_key_end_user_model_max_budget": {"claude-haiku-4-5": {"budget_limit": 2}},
}
with (
@ -1983,12 +1960,8 @@ async def test_model_budget_metadata_propagated_to_summary_call():
)
propagated = mock_call.call_args.kwargs["metadata"]
assert propagated["user_api_key_model_max_budget"] == {
"claude-haiku-4-5": {"budget_limit": 5}
}
assert propagated["user_api_key_end_user_model_max_budget"] == {
"claude-haiku-4-5": {"budget_limit": 2}
}
assert propagated["user_api_key_model_max_budget"] == {"claude-haiku-4-5": {"budget_limit": 5}}
assert propagated["user_api_key_end_user_model_max_budget"] == {"claude-haiku-4-5": {"budget_limit": 2}}
async def test_summary_call_propagates_allowed_model_region():
@ -2460,9 +2433,7 @@ def test_endpoint_returns_anthropic_400_on_context_management_error():
mock_proxy_server.version = "test"
with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
with patch(
"litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
) as mock_cls:
with patch("litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_cls:
mock_instance = MagicMock()
mock_instance.base_process_llm_request = AsyncMock(
side_effect=AnthropicContextManagementError(
@ -2521,9 +2492,7 @@ def test_endpoint_runs_failure_hook_on_500_context_management_error():
mock_proxy_server.version = "test"
with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
with patch(
"litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
) as mock_cls:
with patch("litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_cls:
mock_instance = MagicMock()
mock_instance.base_process_llm_request = AsyncMock(
side_effect=AnthropicContextManagementError(

View file

@ -50,9 +50,7 @@ async def test_unknown_edit_type_is_noop():
messages=messages,
tools=None,
system=None,
context_management_spec={
"edits": [{"type": "totally_not_a_real_edit_20999999"}]
},
context_management_spec={"edits": [{"type": "totally_not_a_real_edit_20999999"}]},
)
assert result.applied_edits == []
assert result.messages == messages

View file

@ -39,9 +39,7 @@ def _text_resp(text: str, model: str = "gpt-4o-mini") -> Dict:
}
def _advisor_call_resp(
question: str = "How do I approach this?", tool_id: str = "tid_01"
) -> Dict:
def _advisor_call_resp(question: str = "How do I approach this?", tool_id: str = "tid_01") -> Dict:
return {
"id": "msg_int_test",
"type": "message",
@ -106,14 +104,10 @@ async def test_full_dispatch_interceptor_fires_and_loop_completes():
assert isinstance(result, dict)
content = result.get("content", [])
text_blocks = [b for b in content if b.get("type") == "text"]
advisor_uses = [
b for b in content if b.get("type") == "tool_use" and b.get("name") == "advisor"
]
advisor_uses = [b for b in content if b.get("type") == "tool_use" and b.get("name") == "advisor"]
assert len(text_blocks) >= 1, "Final response must have text"
assert (
len(advisor_uses) == 0
), "No advisor tool_use blocks must appear in final output"
assert len(advisor_uses) == 0, "No advisor tool_use blocks must appear in final output"
# ---------------------------------------------------------------------------
@ -221,9 +215,7 @@ async def test_named_params_forwarded_into_advisor_executor_subcall():
captured_executor_kwargs: Dict = {}
async def mock_handler(
model, messages, tools, stream, max_tokens, custom_llm_provider, **kwargs
):
async def mock_handler(model, messages, tools, stream, max_tokens, custom_llm_provider, **kwargs):
# First call is the executor sub-call (returns advisor tool_use).
# Capture its kwargs so we can assert the forwarded params.
if not captured_executor_kwargs:
@ -267,8 +259,7 @@ async def test_named_params_forwarded_into_advisor_executor_subcall():
)
assert captured_executor_kwargs["thinking"] == {"type": "adaptive"}, (
"thinking must be forwarded into executor sub-call — see "
"anthropic_messages.handler interceptor invocation."
"thinking must be forwarded into executor sub-call — see anthropic_messages.handler interceptor invocation."
)
# The advisor enriches metadata with `advisor_sub_call` / `parent_request_id`,
# but the original caller fields must survive into the executor sub-call.
@ -304,9 +295,7 @@ async def test_pre_request_hook_override_does_not_collide_with_explicit_kwargs()
captured: Dict = {}
async def mock_handler(
model, messages, tools, stream, max_tokens, custom_llm_provider, **kwargs
):
async def mock_handler(model, messages, tools, stream, max_tokens, custom_llm_provider, **kwargs):
if not captured:
captured.update(
{
@ -320,9 +309,7 @@ async def test_pre_request_hook_override_does_not_collide_with_explicit_kwargs()
return _text_resp("Some advice.", model="claude-opus-4-6")
return _text_resp("Final answer.")
async def fake_pre_request_hooks(
model, messages, tools, stream, custom_llm_provider, **hook_kwargs
):
async def fake_pre_request_hooks(model, messages, tools, stream, custom_llm_provider, **hook_kwargs):
# Simulate a CustomLogger.async_pre_request_hook that overrides several
# named params on its way through. Without the request_kwargs.pop()
# extraction in handler.py, these would collide with the explicit

View file

@ -73,9 +73,7 @@ def _build_simple_text_stream() -> List[bytes]:
},
)
)
chunks.append(
_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
)
chunks.append(_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}))
chunks.append(
_sse_event(
"message_delta",
@ -148,9 +146,7 @@ def _build_tool_use_stream() -> List[bytes]:
},
)
)
chunks.append(
_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
)
chunks.append(_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}))
# tool_use block
chunks.append(
_sse_event(
@ -190,9 +186,7 @@ def _build_tool_use_stream() -> List[bytes]:
},
)
)
chunks.append(
_sse_event("content_block_stop", {"type": "content_block_stop", "index": 1})
)
chunks.append(_sse_event("content_block_stop", {"type": "content_block_stop", "index": 1}))
chunks.append(
_sse_event(
"message_delta",
@ -284,9 +278,7 @@ def _build_hold_back_iterator(
class TestParseSSEEvents:
def test_should_parse_single_event(self):
raw = _sse_event(
"message_start", {"type": "message_start", "message": {"id": "1"}}
)
raw = _sse_event("message_start", {"type": "message_start", "message": {"id": "1"}})
events = _parse_sse_events(raw)
assert len(events) == 1
assert events[0][0] == "message_start"
@ -457,9 +449,7 @@ class TestHandleMessageDelta:
class TestRebuildAnthropicResponse:
def test_should_rebuild_simple_text_response(self):
raw_bytes = _build_simple_text_stream()
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(raw_bytes)
assert result is not None
assert result["id"] == "msg_123"
assert result["model"] == "claude-sonnet-4-20250514"
@ -472,9 +462,7 @@ class TestRebuildAnthropicResponse:
def test_should_rebuild_tool_use_response(self):
raw_bytes = _build_tool_use_stream()
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(raw_bytes)
assert result is not None
assert result["id"] == "msg_tool_456"
assert result["stop_reason"] == "tool_use"
@ -502,23 +490,17 @@ class TestRebuildAnthropicResponse:
},
)
]
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(raw_bytes)
assert result is None
def test_should_handle_empty_bytes(self):
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
[]
)
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse([])
assert result is None
def test_should_handle_multi_event_chunks(self):
"""When multiple SSE events arrive in a single bytes chunk."""
combined = b"".join(_build_simple_text_stream())
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
[combined]
)
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse([combined])
assert result is not None
assert result["content"][0]["text"] == "Hello, world!"
@ -550,9 +532,7 @@ class TestRebuildAnthropicResponse:
),
_sse_event("message_stop", {"type": "message_stop"}),
]
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(raw_bytes)
assert result is not None
assert result["usage"]["cache_creation_input_tokens"] == 50
assert result["usage"]["cache_read_input_tokens"] == 30
@ -593,9 +573,7 @@ class TestRebuildAnthropicResponse:
),
_sse_event("message_stop", {"type": "message_stop"}),
]
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(raw_bytes)
assert result is not None
assert result["content"][0]["type"] == "redacted_thinking"
@ -722,9 +700,7 @@ class TestAgenticStreamingIteratorPhase2:
}
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(
return_value=fake_response
)
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=fake_response)
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
@ -757,9 +733,7 @@ class TestAgenticStreamingIteratorErrorHandling:
mock_stream = MockAsyncStream(chunks)
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(
side_effect=RuntimeError("hook exploded")
)
mock_handler._call_agentic_completion_hooks = AsyncMock(side_effect=RuntimeError("hook exploded"))
mock_logging = MagicMock()
mock_logging.litellm_call_id = "test_call_123"

View file

@ -0,0 +1,118 @@
from pathlib import Path
from typing import Final
import pytest
import litellm
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
AnthropicMessagesConfig,
)
from litellm.llms.anthropic.wif import get_anthropic_wif_token
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,
token_response,
write_token_file,
)
_WIF_PARAMS: Final[dict] = {
"anthropic_federation_rule_id": "fdrl_abc123",
"anthropic_organization_id": "org-uuid-1",
"anthropic_identity_token_file": "/var/run/secrets/identity-token",
}
def test_workload_identity_allowed_for_anthropic() -> None:
assert AnthropicMessagesConfig()._allows_workload_identity is True
def test_workload_identity_blocked_for_minimax() -> None:
assert MinimaxMessagesConfig()._allows_workload_identity is False
def test_workload_identity_blocked_for_tencent() -> None:
assert TencentAnthropicMessagesConfig()._allows_workload_identity is False
def test_minimax_validate_environment_never_attaches_anthropic_wif_credential(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Regression test: before the fix, an Anthropic-WIF-configured proxy would mint a real
Anthropic federation token inside MiniMax's inherited validate_anthropic_messages_environment
and send it as the Authorization header on the MiniMax-routed request. With no MiniMax
credential of its own the deployment must fail closed on the missing key instead."""
monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_prod")
monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-prod-uuid")
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
monkeypatch.delenv("MINIMAX_API_KEY", raising=False)
token_file = write_token_file(tmp_path, "jwt-assertion-value")
litellm_params = {"anthropic_identity_token_file": str(token_file)}
with pytest.raises(litellm.AuthenticationError, match="Missing Anthropic API Key"):
MinimaxMessagesConfig().validate_anthropic_messages_environment(
headers={},
model="MiniMax-M2.1",
messages=[],
optional_params={},
litellm_params=litellm_params,
api_key=None,
api_base="https://api.minimax.io/anthropic",
)
def test_tencent_validate_environment_never_attaches_anthropic_wif_credential(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_prod")
monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-prod-uuid")
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
monkeypatch.delenv("TENCENT_API_KEY", raising=False)
token_file = write_token_file(tmp_path, "jwt-assertion-value")
litellm_params = {"anthropic_identity_token_file": str(token_file)}
monkeypatch.setattr(litellm, "api_key", None)
with pytest.raises(litellm.AuthenticationError, match="Missing Anthropic API Key"):
TencentAnthropicMessagesConfig().validate_anthropic_messages_environment(
headers={},
model="deepseek-v4-pro",
messages=[],
optional_params={},
litellm_params=litellm_params,
api_key=None,
api_base="https://tokenhub-intl.tencentcloudmaas.com",
)
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
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."""
monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_prod")
monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-prod-uuid")
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)
minted: Final = get_anthropic_wif_token(
litellm_params,
"https://api.anthropic.com",
"claude-sonnet-4-5",
engine,
)
assert minted == "sk-ant-oat01-canary"
assert len(poster.requests) == 1
for config in (MinimaxMessagesConfig(), TencentAnthropicMessagesConfig()):
assert config._allows_workload_identity is False
assert len(poster.requests) == 1

View file

@ -93,13 +93,11 @@ def test_messages_drops_speed_for_vertex_opus_with_drop_params(monkeypatch):
"""Regression: a vertex_ai Opus passthrough must drop ``speed`` even though the
prefix-stripped model id maps to a fast-mode-capable direct-Anthropic entry."""
monkeypatch.setattr(litellm, "drop_params", True)
optional_params = (
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
params={"max_tokens": 1024, "speed": "fast"},
model="claude-opus-4-8",
drop_params=False,
custom_llm_provider="vertex_ai",
)
optional_params = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
params={"max_tokens": 1024, "speed": "fast"},
model="claude-opus-4-8",
drop_params=False,
custom_llm_provider="vertex_ai",
)
assert "speed" not in optional_params

View file

@ -13,9 +13,7 @@ def test_output_format_supported_and_transforms_correctly():
config = AnthropicMessagesConfig()
# 1. Verify it's in supported parameters
supported_params = config.get_supported_anthropic_messages_params(
"claude-sonnet-4-5"
)
supported_params = config.get_supported_anthropic_messages_params("claude-sonnet-4-5")
assert "output_format" in supported_params
# 2. Verify transformation preserves output_format and adds beta header

View file

@ -30,26 +30,14 @@ class MockCompletionStreamWithContentAfterStopReason:
self.responses = [
# Initial text content
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content="Hello"), index=0, finish_reason=None
)
],
choices=[StreamingChoices(delta=Delta(content="Hello"), index=0, finish_reason=None)],
),
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content=" world"), index=0, finish_reason=None
)
],
choices=[StreamingChoices(delta=Delta(content=" world"), index=0, finish_reason=None)],
),
# Message delta with stop_reason AND usage (this is how it actually comes from the API)
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content=""), index=0, finish_reason="stop"
)
],
choices=[StreamingChoices(delta=Delta(content=""), index=0, finish_reason="stop")],
usage=Usage(prompt_tokens=230, completion_tokens=65, total_tokens=295),
),
# Additional content after the stop_reason - this simulates the scenario
@ -118,9 +106,9 @@ def test_anthropic_stream_wrapper_content_after_stop_reason():
print(f"Expected chunk types: {expected_types}")
# Verify we have the expected number of chunks
assert len(chunk_types) >= len(
expected_types
), f"Expected at least {len(expected_types)} chunks, got {len(chunk_types)}"
assert len(chunk_types) >= len(expected_types), (
f"Expected at least {len(expected_types)} chunks, got {len(chunk_types)}"
)
# Verify key chunk types are present
assert "message_start" in chunk_types
@ -143,15 +131,9 @@ def test_anthropic_stream_wrapper_content_after_stop_reason():
delta = message_delta_chunk.get("delta", {})
usage = message_delta_chunk.get("usage", {})
assert (
delta.get("stop_reason") == "end_turn"
), f"Expected stop_reason 'end_turn', got {delta.get('stop_reason')}"
assert (
usage.get("input_tokens") == 230
), f"Expected input_tokens 230, got {usage.get('input_tokens')}"
assert (
usage.get("output_tokens") == 65
), f"Expected output_tokens 65, got {usage.get('output_tokens')}"
assert delta.get("stop_reason") == "end_turn", f"Expected stop_reason 'end_turn', got {delta.get('stop_reason')}"
assert usage.get("input_tokens") == 230, f"Expected input_tokens 230, got {usage.get('input_tokens')}"
assert usage.get("output_tokens") == 65, f"Expected output_tokens 65, got {usage.get('output_tokens')}"
# Verify content_block_stop comes before message_delta
content_block_stop_index = None
@ -165,9 +147,7 @@ def test_anthropic_stream_wrapper_content_after_stop_reason():
assert content_block_stop_index is not None, "content_block_stop not found"
assert message_delta_index is not None, "message_delta not found"
assert (
content_block_stop_index < message_delta_index
), "content_block_stop should come before message_delta"
assert content_block_stop_index < message_delta_index, "content_block_stop should come before message_delta"
@pytest.mark.asyncio
@ -210,15 +190,9 @@ async def test_async_anthropic_stream_wrapper_content_after_stop_reason():
delta = message_delta_chunk.get("delta", {})
usage = message_delta_chunk.get("usage", {})
assert (
delta.get("stop_reason") == "end_turn"
), f"Expected stop_reason 'end_turn', got {delta.get('stop_reason')}"
assert (
usage.get("input_tokens") == 230
), f"Expected input_tokens 230, got {usage.get('input_tokens')}"
assert (
usage.get("output_tokens") == 65
), f"Expected output_tokens 65, got {usage.get('output_tokens')}"
assert delta.get("stop_reason") == "end_turn", f"Expected stop_reason 'end_turn', got {delta.get('stop_reason')}"
assert usage.get("input_tokens") == 230, f"Expected input_tokens 230, got {usage.get('input_tokens')}"
assert usage.get("output_tokens") == 65, f"Expected output_tokens 65, got {usage.get('output_tokens')}"
def test_usage_merging_behavior():
@ -234,18 +208,10 @@ def test_usage_merging_behavior():
for chunk in wrapper:
chunks.append(chunk)
# If this is a message_delta with stop_reason, verify it has usage
if (
chunk.get("type") == "message_delta"
and chunk.get("delta", {}).get("stop_reason") is not None
):
if chunk.get("type") == "message_delta" and chunk.get("delta", {}).get("stop_reason") is not None:
usage = chunk.get("usage", {})
assert (
usage.get("input_tokens") is not None
), "Usage should be merged with stop_reason chunk"
assert (
usage.get("output_tokens") is not None
), "Usage should be merged with stop_reason chunk"
assert usage.get("input_tokens") is not None, "Usage should be merged with stop_reason chunk"
assert usage.get("output_tokens") is not None, "Usage should be merged with stop_reason chunk"
break
@ -273,12 +239,8 @@ def test_sse_wrapper_with_content_after_stop_reason():
lines = chunk_str.split("\n")
# Should have event and data lines
assert any(
line.startswith("event: ") for line in lines
), f"Missing event line in: {chunk_str}"
assert any(
line.startswith("data: ") for line in lines
), f"Missing data line in: {chunk_str}"
assert any(line.startswith("event: ") for line in lines), f"Missing event line in: {chunk_str}"
assert any(line.startswith("data: ") for line in lines), f"Missing data line in: {chunk_str}"
@pytest.mark.asyncio
@ -306,12 +268,8 @@ async def test_async_sse_wrapper_with_content_after_stop_reason():
lines = chunk_str.split("\n")
# Should have event and data lines
assert any(
line.startswith("event: ") for line in lines
), f"Missing event line in: {chunk_str}"
assert any(
line.startswith("data: ") for line in lines
), f"Missing data line in: {chunk_str}"
assert any(line.startswith("event: ") for line in lines), f"Missing event line in: {chunk_str}"
assert any(line.startswith("data: ") for line in lines), f"Missing data line in: {chunk_str}"
if __name__ == "__main__":

View file

@ -59,7 +59,7 @@ def test_anthropic_messages_handler_skips_the_gateway_on_recursion():
"litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp",
new=AsyncMock(return_value={"routed": True}),
) as routed:
with pytest.raises(ValueError, match='anthropic_messages_handler is not implemented for sync calls'):
with pytest.raises(ValueError, match="anthropic_messages_handler is not implemented for sync calls"):
anthropic_messages_handler(
max_tokens=100,
messages=[{"role": "user", "content": "hi"}],
@ -78,7 +78,7 @@ def test_anthropic_messages_handler_leaves_native_tools_alone():
"litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp",
new=AsyncMock(return_value={"routed": True}),
) as routed:
with pytest.raises(ValueError, match='anthropic_messages_handler is not implemented for sync calls'):
with pytest.raises(ValueError, match="anthropic_messages_handler is not implemented for sync calls"):
anthropic_messages_handler(
max_tokens=100,
messages=[{"role": "user", "content": "hi"}],
@ -115,9 +115,7 @@ def test_build_tool_result_message_uses_anthropic_tool_result_blocks():
message = _build_tool_result_message([{"tool_call_id": "toolu_1", "result": "9 sections", "name": "read_wiki"}])
assert message["role"] == "user"
assert list(message["content"]) == [
{"type": "tool_result", "tool_use_id": "toolu_1", "content": "9 sections"}
]
assert list(message["content"]) == [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "9 sections"}]
@pytest.mark.asyncio

View file

@ -1,7 +1,6 @@
from typing import List
from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import (
AnthropicStreamWrapper,
)
@ -53,9 +52,7 @@ def construct_text_chunk(text: str) -> ModelResponseStream:
)
def construct_split_tool_call(
id: str, function_name: str, function_arg_parts: List[str]
) -> List[ModelResponseStream]:
def construct_split_tool_call(id: str, function_name: str, function_arg_parts: List[str]) -> List[ModelResponseStream]:
return [
# https://platform.openai.com/docs/guides/function-calling#streaming
ModelResponseStream(
@ -144,10 +141,7 @@ def test_anthropic_stream_wrapper_single_tool_call():
get_weather_calls = 0
for chunk in chunks:
if (
chunk.get("type") == "content_block_start"
and chunk["content_block"]["type"] == "tool_use"
):
if chunk.get("type") == "content_block_start" and chunk["content_block"]["type"] == "tool_use":
if chunk["content_block"]["name"] == "get_weather":
get_weather_calls += 1
@ -203,10 +197,7 @@ def test_anthropic_stream_wrapper_back_to_back_tool_calls():
get_weather_calls = 0
for chunk in chunks:
if (
chunk.get("type") == "content_block_start"
and chunk["content_block"]["type"] == "tool_use"
):
if chunk.get("type") == "content_block_start" and chunk["content_block"]["type"] == "tool_use":
if chunk["content_block"]["name"] == "get_weather":
get_weather_calls += 1
@ -218,9 +209,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text():
*construct_split_tool_call("tooluse_foo", "get_weather", ['{"city":', '"NY"}']),
construct_text_chunk("The weather is nice today."),
*construct_split_tool_call("tooluse_bar", "get_weather", ['{"city":', '"SF"}']),
*construct_split_tool_call(
"tooluse_bar", "get_weather", ['{"city":', '"CHI"}']
),
*construct_split_tool_call("tooluse_bar", "get_weather", ['{"city":', '"CHI"}']),
construct_text_chunk("The weather is not so nice today."),
ModelResponseStream(
choices=[
@ -280,8 +269,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text():
text_deltas = [
chunk["delta"]["text"]
for chunk in chunks
if chunk.get("type") == "content_block_delta"
and chunk["delta"].get("type") == "text_delta"
if chunk.get("type") == "content_block_delta" and chunk["delta"].get("type") == "text_delta"
]
assert text_deltas == [
"The weather is nice today.",
@ -291,10 +279,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text():
get_weather_calls = 0
for chunk in chunks:
if (
chunk.get("type") == "content_block_start"
and chunk["content_block"]["type"] == "tool_use"
):
if chunk.get("type") == "content_block_start" and chunk["content_block"]["type"] == "tool_use":
if chunk["content_block"]["name"] == "get_weather":
get_weather_calls += 1

View file

@ -29,13 +29,12 @@ def _call_handler_and_capture_optional_params(thinking=None, **extra_kwargs):
"""
captured = {}
with patch(
"litellm.llms.anthropic.experimental_pass_through.messages.handler."
"base_llm_http_handler"
) as mock_handler, patch(
"litellm.llms.anthropic.experimental_pass_through.messages.handler."
"ProviderConfigManager"
) as mock_pcm:
with (
patch(
"litellm.llms.anthropic.experimental_pass_through.messages.handler.base_llm_http_handler"
) as mock_handler,
patch("litellm.llms.anthropic.experimental_pass_through.messages.handler.ProviderConfigManager") as mock_pcm,
):
# Make get_provider_anthropic_messages_config return a non-None config
# so the handler takes the native Anthropic path
mock_pcm.get_provider_anthropic_messages_config.return_value = MagicMock()
@ -71,9 +70,7 @@ class TestReasoningAutoSummaryMessages:
def test_adaptive_thinking_gets_display_summarized(self):
"""reasoning_auto_summary=True + thinking.type='adaptive' -> display='summarized'."""
with patch.object(litellm, "reasoning_auto_summary", True):
params = _call_handler_and_capture_optional_params(
thinking={"type": "adaptive", "budget_tokens": 5000}
)
params = _call_handler_and_capture_optional_params(thinking={"type": "adaptive", "budget_tokens": 5000})
thinking = params.get("thinking", {})
assert thinking.get("display") == "summarized"
assert thinking.get("type") == "adaptive"
@ -82,9 +79,7 @@ class TestReasoningAutoSummaryMessages:
def test_enabled_thinking_gets_display_summarized(self):
"""reasoning_auto_summary=True + thinking.type='enabled' -> display='summarized'."""
with patch.object(litellm, "reasoning_auto_summary", True):
params = _call_handler_and_capture_optional_params(
thinking={"type": "enabled", "budget_tokens": 10000}
)
params = _call_handler_and_capture_optional_params(thinking={"type": "enabled", "budget_tokens": 10000})
thinking = params.get("thinking", {})
assert thinking.get("display") == "summarized"
assert thinking.get("type") == "enabled"
@ -92,18 +87,14 @@ class TestReasoningAutoSummaryMessages:
def test_disabled_thinking_no_display(self):
"""reasoning_auto_summary=True + thinking.type='disabled' -> display NOT set."""
with patch.object(litellm, "reasoning_auto_summary", True):
params = _call_handler_and_capture_optional_params(
thinking={"type": "disabled"}
)
params = _call_handler_and_capture_optional_params(thinking={"type": "disabled"})
thinking = params.get("thinking", {})
assert "display" not in thinking
def test_no_injection_when_flag_false(self):
"""reasoning_auto_summary=False + active thinking -> display NOT set."""
with patch.object(litellm, "reasoning_auto_summary", False):
params = _call_handler_and_capture_optional_params(
thinking={"type": "enabled", "budget_tokens": 10000}
)
params = _call_handler_and_capture_optional_params(thinking={"type": "enabled", "budget_tokens": 10000})
thinking = params.get("thinking", {})
assert "display" not in thinking
@ -117,12 +108,11 @@ class TestReasoningAutoSummaryMessages:
def test_env_var_enables_auto_summary(self):
"""LITELLM_REASONING_AUTO_SUMMARY=true env var enables the feature."""
with patch.object(litellm, "reasoning_auto_summary", False), patch.dict(
os.environ, {"LITELLM_REASONING_AUTO_SUMMARY": "true"}
with (
patch.object(litellm, "reasoning_auto_summary", False),
patch.dict(os.environ, {"LITELLM_REASONING_AUTO_SUMMARY": "true"}),
):
params = _call_handler_and_capture_optional_params(
thinking={"type": "adaptive", "budget_tokens": 5000}
)
params = _call_handler_and_capture_optional_params(thinking={"type": "adaptive", "budget_tokens": 5000})
thinking = params.get("thinking", {})
assert thinking.get("display") == "summarized"

View file

@ -23,11 +23,7 @@ def test_optional_param_filtering_unchanged():
"not_a_real_param": "drop me", # invalid key dropped
"stream": True,
}
result = (
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
params
)
)
result = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(params)
assert result == {"temperature": 0.5, "tools": [{"name": "x"}], "stream": True}
assert "top_p" not in result
assert "not_a_real_param" not in result
@ -37,9 +33,7 @@ def test_valid_keys_are_memoized():
_anthropic_messages_optional_param_keys.cache_clear()
first = _anthropic_messages_optional_param_keys()
for _ in range(50):
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
{"temperature": 0.1}
)
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param({"temperature": 0.1})
info = _anthropic_messages_optional_param_keys.cache_info()
# Resolved exactly once despite many calls.
assert info.misses == 1
@ -51,23 +45,16 @@ def test_valid_keys_are_memoized():
def test_empty_params():
assert (
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
{}
)
== {}
)
assert AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param({}) == {}
def test_drop_params_strips_speed_for_unsupported_model():
original = litellm.drop_params
litellm.drop_params = True
try:
result = (
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
params={"speed": "fast", "temperature": 0.5},
model="claude-sonnet-4-6",
)
result = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
params={"speed": "fast", "temperature": 0.5},
model="claude-sonnet-4-6",
)
finally:
litellm.drop_params = original
@ -80,11 +67,9 @@ def test_drop_params_keeps_speed_for_supporting_model():
original = litellm.drop_params
litellm.drop_params = True
try:
result = (
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
params={"speed": "fast"},
model="claude-opus-4-6",
)
result = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
params={"speed": "fast"},
model="claude-opus-4-6",
)
finally:
litellm.drop_params = original

View file

@ -1,4 +1,3 @@
import pytest
from fastapi.testclient import TestClient
@ -14,25 +13,13 @@ class MockCompletionStream:
def __init__(self):
self.responses = [
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content="Hello"), index=0, finish_reason=None
)
],
choices=[StreamingChoices(delta=Delta(content="Hello"), index=0, finish_reason=None)],
),
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content=" World"), index=0, finish_reason=None
)
],
choices=[StreamingChoices(delta=Delta(content=" World"), index=0, finish_reason=None)],
),
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content=""), index=0, finish_reason="stop"
)
],
choices=[StreamingChoices(delta=Delta(content=""), index=0, finish_reason="stop")],
),
]
self.index = 0
@ -50,9 +37,7 @@ class MockCompletionStream:
def test_anthropic_sse_wrapper_format():
"""Test that the SSE wrapper produces proper event and data formatting"""
wrapper = AnthropicStreamWrapper(
completion_stream=MockCompletionStream(), model="claude-3"
)
wrapper = AnthropicStreamWrapper(completion_stream=MockCompletionStream(), model="claude-3")
# Get the first chunk from the SSE wrapper
first_chunk = next(wrapper.anthropic_sse_wrapper())
@ -73,9 +58,7 @@ def test_anthropic_sse_wrapper_format():
def test_anthropic_sse_wrapper_event_types():
"""Test that different chunk types produce correct event types"""
wrapper = AnthropicStreamWrapper(
completion_stream=MockCompletionStream(), model="claude-3"
)
wrapper = AnthropicStreamWrapper(completion_stream=MockCompletionStream(), model="claude-3")
chunks = []
for chunk in wrapper.anthropic_sse_wrapper():
@ -104,18 +87,10 @@ async def test_async_anthropic_sse_wrapper():
def __init__(self):
self.responses = [
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content="Hello"), index=0, finish_reason=None
)
],
choices=[StreamingChoices(delta=Delta(content="Hello"), index=0, finish_reason=None)],
),
ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content=" World"), index=0, finish_reason=None
)
],
choices=[StreamingChoices(delta=Delta(content=" World"), index=0, finish_reason=None)],
),
]
self.index = 0
@ -130,9 +105,7 @@ async def test_async_anthropic_sse_wrapper():
self.index += 1
return response
wrapper = AnthropicStreamWrapper(
completion_stream=AsyncMockCompletionStream(), model="claude-3"
)
wrapper = AnthropicStreamWrapper(completion_stream=AsyncMockCompletionStream(), model="claude-3")
# Get the first chunk from the async SSE wrapper
first_chunk = None

View file

@ -165,7 +165,7 @@ async def test_async_sse_wrapper_treats_message_stop_bytes_as_complete():
def test_is_message_stop_chunk():
assert _is_message_stop_chunk({"type": "message_stop"}) is True
assert _is_message_stop_chunk({"type": "message_delta"}) is False
assert _is_message_stop_chunk(b'event: message_stop\ndata: {}\n\n') is True
assert _is_message_stop_chunk(b"event: message_stop\ndata: {}\n\n") is True
assert _is_message_stop_chunk(b"raw-bytes") is False
assert _is_message_stop_chunk("message_stop") is False
@ -177,7 +177,7 @@ def test_is_message_stop_chunk_ignores_substring_in_payload():
not be treated as a terminal stop event.
"""
delta_frame_with_substring = (
b'event: content_block_delta\n'
b"event: content_block_delta\n"
b'data: {"type": "content_block_delta", "delta": '
b'{"type": "input_json_delta", "partial_json": "\\"message_stop\\""}}\n\n'
)
@ -281,10 +281,11 @@ async def test_async_sse_wrapper_emits_error_when_bytes_stream_only_mentions_mes
payload text contains `message_stop` (but never emits the actual
`event: message_stop` frame) must still be flagged as incomplete.
"""
async def _byte_stream():
yield b'event: message_start\ndata: {"type": "message_start"}\n\n'
yield (
b'event: content_block_delta\n'
b"event: content_block_delta\n"
b'data: {"type": "content_block_delta", "delta": '
b'{"type": "input_json_delta", "partial_json": "\\"message_stop\\""}}\n\n'
)

View file

@ -1481,9 +1481,7 @@ class TestToolResultImages:
},
{
"role": "user",
"content": [
{"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content}
],
"content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content}],
},
]
@ -1857,7 +1855,9 @@ class TestPromptCacheBreakpointToResponses:
]
def test_system_without_breakpoint_still_becomes_instructions(self):
request = _make_request(system=[{"type": "text", "text": "Be concise."}, {"type": "text", "text": "Be helpful."}])
request = _make_request(
system=[{"type": "text", "text": "Be concise."}, {"type": "text", "text": "Be helpful."}]
)
kwargs = _ADAPTER.translate_request(request)
assert kwargs["instructions"] == "Be concise.\nBe helpful."
assert kwargs["input"] == [

View file

@ -5,7 +5,9 @@ Tests the AnthropicFilesConfig class which transforms between
OpenAI-compatible file operations and Anthropic's Files API format.
"""
import asyncio
import io
import threading
import time
import httpx
@ -90,6 +92,38 @@ class TestAnthropicFilesConfig:
api_key=None,
)
@pytest.mark.asyncio
async def test_avalidate_environment_sets_headers(self):
headers = {}
result = await self.config.avalidate_environment(
headers=headers,
model="",
messages=[],
optional_params={},
litellm_params={},
api_key="sk-ant-test-key",
)
assert result["x-api-key"] == "sk-ant-test-key"
assert result["anthropic-version"] == "2023-06-01"
assert result["anthropic-beta"] == ANTHROPIC_FILES_BETA_HEADER
@pytest.mark.asyncio
@patch.dict("os.environ", {}, clear=True)
@patch(
"litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key",
return_value=None,
)
async def test_avalidate_environment_missing_api_key(self, mock_get_key):
with pytest.raises(ValueError, match="Anthropic API key is required"):
await self.config.avalidate_environment(
headers={},
model="",
messages=[],
optional_params={},
litellm_params={},
api_key=None,
)
def test_get_supported_openai_params(self):
params = self.config.get_supported_openai_params(model="")
assert "purpose" in params
@ -187,10 +221,7 @@ class TestAnthropicFilesConfig:
litellm_params={},
)
assert (
url
== f"{ANTHROPIC_FILES_API_BASE}/v1/files/..%2F..%2Fv1%2Fmessages%2Fbatches%3Flimit%3D1%23frag"
)
assert url == f"{ANTHROPIC_FILES_API_BASE}/v1/files/..%2F..%2Fv1%2Fmessages%2Fbatches%3Flimit%3D1%23frag"
assert params == {}
def test_transform_retrieve_file_response(self):
@ -411,6 +442,108 @@ class TestAnthropicFilesConfig:
assert error.message == "Not found"
_WIF_ENV = {
"ANTHROPIC_FEDERATION_RULE_ID": "fdrl_files_seam",
"ANTHROPIC_ORGANIZATION_ID": "org-files-seam",
"ANTHROPIC_IDENTITY_TOKEN": "files-seam-inline-jwt",
}
class _BlockingPoster:
"""A token-endpoint poster 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):
self.thread_ids.append(threading.get_ident())
self.release.wait(timeout=5)
return httpx.Response(
200,
json={
"access_token": "sk-ant-oat01-files-seam",
"token_type": "Bearer",
"expires_in": 3600,
},
)
class TestAnthropicFilesConfigWifAsyncSeam:
"""Regression (Greptile P1): avalidate_environment must resolve workload identity
federation through the async token-exchange facade, never the blocking sync one,
so a cold WIF mint on async file retrieval doesn't freeze the event loop."""
def setup_method(self):
self.config = AnthropicFilesConfig()
@pytest.mark.asyncio
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
for name in (
"ANTHROPIC_API_KEY",
"ANTHROPIC_AUTH_TOKEN",
"ANTHROPIC_API_BASE",
"ANTHROPIC_BASE_URL",
):
monkeypatch.delenv(name, raising=False)
for name, value in _WIF_ENV.items():
monkeypatch.setenv(name, value)
poster = _BlockingPoster()
engine = JwtBearerTokenExchangeEngine(poster=poster)
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)
async def async_shim(litellm_params, api_base, model):
return await aget_anthropic_wif_token(litellm_params, api_base, model, engine)
monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim)
monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim)
ticks = []
async def ticker():
for i in range(20):
await asyncio.sleep(0.005)
ticks.append(i)
ticker_task = asyncio.create_task(ticker())
await asyncio.sleep(0.02)
validate_task = asyncio.create_task(
self.config.avalidate_environment(
headers={},
model="",
messages=[],
optional_params={},
litellm_params={},
api_key=None,
)
)
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.
assert len(ticks) > 0
assert not validate_task.done()
poster.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()
class TestProviderConfigRegistration:
"""Test that AnthropicFilesConfig is properly registered."""

View file

@ -101,9 +101,7 @@ async def test_anthropic_native_interceptor_skipped():
)
h = AdvisorOrchestrationHandler()
assert not h.can_handle(
[ADVISOR_TOOL], "anthropic"
), "Interceptor must NOT trigger for anthropic provider"
assert not h.can_handle([ADVISOR_TOOL], "anthropic"), "Interceptor must NOT trigger for anthropic provider"
# ---------------------------------------------------------------------------
@ -204,9 +202,7 @@ async def test_loop_one_advisor_call():
assert "is_prime" in texts[0]["text"]
# No advisor tool_use blocks in final response
advisor_uses = [
b for b in content if b.get("type") == "tool_use" and b.get("name") == "advisor"
]
advisor_uses = [b for b in content if b.get("type") == "tool_use" and b.get("name") == "advisor"]
assert len(advisor_uses) == 0
@ -366,9 +362,7 @@ async def test_prior_advisor_blocks_replaced_in_history():
# Text block with advisor feedback must be present
text_blocks = [b for b in content if b.get("type") == "text"]
feedback_blocks = [
b for b in text_blocks if "advisor_feedback" in b.get("text", "")
]
feedback_blocks = [b for b in text_blocks if "advisor_feedback" in b.get("text", "")]
assert len(feedback_blocks) >= 1
assert "trial division" in feedback_blocks[0]["text"]
@ -707,11 +701,7 @@ async def test_advisor_ignores_tool_credentials_when_clientside_disabled():
with patch.dict(
sys.modules,
{
"litellm.proxy.proxy_server": _fake_proxy_server(
{"allow_client_side_credentials": False}
)
},
{"litellm.proxy.proxy_server": _fake_proxy_server({"allow_client_side_credentials": False})},
):
captured = await _run_advisor_and_capture_subcall_kwargs()
assert captured["api_key"] is None
@ -726,11 +716,7 @@ async def test_advisor_uses_tool_credentials_when_clientside_enabled():
with patch.dict(
sys.modules,
{
"litellm.proxy.proxy_server": _fake_proxy_server(
{"allow_client_side_credentials": True}
)
},
{"litellm.proxy.proxy_server": _fake_proxy_server({"allow_client_side_credentials": True})},
):
captured = await _run_advisor_and_capture_subcall_kwargs()
assert captured["api_key"] == "sk-other"

View file

@ -1,4 +1,9 @@
import httpx
import pytest
import respx
import litellm
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
from litellm.llms.anthropic.count_tokens.transformation import (
AnthropicCountTokensConfig,
)
@ -88,3 +93,72 @@ def test_transform_no_system_no_tools():
assert "system" not in result
assert "tools" not in result
@pytest.mark.parametrize(
("api_base", "expected"),
[
(None, "https://api.anthropic.com/v1/messages/count_tokens"),
("", "https://api.anthropic.com/v1/messages/count_tokens"),
("https://gateway.example", "https://gateway.example/v1/messages/count_tokens"),
("https://gateway.example/", "https://gateway.example/v1/messages/count_tokens"),
("https://gateway.example/v1", "https://gateway.example/v1/messages/count_tokens"),
("https://gateway.example/anthropic/v1/messages", "https://gateway.example/anthropic/v1/messages/count_tokens"),
],
)
def test_endpoint_appends_count_tokens_path_to_deployment_api_base(api_base, expected, monkeypatch):
monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False)
monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False)
assert AnthropicCountTokensConfig().get_anthropic_count_tokens_endpoint(api_base) == expected
@pytest.mark.parametrize("env_name", ["ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"])
@pytest.mark.parametrize("api_base", [None, ""])
def test_endpoint_without_deployment_api_base_follows_env_base(env_name, api_base, monkeypatch):
"""Chat and the federated exchange resolve an unset deployment base through the environment,
so an env-only gateway must receive the count too, never Anthropic's public host."""
monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False)
monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False)
monkeypatch.setenv(env_name, "https://env-gateway.example/v1/messages/")
assert (
AnthropicCountTokensConfig().get_anthropic_count_tokens_endpoint(api_base)
== "https://env-gateway.example/v1/messages/count_tokens"
)
def test_endpoint_prefers_deployment_api_base_over_env_base(monkeypatch):
monkeypatch.setenv("ANTHROPIC_API_BASE", "https://env-gateway.example")
assert (
AnthropicCountTokensConfig().get_anthropic_count_tokens_endpoint("https://gateway.example/v1")
== "https://gateway.example/v1/messages/count_tokens"
)
@pytest.fixture
def httpx_transport_clients(monkeypatch):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
client_cache = getattr(litellm, "in_memory_llm_clients_cache", None)
if client_cache is not None:
client_cache.flush_cache()
yield
if client_cache is not None:
client_cache.flush_cache()
@pytest.mark.asyncio
async def test_handler_posts_to_count_tokens_path_under_deployment_api_base(httpx_transport_clients):
"""A deployment api_base names the chat host, so a handler that posts to it verbatim lands on
the host root, gets a 404, and the official count silently degrades to the local tokenizer."""
with respx.mock:
route = respx.post("https://gateway.example/v1/messages/count_tokens").mock(
return_value=httpx.Response(200, json={"input_tokens": 7})
)
result = await AnthropicCountTokensHandler().handle_count_tokens_request(
model="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
api_key="sk-ant-api03-test-key",
api_base="https://gateway.example",
)
assert route.called
assert result == {"input_tokens": 7}

View file

@ -70,9 +70,7 @@ class TestAnthropicFilesHandler:
@pytest.fixture
def mock_anthropic_batch_results_canceled(self):
"""Mock Anthropic batch results with canceled status"""
return json.dumps(
{"custom_id": "test-request-3", "result": {"type": "canceled"}}
).encode("utf-8")
return json.dumps({"custom_id": "test-request-3", "result": {"type": "canceled"}}).encode("utf-8")
@pytest.fixture
def mock_anthropic_batch_results_mixed(self):
@ -114,9 +112,7 @@ class TestAnthropicFilesHandler:
return "\n".join(lines).encode("utf-8")
@pytest.mark.asyncio
async def test_afile_content_success(
self, handler, mock_anthropic_batch_results_succeeded
):
async def test_afile_content_success(self, handler, mock_anthropic_batch_results_succeeded):
"""Test successful file content retrieval and transformation"""
file_content_request: FileContentRequest = {
"file_id": "batch_123",
@ -135,16 +131,14 @@ class TestAnthropicFilesHandler:
),
)
with patch(
with patch( # test-quality-ok: the proxy wiring under test is what this patches
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
) as mock_get_client:
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
mock_client = AsyncMock()
mock_client.get = AsyncMock(return_value=mock_response)
mock_get_client.return_value = mock_client
with patch.object(
handler.anthropic_model_info, "get_api_key", return_value="test-api-key"
):
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"):
with patch.object(
handler.anthropic_model_info,
"get_api_base",
@ -161,9 +155,7 @@ class TestAnthropicFilesHandler:
# Verify transformation to OpenAI format
content = result.response.content.decode("utf-8")
lines = [
line for line in content.strip().split("\n") if line.strip()
]
lines = [line for line in content.strip().split("\n") if line.strip()]
assert len(lines) == 1
transformed_result = json.loads(lines[0])
@ -172,18 +164,13 @@ class TestAnthropicFilesHandler:
assert "body" in transformed_result["response"]
# Verify body has required OpenAI format fields
assert "id" in transformed_result["response"]["body"]
assert (
transformed_result["response"]["body"]["object"]
== "chat.completion"
)
assert transformed_result["response"]["body"]["object"] == "chat.completion"
assert "choices" in transformed_result["response"]["body"]
# Verify request_id matches the original message id
assert transformed_result["response"]["request_id"] == "msg_123"
@pytest.mark.asyncio
async def test_afile_content_with_prefix(
self, handler, mock_anthropic_batch_results_succeeded
):
async def test_afile_content_with_prefix(self, handler, mock_anthropic_batch_results_succeeded):
"""Test file content retrieval with anthropic_batch_results: prefix"""
file_content_request: FileContentRequest = {
"file_id": "anthropic_batch_results:batch_123",
@ -203,14 +190,12 @@ class TestAnthropicFilesHandler:
with patch(
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
) as mock_get_client:
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
mock_client = AsyncMock()
mock_client.get = AsyncMock(return_value=mock_response)
mock_get_client.return_value = mock_client
with patch.object(
handler.anthropic_model_info, "get_api_key", return_value="test-api-key"
):
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"):
with patch.object(
handler.anthropic_model_info,
"get_api_base",
@ -228,9 +213,7 @@ class TestAnthropicFilesHandler:
assert "batch_123" in call_url
@pytest.mark.asyncio
async def test_afile_content_errored_result(
self, handler, mock_anthropic_batch_results_errored
):
async def test_afile_content_errored_result(self, handler, mock_anthropic_batch_results_errored):
"""Test transformation of errored batch results"""
file_content_request: FileContentRequest = {
"file_id": "batch_123",
@ -250,14 +233,12 @@ class TestAnthropicFilesHandler:
with patch(
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
) as mock_get_client:
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
mock_client = AsyncMock()
mock_client.get = AsyncMock(return_value=mock_response)
mock_get_client.return_value = mock_client
with patch.object(
handler.anthropic_model_info, "get_api_key", return_value="test-api-key"
):
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"):
with patch.object(
handler.anthropic_model_info,
"get_api_base",
@ -269,29 +250,17 @@ class TestAnthropicFilesHandler:
)
content = result.response.content.decode("utf-8")
lines = [
line for line in content.strip().split("\n") if line.strip()
]
lines = [line for line in content.strip().split("\n") if line.strip()]
assert len(lines) == 1
transformed_result = json.loads(lines[0])
assert transformed_result["custom_id"] == "test-request-2"
assert (
transformed_result["response"]["status_code"] == 400
) # invalid_request_error maps to 400
assert (
transformed_result["response"]["body"]["error"]["type"]
== "invalid_request_error"
)
assert (
transformed_result["response"]["body"]["error"]["message"]
== "Invalid request"
)
assert transformed_result["response"]["status_code"] == 400 # invalid_request_error maps to 400
assert transformed_result["response"]["body"]["error"]["type"] == "invalid_request_error"
assert transformed_result["response"]["body"]["error"]["message"] == "Invalid request"
@pytest.mark.asyncio
async def test_afile_content_canceled_result(
self, handler, mock_anthropic_batch_results_canceled
):
async def test_afile_content_canceled_result(self, handler, mock_anthropic_batch_results_canceled):
"""Test transformation of canceled batch results"""
file_content_request: FileContentRequest = {
"file_id": "batch_123",
@ -311,14 +280,12 @@ class TestAnthropicFilesHandler:
with patch(
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
) as mock_get_client:
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
mock_client = AsyncMock()
mock_client.get = AsyncMock(return_value=mock_response)
mock_get_client.return_value = mock_client
with patch.object(
handler.anthropic_model_info, "get_api_key", return_value="test-api-key"
):
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"):
with patch.object(
handler.anthropic_model_info,
"get_api_base",
@ -330,23 +297,16 @@ class TestAnthropicFilesHandler:
)
content = result.response.content.decode("utf-8")
lines = [
line for line in content.strip().split("\n") if line.strip()
]
lines = [line for line in content.strip().split("\n") if line.strip()]
assert len(lines) == 1
transformed_result = json.loads(lines[0])
assert transformed_result["custom_id"] == "test-request-3"
assert transformed_result["response"]["status_code"] == 400
assert (
"Batch request was canceled"
in transformed_result["response"]["body"]["error"]["message"]
)
assert "Batch request was canceled" in transformed_result["response"]["body"]["error"]["message"]
@pytest.mark.asyncio
async def test_afile_content_mixed_results(
self, handler, mock_anthropic_batch_results_mixed
):
async def test_afile_content_mixed_results(self, handler, mock_anthropic_batch_results_mixed):
"""Test transformation of mixed batch results (succeeded, errored, expired)"""
file_content_request: FileContentRequest = {
"file_id": "batch_123",
@ -366,14 +326,12 @@ class TestAnthropicFilesHandler:
with patch(
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
) as mock_get_client:
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
mock_client = AsyncMock()
mock_client.get = AsyncMock(return_value=mock_response)
mock_get_client.return_value = mock_client
with patch.object(
handler.anthropic_model_info, "get_api_key", return_value="test-api-key"
):
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"):
with patch.object(
handler.anthropic_model_info,
"get_api_base",
@ -385,9 +343,7 @@ class TestAnthropicFilesHandler:
)
content = result.response.content.decode("utf-8")
lines = [
line for line in content.strip().split("\n") if line.strip()
]
lines = [line for line in content.strip().split("\n") if line.strip()]
assert len(lines) == 3
# Check first result (succeeded)
@ -396,9 +352,7 @@ class TestAnthropicFilesHandler:
# Check second result (errored)
result2 = json.loads(lines[1])
assert (
result2["response"]["status_code"] == 429
) # rate_limit_error maps to 429
assert result2["response"]["status_code"] == 429 # rate_limit_error maps to 429
# Check third result (expired)
result3 = json.loads(lines[2])
@ -415,12 +369,12 @@ class TestAnthropicFilesHandler:
}
with patch.object(
handler.anthropic_model_info, "get_auth_header", return_value=None
handler.anthropic_model_info,
"aget_auth_header",
new=AsyncMock(return_value=None),
):
with pytest.raises(ValueError, match="Missing Anthropic API Key"):
await handler.afile_content(
file_content_request=file_content_request, api_key=None
)
await handler.afile_content(file_content_request=file_content_request, api_key=None)
@pytest.mark.asyncio
async def test_afile_content_missing_file_id(self, handler):
@ -432,9 +386,7 @@ class TestAnthropicFilesHandler:
}
with pytest.raises(ValueError, match="file_id is required"):
await handler.afile_content(
file_content_request=file_content_request, api_key="test-api-key"
)
await handler.afile_content(file_content_request=file_content_request, api_key="test-api-key")
@pytest.mark.asyncio
async def test_afile_content_http_error(self, handler):
@ -454,21 +406,17 @@ class TestAnthropicFilesHandler:
),
)
mock_response.raise_for_status = MagicMock(
side_effect=httpx.HTTPStatusError(
"Not Found", request=mock_response.request, response=mock_response
)
side_effect=httpx.HTTPStatusError("Not Found", request=mock_response.request, response=mock_response)
)
with patch(
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
) as mock_get_client:
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
mock_client = AsyncMock()
mock_client.get = AsyncMock(return_value=mock_response)
mock_get_client.return_value = mock_client
with patch.object(
handler.anthropic_model_info, "get_api_key", return_value="test-api-key"
):
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"):
with patch.object(
handler.anthropic_model_info,
"get_api_base",
@ -480,6 +428,86 @@ class TestAnthropicFilesHandler:
api_key="test-api-key",
)
@pytest.mark.asyncio
async def test_afile_content_resolves_wif_via_async_facade(
self, handler, mock_anthropic_batch_results_succeeded, monkeypatch
):
"""Regression: afile_content ran the blocking WIF mint on the event loop
through the sync get_auth_header; it must go through the async facade."""
import threading
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
for name in (
"ANTHROPIC_API_KEY",
"ANTHROPIC_AUTH_TOKEN",
"ANTHROPIC_API_BASE",
"ANTHROPIC_BASE_URL",
):
monkeypatch.delenv(name, raising=False)
monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_files")
monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-files")
monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "files-inline-jwt")
minted = "sk-ant-oat01-files-minted"
thread_ids = []
class ThreadRecordingPoster:
def post(self, url, *, content, headers, timeout):
thread_ids.append(threading.get_ident())
return httpx.Response(
200,
json={"access_token": minted, "token_type": "Bearer", "expires_in": 3600},
)
engine = JwtBearerTokenExchangeEngine(poster=ThreadRecordingPoster())
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)
async def async_shim(litellm_params, api_base, model):
return await aget_anthropic_wif_token(litellm_params, api_base, model, engine)
monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim)
monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim)
mock_response = httpx.Response(
status_code=200,
content=mock_anthropic_batch_results_succeeded,
headers={"content-type": "application/json"},
request=httpx.Request(
method="GET",
url="https://api.anthropic.com/v1/messages/batches/batch_123/results",
),
)
with patch( # test-quality-ok: the proxy wiring under test is what this patches
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
mock_client = AsyncMock()
mock_client.get = AsyncMock(return_value=mock_response)
mock_get_client.return_value = mock_client
await handler.afile_content(
file_content_request={
"file_id": "batch_123",
"extra_headers": None,
"extra_body": None,
},
api_key=None,
)
sent_headers = mock_client.get.call_args.kwargs["headers"]
assert sent_headers["authorization"] == f"Bearer {minted}"
assert "oauth-2025-04-20" in sent_headers["anthropic-beta"]
assert sync_calls == []
assert thread_ids and thread_ids[0] != threading.get_ident()
class TestAnthropicBatchesConfig:
"""Test Anthropic Batches Config for batch retrieval transformation"""
@ -562,15 +590,11 @@ class TestAnthropicBatchesConfig:
)
assert url == "https://api.anthropic.com/v1/messages/batches/batch_123"
def test_transform_retrieve_batch_response_in_progress(
self, config, mock_anthropic_batch_response_in_progress
):
def test_transform_retrieve_batch_response_in_progress(self, config, mock_anthropic_batch_response_in_progress):
"""Test transformation of in_progress batch response"""
mock_response = httpx.Response(
status_code=200,
content=json.dumps(mock_anthropic_batch_response_in_progress).encode(
"utf-8"
),
content=json.dumps(mock_anthropic_batch_response_in_progress).encode("utf-8"),
request=httpx.Request(
method="GET",
url="https://api.anthropic.com/v1/messages/batches/batch_123",
@ -596,9 +620,7 @@ class TestAnthropicBatchesConfig:
assert batch.in_progress_at is not None
assert batch.completed_at is None
def test_transform_retrieve_batch_response_completed(
self, config, mock_anthropic_batch_response_completed
):
def test_transform_retrieve_batch_response_completed(self, config, mock_anthropic_batch_response_completed):
"""Test transformation of completed batch response"""
mock_response = httpx.Response(
status_code=200,
@ -624,9 +646,7 @@ class TestAnthropicBatchesConfig:
assert batch.request_counts.completed == 10
assert batch.request_counts.failed == 0
def test_transform_retrieve_batch_response_canceling(
self, config, mock_anthropic_batch_response_canceling
):
def test_transform_retrieve_batch_response_canceling(self, config, mock_anthropic_batch_response_canceling):
"""Test transformation of canceling batch response"""
mock_response = httpx.Response(
status_code=200,
@ -663,9 +683,7 @@ class TestAnthropicBatchesConfig:
)
logging_obj = MagicMock()
with pytest.raises(
ValueError, match="Failed to parse Anthropic batch response"
):
with pytest.raises(ValueError, match="Failed to parse Anthropic batch response"):
config.transform_retrieve_batch_response(
model="claude-3-5-sonnet-20241022",
raw_response=mock_response,

File diff suppressed because it is too large Load diff

View file

@ -3,8 +3,6 @@ Test that Azure AI Anthropic models have cache pricing configured.
Verifies the fix for issue #19532.
"""
import litellm
from litellm import get_model_info
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
@ -27,16 +25,11 @@ def reload_model_costs():
("claude-sonnet-4-5", 3.75e-06, 3e-07),
],
)
def test_azure_ai_claude_cache_pricing(
model, expected_cache_creation_cost, expected_cache_read_cost
):
def test_azure_ai_claude_cache_pricing(model, expected_cache_creation_cost, expected_cache_read_cost):
"""Test that Azure AI Claude models have correct cache pricing."""
model_info = get_model_info(model=model, custom_llm_provider="azure_ai")
assert model_info.get("cache_creation_input_token_cost") is not None
assert model_info.get("cache_read_input_token_cost") is not None
assert (
model_info.get("cache_creation_input_token_cost")
== expected_cache_creation_cost
)
assert model_info.get("cache_creation_input_token_cost") == expected_cache_creation_cost
assert model_info.get("cache_read_input_token_cost") == expected_cache_read_cost

View file

@ -5,7 +5,6 @@ being either a ``dict`` or a ``ServerToolUse`` pydantic instance.
See https://github.com/BerriAI/litellm/issues/26153.
"""
import pytest
from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests
@ -54,7 +53,8 @@ def test_get_cost_for_anthropic_web_search_with_dict_server_tool_use():
info = _make_model_info(cost_per_query=0.01)
cost = get_cost_for_anthropic_web_search(
model_info=info, usage=usage # type: ignore[arg-type]
model_info=info,
usage=usage, # type: ignore[arg-type]
)
assert cost == pytest.approx(0.03)
@ -65,7 +65,8 @@ def test_get_cost_for_anthropic_web_search_with_pydantic_server_tool_use():
info = _make_model_info(cost_per_query=0.01)
cost = get_cost_for_anthropic_web_search(
model_info=info, usage=usage # type: ignore[arg-type]
model_info=info,
usage=usage, # type: ignore[arg-type]
)
assert cost == pytest.approx(0.03)
@ -76,7 +77,8 @@ def test_get_cost_for_anthropic_web_search_with_none_server_tool_use():
info = _make_model_info(cost_per_query=0.01)
cost = get_cost_for_anthropic_web_search(
model_info=info, usage=usage # type: ignore[arg-type]
model_info=info,
usage=usage, # type: ignore[arg-type]
)
assert cost == 0.0

View file

@ -10,9 +10,9 @@ Regression test for https://github.com/BerriAI/litellm/issues/22040
import os
import sys
sys.path.insert(
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
)
import pytest
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")))
from litellm.llms.anthropic.count_tokens.transformation import (
AnthropicCountTokensConfig,
@ -78,9 +78,177 @@ class TestCountTokensOAuthHeaders:
headers = config.get_required_headers(FAKE_OAUTH_TOKEN)
beta_value = headers.get("anthropic-beta", "")
assert (
"token-counting" in beta_value
), f"token-counting beta missing from OAuth headers: {beta_value}"
assert (
"oauth-2025-04-20" in beta_value
), f"oauth beta missing from OAuth headers: {beta_value}"
assert "token-counting" in beta_value, f"token-counting beta missing from OAuth headers: {beta_value}"
assert "oauth-2025-04-20" in beta_value, f"oauth beta missing from OAuth headers: {beta_value}"
class TestCountTokensUsesWorkloadIdentity:
"""A federated deployment holds no static key. Without minting one, count_tokens returns None
and the caller silently falls back to the local tokenizer, so the number a federated
deployment reports would never come from Anthropic."""
@pytest.mark.asyncio
async def test_a_federated_deployment_mints_and_counts(self, monkeypatch):
from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
minted = "sk-ant-oat01-minted-for-count"
async def fake_mint(_params, _api_base, _model):
return minted
monkeypatch.setattr(token_counter_module, "aget_anthropic_wif_token", fake_mint, raising=False)
monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", fake_mint, raising=False)
seen: dict[str, object] = {}
async def fake_request(**kwargs):
seen.update(kwargs)
return {"input_tokens": 42}
monkeypatch.setattr(
token_counter_module.anthropic_count_tokens_handler,
"handle_count_tokens_request",
fake_request,
raising=False,
)
result = await token_counter_module.AnthropicTokenCounter().count_tokens(
model_to_use="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
contents=None,
deployment={
"litellm_params": {
"model": "anthropic/claude-sonnet-4-5",
"anthropic_federation_rule_id": "fdrl_x",
"anthropic_organization_id": "org-x",
}
},
request_model="claude-sonnet-4-5",
)
assert result is not None
assert result.total_tokens == 42
assert seen["api_key"] == minted
@pytest.mark.asyncio
async def test_an_auth_token_deployment_never_mints(self, monkeypatch):
from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", "bearer-token-for-testing")
mint_calls: list[str] = []
async def fake_mint(_params, _api_base, model):
mint_calls.append(model)
return "sk-ant-oat01-should-not-be-minted"
monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", fake_mint, raising=False)
result = await token_counter_module.AnthropicTokenCounter().count_tokens(
model_to_use="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
contents=None,
deployment={
"litellm_params": {
"model": "anthropic/claude-sonnet-4-5",
"anthropic_federation_rule_id": "fdrl_x",
"anthropic_organization_id": "org-x",
}
},
request_model="claude-sonnet-4-5",
)
assert result is None
assert mint_calls == []
@pytest.mark.asyncio
async def test_a_failed_mint_degrades_like_an_anthropic_error(self, monkeypatch):
import litellm
from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
async def failing_mint(_params, _api_base, model):
raise litellm.AuthenticationError(
message="federation_rule_id is not a well-formed fdrl_ tagged ID",
llm_provider="anthropic",
model=model,
)
monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", failing_mint, raising=False)
result = await token_counter_module.AnthropicTokenCounter().count_tokens(
model_to_use="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
contents=None,
deployment={
"litellm_params": {
"model": "anthropic/claude-sonnet-4-5",
"anthropic_federation_rule_id": "not-a-rule",
"anthropic_organization_id": "org-x",
}
},
request_model="claude-sonnet-4-5",
)
assert result is not None
assert result.error is True
assert result.status_code == 401
assert result.total_tokens == 0
assert "fdrl_" in (result.error_message or "")
@pytest.mark.asyncio
async def test_a_vault_backed_static_key_never_mints(self, monkeypatch):
from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
vault_key = "sk-ant-api03-only-in-the-vault"
def vault_only(secret_name, default_value=None):
return vault_key if secret_name == "ANTHROPIC_API_KEY" else None
monkeypatch.setattr("litellm.secret_managers.main.get_secret_str", vault_only, raising=False)
mint_calls: list[str] = []
async def fake_mint(_params, _api_base, model):
mint_calls.append(model)
return "sk-ant-oat01-should-not-be-minted"
monkeypatch.setattr(token_counter_module, "aget_anthropic_wif_token", fake_mint, raising=False)
monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", fake_mint, raising=False)
seen: dict[str, object] = {}
async def fake_request(**kwargs):
seen.update(kwargs)
return {"input_tokens": 7}
monkeypatch.setattr(
token_counter_module.anthropic_count_tokens_handler,
"handle_count_tokens_request",
fake_request,
raising=False,
)
result = await token_counter_module.AnthropicTokenCounter().count_tokens(
model_to_use="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
contents=None,
deployment={
"litellm_params": {
"model": "anthropic/claude-sonnet-4-5",
"anthropic_federation_rule_id": "fdrl_x",
"anthropic_organization_id": "org-x",
}
},
request_model="claude-sonnet-4-5",
)
assert result is not None
assert result.total_tokens == 7
assert seen["api_key"] == vault_key
assert mint_calls == []

View file

@ -12,9 +12,7 @@ import sys
import os
# Add the parent directory to the path so we can import litellm
sys.path.insert(
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../.."))
)
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")))
import litellm
from litellm.litellm_core_utils.prompt_templates.factory import (
@ -68,10 +66,7 @@ class TestMessageSanitization:
assert sanitized[1]["role"] == "assistant"
assert sanitized[2]["role"] == "tool"
assert sanitized[2]["tool_call_id"] == "toolu_01Kus2cC3ydjBW7UK4GJqBP4"
assert (
"skipped" in sanitized[2]["content"].lower()
or "interrupted" in sanitized[2]["content"].lower()
)
assert "skipped" in sanitized[2]["content"].lower() or "interrupted" in sanitized[2]["content"].lower()
assert "get_weather" in sanitized[2]["content"]
def test_case_a_orphaned_tool_call_multiple(self):
@ -115,12 +110,8 @@ class TestMessageSanitization:
assert len(sanitized) == 4
assert sanitized[0]["role"] == "user"
assert sanitized[1]["role"] == "assistant"
assert (
sanitized[2]["tool_call_id"] == "call_1"
) # Original tool result (first in tool_calls)
assert (
sanitized[3]["tool_call_id"] == "call_2"
) # Dummy added for missing call_2
assert sanitized[2]["tool_call_id"] == "call_1" # Original tool result (first in tool_calls)
assert sanitized[3]["tool_call_id"] == "call_2" # Dummy added for missing call_2
def test_case_b_orphaned_tool_result(self):
"""
@ -188,10 +179,7 @@ class TestMessageSanitization:
assert len(sanitized) == 2
assert sanitized[0]["role"] == "user"
assert (
sanitized[0]["content"]
== "[System: Empty message content sanitised to satisfy protocol]"
)
assert sanitized[0]["content"] == "[System: Empty message content sanitised to satisfy protocol]"
def test_case_c_whitespace_only_content(self):
"""
@ -206,14 +194,8 @@ class TestMessageSanitization:
sanitized = sanitize_messages_for_tool_calling(messages)
assert len(sanitized) == 2
assert (
sanitized[0]["content"]
== "[System: Empty message content sanitised to satisfy protocol]"
)
assert (
sanitized[1]["content"]
== "[System: Empty message content sanitised to satisfy protocol]"
)
assert sanitized[0]["content"] == "[System: Empty message content sanitised to satisfy protocol]"
assert sanitized[1]["content"] == "[System: Empty message content sanitised to satisfy protocol]"
def test_case_c_valid_content_preserved(self):
"""
@ -270,10 +252,7 @@ class TestMessageSanitization:
assert sanitized[2]["role"] == "tool"
assert sanitized[2]["tool_call_id"] == "call_1" # Dummy added
assert sanitized[3]["role"] == "user"
assert (
sanitized[3]["content"]
== "[System: Empty message content sanitised to satisfy protocol]"
)
assert sanitized[3]["content"] == "[System: Empty message content sanitised to satisfy protocol]"
assert sanitized[4]["role"] == "assistant"
def test_modify_params_false_no_sanitization(self):
@ -329,9 +308,7 @@ class TestMessageSanitization:
]
# This should not raise an error and should add dummy tool result
result = anthropic_messages_pt(
messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic"
)
result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic")
# Should have at least 2 messages (user and assistant)
# The tool result will be merged into user content
@ -355,23 +332,17 @@ class TestMessageSanitization:
{"role": "user", "content": ""},
]
result = anthropic_messages_pt(
messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic"
)
result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic")
# All three user messages get merged into one user turn for Anthropic.
assert len(result) == 1
assert result[0]["role"] == "user"
text_blocks = [
b for b in result[0]["content"] if isinstance(b, dict) and b.get("type") == "text"
]
text_blocks = [b for b in result[0]["content"] if isinstance(b, dict) and b.get("type") == "text"]
assert len(text_blocks) == 3
# No text block may be empty — that's the contract Anthropic enforces.
for block in text_blocks:
assert block["text"].strip() != ""
assert text_blocks[2]["text"] == (
"[System: Empty message content sanitised to satisfy protocol]"
)
assert text_blocks[2]["text"] == ("[System: Empty message content sanitised to satisfy protocol]")
def test_empty_text_block_in_list_content_sanitized(self):
"""
@ -392,14 +363,10 @@ class TestMessageSanitization:
},
]
result = anthropic_messages_pt(
messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic"
)
result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic")
assert len(result) == 1
text_blocks = [
b for b in result[0]["content"] if isinstance(b, dict) and b.get("type") == "text"
]
text_blocks = [b for b in result[0]["content"] if isinstance(b, dict) and b.get("type") == "text"]
assert len(text_blocks) == 3
assert text_blocks[0]["text"] == "real content"
for block in text_blocks[1:]:
@ -418,9 +385,7 @@ class TestMessageSanitization:
{"role": "user", "content": "How are you?"},
]
result = anthropic_messages_pt(
messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic"
)
result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic")
# Two user turns + one assistant turn (alternation preserved).
assert len(result) == 3

View file

@ -0,0 +1,484 @@
import base64
import logging
from collections.abc import Mapping
from typing import Final
from urllib.parse import parse_qsl, unquote
import httpx
import pytest
from litellm.llms.base_llm.auth.client_credentials import (
_HttpxSyncKeycloakPoster,
_default_secret_reader,
_new_keycloak_handler,
fetch_keycloak_assertion,
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
TOKEN_URL: Final = "https://keycloak.example/realms/litellm/protocol/openid-connect/token"
CLIENT_ID: Final = "litellm"
CLIENT_SECRET_REF: Final = "oidc/env/KEYCLOAK_CLIENT_SECRET"
CLIENT_SECRET: Final = "s3cr3t-client-value"
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 form_body(self) -> dict[str, str]:
return dict(parse_qsl(self.content.decode()))
class ScriptedPoster:
"""Returns one scripted response per call; records every request it receives."""
def __init__(self, responses: list[httpx.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))
return self._responses.pop(0) if len(self._responses) > 1 else 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
def make_config(
auth_method: str = "client_secret_basic",
scope: str | None = None,
token_url: str = TOKEN_URL,
client_secret_ref: str = CLIENT_SECRET_REF,
client_id: str = CLIENT_ID,
) -> KeycloakSource:
return KeycloakSource(
token_url=token_url,
client_id=client_id,
client_secret_ref=client_secret_ref,
auth_method=auth_method, # pyright: ignore[reportArgumentType] # test-only string widened for parametrization
scope=scope,
)
def secret_reader_returning(secret: str | None):
def reader(ref: str) -> str | None:
assert ref == CLIENT_SECRET_REF
return secret
return reader
DEFAULT_SECRET_READER: Final = secret_reader_returning(CLIENT_SECRET)
def token_response(access_token: str = "keycloak-minted-token") -> httpx.Response:
return httpx.Response(200, json={"access_token": access_token, "token_type": "Bearer", "expires_in": 300})
class TestClientSecretBasic:
def test_sends_basic_auth_header_and_no_secret_in_body(self):
poster = ScriptedPoster([token_response("minted-1")])
token = fetch_keycloak_assertion(
make_config(auth_method="client_secret_basic"), poster=poster, secret_reader=DEFAULT_SECRET_READER
)
assert token == "minted-1"
request = poster.requests[0]
assert request.url == TOKEN_URL
expected_auth = "Basic " + base64.b64encode(f"{CLIENT_ID}:{CLIENT_SECRET}".encode()).decode("ascii")
assert request.headers["authorization"] == expected_auth
assert request.headers["content-type"] == "application/x-www-form-urlencoded"
body = request.form_body()
assert body["grant_type"] == "client_credentials"
assert "client_secret" not in body
assert "client_id" not in body
def test_reserved_characters_are_form_encoded_before_basic(self):
"""RFC 6749 2.3.1 requires the client id and secret be application/x-www-form-urlencoded
(Appendix B) before being base64'd into the Basic header; a raw join lets a reserved
character in either value corrupt the ':'-joined pair Keycloak decodes back out."""
client_id = "id:with+reserved% chars"
client_secret = "secret:with+reserved% chars"
poster = ScriptedPoster([token_response("minted-reserved")])
fetch_keycloak_assertion(
make_config(auth_method="client_secret_basic", client_id=client_id),
poster=poster,
secret_reader=secret_reader_returning(client_secret),
)
header = poster.requests[0].headers["authorization"]
assert header.startswith("Basic ")
decoded = base64.b64decode(header.removeprefix("Basic ")).decode("ascii")
encoded_id, _, encoded_secret = decoded.partition(":")
assert unquote(encoded_id) == client_id
assert unquote(encoded_secret) == client_secret
def test_scope_included_only_when_set(self):
poster = ScriptedPoster([token_response()])
fetch_keycloak_assertion(
make_config(scope="openid profile"), poster=poster, secret_reader=DEFAULT_SECRET_READER
)
assert poster.requests[0].form_body()["scope"] == "openid profile"
poster_no_scope = ScriptedPoster([token_response()])
fetch_keycloak_assertion(make_config(scope=None), poster=poster_no_scope, secret_reader=DEFAULT_SECRET_READER)
assert "scope" not in poster_no_scope.requests[0].form_body()
class TestClientSecretPost:
def test_sends_client_id_and_secret_in_body_with_no_basic_header(self):
poster = ScriptedPoster([token_response("minted-2")])
token = fetch_keycloak_assertion(
make_config(auth_method="client_secret_post"), poster=poster, secret_reader=DEFAULT_SECRET_READER
)
assert token == "minted-2"
request = poster.requests[0]
assert "authorization" not in request.headers
body = request.form_body()
assert body["grant_type"] == "client_credentials"
assert body["client_id"] == CLIENT_ID
assert body["client_secret"] == CLIENT_SECRET
class TestOnePostPerExchange:
def test_exactly_one_post_per_call_no_cache(self):
poster = ScriptedPoster([token_response("first"), token_response("second")])
first = fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
second = fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
assert first == "first"
assert second == "second"
assert len(poster.requests) == 2
class TestInvalidClient:
def test_400_invalid_client_surfaces_redacted_detail(self):
poster = ScriptedPoster(
[httpx.Response(400, json={"error": "invalid_client", "error_description": "unauthorized client"})]
)
with pytest.raises(ValueError, match="invalid_client") as exc_info:
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
assert "unauthorized client" in str(exc_info.value)
assert "400" in str(exc_info.value)
assert CLIENT_SECRET not in str(exc_info.value)
def test_echoed_client_secret_is_never_reflected_into_the_error(self):
"""A misbehaving Keycloak that echoes the submitted client_secret back in its error body
must never leak it into the exception the caller sees."""
long_secret: Final = "reflectable-secret-0123456789"
poster = ScriptedPoster(
[httpx.Response(400, json={"error": "invalid_client", "error_description": f"got {long_secret} in body"})]
)
with pytest.raises(ValueError, match="keycloak") as exc_info:
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=secret_reader_returning(long_secret))
assert long_secret not in str(exc_info.value)
assert "redacted" in str(exc_info.value)
def test_echoed_short_client_secret_is_never_reflected_into_the_error(self):
"""Real Keycloak client secrets are often shorter than a JWT: the reflection probe must
not silently stop protecting a secret just because it is under the probe's usual length."""
short_secret: Final = "hand-set-14ch"
poster = ScriptedPoster(
[httpx.Response(400, json={"error": "invalid_client", "error_description": f"got {short_secret} in body"})]
)
with pytest.raises(ValueError, match="keycloak") as exc_info:
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=secret_reader_returning(short_secret))
assert short_secret not in str(exc_info.value)
assert "redacted" in str(exc_info.value)
class TestUnreachable:
def test_transport_failure_raises_diagnosable_value_error(self):
poster = RaisingPoster(httpx.ConnectError("connection refused"))
with pytest.raises(ValueError, match="ConnectError") as exc_info:
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
assert poster.calls == 1
assert CLIENT_SECRET not in str(exc_info.value)
class TestNon2xx:
def test_500_raises_value_error_with_status_code(self):
poster = ScriptedPoster([httpx.Response(500, json={"error": "server_error"})])
with pytest.raises(ValueError, match="500"):
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
class TestResponseValidation:
def test_missing_access_token_is_a_value_error(self):
poster = ScriptedPoster([httpx.Response(200, json={"token_type": "Bearer"})])
with pytest.raises(ValueError, match="schema validation"):
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
def test_empty_access_token_is_a_value_error(self):
poster = ScriptedPoster([httpx.Response(200, json={"access_token": " "})])
with pytest.raises(ValueError, match="empty access_token"):
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
class TestInsecureTokenUrl:
def test_http_url_is_rejected_before_any_post(self):
poster = ScriptedPoster([token_response()])
with pytest.raises(ValueError, match="https"):
fetch_keycloak_assertion(
make_config(token_url="http://keycloak.example/token"),
poster=poster,
secret_reader=DEFAULT_SECRET_READER,
)
assert poster.requests == []
class TestMissingClientSecret:
def test_unresolvable_secret_ref_raises_value_error_naming_the_ref_not_a_secret(self):
poster = ScriptedPoster([token_response()])
with pytest.raises(ValueError, match=CLIENT_SECRET_REF):
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=secret_reader_returning(None))
assert poster.requests == []
class TestKeycloakAssertionSource:
def test_returns_a_callable_that_fetches_fresh_each_call(self):
poster = ScriptedPoster([token_response("first"), token_response("second")])
source = keycloak_assertion_source(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
assert source() == "first"
assert source() == "second"
assert len(poster.requests) == 2
def test_propagates_the_underlying_fetch_failure(self):
poster = ScriptedPoster([httpx.Response(400, json={"error": "invalid_client"})])
source = keycloak_assertion_source(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
with pytest.raises(ValueError, match="invalid_client"):
source()
class TestClientSecretNeverLeaks:
"""Regression coverage for the load-bearing property: a Keycloak client_secret must never
surface in the assertion_ref, in any error message, or in a log record, however it fails."""
def test_never_in_the_assertion_ref(self):
config = make_config(client_secret_ref=CLIENT_SECRET_REF)
ref = identity_source_ref(config)
assert CLIENT_SECRET not in ref
assert CLIENT_SECRET_REF not in ref
def test_never_in_any_raised_error_message_across_every_failure_mode(self):
config = make_config()
failures = [
lambda: fetch_keycloak_assertion(
config,
poster=ScriptedPoster([httpx.Response(400, json={"error": "invalid_client"})]),
secret_reader=DEFAULT_SECRET_READER,
),
lambda: fetch_keycloak_assertion(
config, poster=RaisingPoster(httpx.ConnectError("boom")), secret_reader=DEFAULT_SECRET_READER
),
lambda: fetch_keycloak_assertion(
config,
poster=ScriptedPoster([httpx.Response(500, json={"error": "server_error"})]),
secret_reader=DEFAULT_SECRET_READER,
),
lambda: fetch_keycloak_assertion(
config, poster=ScriptedPoster([token_response()]), secret_reader=secret_reader_returning(None)
),
]
for fail in failures:
with pytest.raises(ValueError, match="keycloak") as exc_info:
fail()
assert CLIENT_SECRET not in str(exc_info.value)
def test_never_in_a_log_record(self, caplog: pytest.LogCaptureFixture):
with caplog.at_level(logging.DEBUG):
poster = ScriptedPoster(
[httpx.Response(400, json={"error": "invalid_client", "error_description": CLIENT_SECRET})]
)
with pytest.raises(ValueError, match="keycloak"):
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
fetch_keycloak_assertion(
make_config(), poster=ScriptedPoster([token_response()]), secret_reader=DEFAULT_SECRET_READER
)
assert CLIENT_SECRET not in caplog.text
class StubHandler:
"""Stands in for the HTTPHandler the default poster builds, so the poster's own contract
(redirects off, error responses returned rather than raised, no-response guarded) is testable
without a socket."""
def __init__(self, result: httpx.Response | Exception | None) -> None:
self.calls: list[dict[str, object]] = []
self._result = result
def post(self, url: str, *, content: bytes, headers: dict[str, str], timeout: float) -> httpx.Response | None:
self.calls.append({"url": url, "content": content, "headers": headers, "timeout": timeout})
if isinstance(self._result, Exception):
raise self._result
return self._result
class TestDefaultKeycloakPoster:
def test_builds_its_handler_once_with_redirects_disabled(self):
built: list[StubHandler] = []
def factory() -> StubHandler:
handler = StubHandler(httpx.Response(200, json={"access_token": "kc-token"}))
built.append(handler)
return handler
poster: Final = _HttpxSyncKeycloakPoster(handler_factory=factory) # pyright: ignore[reportArgumentType] # StubHandler stands in for the legacy-untyped HTTPHandler
for _ in range(3):
poster.post(TOKEN_URL, content=b"grant_type=client_credentials", headers={}, timeout=1.0)
assert len(built) == 1, "the handler is built once and reused"
assert len(built[0].calls) == 3
def test_the_real_handler_refuses_to_follow_redirects(self):
handler: Final = _new_keycloak_handler()
assert handler.client.follow_redirects is False, (
"a redirected token POST would replay the client secret to whatever host the redirect names"
)
def test_an_http_status_error_becomes_its_response_rather_than_an_exception(self):
response: Final = httpx.Response(
401, json={"error": "invalid_client"}, request=httpx.Request("POST", TOKEN_URL)
)
poster: Final = _HttpxSyncKeycloakPoster(
handler_factory=lambda: StubHandler(
httpx.HTTPStatusError("boom", request=response.request, response=response)
) # pyright: ignore[reportArgumentType] # StubHandler stands in for the legacy-untyped HTTPHandler
)
assert poster.post(TOKEN_URL, content=b"", headers={}, timeout=1.0).status_code == 401
def test_a_missing_response_is_a_transport_error_not_a_none_deref(self):
poster: Final = _HttpxSyncKeycloakPoster(handler_factory=lambda: StubHandler(None)) # pyright: ignore[reportArgumentType] # StubHandler stands in for the legacy-untyped HTTPHandler
with pytest.raises(httpx.TransportError):
poster.post(TOKEN_URL, content=b"", headers={}, timeout=1.0)
class TestDefaultSecretReader:
def test_reads_through_litellm_secret_resolution(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("KEYCLOAK_CLIENT_SECRET_FOR_DEFAULT_READER", CLIENT_SECRET)
assert _default_secret_reader("os.environ/KEYCLOAK_CLIENT_SECRET_FOR_DEFAULT_READER") == CLIENT_SECRET
def test_an_unset_reference_reads_as_none_so_the_caller_raises(self):
assert _default_secret_reader("os.environ/DEFINITELY_NOT_SET_KEYCLOAK_SECRET_REF") is None
class TestOversizedSuccessBody:
def test_a_success_body_over_the_cap_is_refused_before_it_is_parsed(self):
oversized: Final = httpx.Response(200, content=b'{"access_token": "' + b"x" * MAX_RESPONSE_BYTES + b'"}')
with pytest.raises(ValueError, match="exceeded the size cap"):
fetch_keycloak_assertion(
make_config(), poster=ScriptedPoster([oversized]), secret_reader=DEFAULT_SECRET_READER
)
class TestUnresolvedSecretRefIsNotEchoed:
"""An operator who pastes the secret itself into the *_ref field turns that field INTO the
secret, and this error reaches model callers, so it must never echo the value."""
def test_keycloak_ref_value_is_not_in_the_error(self):
from litellm.llms.base_llm.auth.client_credentials import keycloak_assertion_source
from litellm.llms.base_llm.auth.identity_source import KeycloakSource
pasted_secret = "sUp3r-s3cret-value-not-a-pointer"
config = KeycloakSource(
token_url="https://keycloak.example.com/realms/p/protocol/openid-connect/token",
client_id="litellm",
client_secret_ref=pasted_secret,
)
with pytest.raises(ValueError, match="could not be read") as excinfo:
keycloak_assertion_source(config, secret_reader=lambda _ref: None)()
assert pasted_secret not in str(excinfo.value)
assert "withheld" in str(excinfo.value)
def test_internal_issuer_ref_value_is_not_in_the_error(self):
from litellm.llms.base_llm.auth.identity_source import InternalIssuerSource
from litellm.llms.base_llm.auth.internal_issuer import internal_issuer_assertion_source
pasted_pem = "-----BEGIN PRIVATE KEY-----MIGHAgEA-----END PRIVATE KEY-----"
config = InternalIssuerSource(
issuer_url="https://proxy.example.com",
subject="litellm-proxy",
signing_key_ref=pasted_pem,
)
with pytest.raises(ValueError, match="could not be read") as excinfo:
internal_issuer_assertion_source(config, key_reader=lambda _ref: None)()
assert pasted_pem not in str(excinfo.value)
assert "withheld" in str(excinfo.value)
class TestTokenUrlIsNotEchoedWholesale:
"""A token endpoint is configuration and naming it makes the error actionable, but nothing
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
rendered = endpoint_url_for_error_message("https://idp.example/token?client_secret=supersecret")
assert "supersecret" not in rendered
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
rendered = endpoint_url_for_error_message("https://user:pw@idp.example:8443/token")
assert "pw" not in rendered
assert rendered == "https://idp.example:8443/token"
def test_transport_failure_message_carries_no_query_secret(self):
poster = RaisingPoster(httpx.ConnectTimeout("timed out"))
config = make_config(token_url="https://idp.example/token?client_secret=supersecret")
with pytest.raises(ValueError, match="could not reach the keycloak token endpoint") as excinfo:
fetch_keycloak_assertion(config, poster=poster, secret_reader=DEFAULT_SECRET_READER)
assert "supersecret" not in str(excinfo.value)
assert "idp.example/token" in str(excinfo.value)

View file

@ -0,0 +1,239 @@
from types import MappingProxyType
from typing import Final, Literal
import pytest
from pydantic import ValidationError
from litellm.llms.base_llm.auth.identity_source import (
AnthropicIdentitySourceKind,
InternalIssuerSource,
KeycloakSource,
identity_source_config_adapter,
identity_source_ref,
)
SIGNING_KEY_REF: Final = "oidc/env/ISSUER_SIGNING_KEY_PEM"
OTHER_SIGNING_KEY_REF: Final = "oidc/env/OTHER_SIGNING_KEY_PEM"
CLIENT_SECRET_REF: Final = "oidc/env/KEYCLOAK_CLIENT_SECRET"
ISSUER_URL: Final = "https://issuer.internal.example"
SUBJECT: Final = "workload-a"
TOKEN_URL: Final = "https://keycloak.example/realms/litellm/protocol/openid-connect/token"
CLIENT_ID: Final = "litellm"
def make_issuer(
issuer_url: str = ISSUER_URL,
subject: str = SUBJECT,
signing_key_ref: str = SIGNING_KEY_REF,
ttl_seconds: int = 300,
) -> InternalIssuerSource:
return InternalIssuerSource(
issuer_url=issuer_url, subject=subject, signing_key_ref=signing_key_ref, ttl_seconds=ttl_seconds
)
def make_keycloak(
token_url: str = TOKEN_URL,
client_id: str = CLIENT_ID,
client_secret_ref: str = CLIENT_SECRET_REF,
auth_method: Literal["client_secret_basic", "client_secret_post"] = "client_secret_basic",
scope: str | None = None,
) -> KeycloakSource:
return KeycloakSource(
token_url=token_url,
client_id=client_id,
client_secret_ref=client_secret_ref,
auth_method=auth_method,
scope=scope,
)
class TestIdentitySourceRefHashing:
def test_identical_config_hashes_idempotently(self):
assert identity_source_ref(make_issuer()) == identity_source_ref(make_issuer())
def test_ref_is_prefixed_by_kind(self):
assert identity_source_ref(make_issuer()).startswith("oidc/internal_issuer/")
assert identity_source_ref(make_keycloak()).startswith("oidc/keycloak/")
def test_pointer_name_change_changes_ref(self):
"""Two configs differing only in which secret a pointer names must never collide, since a
stale ref would let the token exchange's outer cache key alias two different credentials."""
first: Final = identity_source_ref(make_issuer(signing_key_ref=SIGNING_KEY_REF))
second: Final = identity_source_ref(make_issuer(signing_key_ref=OTHER_SIGNING_KEY_REF))
assert first != second
def test_non_pointer_field_change_changes_ref(self):
first: Final = identity_source_ref(make_keycloak(scope="openid"))
second: Final = identity_source_ref(make_keycloak(scope="openid profile"))
assert first != second
def test_ref_never_contains_the_pointer_field_values(self):
"""The ref is a fixed-width hash, not a serialization of the config, so no field value -
pointer name or otherwise - can leak into the secret-free string echoed into errors."""
ref: Final = identity_source_ref(make_issuer())
assert SIGNING_KEY_REF not in ref
assert "issuer.internal.example" not in ref
def test_different_kinds_with_disjoint_fields_never_collide(self):
assert identity_source_ref(make_issuer()) != identity_source_ref(make_keycloak())
class TestInternalIssuerSourceValidation:
def test_defaults(self):
source: Final = make_issuer()
assert source.kind == AnthropicIdentitySourceKind.internal_issuer
assert source.ttl_seconds == 300
assert source.audience is None
def test_ttl_seconds_over_one_hour_is_rejected(self):
with pytest.raises(ValidationError):
make_issuer(ttl_seconds=3601)
def test_ttl_seconds_at_one_hour_is_accepted(self):
assert make_issuer(ttl_seconds=3600).ttl_seconds == 3600
def test_non_positive_ttl_seconds_is_rejected(self):
with pytest.raises(ValidationError):
make_issuer(ttl_seconds=0)
def test_missing_signing_key_ref_is_rejected(self):
missing_field: Final = MappingProxyType({"issuer_url": ISSUER_URL, "subject": SUBJECT})
with pytest.raises(ValidationError):
InternalIssuerSource.model_validate(missing_field)
def test_keycloak_only_field_is_rejected_as_extra(self):
mixed_variant: Final = MappingProxyType(
{
"issuer_url": ISSUER_URL,
"subject": SUBJECT,
"signing_key_ref": SIGNING_KEY_REF,
"client_secret_ref": CLIENT_SECRET_REF,
}
)
with pytest.raises(ValidationError):
InternalIssuerSource.model_validate(mixed_variant)
def test_is_frozen(self):
source: Final = make_issuer()
with pytest.raises(ValidationError):
source.subject = "workload-b"
def test_secret_pasted_into_wrong_typed_field_is_not_echoed_in_the_error(self):
"""hide_input_in_errors keeps a value the operator pasted into a mistyped field out of the
validation error, so a client_secret headed for the wrong field isn't logged in the raise."""
leaked_secret: Final = "shh-do-not-log-me"
wrong_type: Final = MappingProxyType(
{
"issuer_url": ISSUER_URL,
"subject": SUBJECT,
"signing_key_ref": SIGNING_KEY_REF,
"ttl_seconds": leaked_secret,
}
)
with pytest.raises(ValidationError) as exc_info:
InternalIssuerSource.model_validate(wrong_type)
assert leaked_secret not in str(exc_info.value)
class TestKeycloakSourceValidation:
def test_defaults(self):
source: Final = make_keycloak()
assert source.kind == AnthropicIdentitySourceKind.keycloak
assert source.auth_method == "client_secret_basic"
assert source.scope is None
def test_client_secret_post_is_accepted(self):
assert make_keycloak(auth_method="client_secret_post").auth_method == "client_secret_post"
def test_private_key_jwt_is_not_a_supported_auth_method_yet(self):
unshipped_auth_method: Final = MappingProxyType(
{
"token_url": TOKEN_URL,
"client_id": CLIENT_ID,
"client_secret_ref": CLIENT_SECRET_REF,
"auth_method": "private_key_jwt",
}
)
with pytest.raises(ValidationError):
KeycloakSource.model_validate(unshipped_auth_method)
def test_audience_field_was_dropped(self):
dropped_field: Final = MappingProxyType(
{
"token_url": TOKEN_URL,
"client_id": CLIENT_ID,
"client_secret_ref": CLIENT_SECRET_REF,
"audience": "https://anthropic.example",
}
)
with pytest.raises(ValidationError):
KeycloakSource.model_validate(dropped_field)
def test_missing_client_secret_ref_is_rejected(self):
missing_field: Final = MappingProxyType({"token_url": TOKEN_URL, "client_id": CLIENT_ID})
with pytest.raises(ValidationError):
KeycloakSource.model_validate(missing_field)
class TestDiscriminatedUnionParsing:
def test_parses_internal_issuer_variant(self):
parsed: Final = identity_source_config_adapter.validate_python(
MappingProxyType(
{
"kind": "internal_issuer",
"issuer_url": ISSUER_URL,
"subject": SUBJECT,
"signing_key_ref": SIGNING_KEY_REF,
}
)
)
assert isinstance(parsed, InternalIssuerSource)
def test_parses_keycloak_variant(self):
parsed: Final = identity_source_config_adapter.validate_python(
MappingProxyType(
{
"kind": "keycloak",
"token_url": TOKEN_URL,
"client_id": CLIENT_ID,
"client_secret_ref": CLIENT_SECRET_REF,
}
)
)
assert isinstance(parsed, KeycloakSource)
def test_unknown_kind_is_a_hard_error(self):
with pytest.raises(ValidationError):
identity_source_config_adapter.validate_python(MappingProxyType({"kind": "token_file"}))
def test_mixed_variant_fields_are_a_hard_error(self):
"""A keycloak field on an internal_issuer-tagged payload must fail closed rather than be
silently dropped or silently accepted as if it selected the other variant."""
with pytest.raises(ValidationError):
identity_source_config_adapter.validate_python(
MappingProxyType(
{
"kind": "internal_issuer",
"issuer_url": ISSUER_URL,
"subject": SUBJECT,
"signing_key_ref": SIGNING_KEY_REF,
"client_secret_ref": CLIENT_SECRET_REF,
}
)
)

View file

@ -0,0 +1,188 @@
import json
from typing import Final
import jwt
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec
from litellm.llms.base_llm.auth.identity_source import InternalIssuerSource
from litellm.llms.base_llm.auth.internal_issuer import (
internal_issuer_assertion_source,
internal_issuer_jwks_document,
mint_internal_issuer_assertion,
)
from litellm.llms.base_llm.auth.jwt_signing import build_jwks, rfc7638_thumbprint
SIGNING_KEY_REF: Final = "oidc/env/ISSUER_SIGNING_KEY_PEM"
ISSUER_URL: Final = "https://issuer.internal.example"
SUBJECT: Final = "workload-a"
_PRIVATE_VALUE: Final = 90123456789012345678901234567890123456789012345678901234567890
def signing_key() -> ec.EllipticCurvePrivateKey:
return ec.derive_private_key(_PRIVATE_VALUE, ec.SECP256R1())
def pem_of(key: ec.EllipticCurvePrivateKey) -> str:
return key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption(),
).decode()
def make_config(
issuer_url: str = ISSUER_URL,
subject: str = SUBJECT,
audience: str | None = None,
ttl_seconds: int = 300,
signing_key_ref: str = SIGNING_KEY_REF,
) -> InternalIssuerSource:
return InternalIssuerSource(
issuer_url=issuer_url,
subject=subject,
audience=audience,
ttl_seconds=ttl_seconds,
signing_key_ref=signing_key_ref,
)
def key_reader_returning(pem: str | None):
def reader(ref: str) -> str | None:
assert ref == SIGNING_KEY_REF
return pem
return reader
class FakeClock:
def __init__(self, value: float) -> None:
self._value: Final = value
def __call__(self) -> float:
return self._value
def decode_ignoring_wall_clock(token: str, public_key: ec.EllipticCurvePublicKey) -> dict:
"""Tests mint with a fixed past ``FakeClock`` and no expected audience, so PyJWT's
real-wall-clock ``exp``/``aud`` checks (irrelevant to what these tests verify) are disabled."""
return jwt.decode(token, public_key, algorithms=["ES256"], options={"verify_exp": False, "verify_aud": False})
class TestMintInternalIssuerAssertion:
def test_required_claims_and_asymmetric_alg(self):
key: Final = signing_key()
config: Final = make_config(ttl_seconds=300)
token: Final = mint_internal_issuer_assertion(
config, key_reader=key_reader_returning(pem_of(key)), clock=FakeClock(1_700_000_000.0)
)
header: Final = jwt.get_unverified_header(token)
claims: Final = decode_ignoring_wall_clock(token, key.public_key())
assert header["alg"] == "ES256"
assert claims["sub"] == SUBJECT
assert claims["iss"] == ISSUER_URL
assert claims["iat"] == 1_700_000_000
assert claims["exp"] == 1_700_000_300
def test_kid_matches_the_published_jwks(self):
key: Final = signing_key()
config: Final = make_config()
token: Final = mint_internal_issuer_assertion(
config, key_reader=key_reader_returning(pem_of(key)), clock=FakeClock(1_700_000_000.0)
)
header_kid: Final = jwt.get_unverified_header(token)["kid"]
published_kid: Final = build_jwks(key.public_key())["keys"][0]["kid"]
assert header_kid == published_kid == rfc7638_thumbprint(key.public_key())
def test_ttl_bounds_exp_minus_iat(self):
key: Final = signing_key()
config: Final = make_config(ttl_seconds=120)
token: Final = mint_internal_issuer_assertion(
config, key_reader=key_reader_returning(pem_of(key)), clock=FakeClock(1_700_000_000.0)
)
claims: Final = decode_ignoring_wall_clock(token, key.public_key())
assert claims["exp"] - claims["iat"] == 120
def test_audience_included_only_when_set(self):
key: Final = signing_key()
without_audience: Final = mint_internal_issuer_assertion(
make_config(audience=None), key_reader=key_reader_returning(pem_of(key)), clock=FakeClock(1_700_000_000.0)
)
with_audience: Final = mint_internal_issuer_assertion(
make_config(audience="urn:anthropic:federation"),
key_reader=key_reader_returning(pem_of(key)),
clock=FakeClock(1_700_000_000.0),
)
claims_without: Final = decode_ignoring_wall_clock(without_audience, key.public_key())
claims_with: Final = decode_ignoring_wall_clock(with_audience, key.public_key())
assert "aud" not in claims_without
assert claims_with["aud"] == "urn:anthropic:federation"
def test_jti_is_present_and_fresh_on_every_mint(self):
key: Final = signing_key()
config: Final = make_config()
reader: Final = key_reader_returning(pem_of(key))
first: Final = decode_ignoring_wall_clock(
mint_internal_issuer_assertion(config, key_reader=reader, clock=FakeClock(1_700_000_000.0)),
key.public_key(),
)
second: Final = decode_ignoring_wall_clock(
mint_internal_issuer_assertion(config, key_reader=reader, clock=FakeClock(1_700_000_000.0)),
key.public_key(),
)
assert first["jti"] and second["jti"]
assert first["jti"] != second["jti"]
def test_missing_signing_key_raises_value_error_naming_the_ref_not_a_secret(self):
with pytest.raises(ValueError, match=SIGNING_KEY_REF):
mint_internal_issuer_assertion(make_config(), key_reader=key_reader_returning(None))
def test_malformed_signing_key_raises_value_error(self):
with pytest.raises(ValueError, match="not a valid unencrypted PEM"):
mint_internal_issuer_assertion(make_config(), key_reader=key_reader_returning("not-a-pem"))
class TestInternalIssuerAssertionSource:
def test_returns_a_callable_that_mints_fresh_each_call(self):
key: Final = signing_key()
source: Final = internal_issuer_assertion_source(make_config(), key_reader=key_reader_returning(pem_of(key)))
first: Final = jwt.decode(source(), key.public_key(), algorithms=["ES256"])
second: Final = jwt.decode(source(), key.public_key(), algorithms=["ES256"])
assert first["jti"] != second["jti"]
def test_propagates_the_underlying_mint_failure(self):
source: Final = internal_issuer_assertion_source(make_config(), key_reader=key_reader_returning(None))
with pytest.raises(ValueError, match=SIGNING_KEY_REF):
source()
class TestInternalIssuerJwksDocument:
def test_matches_the_key_used_to_mint(self):
key: Final = signing_key()
config: Final = make_config()
reader: Final = key_reader_returning(pem_of(key))
document: Final = json.loads(internal_issuer_jwks_document(config, key_reader=reader))
token: Final = mint_internal_issuer_assertion(config, key_reader=reader, clock=FakeClock(1_700_000_000.0))
assert document["keys"][0]["kid"] == jwt.get_unverified_header(token)["kid"]
assert decode_ignoring_wall_clock(token, key.public_key())
def test_missing_signing_key_raises_value_error(self):
with pytest.raises(ValueError, match=SIGNING_KEY_REF):
internal_issuer_jwks_document(make_config(), key_reader=key_reader_returning(None))

View file

@ -0,0 +1,213 @@
import base64
import hashlib
import json
import subprocess
import sys
import textwrap
import time
from typing import Final
import jwt
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec, rsa
from litellm.llms.base_llm.auth.jwt_signing import (
MISSING_SIGNING_DEPENDENCIES_MESSAGE,
build_jwk,
build_jwks,
jwks_document_json,
load_es256_private_key,
rfc7638_thumbprint,
sign_es256_jwt,
)
_FIXED_PRIVATE_VALUE: Final = 55090612345678901234567890123456789012345678901234567890123456
_OTHER_PRIVATE_VALUE: Final = 1
def fixed_private_key(value: int = _FIXED_PRIVATE_VALUE) -> ec.EllipticCurvePrivateKey:
return ec.derive_private_key(value, ec.SECP256R1())
def pem_of(key: ec.EllipticCurvePrivateKey) -> str:
return key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption(),
).decode()
def independent_thumbprint(public_key: ec.EllipticCurvePublicKey) -> str:
"""Recomputes RFC 7638 by hand, deliberately not sharing a single line of code with
``jwt_signing.rfc7638_thumbprint`` -- a mutation that broke the real implementation must not
also break this reference, or the two would trivially agree by sharing the bug."""
numbers: Final = public_key.public_numbers()
x: Final = base64.urlsafe_b64encode(numbers.x.to_bytes(32, "big")).rstrip(b"=").decode()
y: Final = base64.urlsafe_b64encode(numbers.y.to_bytes(32, "big")).rstrip(b"=").decode()
canonical: Final = f'{{"crv":"P-256","kty":"EC","x":"{x}","y":"{y}"}}'
return base64.urlsafe_b64encode(hashlib.sha256(canonical.encode()).digest()).rstrip(b"=").decode()
class TestLoadEs256PrivateKey:
def test_valid_ec_p256_pem_loads(self):
key: Final = load_es256_private_key(pem_of(fixed_private_key()))
assert isinstance(key, ec.EllipticCurvePrivateKey)
assert isinstance(key.curve, ec.SECP256R1)
def test_garbage_pem_is_rejected(self):
with pytest.raises(ValueError, match="not a valid unencrypted PEM"):
load_es256_private_key("not a pem")
def test_rsa_key_is_rejected(self):
rsa_pem: Final = (
rsa.generate_private_key(public_exponent=65537, key_size=2048)
.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption(),
)
.decode()
)
with pytest.raises(ValueError, match="P-256"):
load_es256_private_key(rsa_pem)
def test_non_p256_curve_is_rejected(self):
secp384_pem: Final = (
ec.generate_private_key(ec.SECP384R1())
.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption(),
)
.decode()
)
with pytest.raises(ValueError, match="P-256"):
load_es256_private_key(secp384_pem)
def test_error_never_echoes_key_material(self):
pem: Final = pem_of(fixed_private_key())
with pytest.raises(ValueError, match="P-256"):
load_es256_private_key(pem_of(ec.generate_private_key(ec.SECP384R1())))
with pytest.raises(ValueError, match="not a valid unencrypted PEM") as exc_info:
load_es256_private_key("garbage-not-a-pem")
assert pem not in str(exc_info.value)
class TestRfc7638Thumbprint:
def test_matches_independent_recomputation(self):
public_key: Final = fixed_private_key().public_key()
assert rfc7638_thumbprint(public_key) == independent_thumbprint(public_key)
def test_different_keys_have_different_thumbprints(self):
first: Final = fixed_private_key(_FIXED_PRIVATE_VALUE).public_key()
second: Final = fixed_private_key(_OTHER_PRIVATE_VALUE).public_key()
assert rfc7638_thumbprint(first) != rfc7638_thumbprint(second)
def test_thumbprint_is_deterministic(self):
public_key: Final = fixed_private_key().public_key()
assert rfc7638_thumbprint(public_key) == rfc7638_thumbprint(public_key)
class TestBuildJwks:
def test_jwks_contains_one_key_matching_the_thumbprint(self):
public_key: Final = fixed_private_key().public_key()
jwks: Final = build_jwks(public_key)
assert len(jwks["keys"]) == 1
assert jwks["keys"][0]["kid"] == rfc7638_thumbprint(public_key)
assert jwks["keys"][0]["kty"] == "EC"
assert jwks["keys"][0]["crv"] == "P-256"
assert jwks["keys"][0]["alg"] == "ES256"
def test_build_jwk_stamps_the_given_kid_verbatim(self):
jwk: Final = build_jwk(fixed_private_key().public_key(), kid="caller-supplied-kid")
assert jwk["kid"] == "caller-supplied-kid"
def test_jwks_document_json_round_trips_through_build_jwks(self):
key: Final = fixed_private_key()
document: Final = json.loads(jwks_document_json(pem_of(key)))
jwks: Final = build_jwks(key.public_key())
assert document == {"keys": [dict(jwk) for jwk in jwks["keys"]]}
class TestSignEs256Jwt:
def test_minted_token_verifies_against_the_matching_public_key(self):
key: Final = fixed_private_key()
now: Final = int(time.time())
claims: Final = {"sub": "workload-a", "iss": "https://issuer.example", "iat": now, "exp": now + 300}
token: Final = sign_es256_jwt(pem_of(key), claims)
decoded: Final = jwt.decode(token, key.public_key(), algorithms=["ES256"])
assert decoded == claims
def test_header_alg_is_es256(self):
token: Final = sign_es256_jwt(pem_of(fixed_private_key()), {"sub": "x"})
assert jwt.get_unverified_header(token)["alg"] == "ES256"
def test_header_kid_matches_the_published_jwks(self):
key: Final = fixed_private_key()
token: Final = sign_es256_jwt(pem_of(key), {"sub": "x"})
header_kid: Final = jwt.get_unverified_header(token)["kid"]
published_kid: Final = build_jwks(key.public_key())["keys"][0]["kid"]
assert header_kid == published_kid == rfc7638_thumbprint(key.public_key())
def test_wrong_key_fails_verification(self):
signing_key: Final = fixed_private_key(_FIXED_PRIVATE_VALUE)
other_key: Final = fixed_private_key(_OTHER_PRIVATE_VALUE)
token: Final = sign_es256_jwt(pem_of(signing_key), {"sub": "x"})
with pytest.raises(jwt.exceptions.InvalidSignatureError):
jwt.decode(token, other_key.public_key(), algorithms=["ES256"])
class TestBaseSdkImport:
"""A base ``pip install litellm`` has neither PyJWT nor cryptography (both are proxy extras),
and ``litellm/__init__`` reaches this module through the Anthropic provider, so a
module-level import of either would break ``import litellm`` for every base SDK user."""
def test_module_imports_with_pyjwt_and_cryptography_absent(self):
script: Final = textwrap.dedent(
"""
import sys
class Blocker:
def find_spec(self, name, path=None, target=None):
if name.split(".")[0] in {"jwt", "cryptography"}:
raise ModuleNotFoundError(f"No module named {name!r}")
sys.meta_path.insert(0, Blocker())
import litellm
from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json
try:
jwks_document_json("not a key")
except ImportError as e:
print(e)
"""
)
result: Final = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=False)
assert result.returncode == 0, result.stderr[-2000:]
assert result.stdout.strip() == MISSING_SIGNING_DEPENDENCIES_MESSAGE
def test_signing_reports_the_missing_extra(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setitem(sys.modules, "jwt", None)
with pytest.raises(ImportError, match="litellm\\[proxy\\]"):
sign_es256_jwt(pem_of(fixed_private_key()), {"sub": "x"})

File diff suppressed because it is too large Load diff

View file

@ -1,6 +1,7 @@
import asyncio
import json
import logging
import threading
import time
from unittest.mock import AsyncMock, Mock, patch
@ -17,7 +18,9 @@ from litellm.llms.base_llm.audio_transcription.transformation import (
AudioTranscriptionRequestData,
BaseAudioTranscriptionConfig,
)
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.files.transformation import BaseFilesConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import (
BaseLLMHTTPHandler,
@ -455,7 +458,7 @@ async def test_async_anthropic_messages_handler_extra_headers():
# Mock the config
mock_config = Mock()
mock_config.validate_anthropic_messages_environment = Mock(
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com")
)
mock_config.transform_anthropic_messages_request = Mock(
@ -502,7 +505,7 @@ async def test_async_anthropic_messages_handler_extra_headers():
captured_headers.update(kwargs.get("headers", {}))
return ({"x-api-key": "test-key"}, "https://api.anthropic.com")
mock_config.validate_anthropic_messages_environment = capture_validate
mock_config.avalidate_anthropic_messages_environment = AsyncMock(side_effect=capture_validate)
try:
await handler.async_anthropic_messages_handler(
@ -758,7 +761,7 @@ async def test_async_anthropic_messages_handler_passes_litellm_metadata():
handler = BaseLLMHTTPHandler()
mock_config = Mock()
mock_config.validate_anthropic_messages_environment = Mock(
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com")
)
mock_config.transform_anthropic_messages_request = Mock(
@ -837,7 +840,7 @@ async def test_async_anthropic_messages_handler_forwards_router_model_info():
handler = BaseLLMHTTPHandler()
mock_config = Mock()
mock_config.validate_anthropic_messages_environment = Mock(
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com")
)
mock_config.transform_anthropic_messages_request = Mock(
@ -929,7 +932,7 @@ async def test_async_anthropic_messages_handler_header_priority():
captured_headers.update(kwargs.get("headers", {}))
return ({"x-api-key": "test-key"}, "https://api.anthropic.com")
mock_config.validate_anthropic_messages_environment = capture_validate
mock_config.avalidate_anthropic_messages_environment = AsyncMock(side_effect=capture_validate)
mock_config.transform_anthropic_messages_request = Mock(
return_value={"model": "claude-3-opus-20240229", "messages": []}
)
@ -968,7 +971,7 @@ async def test_async_anthropic_messages_handler_drops_top_level_and_nested_param
handler = BaseLLMHTTPHandler()
mock_config = Mock()
mock_config.validate_anthropic_messages_environment = Mock(
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com")
)
@ -1181,9 +1184,7 @@ def test_sync_delete_responses_sets_json_content_type():
({}, True, None, None),
],
)
def test_resolve_anthropic_messages_timeout(
monkeypatch, litellm_params_kwargs, stream, global_timeout, expected
):
def test_resolve_anthropic_messages_timeout(monkeypatch, litellm_params_kwargs, stream, global_timeout, expected):
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS
if global_timeout is None:
@ -1199,9 +1200,7 @@ def test_resolve_anthropic_messages_timeout(
)
else:
monkeypatch.setattr("litellm.request_timeout", global_timeout, raising=False)
monkeypatch.setattr(
"litellm.request_timeout_explicitly_set", True, raising=False
)
monkeypatch.setattr("litellm.request_timeout_explicitly_set", True, raising=False)
resolved = BaseLLMHTTPHandler._resolve_anthropic_messages_timeout(
litellm_params=GenericLiteLLMParams(**litellm_params_kwargs),
@ -1222,13 +1221,11 @@ async def test_async_anthropic_messages_handler_forwards_request_timeout(monkeyp
handler = BaseLLMHTTPHandler()
mock_config = Mock()
mock_config.validate_anthropic_messages_environment = Mock(
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
return_value=({"x-api-key": "k"}, "https://api.anthropic.com")
)
mock_config.should_filter_anthropic_beta_headers = Mock(return_value=False)
mock_config.transform_anthropic_messages_request = Mock(
return_value={"model": "claude", "messages": []}
)
mock_config.transform_anthropic_messages_request = Mock(return_value={"model": "claude", "messages": []})
mock_config.get_complete_url = Mock(return_value="https://api.anthropic.com/v1/messages")
mock_config.sign_request = Mock(return_value=({"x-api-key": "k"}, None))
mock_config.max_retry_on_anthropic_messages_http_error = 1
@ -1270,13 +1267,11 @@ async def test_async_anthropic_messages_handler_forwards_stream_timeout(monkeypa
handler = BaseLLMHTTPHandler()
mock_config = Mock()
mock_config.validate_anthropic_messages_environment = Mock(
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
return_value=({"x-api-key": "k"}, "https://api.anthropic.com")
)
mock_config.should_filter_anthropic_beta_headers = Mock(return_value=False)
mock_config.transform_anthropic_messages_request = Mock(
return_value={"model": "claude", "messages": []}
)
mock_config.transform_anthropic_messages_request = Mock(return_value={"model": "claude", "messages": []})
mock_config.get_complete_url = Mock(return_value="https://api.anthropic.com/v1/messages")
mock_config.sign_request = Mock(return_value=({"x-api-key": "k"}, None))
mock_config.max_retry_on_anthropic_messages_http_error = 1
@ -1678,7 +1673,7 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks(
handler = BaseLLMHTTPHandler()
mock_config = Mock()
mock_config.validate_anthropic_messages_environment = Mock(
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
return_value=({"x-api-key": "sk-test"}, "https://api.anthropic.com")
)
mock_config.transform_anthropic_messages_request = Mock(
@ -1686,7 +1681,13 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks(
)
mock_config.sign_request = Mock(return_value=({}, None))
fake_raw_response = {"id": "msg_1", "type": "message", "role": "assistant", "content": [], "stop_reason": "end_turn"}
fake_raw_response = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"content": [],
"stop_reason": "end_turn",
}
mock_config.transform_anthropic_messages_response = Mock(return_value=fake_raw_response)
mock_logging_obj = Mock()
@ -1706,10 +1707,17 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks(
mock_httpx_response.status_code = 200
with (
patch.object(handler, "_async_post_anthropic_messages_with_http_error_retry", new=AsyncMock(return_value=mock_httpx_response)),
patch.object(
handler,
"_async_post_anthropic_messages_with_http_error_retry",
new=AsyncMock(return_value=mock_httpx_response),
),
patch.object(handler, "_call_agentic_completion_hooks", side_effect=fake_agentic_hooks),
patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client"),
patch("litellm.litellm_core_utils.get_provider_specific_headers.ProviderSpecificHeaderUtils.get_provider_specific_headers", return_value=None),
patch( # test-quality-ok: the proxy wiring under test is what this patches
"litellm.litellm_core_utils.get_provider_specific_headers.ProviderSpecificHeaderUtils.get_provider_specific_headers",
return_value=None,
),
):
result = await handler.async_anthropic_messages_handler(
model="claude-haiku",
@ -1951,7 +1959,9 @@ def test_audio_transcriptions_sends_dict_data_as_json_body():
form-encodes it and silently ignores json=; JSON-body providers (e.g.
Google Speech-to-Text) need an application/json body."""
captured = {}
client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_json_transcription_request(captured))))
client = HTTPHandler(
client=httpx.Client(transport=httpx.MockTransport(_capture_json_transcription_request(captured)))
)
response = BaseLLMHTTPHandler().audio_transcriptions(
client=client,
@ -2110,6 +2120,105 @@ def test_sync_retrieve_file_content_raises_on_http_error():
assert exc_info.value.status_code == 404
_FILE_CONTENT_WIF_ENV = {
"ANTHROPIC_FEDERATION_RULE_ID": "fdrl_llm_http_handler_seam",
"ANTHROPIC_ORGANIZATION_ID": "org-llm-http-handler-seam",
"ANTHROPIC_IDENTITY_TOKEN": "llm-http-handler-seam-inline-jwt",
}
class _BlockingWifPoster:
"""A token-endpoint poster 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):
self.thread_ids.append(threading.get_ident())
self.release.wait(timeout=5)
return httpx.Response(
200,
json={
"access_token": "sk-ant-oat01-llm-http-handler-seam",
"token_type": "Bearer",
"expires_in": 3600,
},
)
@pytest.mark.asyncio
async def test_async_retrieve_file_content_wif_exchange_does_not_block_event_loop(monkeypatch):
"""Regression (Greptile P1): async_retrieve_file_content called the synchronous
validate_environment directly, so a cold WIF mint on this call site froze the
event loop until the exchange finished. It must resolve credentials through the
async facade instead."""
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
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)
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)
async def async_shim(litellm_params, api_base, model):
return await aget_anthropic_wif_token(litellm_params, api_base, model, engine)
monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim)
monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim)
handler = BaseLLMHTTPHandler()
client = Mock(spec=AsyncHTTPHandler)
client.get = AsyncMock(return_value=httpx.Response(status_code=200, content=b"file bytes"))
ticks = []
async def ticker():
for i in range(20):
await asyncio.sleep(0.005)
ticks.append(i)
ticker_task = asyncio.create_task(ticker())
await asyncio.sleep(0.02)
retrieve_task = asyncio.create_task(
handler.async_retrieve_file_content(
file_content_request={"file_id": "file-abc"},
provider_config=AnthropicFilesConfig(),
litellm_params={},
headers={},
logging_obj=Mock(),
client=client,
)
)
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.
assert len(ticks) > 0
assert not retrieve_task.done()
poster.release.set()
await retrieve_task
await ticker_task
assert sync_calls == []
assert poster.thread_ids
assert poster.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"
_UPSTREAM_NOT_FOUND_BODY = {
"error": {
"message": "Response with id 'resp_abc' not found.",
@ -2257,9 +2366,7 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques
ok_response = httpx.Response(200, json={"id": "msg_1"}, request=httpx.Request("POST", request_url))
class FakeAsyncClient:
async def post(
self, url, headers, data, stream=False, logging_obj=None, timeout=None
):
async def post(self, url, headers, data, stream=False, logging_obj=None, timeout=None):
posts.append({"headers": dict(headers), "data": data})
return invalid_signature_response if len(posts) == 1 else ok_response
@ -2576,7 +2683,7 @@ async def test_async_anthropic_messages_handler_carries_deployment_vertex_locati
custom_llm_provider="vertex_ai",
)
mock_config = Mock()
mock_config.validate_anthropic_messages_environment = Mock(
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
return_value=({"authorization": "Bearer t"}, "https://us-east5-aiplatform.googleapis.com")
)
mock_config.transform_anthropic_messages_request = Mock(
@ -3100,6 +3207,70 @@ async def test_a_provider_that_keeps_rejecting_is_not_retried_forever_on_the_asy
assert len(recorder.bodies) == 2
def _async_client_returning(response: Mock) -> AsyncMock:
client = AsyncMock(spec=AsyncHTTPHandler)
client.post.return_value = response
return client
@pytest.mark.asyncio
async def test_create_file_async_awaits_the_provider_credential_hook_instead_of_blocking():
provider_config = Mock(spec=BaseFilesConfig)
provider_config.validate_environment.side_effect = AssertionError("sync validate_environment ran on the event loop")
provider_config.avalidate_environment = AsyncMock(return_value={"x-api-key": "federated"})
provider_config.get_complete_file_url.return_value = "https://files.example/v1/files"
provider_config.transform_create_file_request.return_value = {"file": ("batch.jsonl", b"{}", "application/jsonl")}
file_object = object()
provider_config.transform_create_file_response.return_value = file_object
client = _async_client_returning(Mock(spec=httpx.Response))
result = await BaseLLMHTTPHandler().create_file(
create_file_data={"file": b"{}", "purpose": "batch"},
litellm_params={},
provider_config=provider_config,
headers={},
api_base=None,
api_key=None,
logging_obj=Mock(),
_is_async=True,
client=client,
)
assert result is file_object
provider_config.validate_environment.assert_not_called()
provider_config.avalidate_environment.assert_awaited_once()
assert client.post.call_args.kwargs["headers"] == {"x-api-key": "federated"}
assert client.post.call_args.kwargs["url"] == "https://files.example/v1/files"
@pytest.mark.asyncio
async def test_create_batch_async_validates_credentials_off_the_event_loop():
provider_config = Mock(spec=BaseBatchesConfig)
provider_config.validate_environment.side_effect = lambda **_: {"x-validated-on": str(threading.get_ident())}
provider_config.get_complete_batch_url.return_value = "https://batches.example/v1/messages/batches"
provider_config.transform_create_batch_request.return_value = {"requests": []}
batch = object()
provider_config.transform_create_batch_response.return_value = batch
client = _async_client_returning(Mock(spec=httpx.Response))
result = await BaseLLMHTTPHandler().create_batch(
create_batch_data={"input_file_id": "file_1", "endpoint": "/v1/chat/completions", "completion_window": "24h"},
litellm_params={},
provider_config=provider_config,
headers={},
api_base=None,
api_key=None,
logging_obj=Mock(),
_is_async=True,
client=client,
model="claude-sonnet-4-5",
)
assert result is batch
validated_on = client.post.call_args.kwargs["headers"]["x-validated-on"]
assert validated_on != str(threading.get_ident())
assert client.post.call_args.kwargs["url"] == "https://batches.example/v1/messages/batches"
CONTAINER_NOT_FOUND_BODY = {
"error": {
"message": "Container with id 'cntr_gone' not found.",

View file

@ -10,6 +10,7 @@ from openai import AsyncOpenAI, OpenAI
import litellm
from litellm.llms.litellm_proxy.responses.transformation import LiteLLMProxyResponsesAPIConfig
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIError
from litellm.llms.openai.openai import OpenAIChatCompletion
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
@ -22,6 +23,17 @@ from litellm.llms.openai.workload_identity import (
from litellm.types.router import GenericLiteLLMParams
TOKEN_EXCHANGE_URL: Final = "https://auth.openai.com/oauth/token"
CHAT_COMPLETIONS_URL: Final = "https://api.openai.com/v1/chat/completions"
EMBEDDINGS_URL: Final = "https://api.openai.com/v1/embeddings"
MODELS_URL: Final = "https://api.openai.com/v1/models"
CHAT_COMPLETION_BODY: Final = {
"id": "chatcmpl-wif",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
@pytest.fixture
@ -279,3 +291,342 @@ class TestResponsesValidateEnvironment:
headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams()
)
assert headers["Authorization"] == "Bearer None"
@pytest.fixture
def deployment_wif(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> dict[str, str]:
token_file: Final = tmp_path / "deployment_subject_token.jwt"
token_file.write_text("subject-token-from-deployment-file")
for name in (
"OPENAI_API_KEY",
"OPENAI_BASE_URL",
"OPENAI_API_BASE",
"OPENAI_IDENTITY_PROVIDER_ID",
"OPENAI_SERVICE_ACCOUNT_ID",
"OPENAI_IDENTITY_TOKEN_FILE",
):
monkeypatch.delenv(name, raising=False)
monkeypatch.setattr(litellm, "api_base", None)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
_workload_identity_auth.cache_clear()
litellm.in_memory_llm_clients_cache.flush_cache()
return {
"openai_identity_provider_id": "idp_deployment",
"openai_service_account_id": "user-deployment",
"openai_identity_token_file": str(token_file),
}
def deployment_config(deployment_wif: dict[str, str]) -> OpenAIWorkloadIdentityConfig:
return OpenAIWorkloadIdentityConfig(
identity_provider_id="idp_deployment",
service_account_id="user-deployment",
token_file=deployment_wif["openai_identity_token_file"],
)
def mock_chat_completions() -> respx.Route:
return respx.post(CHAT_COMPLETIONS_URL).mock(return_value=httpx.Response(200, json=CHAT_COMPLETION_BODY))
def mock_streaming_chat_completions() -> respx.Route:
chunk: Final = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"}
events: Final = (
{**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": None}]},
{**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
)
body: Final = "".join(f"data: {json.dumps(event)}\n\n" for event in events) + "data: [DONE]\n\n"
return respx.post(CHAT_COMPLETIONS_URL).mock(
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=body)
)
class TestResolveConfigFromDeployment:
def test_resolves_from_litellm_params_without_env(self, deployment_wif: dict[str, str]) -> None:
assert resolve_openai_workload_identity_config(
api_key=None, api_base=None, litellm_params=deployment_wif
) == deployment_config(deployment_wif)
def test_env_alone_disables_nothing_when_params_are_absent(self, deployment_wif: dict[str, str]) -> None:
assert resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=None) is None
def test_unrelated_litellm_params_do_not_resolve(self, deployment_wif: dict[str, str]) -> None:
assert (
resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params={"model": "gpt-4o"})
is None
)
def test_litellm_params_beat_env(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
config: Final = resolve_openai_workload_identity_config(
api_key=None,
api_base=None,
litellm_params={
"openai_identity_provider_id": "idp_deployment",
"openai_service_account_id": "user-deployment",
"openai_identity_token_file": wif_env.token_file,
},
)
assert config == OpenAIWorkloadIdentityConfig(
identity_provider_id="idp_deployment",
service_account_id="user-deployment",
token_file=wif_env.token_file,
)
def test_partial_litellm_params_fill_from_env_per_field(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
config: Final = resolve_openai_workload_identity_config(
api_key=None, api_base=None, litellm_params={"openai_identity_provider_id": "idp_deployment"}
)
assert config == OpenAIWorkloadIdentityConfig(
identity_provider_id="idp_deployment",
service_account_id=wif_env.service_account_id,
token_file=wif_env.token_file,
)
@pytest.mark.parametrize("blank", ["", None, 7])
def test_blank_or_non_string_param_falls_back_to_env(
self, wif_env: OpenAIWorkloadIdentityConfig, blank: object
) -> None:
config: Final = resolve_openai_workload_identity_config(
api_key=None, api_base=None, litellm_params={"openai_identity_provider_id": blank}
)
assert config == wif_env
def test_partial_litellm_params_without_env_disable(self, deployment_wif: dict[str, str]) -> None:
partial: Final = {key: value for key, value in deployment_wif.items() if key != "openai_identity_token_file"}
assert resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=partial) is None
def test_static_api_key_beats_litellm_params(self, deployment_wif: dict[str, str]) -> None:
assert (
resolve_openai_workload_identity_config(api_key="sk-static", api_base=None, litellm_params=deployment_wif)
is None
)
def test_env_openai_api_key_beats_litellm_params(
self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env")
assert (
resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=deployment_wif) is None
)
def test_foreign_api_base_disables_deployment_wif(self, deployment_wif: dict[str, str]) -> None:
assert (
resolve_openai_workload_identity_config(
api_key=None, api_base="https://my-vllm.internal/v1", litellm_params=deployment_wif
)
is None
)
class TestDeploymentClientConstruction:
def test_sync_client_from_deployment_params(self, deployment_wif: dict[str, str]) -> None:
client: Final = OpenAIChatCompletion()._get_openai_client(
is_async=False, api_key=None, api_base=None, litellm_params=deployment_wif
)
assert isinstance(client, OpenAI)
assert client.api_key == "workload-identity-auth"
assert client._workload_identity_auth is not None
def test_async_client_from_deployment_params(self, deployment_wif: dict[str, str]) -> None:
client: Final = OpenAIChatCompletion()._get_openai_client(
is_async=True, api_key=None, api_base=None, litellm_params=deployment_wif
)
assert isinstance(client, AsyncOpenAI)
assert client._workload_identity_auth is not None
def test_distinct_deployments_get_distinct_cached_clients(self, deployment_wif: dict[str, str]) -> None:
other_deployment: Final = {**deployment_wif, "openai_service_account_id": "user-other"}
handler: Final = OpenAIChatCompletion()
first: Final = handler._get_openai_client(
is_async=False, api_key=None, api_base=None, litellm_params=deployment_wif
)
second: Final = handler._get_openai_client(
is_async=False, api_key=None, api_base=None, litellm_params=other_deployment
)
again: Final = handler._get_openai_client(
is_async=False, api_key=None, api_base=None, litellm_params=dict(deployment_wif)
)
assert first is not second
assert again is first
@respx.mock
def test_completion_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None:
mock_token_exchange("deployment-bearer")
completion_route: Final = mock_chat_completions()
response: Final = litellm.completion(
model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], **deployment_wif
)
assert response.choices[0].message.content == "ok"
request: Final = completion_route.calls.last.request
assert request.headers["Authorization"] == "Bearer deployment-bearer"
assert not any(key.startswith("openai_") for key in json.loads(request.content))
@respx.mock
def test_streaming_completion_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None:
mock_token_exchange("stream-bearer")
stream_route: Final = mock_streaming_chat_completions()
chunks: Final = tuple(
litellm.completion(
model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **deployment_wif
)
)
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok"
assert stream_route.calls.last.request.headers["Authorization"] == "Bearer stream-bearer"
@respx.mock
@pytest.mark.asyncio
async def test_async_streaming_completion_kwargs_carry_exchanged_bearer(
self, deployment_wif: dict[str, str]
) -> None:
mock_token_exchange("async-stream-bearer")
stream_route: Final = mock_streaming_chat_completions()
stream: Final = await litellm.acompletion(
model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **deployment_wif
)
chunks: Final = tuple([chunk async for chunk in stream])
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok"
assert stream_route.calls.last.request.headers["Authorization"] == "Bearer async-stream-bearer"
@respx.mock
def test_router_deployment_without_api_key_authenticates_via_token_exchange(
self, deployment_wif: dict[str, str]
) -> None:
exchange_route: Final = mock_token_exchange("router-bearer")
completion_route: Final = mock_chat_completions()
router: Final = litellm.Router(
model_list=[{"model_name": "wif-gpt", "litellm_params": {"model": "openai/gpt-4o-mini", **deployment_wif}}]
)
response: Final = router.completion(model="wif-gpt", messages=[{"role": "user", "content": "hi"}])
assert response.choices[0].message.content == "ok"
assert exchange_route.called
assert completion_route.calls.last.request.headers["Authorization"] == "Bearer router-bearer"
@respx.mock
def test_embedding_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None:
mock_token_exchange("embedding-bearer")
embeddings_route: Final = respx.post(EMBEDDINGS_URL).mock(
return_value=httpx.Response(
200,
json={
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}],
"model": "text-embedding-3-small",
"usage": {"prompt_tokens": 1, "total_tokens": 1},
},
)
)
litellm.embedding(model="openai/text-embedding-3-small", input=["hi"], **deployment_wif)
assert embeddings_route.calls.last.request.headers["Authorization"] == "Bearer embedding-bearer"
class TestResponsesValidateEnvironmentFromDeployment:
@respx.mock
def test_mints_bearer_from_litellm_params(self, deployment_wif: dict[str, str]) -> None:
mock_token_exchange("responses-bearer")
headers: Final = OpenAIResponsesAPIConfig().validate_environment(
headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams(**deployment_wif)
)
assert headers["Authorization"] == "Bearer responses-bearer"
def test_static_key_in_litellm_params_wins(self, deployment_wif: dict[str, str]) -> None:
headers: Final = OpenAIResponsesAPIConfig().validate_environment(
headers={},
model="gpt-4o-mini",
litellm_params=GenericLiteLLMParams(api_key="sk-responses", **deployment_wif),
)
assert headers["Authorization"] == "Bearer sk-responses"
class TestDiscoverModels:
@staticmethod
def mock_models() -> respx.Route:
return respx.get(MODELS_URL).mock(
return_value=httpx.Response(200, json={"data": [{"id": "gpt-4o-mini"}, {"id": "gpt-4.1"}]})
)
@respx.mock
def test_discovers_with_exchanged_bearer_from_litellm_params(self, deployment_wif: dict[str, str]) -> None:
mock_token_exchange("discovery-bearer")
models_route: Final = self.mock_models()
assert OpenAIGPTConfig().discover_models(deployment_wif) == ["gpt-4o-mini", "gpt-4.1"]
assert models_route.calls.last.request.headers["Authorization"] == "Bearer discovery-bearer"
@respx.mock
def test_discovers_with_env_wif_when_params_carry_no_key(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
mock_token_exchange("env-discovery-bearer")
models_route: Final = self.mock_models()
OpenAIGPTConfig().discover_models({})
assert models_route.calls.last.request.headers["Authorization"] == "Bearer env-discovery-bearer"
@respx.mock
def test_static_api_key_in_params_skips_token_exchange(self, deployment_wif: dict[str, str]) -> None:
exchange_route: Final = mock_token_exchange()
models_route: Final = self.mock_models()
OpenAIGPTConfig().discover_models({**deployment_wif, "api_key": "sk-discovery"})
assert models_route.calls.last.request.headers["Authorization"] == "Bearer sk-discovery"
assert not exchange_route.called
@respx.mock
def test_blank_api_base_in_params_discovers_from_openai(self, deployment_wif: dict[str, str]) -> None:
mock_token_exchange("blank-base-bearer")
models_route: Final = self.mock_models()
assert OpenAIGPTConfig().discover_models({**deployment_wif, "api_base": ""}) == ["gpt-4o-mini", "gpt-4.1"]
assert models_route.calls.last.request.headers["Authorization"] == "Bearer blank-base-bearer"
@respx.mock
def test_openai_compatible_subclass_never_mints_wif(self, deployment_wif: dict[str, str]) -> None:
exchange_route: Final = mock_token_exchange()
models_route: Final = self.mock_models()
class CompatibleConfig(OpenAIGPTConfig):
pass
CompatibleConfig().discover_models(deployment_wif)
assert models_route.calls.last.request.headers["Authorization"] == "Bearer None"
assert not exchange_route.called
@respx.mock
def test_empty_static_key_never_borrows_the_env_key(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "sk-env-key-that-must-stay-home")
foreign_models: Final = respx.get("https://third-party.example/v1/models").mock(
return_value=httpx.Response(200, json={"data": [{"id": "other-model"}]})
)
assert OpenAIGPTConfig().get_models(api_key="", api_base="https://third-party.example") == ["other-model"]
assert foreign_models.calls.last.request.headers["Authorization"] == "Bearer "
class TestClientsideBaseOverride:
def test_client_api_base_override_clears_deployment_wif(self, deployment_wif: dict[str, str]) -> None:
from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params
redirected: Final = get_dynamic_litellm_params(
litellm_params={"model": "openai/gpt-4o-mini", **deployment_wif},
request_kwargs={"api_base": "https://not-openai.example/v1"},
)
assert not any(key in redirected for key in deployment_wif)
assert (
resolve_openai_workload_identity_config(
api_key=None, api_base=redirected["api_base"], litellm_params=redirected
)
is None
)

Some files were not shown because too many files have changed in this diff Show more