scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

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

Browse source 
* 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>

* fix(anthropic): let batch-result downloads mint from deployment params and accept host:port allowlist entries

The files handler enabled workload identity on batch-result downloads but never received the
deployment's litellm_params, so a deployment authenticating through a named credential could only
mint from process-wide env vars. It now threads litellm_params through to the auth header the way
the batch retrieve path already does.

LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS entries written as host:port were read by urlsplit as a scheme,
so the allowlist kept the raw entry while the exchange compared bare hostnames and refused the
gateway. Entries are now parsed as network locations whether or not they carry a scheme.

* fix(types): move the WIF kwargs key sets to a leaf module so the kwargs funnel imports without a cycle

* test(anthropic): pin case-insensitive matching of WIF exchange-host allowlist entries

* fix(anthropic): end workload identity federation errors without a period so the router suffix reads cleanly

* fix(proxy): decrypt stored litellm_params before the WIF write gate

* fix(proxy): hide WIF secret references from /health output

* fix(proxy): keep the proxy error shape on credential endpoint refusals

* fix(proxy): hide identity token file paths from /health output

* fix(anthropic): rename the federation workspace param so Bedrock's anthropic_workspace_id keeps working

The Bedrock Claude Platform route already reads anthropic_workspace_id from
optional_params, so banning that spelling as a server-owned federation
parameter broke a pre-existing client capability. The federation field is now
anthropic_federation_workspace_id (env ANTHROPIC_FEDERATION_WORKSPACE_ID),
which restores the base branch's behavior for Bedrock callers, drops the
Bedrock-specific hint from the refusal message, and deletes the unconditional
ban constant that no longer had a reader

* fix(auth): share one exchanged token across workers reading the same assertion

Anthropic accepts each identity assertion exactly once, so two uvicorn
workers reading the same token file both minting from it means the second
exchange is denied with jti_reused. Minted tokens now land in a per-user
0700 cache directory guarded by a file lock, so workers on the same host
reuse one exchange until the token expires or the assertion rotates. A 401
is only retried when the re-read assertion actually differs, and the denial
hint explains jti_reused. LITELLM_TOKEN_EXCHANGE_CACHE_DIR moves the cache
and an empty value disables it

* fix: keep anthropic federation from being shadowed or leaked

An empty or whitespace-only ANTHROPIC_API_KEY counted as set, so a federated
deployment sent an empty x-api-key on every call instead of minting a token.
Blank values now read as unset, and a real static key on a federated deployment
logs once that it outranks federation and nothing is being federated.

The exchange-host allowlist matched hostnames only, so a second process on
another port of an allowed host was trusted with the workload's identity token.
An entry that names a port now trusts that port alone, while a bare host still
trusts every port.

The shared token store exists so the workers reading one projected token file do
not each spend its single-use jti. A source that mints its own assertion per
exchange shares nothing with another worker, so it no longer writes a live token
to disk for a lookup that can never hit.

* fix: unlink a staged token file a failed write leaves behind

The 401 denial hint now also says federation ignores ANTHROPIC_WORKSPACE_ID, which the Bedrock Claude platform provider already reads.

* refactor: move anthropic jwks derivation behind a provider-owned tagged union

* fix: unlink the staged token file when its write fails at close

A buffered write only reaches the disk when the handle closes, so a full disk surfaces at close and left the staging file behind holding a usable token.

* fix(anthropic): close the staging descriptor before writing the shared token file

* fix(wif): judge federation writes by what they set, not what is stored

The admin gate read the stored deployment, so a team admin lost edit, delete
and Test Connection on any deployment carrying federation params. It now
returns early unless the submitted fields touch the federation surface, and a
Test Connection probe that points the deployment at its own api_base is still
refused, with the 403 no longer wrapped into a 500

The rest of the same review pass: POST /model/new refuses only a blocking
value of `blocked`, so a client that always sends `blocked: false` is not
turned away; a request body can no longer pick which federated identity to
mint as by naming a stored credential; an advisory refresh the executor
refuses disarms the entry instead of wedging the identity until the follower
timeout; the static-key shadow warning resolves its env fallback inside the
cache instead of once per request; credential writes drop nulls before
storing them; the token exchange validates the endpoint URL before reading an
assertion and keeps refusing redirects across a client heal; /health hides
every server-owned federation field from non-admins; and the async create_file
and create_batch paths say which setting is missing when the provider resolves
no URL

* fix(proxy): let a deployment write name a federated credential

reject_federated_credential_reference runs from is_request_body_safe, which
pre_db_read_auth_checks calls on every route, so it also fired on POST
/model/new, /model/update, /model/{id}/update and /health/test_connection. A
proxy admin could no longer attach a federated credential to a deployment over
the API or the Admin UI, leaving a static config.yaml entry as the only way to
configure the feature the rejection told the caller to go configure, and
_reject_non_admin_wif_write never got to make the call it exists to make.

is_request_body_safe now takes the route and skips only the credential-reference
check on the routes that reach can_user_make_model_call. Federation fields typed
inline into a body stay refused everywhere, and a call naming a federated
credential still cannot pick the identity it mints as.

* refactor(proxy): derive health display policy from the federation key sets

The health check module hand-copied the five workload identity fields whose
value is a credential, so a shared proxy surface named provider-specific
parameters and a newly added secret-bearing field would have gone on being
displayed until someone remembered both places

WIF_SECRET_BEARING_KEYS now sits beside the key sets it splits out of,
types/utils derives secret_bearing_wif_litellm_params from it, and the health
layer splats that tuple the same way it already splats the admin-only one

* fix(anthropic_wif): treat blank identity-source fields as unset

* test(proxy): classify the federation params in the credential slot registry

main's registry test (#43298) now fails the build for any credential-named
deployment param without a classification. The five federation fields that
carry a token, a token file path, or a signing or client secret reference are
Unplanted, matching WIF_SECRET_BEARING_KEYS; the four remaining Keycloak
settings name a URL, a client id, an auth method, or a scope and are NotSecret

* fix(anthropic_wif): declare federation params as owned connection leaves and chart their metrics

Register the 18 Anthropic and 3 OpenAI federation params as frozen
ConnectionSettings leaves so the owned-kwarg registry, the kwargs funnel
and the request-body ban list read one declaration. Pass the deployment
api_base through to the count-tokens handler instead of a pre-suffixed
URL, which doubled the /count_tokens path on main's prompt-cache
predictor. Add the five litellm_anthropic_wif_* families to the
all-metrics Grafana dashboard.

* fix(credentials): gate PATCH on WIF fields resolved from model_id

The credential PATCH handler checked server-owned workload identity
federation fields only on the values the caller sent, while a body that
named a deployment through model_id had its credential values resolved
after that check. A non-admin could therefore copy a federated
deployment's WIF fields onto an ordinary credential. Resolve the incoming
values first and run the non-admin gate on them, matching the POST path

* fix(anthropic): count tokens with ANTHROPIC_AUTH_TOKEN through the shared auth header

Count-tokens walked its own credential ladder: a static key, else skip minting when
ANTHROPIC_AUTH_TOKEN is set, else mint a federated token. With only the auth token set it
forwarded nothing and the proxy silently fell back to its local tokenizer while chat on the
same deployment authenticated with that token. The handler now takes the auth header that
AnthropicModelInfo.aget_auth_header resolves, the same ladder chat, files, batches and skills
use, and merges the oauth beta a minted or consumer token carries with the token-counting beta

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: derhornspieler <15236687+derhornspieler@users.noreply.github.com>
Co-authored-by: mateo-berri <happymvw@gmail.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-03 17:08:30 -07:00 • committed by GitHub
parent dfdd496db8
commit 0ed1c08f02
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
122 changed files with 16215 additions and 2158 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

127
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json
View file

@ -5710,6 +5710,17 @@
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "histogram_quantile(0.95, sum(rate(litellm_anthropic_wif_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "anthropic_wif",
"range": true,
"refId": "A"
},
{
"datasource": {
"type": "prometheus",
@ -5719,7 +5730,7 @@
"expr": "histogram_quantile(0.95, sum(rate(litellm_auth_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "auth",
"range": true,
"refId": "A"
"refId": "B"
},
{
"datasource": {
@ -5730,7 +5741,7 @@
"expr": "histogram_quantile(0.95, sum(rate(litellm_batch_write_to_db_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "batch_write_to_db",
"range": true,
"refId": "B"
"refId": "C"
},
{
"datasource": {
@ -5741,7 +5752,7 @@
"expr": "histogram_quantile(0.95, sum(rate(litellm_postgres_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "postgres",
"range": true,
"refId": "C"
"refId": "D"
},
{
"datasource": {
@ -5752,7 +5763,7 @@
"expr": "histogram_quantile(0.95, sum(rate(litellm_proxy_pre_call_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "proxy_pre_call",
"range": true,
"refId": "D"
"refId": "E"
},
{
"datasource": {
@ -5763,7 +5774,7 @@
"expr": "histogram_quantile(0.95, sum(rate(litellm_redis_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "redis",
"range": true,
"refId": "E"
"refId": "F"
},
{
"datasource": {
@ -5774,7 +5785,7 @@
"expr": "histogram_quantile(0.95, sum(rate(litellm_redis_daily_org_spend_update_queue_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "redis_daily_org_spend_update_queue",
"range": true,
"refId": "F"
"refId": "G"
},
{
"datasource": {
@ -5785,7 +5796,7 @@
"expr": "histogram_quantile(0.95, sum(rate(litellm_redis_daily_tag_spend_update_queue_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "redis_daily_tag_spend_update_queue",
"range": true,
"refId": "G"
"refId": "H"
},
{
"datasource": {
@ -5796,7 +5807,7 @@
"expr": "histogram_quantile(0.95, sum(rate(litellm_redis_daily_team_spend_update_queue_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "redis_daily_team_spend_update_queue",
"range": true,
"refId": "H"
"refId": "I"
},
{
"datasource": {
@ -5807,7 +5818,7 @@
"expr": "histogram_quantile(0.95, sum(rate(litellm_redis_window_spend_update_queue_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "redis_window_spend_update_queue",
"range": true,
"refId": "I"
"refId": "J"
},
{
"datasource": {
@ -5818,7 +5829,7 @@
"expr": "histogram_quantile(0.95, sum(rate(litellm_reset_budget_job_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "reset_budget_job",
"range": true,
"refId": "J"
"refId": "K"
},
{
"datasource": {
@ -5829,7 +5840,7 @@
"expr": "histogram_quantile(0.95, sum(rate(litellm_router_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "router",
"range": true,
"refId": "K"
"refId": "L"
},
{
"datasource": {
@ -5840,7 +5851,7 @@
"expr": "histogram_quantile(0.95, sum(rate(litellm_self_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "self",
"range": true,
"refId": "L"
"refId": "M"
}
],
"title": "Service latency p95 (litellm_<service>_latency)",
@ -5888,6 +5899,28 @@
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(rate(litellm_anthropic_wif_total_requests_total[$__rate_interval]))",
"legendFormat": "anthropic_wif",
"range": true,
"refId": "A"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(rate(litellm_anthropic_wif_cache_total_requests_total[$__rate_interval]))",
"legendFormat": "anthropic_wif_cache",
"range": true,
"refId": "B"
},
{
"datasource": {
"type": "prometheus",
@ -5897,7 +5930,7 @@
"expr": "sum(rate(litellm_auth_total_requests_total[$__rate_interval]))",
"legendFormat": "auth",
"range": true,
"refId": "A"
"refId": "C"
},
{
"datasource": {
@ -5908,7 +5941,7 @@
"expr": "sum(rate(litellm_batch_write_to_db_total_requests_total[$__rate_interval]))",
"legendFormat": "batch_write_to_db",
"range": true,
"refId": "B"
"refId": "D"
},
{
"datasource": {
@ -5919,7 +5952,7 @@
"expr": "sum(rate(litellm_postgres_total_requests_total[$__rate_interval]))",
"legendFormat": "postgres",
"range": true,
"refId": "C"
"refId": "E"
},
{
"datasource": {
@ -5930,7 +5963,7 @@
"expr": "sum(rate(litellm_proxy_pre_call_total_requests_total[$__rate_interval]))",
"legendFormat": "proxy_pre_call",
"range": true,
"refId": "D"
"refId": "F"
},
{
"datasource": {
@ -5941,7 +5974,7 @@
"expr": "sum(rate(litellm_redis_total_requests_total[$__rate_interval]))",
"legendFormat": "redis",
"range": true,
"refId": "E"
"refId": "G"
},
{
"datasource": {
@ -5952,7 +5985,7 @@
"expr": "sum(rate(litellm_redis_daily_org_spend_update_queue_total_requests_total[$__rate_interval]))",
"legendFormat": "redis_daily_org_spend_update_queue",
"range": true,
"refId": "F"
"refId": "H"
},
{
"datasource": {
@ -5963,7 +5996,7 @@
"expr": "sum(rate(litellm_redis_daily_tag_spend_update_queue_total_requests_total[$__rate_interval]))",
"legendFormat": "redis_daily_tag_spend_update_queue",
"range": true,
"refId": "G"
"refId": "I"
},
{
"datasource": {
@ -5974,7 +6007,7 @@
"expr": "sum(rate(litellm_redis_daily_team_spend_update_queue_total_requests_total[$__rate_interval]))",
"legendFormat": "redis_daily_team_spend_update_queue",
"range": true,
"refId": "H"
"refId": "J"
},
{
"datasource": {
@ -5985,7 +6018,7 @@
"expr": "sum(rate(litellm_redis_window_spend_update_queue_total_requests_total[$__rate_interval]))",
"legendFormat": "redis_window_spend_update_queue",
"range": true,
"refId": "I"
"refId": "K"
},
{
"datasource": {
@ -5996,7 +6029,7 @@
"expr": "sum(rate(litellm_reset_budget_job_total_requests_total[$__rate_interval]))",
"legendFormat": "reset_budget_job",
"range": true,
"refId": "J"
"refId": "L"
},
{
"datasource": {
@ -6007,7 +6040,7 @@
"expr": "sum(rate(litellm_router_total_requests_total[$__rate_interval]))",
"legendFormat": "router",
"range": true,
"refId": "K"
"refId": "M"
},
{
"datasource": {
@ -6018,7 +6051,7 @@
"expr": "sum(rate(litellm_self_total_requests_total[$__rate_interval]))",
"legendFormat": "self",
"range": true,
"refId": "L"
"refId": "N"
}
],
"title": "Service request rate (litellm_<service>_total_requests)",
@ -6066,6 +6099,28 @@
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(rate(litellm_anthropic_wif_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "anthropic_wif / {{error_class}}",
"range": true,
"refId": "A"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(rate(litellm_anthropic_wif_cache_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "anthropic_wif_cache / {{error_class}}",
"range": true,
"refId": "B"
},
{
"datasource": {
"type": "prometheus",
@ -6075,7 +6130,7 @@
"expr": "sum(rate(litellm_auth_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "auth / {{error_class}}",
"range": true,
"refId": "A"
"refId": "C"
},
{
"datasource": {
@ -6086,7 +6141,7 @@
"expr": "sum(rate(litellm_batch_write_to_db_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "batch_write_to_db / {{error_class}}",
"range": true,
"refId": "B"
"refId": "D"
},
{
"datasource": {
@ -6097,7 +6152,7 @@
"expr": "sum(rate(litellm_postgres_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "postgres / {{error_class}}",
"range": true,
"refId": "C"
"refId": "E"
},
{
"datasource": {
@ -6108,7 +6163,7 @@
"expr": "sum(rate(litellm_proxy_pre_call_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "proxy_pre_call / {{error_class}}",
"range": true,
"refId": "D"
"refId": "F"
},
{
"datasource": {
@ -6119,7 +6174,7 @@
"expr": "sum(rate(litellm_redis_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "redis / {{error_class}}",
"range": true,
"refId": "E"
"refId": "G"
},
{
"datasource": {
@ -6130,7 +6185,7 @@
"expr": "sum(rate(litellm_redis_daily_org_spend_update_queue_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "redis_daily_org_spend_update_queue / {{error_class}}",
"range": true,
"refId": "F"
"refId": "H"
},
{
"datasource": {
@ -6141,7 +6196,7 @@
"expr": "sum(rate(litellm_redis_daily_tag_spend_update_queue_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "redis_daily_tag_spend_update_queue / {{error_class}}",
"range": true,
"refId": "G"
"refId": "I"
},
{
"datasource": {
@ -6152,7 +6207,7 @@
"expr": "sum(rate(litellm_redis_daily_team_spend_update_queue_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "redis_daily_team_spend_update_queue / {{error_class}}",
"range": true,
"refId": "H"
"refId": "J"
},
{
"datasource": {
@ -6163,7 +6218,7 @@
"expr": "sum(rate(litellm_redis_window_spend_update_queue_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "redis_window_spend_update_queue / {{error_class}}",
"range": true,
"refId": "I"
"refId": "K"
},
{
"datasource": {
@ -6174,7 +6229,7 @@
"expr": "sum(rate(litellm_reset_budget_job_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "reset_budget_job / {{error_class}}",
"range": true,
"refId": "J"
"refId": "L"
},
{
"datasource": {
@ -6185,7 +6240,7 @@
"expr": "sum(rate(litellm_router_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "router / {{error_class}}",
"range": true,
"refId": "K"
"refId": "M"
},
{
"datasource": {
@ -6196,7 +6251,7 @@
"expr": "sum(rate(litellm_self_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "self / {{error_class}}",
"range": true,
"refId": "L"
"refId": "N"
}
],
"title": "Service failure rate (litellm_<service>_failed_requests)",

2
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md
View file

@ -1,6 +1,6 @@
# LiteLLM All Prometheus Metrics dashboard
Every `litellm_*` metric family the proxy can expose on `/metrics` (136 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
Every `litellm_*` metric family the proxy can expose on `/metrics` (141 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/unit/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard

4
litellm/batches/batch_utils.py
View file

@ -17,6 +17,7 @@ from litellm.llms.vertex_ai.batches.transformation import (
)
from litellm.types.llms.openai import Batch
from litellm.types.utils import ModelInfo, Usage
from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS
from litellm.utils import token_counter
@ -543,6 +544,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:

12
litellm/batches/main.py
View file

@ -200,7 +200,7 @@ def create_batch(
LiteLLM Equivalent of POST: https://api.openai.com/v1/batches
"""
try:
optional_params: Final = GenericLiteLLMParams(**kwargs)
optional_params: Final = GenericLiteLLMParams.model_validate(kwargs)
litellm_call_id: Final = kwargs.get("litellm_call_id", None)
proxy_server_request: Final = kwargs.get("proxy_server_request", None)
model_info: Final = kwargs.get("model_info", None)
@ -217,7 +217,7 @@ def create_batch(
)
_is_async: Final = kwargs.pop("acreate_batch", False) is True
litellm_params: Final = dict(GenericLiteLLMParams(**kwargs))
litellm_params: Final = dict(GenericLiteLLMParams.model_validate(kwargs))
litellm_logging_obj: Final[LiteLLMLoggingObj] = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None))
### TIMEOUT LOGIC ###
timeout: Final = _resolve_timeout(optional_params, kwargs, custom_llm_provider)
@ -530,6 +530,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)
response = anthropic_batches_instance.retrieve_batch(
_is_async=_is_async,
batch_id=batch_id,
@ -537,6 +538,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(
@ -573,7 +575,7 @@ def retrieve_batch(
LiteLLM Equivalent of GET https://api.openai.com/v1/batches/{batch_id}
"""
try:
optional_params: Final = GenericLiteLLMParams(**kwargs)
optional_params: Final = GenericLiteLLMParams.model_validate(kwargs)
litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj", None)
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
@ -755,7 +757,7 @@ def list_batches(
"""
try:
# set API KEY
optional_params: Final = GenericLiteLLMParams(**kwargs)
optional_params: Final = GenericLiteLLMParams.model_validate(kwargs)
litellm_params: Final = get_litellm_params(
custom_llm_provider=custom_llm_provider,
**kwargs,
@ -956,7 +958,7 @@ def cancel_batch(
verbose_logger.exception(
"litellm.batches.main.py::cancel_batch() - Error inferring custom_llm_provider - %s", e
)
optional_params: Final = GenericLiteLLMParams(**kwargs)
optional_params: Final = GenericLiteLLMParams.model_validate(kwargs)
litellm_params: Final = get_litellm_params(
custom_llm_provider=custom_llm_provider,
**kwargs,

3
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -11,6 +11,7 @@ from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.llms.openai.data_residency import infer_openai_data_residency
from litellm.types.litellm_params import MAX_CONTROL_INT_DIGITS, ControlOptions
from litellm.types.router import CustomPricingLiteLLMParams
from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS
AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset(
{
@ -70,6 +71,8 @@ OPTIONAL_KWARGS_KEYS: Final = (
}
)
| AWS_CREDENTIAL_KWARGS_KEYS
| ANTHROPIC_WIF_KWARGS_KEYS
| OPENAI_WIF_KWARGS_KEYS
| frozenset(CustomPricingLiteLLMParams.model_fields)
)

18
litellm/llms/anthropic/batches/handler.py
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,
)
)

28
litellm/llms/anthropic/batches/transformation.py
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:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
@ -69,24 +71,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,

4
litellm/llms/anthropic/chat/transformation.py
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 BaseModel, ValidationError
@ -296,6 +296,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

353
litellm/llms/anthropic/common_utils.py
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, TypeVar
from typing import Any, ClassVar, Final, Literal, TypeVar
from urllib.parse import quote
import httpx
from pydantic import BaseModel, ConfigDict, StrictBool, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, Field, StrictBool, TypeAdapter, ValidationError
import litellm
from litellm.constants import (
@ -28,8 +29,15 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
THOUGHT_SIGNATURE_SEPARATOR,
)
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of
from litellm.llms.anthropic.wif import (
aget_anthropic_wif_token,
anthropic_base_without_chat_suffix,
get_anthropic_wif_token,
warn_if_static_credential_shadows_federation,
)
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_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER,
@ -222,6 +230,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:
@ -240,12 +271,22 @@ def resolve_used_client_oauth_token(client_sent_oauth_token: object, custom_llm_
return client_sent_oauth_token and custom_llm_provider in ANTHROPIC_OAUTH_FORWARD_PROVIDERS
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 _beta_header_values(side: str | Sequence[str] | None) -> tuple[str, ...]:
if not side:
return ()
if isinstance(side, str):
return (side,)
return tuple(entry for entry in side if isinstance(entry, str))
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.
"""
joined: Final = ",".join(_beta_header_values(existing) + _beta_header_values(new_beta))
betas: Final = frozenset(b.strip() for b in joined.split(",") if b.strip())
return ",".join(sorted(betas))
@ -272,7 +313,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)
@ -280,7 +323,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
@ -316,7 +361,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)
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
@ -940,7 +1057,9 @@ class AnthropicModelInfo(BaseLLMModelInfo):
return list(set(betas).union(thinking_display_betas, tool_change_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-")
):
@ -948,6 +1067,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,
@ -972,6 +1118,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
is_mid_conversation_output_config_used: bool = False,
is_thinking_display_updates_used: bool = False,
is_mid_conversation_tool_change_used: bool = False,
wif_minted: bool = False,
) -> dict:
betas: Final = set()
# Anthropic no longer requires the prompt-caching beta header
@ -1010,20 +1157,21 @@ class AnthropicModelInfo(BaseLLMModelInfo):
if is_mid_conversation_output_config_used:
betas.add(ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER)
_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)
@ -1055,10 +1203,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)
@ -1067,9 +1216,25 @@ 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:
if (api_key is not None or auth_token is not None) and config_allows_workload_identity(self):
warn_if_static_credential_shadows_federation(params_mapping, model)
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,
)
@ -1095,7 +1260,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,
is_mid_conversation_output_config_used=is_mid_conversation_output_config_used,
@ -1113,11 +1278,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:
@ -1132,9 +1298,13 @@ class AnthropicModelInfo(BaseLLMModelInfo):
@staticmethod
def get_api_key(api_key: str | None = None) -> str | None:
from litellm.secret_managers.main import get_secret_str
"""An empty or whitespace-only key counts as unset: it can never authenticate anything, and
treating it as set would silently outrank workload identity federation."""
from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str
return api_key or get_secret_str("ANTHROPIC_API_KEY")
return normalize_nonempty_secret_str(api_key) or normalize_nonempty_secret_str(
get_secret_str("ANTHROPIC_API_KEY")
)
@staticmethod
def get_auth_token(auth_token: str | None = None) -> str | None:
@ -1143,61 +1313,130 @@ class AnthropicModelInfo(BaseLLMModelInfo):
Unlike api_key (which uses X-Api-Key header), auth_token uses
Authorization: Bearer header, matching the official Anthropic SDK behavior.
"""
from litellm.secret_managers.main import get_secret_str
from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str
return auth_token or get_secret_str("ANTHROPIC_AUTH_TOKEN")
return normalize_nonempty_secret_str(auth_token) or normalize_nonempty_secret_str(
get_secret_str("ANTHROPIC_AUTH_TOKEN")
)
@staticmethod
def get_auth_header(
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 ["anthropic/" + model_id for model_id in model_ids]
def get_token_counter(self) -> BaseTokenCounter | None:
"""

10
litellm/llms/anthropic/count_tokens/handler.py
View file

@ -32,7 +32,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
self,
model: str,
messages: list[dict[str, JsonValue]],
api_key: str,
auth_header: Mapping[str, str],
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
tools: list[dict[str, JsonValue]] | None = None,
@ -45,8 +45,8 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
Args:
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
auth_header: The resolved Anthropic auth header (``AnthropicModelInfo.get_auth_header``)
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:
@ -73,12 +73,12 @@ 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)
# Get required headers
headers: Final = self.get_required_headers(api_key)
headers: Final = self.get_count_tokens_headers(auth_header)
# Use LiteLLM's async httpx client
async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.ANTHROPIC)

31
litellm/llms/anthropic/count_tokens/token_counter.py
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,31 @@ 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
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")
try:
auth_header: Final = await AnthropicModelInfo.aget_auth_header(
api_key=litellm_params.get("api_key"),
api_base=api_base,
litellm_params=litellm_params,
allow_workload_identity=True,
)
if auth_header is None:
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,
auth_header=auth_header,
api_base=api_base,
tools=tools,
system=system,
)
@ -80,8 +83,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,

42
litellm/llms/anthropic/count_tokens/transformation.py
View file

@ -11,6 +11,8 @@ from typing import Final
from pydantic import JsonValue, TypeAdapter
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers
from litellm.llms.anthropic.wif import resolve_anthropic_base
_COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue])
COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config")
@ -26,14 +28,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,
@ -64,28 +73,19 @@ class AnthropicCountTokensConfig:
)
)
def get_required_headers(self, api_key: str) -> dict[str, str]:
"""
Get the required headers for the CountTokens API.
Args:
api_key: The Anthropic API key
Returns:
Dictionary of required headers
"""
from litellm.llms.anthropic.common_utils import (
optionally_handle_anthropic_oauth,
)
headers: dict[str, str] = {
def get_count_tokens_headers(self, auth_header: Mapping[str, str]) -> dict[str, str]:
"""The count-tokens headers around a resolved Anthropic auth header
(``AnthropicModelInfo.get_auth_header``): x-api-key for a static key, an Authorization
bearer for ``ANTHROPIC_AUTH_TOKEN`` and for sk-ant-oat tokens, whose mandatory oauth beta
merges with the token-counting beta instead of replacing it."""
return {
"Content-Type": "application/json",
"x-api-key": api_key,
"anthropic-version": "2023-06-01",
"anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION,
**auth_header,
"anthropic-beta": merge_anthropic_beta_headers(
auth_header.get("anthropic-beta"), ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
),
}
headers, _ = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key)
return headers
def validate_request(
self,

13
litellm/llms/anthropic/files/handler.py
View file

@ -1,7 +1,7 @@
import asyncio
import json
import time
from collections.abc import Coroutine
from collections.abc import Coroutine, Mapping
from typing import Final
import httpx
@ -43,6 +43,7 @@ class AnthropicFilesHandler:
api_key: str | None = None,
timeout: float | httpx.Timeout = 600.0,
max_retries: int | None = None,
litellm_params: Mapping[str, object] | None = None,
) -> HttpxBinaryResponseContent:
"""
Async: Retrieve file content from Anthropic.
@ -56,6 +57,7 @@ class AnthropicFilesHandler:
api_key: Anthropic API key
timeout: Request timeout
max_retries: Max retry attempts (unused for now)
litellm_params: Deployment params, so a named credential's federation settings reach the mint
Returns:
HttpxBinaryResponseContent: Binary content wrapped in compatible response format
@ -73,7 +75,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, litellm_params=litellm_params, allow_workload_identity=True
)
if auth_header is None:
raise ValueError("Missing Anthropic API Key")
@ -116,6 +120,7 @@ class AnthropicFilesHandler:
api_key: str | None = None,
timeout: float | httpx.Timeout = 600.0,
max_retries: int | None = None,
litellm_params: Mapping[str, object] | None = None,
) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]:
"""
Retrieve file content from Anthropic.
@ -130,6 +135,7 @@ class AnthropicFilesHandler:
api_key: Anthropic API key
timeout: Request timeout
max_retries: Max retry attempts (unused for now)
litellm_params: Deployment params, so a named credential's federation settings reach the mint
Returns:
HttpxBinaryResponseContent or Coroutine: Binary content wrapped in compatible response format
@ -139,7 +145,9 @@ class AnthropicFilesHandler:
file_content_request=file_content_request,
api_base=api_base,
api_key=api_key,
timeout=timeout,
max_retries=max_retries,
litellm_params=litellm_params,
)
else:
return asyncio.run(
@ -149,6 +157,7 @@ class AnthropicFilesHandler:
api_key=api_key,
timeout=timeout,
max_retries=max_retries,
litellm_params=litellm_params,
)
)

62
litellm/llms/anthropic/files/transformation.py
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"]

124
litellm/llms/anthropic/pass_through/messages/transformation.py
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
@ -25,6 +25,7 @@ from litellm.types.router import GenericLiteLLMParams
from ...common_utils import (
AnthropicError,
AnthropicModelInfo,
merge_anthropic_beta_headers,
optionally_handle_anthropic_oauth,
requires_native_compaction_beta,
strip_advisor_blocks_from_messages,
@ -38,6 +39,17 @@ from .mid_conversation_system import (
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 "
@ -55,6 +67,8 @@ def _messages_carry_output_config(messages: Sequence[object]) -> bool:
class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
_workload_identity_eligible: ClassVar[bool] = True
@property
def custom_llm_provider(self) -> str | None:
return "anthropic"
@ -256,33 +270,109 @@ 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, messages), 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, messages), 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, # mutable-ok: out-param
optional_params: dict, # mutable-ok: out-param
messages: list[Any], # mutable-ok: mirrors the validate_anthropic_messages_environment contract
) -> 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,
messages=messages,
)
return headers, api_base
@staticmethod
def _translate_reasoning_effort_to_anthropic(
model: str, optional_params: dict, max_tokens: int | None, custom_llm_provider: str

8
litellm/llms/anthropic/prompt_cache_prediction.py
View file

@ -524,17 +524,19 @@ async def count_prompt_tokens(
body: Mapping[str, JsonValue],
api_base: str | None = None,
) -> int | None:
auth_header: Final = AnthropicModelInfo.get_auth_header(api_key=api_key, api_base=api_base)
if auth_header is None:
return None
try:
native: Final = _CountBody.model_validate(body)
count_url: Final = _messages_url(model, api_key, api_base) + "/count_tokens"
result: Final = _CountResult.model_validate(
await _counter.handle_count_tokens_request(
model=model,
messages=_count_objects(native.messages),
tools=_count_objects(native.tools) if native.tools is not None else None,
system=_JSON_OBJECT.validate_python(MappingProxyType({"system": native.system}))["system"],
api_key=api_key,
api_base=count_url,
auth_header=auth_header,
api_base=api_base,
optional_params=_JSON_OBJECT.validate_python(
MappingProxyType({key: body[key] for key in COUNT_TOKEN_OPTION_NAMES if key in body})
),

56
litellm/llms/anthropic/skills/transformation.py
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 {
**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,

604
litellm/llms/anthropic/wif.py Normal file
View file

@ -0,0 +1,604 @@
"""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 dataclasses import dataclass
from functools import lru_cache
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._logging import verbose_logger
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,
internal_issuer_jwks_document,
)
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/"
_SHADOWED_DEPLOYMENT_WARNING_CAP: Final = 512
_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"
_SCHEME_DEFAULT_PORTS: Final[Mapping[str, int]] = MappingProxyType({"http": 80, "https": 443})
_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 (types/workload_identity.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. jti_reused means this"
" identity token was already exchanged once: Anthropic accepts each assertion a single time, so a"
" token file or env var has to rotate before the minted token expires (the rule's"
" token_lifetime_seconds), or switch to the internal issuer or Keycloak source, which mint a"
" fresh assertion per exchange"
)
_WORKSPACE_HINT: Final = (
"If the federation rule is enabled in more than one workspace, set anthropic_federation_workspace_id"
" (or ANTHROPIC_FEDERATION_WORKSPACE_ID) to the wrkspc_ id of the workspace to mint tokens for, or to 'default'."
" Federation does not read ANTHROPIC_WORKSPACE_ID, which the Bedrock Claude platform provider already uses"
)
_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_federation_workspace_id", "ANTHROPIC_FEDERATION_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 _is_set(value)}
)
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 and _is_set(value)}
)
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
@dataclass(frozen=True, slots=True)
class ExportedJwks:
document: str
@dataclass(frozen=True, slots=True)
class NotAnInternalIssuerCredential:
required_param: str
required_value: str
@dataclass(frozen=True, slots=True)
class UnbuildableIdentitySource:
message: str
AnthropicJwksExport = ExportedJwks | NotAnInternalIssuerCredential | UnbuildableIdentitySource
def anthropic_internal_issuer_jwks(credential_values: Mapping[str, object]) -> AnthropicJwksExport:
"""Derive the public JWKS a stored anthropic credential publishes to its federation issuer.
The private signing key stays in this process; only the derived public document comes back."""
if credential_values.get(_IDENTITY_SOURCE_PARAM) != AnthropicIdentitySourceKind.internal_issuer.value:
return NotAnInternalIssuerCredential(
required_param=_IDENTITY_SOURCE_PARAM,
required_value=AnthropicIdentitySourceKind.internal_issuer.value,
)
try:
issuer_source: Final = _build_variant(InternalIssuerSource, credential_values, _INTERNAL_ISSUER_FIELD_MAP)
return ExportedJwks(internal_issuer_jwks_document(issuer_source))
except (litellm.AuthenticationError, ValueError) as e:
return UnbuildableIdentitySource(str(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 _allowlisted_authority(entry: str) -> tuple[str, int | None] | None:
"""One allowlist entry as ``(host, port)``. The port stays ``None`` unless the entry spells one
out, so ``gateway.internal`` trusts that host on every port while ``gateway.internal:8443``
trusts only 8443."""
parts: Final = urlsplit(entry if "://" in entry else f"//{entry}")
try:
port: Final = parts.port
except ValueError:
return None
return (parts.hostname, port) if parts.hostname else None
def _exchange_authority(exchange_base: str) -> tuple[str, int | None]:
"""The host and port an exchange would actually reach, filling in the scheme's default port so
an operator who wrote ``api.anthropic.com:443`` still matches ``https://api.anthropic.com``."""
parts: Final = urlsplit(exchange_base)
try:
port: Final = parts.port
except ValueError:
return "", None
return (parts.hostname or "").lower(), port if port is not None else _SCHEME_DEFAULT_PORTS.get(parts.scheme)
def _trusted_exchange_authorities() -> frozenset[tuple[str, int | None]]:
"""Authorities a federated exchange may reach: Anthropic's own, plus whatever the operator put in
the environment. Comma separated, case folded, each entry a URL, a bare host, or ``host:port``."""
configured: Final = os.getenv(_TRUSTED_EXCHANGE_HOSTS_ENV) or ""
entries: Final = (_allowlisted_authority(entry.strip()) for entry in configured.split(",") if entry.strip())
return frozenset(chain(((_DEFAULT_TRUSTED_EXCHANGE_HOST, None),), (entry for entry in entries if entry)))
def _raise_if_exchange_host_untrusted(exchange_base: str, model: str) -> None:
"""The federated exchange refuses any authority the operator has not vouched for, whatever wrote
the deployment's api_base. Exact host match, never a substring: ``api.anthropic.com.evil.test``
contains the real host and must not pass. An entry naming a port trusts that port alone, so a
second process on another port of an allowed host is refused."""
host, port = _exchange_authority(exchange_base)
if host and any(
host == allowed_host and allowed_port in (None, port)
for allowed_host, allowed_port in _trusted_exchange_authorities()
):
return
refused: Final = f"{host}:{port}" if host and port is not None else host
raise litellm.AuthenticationError(
message=(
f"Anthropic workload identity federation refused to use host {refused 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 host, or host:port to pin the port, "
f"to the {_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,
)
@lru_cache(maxsize=_SHADOWED_DEPLOYMENT_WARNING_CAP)
def _warn_static_credential_shadows_federation(model: str, configured_rule_id: str | None) -> None:
"""Memoized so a shadowed deployment says this once rather than once per request.
The environment fallback is resolved in here rather than by the caller so it too costs one
secret-manager read per deployment: every static-key Anthropic call reaches this, and a
per-request read of a rule id almost nobody sets is an ERROR log with a traceback per call on
the deployments that shadow nothing.
"""
if configured_rule_id is None and _env_str("ANTHROPIC_FEDERATION_RULE_ID") is None:
return
verbose_logger.warning(
"Anthropic deployment %s is configured for workload identity federation, but a static "
"ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN is set and takes precedence, so every call bills "
"that credential and no federated token is minted. Unset it to federate.",
model or "(unnamed)",
)
def warn_if_static_credential_shadows_federation(litellm_params: Mapping[str, object] | None, model: str) -> None:
"""A process-wide static credential outranks federation everywhere in the provider, which is the
Anthropic SDK's own precedence. An operator who configured federation and left a key behind would
otherwise get no signal at all that none of their calls are federated."""
if litellm_params is not None and litellm_params.get(_DISABLE_WIF_PARAM) is True:
return
_warn_static_credential_shadows_federation(model, _param_str(litellm_params, "anthropic_federation_rule_id"))
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(
(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 _is_set(value: object) -> bool:
return value is not None and value != ""
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:
hints: Final = (
_DENIAL_HINT,
"" if workspace_id_set else _WORKSPACE_HINT,
"" if service_account_id_set else _SERVICE_ACCOUNT_HINT,
)
return " " + ". ".join(hint for hint in hints if 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)

2
litellm/llms/azure_ai/embed/handler.py
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

23
litellm/llms/base_llm/anthropic_messages/transformation.py
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,

99
litellm/llms/base_llm/auth/__init__.py Normal file
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",
)

225
litellm/llms/base_llm/auth/client_credentials.py Normal file
View file

@ -0,0 +1,225 @@
"""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]
_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
return HTTPHandler(timeout=httpx.Timeout(timeout=30.0, connect=5.0), follow_redirects=False)
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),
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)

76
litellm/llms/base_llm/auth/identity_source.py Normal file
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>"

86
litellm/llms/base_llm/auth/internal_issuer.py Normal file
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]
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))

115
litellm/llms/base_llm/auth/jwt_signing.py Normal file
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())),
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"]]},
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}
return jwt.encode(dict(claims), key, algorithm=ALG, headers=headers)

181
litellm/llms/base_llm/auth/shared_token_store.py Normal file
View file

@ -0,0 +1,181 @@
"""Same-host token store for the JWT-bearer exchange engine.
Anthropic accepts an assertion carrying a ``jti`` once per issuer, so every uvicorn worker that
reads the same projected token file must share the token the first exchange minted instead of
re-sending the same assertion. The engine keys the store by cache key and only reuses a stored
token minted from the assertion it currently holds; a rotated assertion always buys a fresh token.
"""
import contextlib
import os
import sys
import tempfile
import threading
from collections.abc import Generator
from dataclasses import dataclass
from pathlib import Path
from typing import Final, Protocol
from pydantic import BaseModel, SecretStr, ValidationError
from litellm._logging import verbose_logger
CACHE_DIR_ENV: Final = "LITELLM_TOKEN_EXCHANGE_CACHE_DIR"
@dataclass(frozen=True, slots=True)
class StoredToken:
access_token: SecretStr
expires_at_epoch: float | None
assertion_sha256: str
class SharedTokenStore(Protocol):
"""Every method is best-effort: a store that cannot read, write, or lock degrades to a per-process
cache and never raises into the mint path."""
def load(self, key: str) -> StoredToken | None: ...
def save(self, key: str, token: StoredToken) -> None: ...
def delete(self, key: str) -> None: ...
def lock(self, key: str) -> contextlib.AbstractContextManager[None]: ...
class _StoredTokenFile(BaseModel):
access_token: str
expires_at_epoch: float | None
assertion_sha256: str
def _directory_is_private(directory: Path) -> bool:
try:
directory.mkdir(mode=0o700, exist_ok=True)
stat: Final = directory.stat()
except OSError as e:
verbose_logger.warning("Token exchange cache directory %s is unusable (%s); caching per process", directory, e)
return False
if stat.st_uid != os.getuid() or stat.st_mode & 0o077:
verbose_logger.warning(
"Token exchange cache directory %s must be owned by uid %d with mode 0700; caching per process",
directory,
os.getuid(),
)
return False
return True
def _unlink(path: Path) -> None:
with contextlib.suppress(OSError):
path.unlink()
def _write_token_file(directory: Path, key: str, body: bytes) -> None:
"""The token is staged in its own file and renamed over the entry, so a reader never sees a
half-written one. Every failure unlinks the staging file, including the buffered write that only
reaches the disk when the handle closes: nothing else sweeps this directory, and that file holds
a token that still works. The rename leaves nothing behind for the unlink to find."""
descriptor, name = tempfile.mkstemp(dir=directory, prefix=f"{key}.")
os.close(descriptor)
staged: Final = Path(name)
try:
staged.write_bytes(body)
os.replace(staged, directory / f"{key}.json")
finally:
_unlink(staged)
class FileTokenStore:
"""One ``<cache key>.json`` (mode 0600) and one ``<cache key>.lock`` (flock) per identity under a
directory only the proxy's uid can enter; the directory is checked on first use, not at import."""
def __init__(self, directory: Path) -> None:
self._directory: Final = directory
self._ready_lock: Final = threading.Lock()
self._ready: bool | None = None
@property
def directory(self) -> Path:
return self._directory
def _usable(self) -> bool:
with self._ready_lock:
if self._ready is None:
self._ready = _directory_is_private(self._directory)
return self._ready
def load(self, key: str) -> StoredToken | None:
if not self._usable():
return None
try:
raw: Final = (self._directory / f"{key}.json").read_bytes()
parsed: Final = _StoredTokenFile.model_validate_json(raw)
except FileNotFoundError:
return None
except (OSError, ValidationError) as e:
verbose_logger.debug("Ignoring unreadable token exchange cache entry: %s", e)
return None
return StoredToken(
access_token=SecretStr(parsed.access_token),
expires_at_epoch=parsed.expires_at_epoch,
assertion_sha256=parsed.assertion_sha256,
)
def save(self, key: str, token: StoredToken) -> None:
if not self._usable():
return
body: Final = (
_StoredTokenFile(
access_token=token.access_token.get_secret_value(),
expires_at_epoch=token.expires_at_epoch,
assertion_sha256=token.assertion_sha256,
)
.model_dump_json()
.encode()
)
try:
_write_token_file(self._directory, key, body)
except OSError as e:
verbose_logger.debug("Token exchange cache entry not written: %s", e)
def delete(self, key: str) -> None:
if not self._usable():
return
with contextlib.suppress(FileNotFoundError, OSError):
(self._directory / f"{key}.json").unlink()
@contextlib.contextmanager
def lock(self, key: str) -> Generator[None]:
if sys.platform == "win32" or not self._usable():
yield
return
import fcntl
try:
fd: Final = os.open(self._directory / f"{key}.lock", os.O_RDWR | os.O_CREAT, 0o600)
except OSError as e:
verbose_logger.debug("Token exchange cache lock unavailable (%s); minting without it", e)
yield
return
try:
fcntl.flock(fd, fcntl.LOCK_EX)
yield
finally:
with contextlib.suppress(OSError):
fcntl.flock(fd, fcntl.LOCK_UN)
os.close(fd)
def default_shared_token_store() -> SharedTokenStore | None:
"""``LITELLM_TOKEN_EXCHANGE_CACHE_DIR`` relocates the store; setting it empty disables it. Without
it the store lives under the temp directory, keyed by uid, so the workers of one proxy share it and
other users on the host cannot read it. Windows has no ``flock``, so it caches per process there."""
if sys.platform == "win32":
return None
configured: Final = os.environ.get(CACHE_DIR_ENV)
if configured == "":
return None
if configured is not None:
return FileTokenStore(Path(configured))
return FileTokenStore(Path(tempfile.gettempdir()) / f"litellm-token-exchange-{os.getuid()}")

941
litellm/llms/base_llm/auth/token_exchange.py Normal file
View file

@ -0,0 +1,941 @@
"""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, Iterator, Mapping, Sequence
from concurrent.futures import Executor, ThreadPoolExecutor
from dataclasses import dataclass
from itertools import chain
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.shared_token_store import SharedTokenStore, StoredToken, default_shared_token_store
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))
)
windows: Final = chain.from_iterable(_character_runs(candidate, run) for candidate in compacted_candidates)
return any(window in compacted_secret for window in windows)
def _character_runs(compacted: str, run: int) -> Iterator[str]:
return (compacted[start : start + run] 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(body[name])}" for name in _OAUTH_ERROR_FIELDS if body.get(name) 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(
str(value[key])[:_REDACTION_CAP] for key in _NESTED_ERROR_FIELDS if value.get(key) 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 _shares_one_assertion_across_workers(spec: TokenExchangeSpec) -> bool:
"""The store exists so the workers reading one projected token file don't each spend that file's
single-use ``jti``. A source that mints its own assertion per exchange shares nothing with another
worker, so it never reads the store, never finds a hit there, and keeps its minted token off disk."""
return spec.assertion_source is None
def _assertion_digest(assertion: SecretStr) -> str:
return hashlib.sha256(assertion.get_secret_value().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(
{
**spec.static_body,
spec.assertion_field: assertion.get_secret_value(),
}
).encode()
return urlencode(
{
**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
return HTTPHandler(timeout=httpx.Timeout(timeout=30.0, connect=5.0), follow_redirects=False)
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),
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],
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
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 disarm(self, now: float) -> None:
"""Undo ``arm`` for a refresh that never started. Nothing is on its way to publish, so the
entry must stop reading as in-flight, and the backoff keeps every later caller from
re-attempting a schedule that just failed."""
self.backoff_until = now + ADVISORY_REFRESH_BACKOFF_SECONDS
self.in_flight = False
self.done.set()
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
def _denied(attempt: _Unauthorized) -> TokenEndpointError:
return redact_oauth_error_body(attempt.response.status_code, _capped_body_text(attempt.response), attempt.assertion)
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,
shared_store: SharedTokenStore | None = None,
wall_clock: Callable[[], float] = time.time,
) -> 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._shared_store: Final = shared_store
self._wall_clock: Final = wall_clock
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)
decision = self._classify_and_arm_locked(entry)
match decision:
case _Serve(token=token):
self._report_cache_hit()
return token
case _ServeAndRefresh(token=token):
self._report_cache_hit()
self._submit_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)
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)
if self._shared_store is not None:
self._shared_store.delete(key)
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(thread_name_prefix="litellm-token-exchange-refresh")
return self._refresh_executor
def _submit_advisory_refresh(self, spec: TokenExchangeSpec, entry: _Entry) -> None:
"""The entry is already armed, so an executor that refuses the work would leave it reading
as in-flight with nothing on its way to publish, and every later caller would wait out the
follower timeout and fail. A refused submit disarms it and the cached token keeps serving."""
try:
self._executor_instance().submit(self._advisory_refresh, spec, entry)
except RuntimeError as e:
verbose_logger.debug("token exchange advisory refresh could not be scheduled: %s", e)
with self._lock:
entry.disarm(self._clock())
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:
url_check: Final = validate_token_endpoint_url(spec.token_url)
if isinstance(url_check, InsecureTokenUrl):
return url_check
fetch: Final = _assertion_fetch(self._assertion_reader, spec)
assertion: Final = _read_assertion(fetch, spec.assertion_ref)
if isinstance(assertion, AssertionSourceError):
return assertion
if self._shared_store is None or not _shares_one_assertion_across_workers(spec):
return self._mint(spec, fetch, assertion)
key: Final = _cache_key(spec)
with self._shared_store.lock(key):
shared: Final = self._shared_token(self._shared_store.load(key), _assertion_digest(assertion))
if shared is not None:
return shared
minted: Final = self._mint(spec, fetch, assertion)
if isinstance(minted, MintedToken):
self._shared_store.save(key, self._stored_token(minted))
return minted
def _shared_token(self, stored: StoredToken | None, assertion_sha256: str) -> MintedToken | None:
"""A stored token minted from the very assertion this process holds is the token that assertion
bought: another worker sharing the token file already exchanged it, and an issuer enforcing
single-use ``jti`` would only deny a second exchange."""
if stored is None or stored.assertion_sha256 != assertion_sha256:
return None
if stored.expires_at_epoch is None:
return MintedToken(access_token=stored.access_token, expires_at=None, assertion_sha256=assertion_sha256)
remaining: Final = stored.expires_at_epoch - self._wall_clock()
if remaining <= 0.0:
return None
return MintedToken(
access_token=stored.access_token,
expires_at=self._clock() + remaining,
assertion_sha256=assertion_sha256,
)
def _stored_token(self, token: MintedToken) -> StoredToken:
return StoredToken(
access_token=token.access_token,
expires_at_epoch=(
None if token.expires_at is None else self._wall_clock() + (token.expires_at - self._clock())
),
assertion_sha256=token.assertion_sha256,
)
def _mint(self, spec: TokenExchangeSpec, fetch: AssertionSource, assertion: SecretStr) -> ExchangeResult:
"""One 401 earns one retry, and only with an assertion that changed since the first attempt: a
token file rotated between the read and the POST is worth resending, the same assertion is not,
since an issuer that already consumed its ``jti`` denies it again."""
first: Final = self._post_assertion(spec, assertion)
if not isinstance(first, _Unauthorized):
return first
reread: Final = _read_assertion(fetch, spec.assertion_ref)
if isinstance(reread, AssertionSourceError):
return reread
if reread.get_secret_value() == assertion.get_secret_value():
return _denied(first)
second: Final = self._post_assertion(spec, reread)
if isinstance(second, _Unauthorized):
return _denied(second)
return second
def _post_assertion(self, spec: TokenExchangeSpec, assertion: SecretStr) -> "ExchangeResult | _Unauthorized":
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) -> 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),
assertion_sha256=_assertion_digest(assertion),
)
default_token_exchange_engine: Final = JwtBearerTokenExchangeEngine(shared_store=default_shared_token_store())

100
litellm/llms/base_llm/auth/types.py Normal file
View file

@ -0,0 +1,100 @@
"""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]
AssertionSource: TypeAlias = Callable[[], str | None]
@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
assertion_sha256: str
@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: ...

17
litellm/llms/base_llm/base_utils.py
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
@ -65,6 +66,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:

6
litellm/llms/custom_httpx/http_handler.py
View file

@ -1382,10 +1382,12 @@ class HTTPHandler:
ssl_verify: bool | str | None = None,
disable_default_headers: bool
| None = False, # arize phoenix returns different API responses when user agent header in request
follow_redirects: bool = True,
):
self.timeout = timeout
self.ssl_verify = ssl_verify
self.disable_default_headers = disable_default_headers
self.follow_redirects = follow_redirects
self._owns_client = client is None
self._heal_lock = threading.Lock()
self._client = self.create_client() if client is None else client
@ -1410,7 +1412,7 @@ class HTTPHandler:
cert=cert,
headers=default_headers,
cookies=blocked_cookie_jar(),
follow_redirects=True,
follow_redirects=self.follow_redirects,
http2=http2_enabled(),
)
@ -1436,7 +1438,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,
):

264
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -1,4 +1,5 @@
import asyncio
import inspect
import json
import ssl
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence
@ -18,6 +19,7 @@ from typing import (
Union,
cast,
get_type_hints,
runtime_checkable,
)
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
@ -277,6 +279,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,
)
class _SignedBodyKwargs(TypedDict, total=False):
data: ReadOnly[bytes]
json: ReadOnly[dict[str, object]]
@ -387,6 +438,21 @@ class _PreparedFileContentRequest(NamedTuple):
headers: dict
def _logged_file_content_request(
url: str,
params: dict,
request_headers: dict,
file_content_request: "FileContentRequest",
logging_obj: LiteLLMLoggingObj,
) -> _PreparedFileContentRequest:
logging_obj.pre_call(
input="",
api_key="",
additional_args={"api_base": url, "headers": request_headers, "file_id": file_content_request.get("file_id")},
)
return _PreparedFileContentRequest(url=url, params=params, headers=request_headers)
async def _aiter_bytes_then_close(response: httpx.Response, *, chunk_size: int) -> AsyncGenerator[bytes, None]:
try:
async for chunk in response.aiter_bytes(chunk_size=chunk_size):
@ -2027,7 +2093,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,
@ -3324,6 +3390,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,
@ -3353,18 +3432,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:
@ -3492,6 +3559,54 @@ 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=[],
optional_params={},
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={},
litellm_params=litellm_params,
data=create_file_data,
)
if not complete_api_base:
raise ValueError("api_base is required for create_file")
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={},
),
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"],
@ -3742,6 +3857,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,
@ -3770,19 +3899,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:
@ -3920,6 +4036,56 @@ 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=[],
optional_params={},
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={},
litellm_params=litellm_params,
data=create_batch_data,
)
if not complete_api_base:
raise ValueError("api_base is required for create_batch")
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={},
),
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,
@ -4521,7 +4687,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="",
@ -4646,7 +4813,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="",
@ -4770,7 +4938,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="",
@ -4958,7 +5127,7 @@ class BaseLLMHTTPHandler:
else:
async_httpx_client = client
prepared: Final = self._prepare_file_content_request(
prepared: Final = await self._aprepare_file_content_request(
file_content_request=file_content_request,
provider_config=provider_config,
litellm_params=litellm_params,
@ -5004,7 +5173,7 @@ class BaseLLMHTTPHandler:
client if client is not None else get_async_httpx_client(llm_provider=provider_config.custom_llm_provider)
)
prepared: Final = self._prepare_file_content_request(
prepared: Final = await self._aprepare_file_content_request(
file_content_request=file_content_request,
provider_config=provider_config,
litellm_params=litellm_params,
@ -5062,16 +5231,31 @@ class BaseLLMHTTPHandler:
optional_params={},
litellm_params=litellm_params,
)
logging_obj.pre_call(
input="",
api_key="",
additional_args={
"api_base": url,
"headers": request_headers,
"file_id": file_content_request.get("file_id"),
},
return _logged_file_content_request(url, params, request_headers, file_content_request, logging_obj)
@staticmethod
async def _aprepare_file_content_request(
file_content_request: "FileContentRequest",
provider_config: BaseFilesConfig,
litellm_params: dict,
headers: dict,
logging_obj: LiteLLMLoggingObj,
) -> "_PreparedFileContentRequest":
url, params = provider_config.transform_file_content_request(
file_content_request=file_content_request,
optional_params={},
litellm_params=litellm_params,
)
return _PreparedFileContentRequest(url=url, params=params, headers=request_headers)
request_headers: Final = await _avalidate_files_environment(
provider_config,
api_key=litellm_params.get("api_key"),
headers=headers,
model="",
messages=[],
optional_params={},
litellm_params=litellm_params,
)
return _logged_file_content_request(url, params, request_headers, file_content_request, logging_obj)
def _prepare_fake_stream_request(
self,

53
litellm/llms/openai/chat/gpt_transformation.py
View file

@ -58,6 +58,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:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
@ -73,6 +74,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
@ -765,28 +771,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:

16
litellm/llms/openai/openai.py
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
@ -928,6 +933,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
organization=organization,
client=client,
shared_session=shared_session,
litellm_params=litellm_params,
)
## LOGGING
@ -1024,6 +1030,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))
@ -1037,6 +1044,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
max_retries=max_retries,
organization=organization,
client=client,
litellm_params=litellm_params,
)
## LOGGING
logging_obj.pre_call(
@ -1109,6 +1117,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
organization=organization,
client=client,
shared_session=shared_session,
litellm_params=litellm_params,
)
## LOGGING
logging_obj.pre_call(
@ -1243,6 +1252,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(
@ -1253,6 +1263,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,
@ -1316,6 +1327,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:
@ -1342,6 +1354,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(
@ -1351,6 +1364,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
timeout=timeout,
max_retries=max_retries,
client=client,
litellm_params=litellm_params,
)
## embedding CALL

7
litellm/llms/openai/responses/transformation.py
View file

@ -30,6 +30,7 @@ from litellm.types.llms.openai import *
from litellm.types.responses.main import *
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
from litellm.types.workload_identity import OPENAI_WIF_KWARGS_KEYS
from ..common_utils import OpenAIError
from ..workload_identity import get_workload_identity_bearer_token, resolve_openai_workload_identity_config
@ -599,7 +600,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
)

19
litellm/llms/openai/workload_identity.py
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,
@ -77,6 +81,13 @@ async def get_workload_identity_bearer_token_for_api_base(api_base: str) -> str
return await _workload_identity_auth(config).get_token_async()
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

17
litellm/main.py
View file

@ -144,6 +144,7 @@ from litellm.types.utils import (
RawRequestTypedDict,
StreamingChoices,
)
from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS
from litellm.utils import (
Choices,
CustomStreamWrapper,
@ -5489,7 +5490,9 @@ def completion(
api_base=api_base,
api_key=api_key,
litellm_params=(
GenericLiteLLMParams(**_supplemental_provider_params) if _supplemental_provider_params else None
GenericLiteLLMParams.model_validate(_supplemental_provider_params)
if _supplemental_provider_params
else None
),
)
@ -5716,7 +5719,12 @@ def completion(
gigachat_access_token=kwargs.get("gigachat_access_token"),
**{
key: kwargs[key]
for key in (*AWS_CREDENTIAL_KWARGS_KEYS, PROVIDER_AFFINITY_HEADER_KWARG_KEY)
for key in (
*AWS_CREDENTIAL_KWARGS_KEYS,
*ANTHROPIC_WIF_KWARGS_KEYS,
*OPENAI_WIF_KWARGS_KEYS,
PROVIDER_AFFINITY_HEADER_KWARG_KEY,
)
if key in kwargs
},
)
@ -6553,6 +6561,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")
@ -7802,7 +7811,7 @@ async def amoderation(
# only supports open ai for now
api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
optional_params: Final = GenericLiteLLMParams(**kwargs)
optional_params: Final = GenericLiteLLMParams.model_validate(kwargs)
litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj", None)
_dynamic_api_base = None
try:
@ -8519,7 +8528,7 @@ def speech(
VertexAITextToSpeechConfig,
)
generic_optional_params: Final = GenericLiteLLMParams(**kwargs)
generic_optional_params: Final = GenericLiteLLMParams.model_validate(kwargs)
# Handle Gemini models separately (they use speech_to_completion_bridge)
if "gemini" in model:

7
litellm/models/credentials.py
View file

@ -7,7 +7,7 @@ layer; ``litellm.types.utils`` re-exports them for backwards compatibility.
from collections.abc import Mapping
from pydantic import BaseModel, model_validator
from pydantic import BaseModel, Field, model_validator
class CredentialBase(BaseModel):
@ -17,6 +17,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):
@ -36,3 +40,4 @@ class UpdateCredentialItem(BaseModel):
credential_info: Mapping[str, object]
credential_values: Mapping[str, object] | None = None
model_id: str | None = None
credential_values_to_delete: tuple[str, ...] | None = None

90
litellm/proxy/auth/auth_utils.py
View file

@ -20,6 +20,7 @@ from litellm.constants import (
MINIMUM_CUSTOM_KEY_LENGTH,
STANDARD_CUSTOMER_ID_HEADERS,
)
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.url_utils import (
SSRFError,
@ -34,7 +35,8 @@ from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_me
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
)
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment, server_owned_wif_fields_present
from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params
from litellm.types.utils import CustomPricingLiteLLMParams
@ -228,6 +230,63 @@ 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``.
# 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
_CREDENTIAL_VALUES: Final = TypeAdapter(dict[str, object])
def reject_federated_credential_reference(body: Mapping[str, object]) -> None:
"""Raise ``ValueError`` if a request body picks the federated identity by credential name.
``load_credentials_from_list`` merges a named credential's values into the call, so a body
naming a federated credential moves the token exchange onto that credential's federation rule
and organization exactly as sending the federation fields inline would, which
``reject_server_owned_wif_params`` already refuses. Attaching one is a deployment decision, so
it is refused with no client-side opt-in to relax it, the same as the inline form.
Only credentials already loaded into memory can be resolved here, which is every credential the
proxy would resolve for the call itself: ``load_credentials_from_list`` reads the same list.
``route_manages_deployments`` names the routes this is not applied to, where choosing the
credential a deployment federates through is the point of the call.
"""
named: Final = body.get("litellm_credential_name")
if not isinstance(named, str) or not named:
return
wif_fields: Final = server_owned_wif_fields_present(
_CREDENTIAL_VALUES.validate_python(CredentialAccessor.get_credential_values(named))
)
if wif_fields:
raise ValueError(
f"Rejected Request: litellm_credential_name={named!r} names a credential configured for "
f"workload identity federation ({wif_fields[0]}), which a request body cannot choose. "
"A proxy admin attaches it to a deployment."
)
_DEPLOYMENT_MANAGEMENT_ROUTES: Final[frozenset[str]] = frozenset(
("/model/new", "/model/update", "/model/delete", "/health/test_connection")
)
_DEPLOYMENT_ID_UPDATE_ROUTE: Final = re.compile(r"^/model/[^/]+/update$")
def route_manages_deployments(route: str | None) -> bool:
"""Whether ``route`` configures a deployment instead of calling one.
These are the routes that reach ``ModelManagementAuthChecks.can_user_make_model_call``, where
a federated write is judged by ``_reject_non_admin_wif_write`` against what the write sets and
what the deployment already stores: a proxy admin goes through, anyone else is refused with a
403 naming the field. ``reject_federated_credential_reference`` runs ahead of that gate on
every route, so without this exemption a proxy admin could not attach a federated credential
to a deployment over the API or the Admin UI at all, leaving a static ``config.yaml`` entry as
the only way to configure the feature the rejection tells the caller to go configure.
"""
return route is not None and (
route in _DEPLOYMENT_MANAGEMENT_ROUTES or _DEPLOYMENT_ID_UPDATE_ROUTE.fullmatch(route) is not None
)
_NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body")
# Metadata containers that carry per-request configuration consumed by the
@ -380,12 +439,17 @@ def _check_banned_params(
general_settings: dict,
llm_router: Router | None,
model: str,
*,
manages_deployments: bool = False,
) -> None:
"""Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in.
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)
if not manages_deployments:
reject_federated_credential_reference(body)
for param in _BANNED_REQUEST_BODY_PARAMS:
if param not in body:
continue
@ -472,7 +536,14 @@ def _reject_url_valued_fallback_target(value: str) -> None:
)
def is_request_body_safe(request_body: dict, general_settings: dict, llm_router: Router | None, model: str) -> bool:
def is_request_body_safe(
request_body: dict,
general_settings: dict,
llm_router: Router | None,
model: str,
*,
route: str | None = None,
) -> bool:
"""
Check if the request body is safe.
@ -500,25 +571,27 @@ def is_request_body_safe(request_body: dict, general_settings: dict, llm_router:
"""
if "model_list" in request_body:
raise ValueError("Rejected Request: model_list is not allowed in the request body.")
_check_banned_params(request_body, general_settings, llm_router, model)
manages_deployments: Final = route_manages_deployments(route)
_check_banned_params(request_body, general_settings, llm_router, model, manages_deployments=manages_deployments)
for nested_key in _NESTED_CONFIG_KEYS:
nested = _coerce_metadata_to_dict(request_body.get(nested_key))
if nested is not None:
_check_banned_params(nested, general_settings, llm_router, model)
_check_banned_params(nested, general_settings, llm_router, model, manages_deployments=manages_deployments)
for metadata_key in _NESTED_METADATA_KEYS:
metadata = _coerce_metadata_to_dict(request_body.get(metadata_key))
if metadata is not None:
_check_banned_params(metadata, general_settings, llm_router, model)
_check_banned_params(metadata, general_settings, llm_router, model, manages_deployments=manages_deployments)
if any(isinstance(key, str) and key.startswith(f"{metadata_key}[") for key in request_body):
_check_banned_params(
extract_nested_form_metadata(form_data=request_body, prefix=f"{metadata_key}["),
general_settings,
llm_router,
model,
manages_deployments=manages_deployments,
)
for target in iter_request_fallback_targets(request_body):
if isinstance(target, dict):
_check_banned_params(target, general_settings, llm_router, model)
_check_banned_params(target, general_settings, llm_router, model, manages_deployments=manages_deployments)
target_model = target.get("model")
if isinstance(target_model, str):
_reject_url_valued_fallback_target(target_model)
@ -526,6 +599,9 @@ 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)
if not manages_deployments:
reject_federated_credential_reference(litellm_params)
litellm_params_metadata: Final = _coerce_metadata_to_dict(litellm_params.get("metadata"))
if litellm_params_metadata is not None:
_check_banned_params(
@ -533,6 +609,7 @@ def is_request_body_safe(request_body: dict, general_settings: dict, llm_router:
general_settings,
llm_router,
model,
manages_deployments=manages_deployments,
)
return True
@ -585,6 +662,7 @@ async def pre_db_read_auth_checks(
general_settings=general_settings,
llm_router=llm_router,
model=request_data.get("model", ""), # [TODO] use model passed in url as well (azure openai routes)
route=route,
)
# Check 3. Check if IP address is allowed

189
litellm/proxy/common_utils/credential_hydration.py Normal file
View file

@ -0,0 +1,189 @@
"""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.router_utils.clientside_credential_handler import clientside_credential_keys
from litellm.types.router import (
GenericLiteLLMParams,
server_owned_wif_fields_named,
server_owned_wif_fields_present,
)
from litellm.types.utils import CredentialItem, LlmProviders, server_owned_wif_litellm_params
_LITELLM_PROVIDER_IDS: Final = frozenset(provider.value for provider in LlmProviders)
_FEDERATION_SURFACE_FIELDS: Final = frozenset(
(
*clientside_credential_keys,
"configurable_clientside_auth_params",
"litellm_credential_name",
*server_owned_wif_litellm_params,
)
)
def write_touches_federation_surface(incoming: Mapping[str, object] | None) -> bool:
"""Whether this write can move or re-scope the token a federated deployment mints.
Three groups of fields can. The federation parameters choose which server-side secret is read
and what the minted token is scoped to. ``litellm_credential_name`` resolves to those same
parameters by reference. ``api_key``, ``api_base``, ``base_url``, and the
``configurable_clientside_auth_params`` that let a caller override them decide where the
resulting token is sent. A write setting none of them leaves the federation configuration
exactly as the proxy admin left it, so renaming a federated deployment or changing its rpm
stays an ordinary team-admin edit.
"""
return incoming is not None and not _FEDERATION_SURFACE_FIELDS.isdisjoint(incoming.keys())
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.
"""
matching: Final = tuple(c for c in litellm.credential_list if c.credential_name == credential_name)
in_memory: Final = tuple(chain.from_iterable(server_owned_wif_fields_named(c.credential_values) for c in matching))
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))
def submitted_litellm_params(params: GenericLiteLLMParams | None) -> Mapping[str, object] | None:
"""The fields a pydantic write actually set, as the mapping the federation gate reads.
Only the set fields belong here: ``GenericLiteLLMParams`` declares every federation field, so
the whole model would report every write as touching all of them.
"""
if params is None:
return None
return MappingProxyType({name: getattr(params, name, None) for name in params.model_fields_set})
async def effective_server_owned_wif_fields(
stored: Mapping[str, object] | None,
incoming: Mapping[str, object] | 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, since it is
a full deployment and a declared-but-unset field is not a federation field it carries.
``incoming`` holds only 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.keys())
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: Mapping[str, object] | 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.get("litellm_credential_name")
return tuple(dict.fromkeys(name for name in (from_stored, from_incoming) if isinstance(name, str)))

232
litellm/proxy/credential_endpoints/endpoints.py
View file

@ -9,26 +9,134 @@ from typing import (
cast, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict
)
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response, status
from pydantic import TypeAdapter
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.llms.anthropic.wif import (
ExportedJwks,
NotAnInternalIssuerCredential,
UnbuildableIdentitySource,
anthropic_internal_issuer_jwks,
)
from litellm.models.credentials import UpdateCredentialItem
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy._types import (
CommonProxyErrors,
LitellmUserRoles,
ProxyErrorTypes,
ProxyException,
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.base_repository import is_unique_violation
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()
_CREDENTIAL_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
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 ProxyException(
message=(
f"Only proxy admins can change {wif_fields[0]!r}, a server-owned workload identity federation parameter."
),
type=ProxyErrorTypes.auth_error.value,
code=status.HTTP_403_FORBIDDEN,
param=wif_fields[0],
)
def _incoming_wif_fields(incoming_values: Mapping[str, object], credential: UpdateCredentialItem) -> tuple[str, ...]:
"""WIF fields the request touches: the ones its values set (to any value, ``None`` included,
since the key alone is what the federation resolver reacts to), whether the caller sent them
or named a deployment through ``model_id`` for the proxy to copy them from, 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(incoming_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: UpdateCredentialItem) -> None:
overlap: Final = frozenset(credential.credential_values or ()) & 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 _without_null_values(credential_values: Mapping[str, object]) -> dict[str, object]:
"""A null carries no credential, and the federation resolver refuses a foreign variant's field by
KEY, so a stored ``{"anthropic_issuer_url": null}`` wedges every deployment that names this
credential. ``model_dump(exclude_none=True)`` cannot do this: it drops the model's own null
fields, and ``credential_values`` is a mapping inside one of them.
"""
return {key: value for key, value in credential_values.items() if value is not None}
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(_without_null_values(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:
@ -108,13 +216,17 @@ 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_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_DICT_ADAPTER.validate_python(credential_values),
credential_values=_without_null_values(_CREDENTIAL_DICT_ADAPTER.validate_python(credential_values)),
credential_info=credential.credential_info,
)
encrypted_credential: Final = CredentialHelperUtils.encrypt_credential_values(processed_credential)
credentials_dict: Final = encrypted_credential.model_dump()
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)
)
@ -204,6 +316,65 @@ 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"],
)
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={"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={"error": f"No anthropic credential named {credential_name!r}."},
)
match anthropic_internal_issuer_jwks(credential.credential_values):
case ExportedJwks(document):
return Response(content=document, media_type="application/json")
case NotAnInternalIssuerCredential(required_param, required_value):
raise HTTPException(
status_code=404,
detail={
"error": (
f"Credential {credential_name!r} is not configured with "
f"{required_param}={required_value!r}."
)
},
)
case UnbuildableIdentitySource(message):
raise HTTPException(
status_code=400,
detail={"error": message},
)
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)],
@ -268,6 +439,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(
@ -307,9 +481,10 @@ def update_db_credential(
# update litellm params
if encrypted_credential.credential_values:
# Encrypt any sensitive values
encrypted_params: Final = {k: v for k, v in encrypted_credential.credential_values.items()}
merged_credential.credential_values.update(_without_null_values(encrypted_credential.credential_values))
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:
@ -340,6 +515,13 @@ async def update_credential(
from litellm.proxy.proxy_server import prisma_client
try:
_reject_overlapping_credential_values(credential)
incoming_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
_resolve_deployment_credentials(llm_router, credential.model_id)
if credential.model_id
else credential.credential_values or {}
)
_reject_non_admin_wif_fields(_incoming_wif_fields(incoming_values, credential), user_api_key_dict)
if prisma_client is None:
raise HTTPException(
status_code=500,
@ -349,18 +531,20 @@ 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)
patch: Final = CredentialItem(
credential_name=credential.credential_name,
credential_info=_CREDENTIAL_DICT_ADAPTER.validate_python(credential.credential_info),
credential_values=_CREDENTIAL_DICT_ADAPTER.validate_python(
_resolve_deployment_credentials(llm_router, credential.model_id)
if credential.model_id
else credential.credential_values or {}
),
credential_values=incoming_values,
credential_values_to_delete=credential.credential_values_to_delete,
)
merged_credential: Final = update_db_credential(db_credential, patch)
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,
@ -371,29 +555,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 patch.credential_values:
in_memory_values.update(patch.credential_values)
in_memory_info: Final = dict(existing_in_memory.credential_info or {})
if patch.credential_info:
in_memory_info.update(patch.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(patch, credential_name, merged_credential.credential_name)
return {"success": True, "message": "Credential updated successfully"}
except Exception as e:

14
litellm/proxy/health_check.py
View file

@ -32,11 +32,17 @@ from litellm.router_utils.auto_router_model_naming import (
classify_strategy_router_model,
strategy_router_dependencies,
)
from litellm.types.utils import secret_bearing_wif_litellm_params, server_owned_wif_litellm_params
# Provider routing fields. Allowed for proxy admins so they can see which
# region/version a deployment is checking; gated at the endpoint layer for
# non-admin callers (see _strip_admin_only_fields_from_health_result).
ADMIN_ONLY_HEALTH_DISPLAY_PARAMS: Final = ("api_base", "api_version", "aws_bedrock_runtime_endpoint")
# Provider routing and workload identity federation fields. Allowed for proxy admins so they can
# see which region/version a deployment is checking and which identity it federates as; gated at
# the endpoint layer for non-admin callers (see _strip_admin_only_fields_from_health_result).
ADMIN_ONLY_HEALTH_DISPLAY_PARAMS: Final = (
"api_base",
"api_version",
"aws_bedrock_runtime_endpoint",
*(name for name in server_owned_wif_litellm_params if name not in secret_bearing_wif_litellm_params),
)
MINIMAL_DISPLAY_PARAMS: Final = frozenset({"model", "mode_error"})

16
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -42,6 +42,7 @@ from litellm.proxy.auth.auth_checks 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.model_checks import get_key_models
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -956,9 +957,9 @@ def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
def _strip_admin_only_fields_from_health_result(result: dict) -> dict:
"""
Return a copy of the /health response with provider routing fields
(``ADMIN_ONLY_HEALTH_DISPLAY_PARAMS``) removed from each healthy/unhealthy
endpoint entry. Used to hide those fields from non-admin callers while
Return a copy of the /health response with the admin-only fields (provider routing plus the
workload identity federation params naming the identity a deployment mints as) removed from
each healthy/unhealthy endpoint entry. Used to hide those fields from non-admin callers while
still showing them which deployments they own and whether each one is
healthy. Proxy admins receive the unmodified result.
"""
@ -2203,6 +2204,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(
@ -2228,12 +2230,16 @@ async def test_model_connection(
await ModelManagementAuthChecks.can_user_make_model_call(
model_params=Deployment(
model_name="test_model",
litellm_params=LiteLLM_Params(**litellm_params),
litellm_params=LiteLLM_Params.model_validate(litellm_params),
model_info=resolved_model_info,
),
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
premium_user=premium_user,
# The probe is a write of the caller's own params onto the stored deployment, so the
# caller's params are the incoming side: a probe that redirects a federated
# deployment's api_base is an admin's action, an unmodified probe of it is not.
incoming_params=request_litellm_params,
)
raw_params_mode: Final[object] = litellm_params.pop("mode", None)
probe_mode: Final = (
@ -2260,7 +2266,7 @@ async def test_model_connection(
"result": cleaned_result,
}
except HTTPException as e:
except (HTTPException, ProxyException) as e:
raise e
except Exception as e:
verbose_proxy_logger.debug("litellm.proxy.health_endpoints.test_model_connection(): Exception occurred - %s", e)

102
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -63,6 +63,12 @@ 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,
submitted_litellm_params,
write_touches_federation_surface,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
@ -362,6 +368,27 @@ 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 out
from under the proxy admin.
Only a blocking value is refused. A create that sends ``blocked: false`` asks for the state
every create already lands in, and dashboards and SDKs send the whole model shape on every
create, so refusing the flag's presence would turn a working non-admin create into a 403.
"""
if blocked 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",
)
def _stored_credential_name(existing_litellm_params: GenericLiteLLMParams | None) -> str | None:
if existing_litellm_params is None or existing_litellm_params.litellm_credential_name is None:
return None
@ -437,6 +464,20 @@ def _effective_complexity_router_config(
).effective
def _decrypted_litellm_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, object]:
dumped: Final[Mapping[str, object]] = litellm_params.model_dump(exclude_none=True)
return MappingProxyType(
{
name: (
decrypt_value_helper(value=value, key=name, exception_type="debug", return_original_value=True)
if isinstance(value, str)
else value
)
for name, value in dumped.items()
}
)
def _effective_model(
incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None
) -> str | None:
@ -1161,6 +1202,7 @@ async def patch_model(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
premium_user=premium_user,
incoming_params=submitted_litellm_params(patch_data.litellm_params),
member_operation="update",
incoming_model_params=patch_data,
)
@ -1524,6 +1566,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)
@ -2105,12 +2149,51 @@ class ModelManagementAuthChecks:
)
return True
@staticmethod
async def _reject_non_admin_wif_write(
*,
model_params: Deployment,
incoming_params: Mapping[str, object] | None,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
) -> None:
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return
if not write_touches_federation_surface(incoming_params):
return
stored: Final = _decrypted_litellm_params(model_params.litellm_params)
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 change the credentials of a deployment configured for "
f"workload identity 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.
named: Final = None if incoming_params is None else incoming_params.get("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: Mapping[str, object] | None,
allow_missing_team: bool = False,
member_operation: Literal["create", "update"] | None = None,
incoming_model_params: updateDeployment | None = None,
@ -2120,6 +2203,19 @@ class ModelManagementAuthChecks:
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
):
raise HTTPException(status_code=403, detail="View-only users cannot manage models.")
# 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 point a federated deployment somewhere else.
# Evaluated on the RESULTING deployment: a patch attaching a credential by name inherits
# whatever that credential holds. `incoming_params` carries only the fields the write set
# and 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.team_id is not None:
team_obj_row: Final = await _repo_team_table(prisma_client).find_unique(
@ -2233,6 +2329,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,
)
@ -2355,7 +2452,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",
@ -2424,10 +2520,13 @@ async def add_new_model(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
premium_user=premium_user,
incoming_params=submitted_litellm_params(model_params.litellm_params),
member_operation="create",
)
member_write: Final = write_authorization if isinstance(write_authorization, MemberAutoRouterWrite) else None
_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,
@ -2625,6 +2724,7 @@ async def update_model(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
premium_user=premium_user,
incoming_params=submitted_litellm_params(model_params.litellm_params),
member_operation="update",
incoming_model_params=model_params,
)

23
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -45,7 +45,7 @@ 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 AnthropicModelInfo, merge_anthropic_beta_headers
from litellm.llms.azure.passthrough.transformation import (
foreign_azure_deployment,
is_azure_body_model_inference_endpoint,
@ -914,7 +914,9 @@ 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
)
endpoint_func: Final = create_pass_through_route(
endpoint=endpoint,
target=str(updated_url),
@ -2645,9 +2647,20 @@ def _upstream_headers_for_anthropic_route(
caller_headers: Final = _caller_headers_without_litellm_secrets(
request, user_api_key_dict, _HEADERS_NEVER_FORWARDED_TO_ANTHROPIC
)
if proxy_auth_header is None and _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS.isdisjoint(caller_headers):
raise HTTPException(status_code=401, detail=_CREDENTIALLESS_ANTHROPIC_MISSING_CREDENTIAL_DETAIL)
return MappingProxyType({**caller_headers, **(proxy_auth_header or {})})
if proxy_auth_header is None:
if _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS.isdisjoint(caller_headers):
raise HTTPException(status_code=401, detail=_CREDENTIALLESS_ANTHROPIC_MISSING_CREDENTIAL_DETAIL)
return caller_headers
forwarded: Final = MappingProxyType(
{name: value for name, value in caller_headers.items() if name not in _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS}
)
caller_beta, credential_beta = caller_headers.get("anthropic-beta"), proxy_auth_header.get("anthropic-beta")
merged_beta: Final = (
{"anthropic-beta": merge_anthropic_beta_headers(caller_beta, credential_beta)}
if caller_beta and credential_beta
else {}
)
return MappingProxyType({**forwarded, **proxy_auth_header, **merged_beta})
def _upstream_headers_for_bedrock_agent_runtime_route(

20
litellm/proxy/proxy_server.py
View file

@ -413,6 +413,7 @@ from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_pr
from litellm.proxy.common_utils.codex_model_catalog import codex_model_list_body
from litellm.proxy.common_utils.config_includes import resolve_include_file_path, resolve_includes
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.discoverable_model_filter import discoverable_rows, undiscoverable_model_names
@ -901,6 +902,7 @@ from litellm.types.router import (
RoutingGroup,
RoutingPlugin,
SearchToolTypedDict,
holds_secret_pointer,
updateDeployment,
)
from litellm.types.router import ModelInfo as RouterModelInfo
@ -5793,11 +5795,11 @@ class ProxyConfig:
return config
return {
key: self._resolved_config_value(value=value, depth=depth, max_depth=max_depth)
key: self._resolved_config_value(key=key, value=value, depth=depth, max_depth=max_depth)
for key, value in config.items()
}
def _resolved_config_value(self, value: object, depth: int, max_depth: int) -> object:
def _resolved_config_value(self, key: str, value: object, depth: int, max_depth: int) -> object:
if isinstance(value, dict):
return self._check_for_os_environ_vars(config=value, depth=depth + 1, max_depth=max_depth)
if isinstance(value, list):
@ -5807,7 +5809,7 @@ class ProxyConfig:
else item
for item in value
]
if isinstance(value, str) and value.startswith("os.environ/"):
if isinstance(value, str) and value.startswith("os.environ/") and not holds_secret_pointer(key):
resolved: Final = 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)
@ -6938,7 +6940,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)
@ -7349,7 +7351,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 ##
@ -7394,7 +7396,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
@ -9118,7 +9124,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

22
litellm/router.py
View file

@ -307,6 +307,7 @@ from litellm.types.router import (
RoutingStrategy,
SearchToolTypedDict,
TaggedPreRoutingStrategy,
holds_secret_pointer,
)
from litellm.types.services import ServiceTypes
from litellm.types.utils import (
@ -3987,7 +3988,7 @@ class Router:
model_info["original_model_id"] = original_model_id
deployment_pydantic_obj: Final = Deployment(
model_name=model_group,
litellm_params=LiteLLM_Params(**dynamic_litellm_params),
litellm_params=LiteLLM_Params.model_validate(dynamic_litellm_params),
model_info=model_info,
)
Router._register_deployment_pricing(deployment=deployment_pydantic_obj)
@ -9000,13 +9001,10 @@ class Router:
if access_windows_error is not None:
raise ValueError(access_windows_error)
zeroed_pricing: Final = zeroed_ptu_pricing(_model_info, _litellm_params) if config_sourced else None
litellm_params: Final[LiteLLM_Params] = LiteLLM_Params(
**( # pyright: ignore[reportArgumentType] # untyped merged dict; already true for every field here
_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,
@ -9362,7 +9360,7 @@ class Router:
continue
deployment = Deployment(
model_name=model_name,
litellm_params=(lp if not isinstance(lp, dict) else LiteLLM_Params(**lp)),
litellm_params=(lp if not isinstance(lp, dict) else LiteLLM_Params.model_validate(lp)),
model_info=(entry.get("model_info") if isinstance(entry, dict) else entry.model_info),
)
if self._has_registered_strategy(self.adaptive_routers, model_name, self._deployment_tags(deployment)):
@ -9584,7 +9582,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", {})
@ -10791,7 +10789,7 @@ class Router:
if isinstance(litellm_params_data, LiteLLM_Params):
litellm_params = litellm_params_data
elif isinstance(litellm_params_data, dict) and "model" in litellm_params_data:
litellm_params = LiteLLM_Params(**litellm_params_data)
litellm_params = LiteLLM_Params.model_validate(litellm_params_data)
else:
raise ValueError(
f"Deployment missing valid litellm_params. "
@ -12633,7 +12631,7 @@ class Router:
if allowed_model_region is not None:
if not is_region_allowed(
litellm_params=LiteLLM_Params(**_litellm_params),
litellm_params=LiteLLM_Params.model_validate(_litellm_params),
allowed_model_region=allowed_model_region,
):
invalid_model_indices.add(idx)
@ -12651,7 +12649,7 @@ class Router:
_,
) = litellm.get_llm_provider(
model=_dep_model_for_params,
litellm_params=LiteLLM_Params(**_litellm_params),
litellm_params=LiteLLM_Params.model_validate(_litellm_params),
)
except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not fail the request
verbose_router_logger.debug(

17
litellm/router_utils/clientside_credential_handler.py
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

10
litellm/router_utils/fallback_event_handlers.py
View file

@ -28,7 +28,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
@ -640,6 +640,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

8
litellm/secret_managers/main.py
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."
)

31
litellm/types/litellm_params.py
View file

@ -104,10 +104,41 @@ class BedrockBatchConnection:
bedrock_tags: Sequence[Mapping[str, str]] | None = None
@dataclass(frozen=True, slots=True, kw_only=True)
class AnthropicFederationConnection:
anthropic_federation_rule_id: str | None = None
anthropic_organization_id: str | None = None
anthropic_service_account_id: str | None = None
anthropic_federation_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
anthropic_disable_workload_identity_federation: bool | None = None
@dataclass(frozen=True, slots=True, kw_only=True)
class OpenAIFederationConnection:
openai_identity_provider_id: str | None = None
openai_service_account_id: str | None = None
openai_identity_token_file: str | None = None
@dataclass(frozen=True, slots=True, kw_only=True)
class ConnectionSettings:
provider: ProviderConnection
bedrock_batch: BedrockBatchConnection
anthropic_federation: AnthropicFederationConnection
openai_federation: OpenAIFederationConnection
@dataclass(frozen=True, slots=True, kw_only=True)

1
litellm/types/llms/anthropic.py
View file

@ -784,5 +784,6 @@ ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER: Final = "fine-grained-tool-st
# 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"

91
litellm/types/router.py
View file

@ -4,7 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc
import datetime
import enum
from collections.abc import Mapping, Sequence
from collections.abc import Container, Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
@ -34,6 +34,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):
@ -374,6 +378,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_federation_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__"})
@ -648,6 +715,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=())
@ -1243,3 +1313,22 @@ class AdaptiveRouterPreferences(BaseModel):
quality_tier: int = Field(ge=1, le=3)
strengths: list[RequestType] = Field(default_factory=list)
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 settable inline on a client surface, with or without a client-side credential
opt-in. Naming a stored credential that already holds them is the other way in and has its own
gate: ``_check_banned_params`` resolves ``litellm_credential_name`` and refuses a federated one.
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. A proxy admin configures it on the deployment "
"or on a stored credential."
)

5
litellm/types/services.py
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,9 @@ 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: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]},
# cache hits are counter-only: no HTTP call happens, so observing a latency would be a lie
ServiceTypes.ANTHROPIC_WIF_CACHE.value: {"metrics": [ServiceMetrics.COUNTER]},
# 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]},

10
litellm/types/utils.py
View file

@ -55,6 +55,11 @@ from litellm.types.llms.base import (
LiteLLMPydanticObjectBase,
)
from litellm.types.mcp import MCPServerCostInfo
from litellm.types.workload_identity import (
ANTHROPIC_WIF_KWARGS_KEYS,
OPENAI_WIF_KWARGS_KEYS,
WIF_SECRET_BEARING_KEYS,
)
from ..litellm_core_utils.core_helpers import map_finish_reason, process_response_headers
from . import litellm_params as _litellm_params
@ -3958,6 +3963,11 @@ bedrock_batch_litellm_params: Final = BEDROCK_BATCH_KWARG_NAMES
TRUSTED_CALLBACK_VARS_FIELD: Final = _litellm_params.TRUSTED_CALLBACK_VARS_FIELD
ADDRESSED_RESPONSE_ID_FIELD: Final = _litellm_params.ADDRESSED_RESPONSE_ID_FIELD
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
secret_bearing_wif_litellm_params: Final = tuple(sorted(WIF_SECRET_BEARING_KEYS))
all_litellm_params = [ # rebind-ok: two star imports in litellm/__init__.py re-bind it
*OWNED_KWARG_NAMES,
*KWARG_ARTIFACTS,

35
litellm/types/workload_identity.py Normal file
View file

@ -0,0 +1,35 @@
"""litellm_params keys that configure workload identity federation.
The key sets derive from the federation connection leaves in ``types.litellm_params``, so the
kwargs funnel (``litellm_core_utils.get_litellm_params``) and the request-body ban list
(``types.utils.all_litellm_params``) read one declaration. Every key rides the funnel into
``litellm_params`` and is banned from request bodies, which also covers
``anthropic_disable_workload_identity_federation``: the proxy sets it when a client redirects
``api_base`` so a federated deployment stops minting for a base the caller chose, and a caller
must not be able to set it in either direction.
"""
from typing import Final
from litellm.types.litellm_params import AnthropicFederationConnection, OpenAIFederationConnection, wire_names
ANTHROPIC_WIF_KWARGS_KEYS: Final = frozenset(wire_names(AnthropicFederationConnection))
OPENAI_WIF_KWARGS_KEYS: Final = frozenset(wire_names(OpenAIFederationConnection))
WIF_SECRET_BEARING_KEYS: Final = frozenset(
{
"anthropic_identity_token",
"anthropic_identity_token_file",
"anthropic_issuer_signing_key_ref",
"anthropic_keycloak_client_secret_ref",
"openai_identity_token_file",
}
)
"""The federation keys whose value is a credential, or the path or reference that reaches one.
The rest of the sets above name a rule, an organization, a workspace, or a URL: an operator has
to be able to read those back to tell what a deployment federates as. These carry the secret
itself, so no surface displays them to anyone. Splitting the sensitivity out here keeps the
callers that redact them from having to know which provider a field belongs to.
"""

5
litellm/utils.py
View file

@ -7783,9 +7783,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)

1
terraform/provider/tools/endpointaudit/coverage_allowlist.txt
View file

@ -65,6 +65,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}

34
tests/unit/batches/test_batch_utils.py
View file

@ -1890,6 +1890,40 @@ def test_unparsable_bedrock_batch_usage_warns(caplog):
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)
def test_total_cost_bills_cached_tokens_per_line_at_the_batch_cached_rate():
responses_row = _success_row(
usage={

25
tests/unit/integrations/test_prometheus_services.py
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",
}

145
tests/unit/litellm_core_utils/test_get_litellm_params.py
View file

@ -20,6 +20,28 @@ from litellm.litellm_core_utils.get_litellm_params import (
)
from litellm.types.litellm_params import ControlOptions
def _funnel_kwargs_completion_forwards(monkeypatch, wif_kwargs: dict[str, object]) -> dict[str, object]:
"""completion() names its get_litellm_params arguments one by one, so a key the funnel knows
is still dropped unless that call site forwards it from its own kwargs."""
from unittest.mock import MagicMock
import litellm
import litellm.main as litellm_main
spy = MagicMock(wraps=litellm_main.get_litellm_params)
monkeypatch.setattr( # test-quality-ok: completion() has no injection seam for its kwargs funnel
litellm_main, "get_litellm_params", spy
)
litellm.completion(
model="anthropic/claude-sonnet-5",
messages=[{"role": "user", "content": "hi"}],
mock_response="ok",
**wif_kwargs,
)
return spy.call_args.kwargs
NAMED_PRICE_PARAMS: Final = frozenset(
{
"input_cost_per_token",
@ -357,3 +379,126 @@ def test_drop_params_strings_reach_litellm_params_as_flags(
value: str | bool | None, expected: bool | None
) -> None:
assert get_litellm_params(drop_params=value)["drop_params"] is expected
class TestAnthropicWifKeys:
"""The six anthropic_* WIF keys need dual registration: 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_federation_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, monkeypatch):
forwarded = _funnel_kwargs_completion_forwards(monkeypatch, self.SIX_KEYS)
assert {key: forwarded[key] for key in self.SIX_KEYS} == self.SIX_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.types.workload_identity 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, monkeypatch):
forwarded = _funnel_kwargs_completion_forwards(monkeypatch, self.NEW_KEYS)
assert {key: forwarded[key] for key in self.NEW_KEYS} == self.NEW_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.types.workload_identity 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, monkeypatch):
forwarded = _funnel_kwargs_completion_forwards(monkeypatch, self.THREE_KEYS)
assert {key: forwarded[key] for key in self.THREE_KEYS} == self.THREE_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

146
tests/unit/llms/anthropic/batches/test_handler.py
View file

@ -13,6 +13,8 @@ The sync ``retrieve_batch`` dispatch (``_is_async`` true -> coroutine, false ->
asyncio.run) is exercised directly.
"""
import asyncio
import threading
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -34,9 +36,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"),
)
@ -58,9 +58,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(
@ -75,9 +73,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"
@ -92,9 +88,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
@ -110,14 +104,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",
@ -133,9 +123,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",
@ -164,9 +152,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
@ -175,9 +161,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)
@ -212,21 +196,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",
@ -280,3 +258,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"

47
tests/unit/llms/anthropic/batches/test_transformation.py
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(

46
tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py
View file

@ -49,9 +49,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
@ -198,9 +196,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):
@ -212,9 +208,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,
@ -226,9 +220,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:
@ -527,17 +519,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"}
@ -1439,9 +1425,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"]
@ -1472,9 +1456,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",
@ -1637,18 +1619,14 @@ 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"]] == [
"You are a helpful geography assistant.",
"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.",
@ -2087,9 +2065,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"]

321
tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py
View file

@ -26,9 +26,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
@ -145,9 +143,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
@ -155,19 +151,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",
@ -191,17 +182,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. "},
@ -285,9 +271,7 @@ def test_streamed_signed_thinking_round_trips_to_the_next_turn_once():
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",
@ -307,17 +291,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."
@ -328,9 +307,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",
@ -349,17 +326,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."
@ -370,9 +342,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",
@ -383,9 +353,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")
@ -394,15 +362,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,
)
@ -416,17 +380,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,
)
@ -434,9 +394,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,
)
@ -444,9 +402,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,
)
@ -477,9 +433,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(
@ -493,9 +447,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,
)
@ -520,27 +472,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 = {
@ -559,25 +503,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 = {
@ -637,9 +575,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_message_delta_without_usage_returns_chunk_with_no_usage():
@ -830,9 +766,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
@ -913,13 +847,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 == [
{
@ -1011,9 +941,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():
@ -1066,15 +996,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,
@ -1099,9 +1025,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"}}'
@ -1109,31 +1033,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"}}'
@ -1141,18 +1055,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'
@ -1320,9 +1230,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)
@ -1397,9 +1305,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"
@ -1415,9 +1321,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
@ -1516,9 +1420,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 = [
@ -1589,23 +1491,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():
@ -1619,9 +1513,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 = [
@ -1692,25 +1584,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():
@ -1723,9 +1605,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 = [
@ -1769,9 +1649,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():
@ -1781,9 +1661,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 = [
@ -1851,20 +1729,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"
@ -1877,9 +1747,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 = {
@ -1915,21 +1783,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():
@ -1938,9 +1798,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 = {
@ -1959,9 +1817,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():
@ -2157,8 +2015,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"
@ -2322,9 +2179,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():
@ -2387,12 +2242,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 TestAnthropicChatCompletionPreCallLogging:

18
tests/unit/llms/anthropic/chat/test_code_interpreter_results_extraction.py
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

141
tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py
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."""

24
tests/unit/llms/anthropic/messages/test_advisor_orchestration.py
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"

8
tests/unit/llms/anthropic/pass_through/adapters/test_handler_output_config_passthrough.py
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.pass_through.adapters.handler import (
ANTHROPIC_ONLY_REQUEST_KEYS,
@ -187,9 +185,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},
)

26
tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py
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)

34
tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_compaction.py
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

8
tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_empty_choices.py
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"

14
tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_first_delta.py
View file

@ -637,11 +637,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."]
@ -1110,9 +1106,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
@ -1127,9 +1121,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])

42
tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_tool_args.py
View file

@ -129,15 +129,13 @@ 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]
@ -230,8 +228,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 +288,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 json.loads(events[input_json_delta_idx]["delta"]["partial_json"]) == {"location": "Boston"}
@ -343,9 +338,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 +367,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"}'

8
tests/unit/llms/anthropic/pass_through/context_management/test_clear_tool_uses.py
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",

73
tests/unit/llms/anthropic/pass_through/context_management/test_compact.py
View file

@ -196,10 +196,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():
@ -428,9 +425,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
@ -515,9 +510,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(
@ -1071,13 +1064,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():
@ -1160,9 +1151,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():
@ -1287,9 +1278,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()
@ -1318,9 +1307,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()
@ -1575,9 +1562,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 (
@ -1628,9 +1613,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 (
@ -1671,9 +1654,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"
@ -1695,9 +1676,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 (
@ -2188,9 +2167,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 (
@ -2215,12 +2192,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():
@ -2692,9 +2665,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(
@ -2753,9 +2724,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(

4
tests/unit/llms/anthropic/pass_through/context_management/test_dispatcher.py
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

27
tests/unit/llms/anthropic/pass_through/messages/test_advisor_integration.py
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

52
tests/unit/llms/anthropic/pass_through/messages/test_agentic_streaming_iterator.py
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"

118
tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py Normal file
View file

@ -0,0 +1,118 @@
from pathlib import Path
from typing import Final
import pytest
import litellm
from litellm.llms.anthropic.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.unit.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

12
tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_speed.py
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

4
tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_structured_outputs.py
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

82
tests/unit/llms/anthropic/pass_through/messages/test_content_after_stop_reason.py
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__":

26
tests/unit/llms/anthropic/pass_through/messages/test_parallel_tool_calls.py
View file

@ -53,9 +53,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 +142,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 +198,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 +210,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 +270,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 +280,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

25
tests/unit/llms/anthropic/pass_through/messages/test_reasoning_auto_summary_messages.py
View file

@ -71,9 +71,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 +80,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 +88,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 +109,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"

33
tests/unit/llms/anthropic/pass_through/messages/test_request_optional_param_utils.py
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

43
tests/unit/llms/anthropic/pass_through/messages/test_sse_wrapper.py
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

7
tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py
View file

@ -167,7 +167,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
@ -198,7 +198,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'
)
@ -302,10 +302,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'
)

1625
tests/unit/llms/anthropic/test_anthropic_common_utils.py
View file

File diff suppressed because it is too large Load diff

74
tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py
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"}],
auth_header={"x-api-key": "sk-ant-api03-test-key"},
api_base="https://gateway.example",
)
assert route.called
assert result == {"input_tokens": 7}

296
tests/unit/llms/anthropic/test_anthropic_files_and_batches.py
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,160 @@ 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()
@pytest.mark.asyncio
async def test_afile_content_mints_from_the_deployment_litellm_params(
self, handler, mock_anthropic_batch_results_succeeded, monkeypatch
):
"""Regression: a deployment that authenticates through a named credential carries its
federation settings in litellm_params, and afile_content dropped them, so only
process-wide env vars could ever mint on a batch-result download."""
from litellm.llms.anthropic import common_utils as anthropic_common_utils
from litellm.llms.anthropic.wif import aget_anthropic_wif_token
from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine
for name in (
"ANTHROPIC_API_KEY",
"ANTHROPIC_AUTH_TOKEN",
"ANTHROPIC_API_BASE",
"ANTHROPIC_BASE_URL",
"ANTHROPIC_FEDERATION_RULE_ID",
"ANTHROPIC_ORGANIZATION_ID",
"ANTHROPIC_IDENTITY_TOKEN",
):
monkeypatch.delenv(name, raising=False)
monkeypatch.setenv("CREDENTIAL_IDENTITY_JWT", "credential-inline-jwt")
minted = "sk-ant-oat01-credential-minted"
class Poster:
def post(self, url, *, content, headers, timeout):
return httpx.Response(
200,
json={"access_token": minted, "token_type": "Bearer", "expires_in": 3600},
)
engine = JwtBearerTokenExchangeEngine(poster=Poster())
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, "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,
litellm_params={
"anthropic_federation_rule_id": "fdrl_credential",
"anthropic_organization_id": "org-credential",
"anthropic_identity_token": "oidc/env/CREDENTIAL_IDENTITY_JWT",
},
)
sent_headers = mock_client.get.call_args.kwargs["headers"]
assert sent_headers["authorization"] == f"Bearer {minted}"
class TestAnthropicBatchesConfig:
"""Test Anthropic Batches Config for batch retrieval transformation"""
@ -562,15 +664,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 +694,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 +720,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 +757,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,

15
tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py
View file

@ -115,7 +115,16 @@ def test_prediction_header_eligibility(headers: Mapping[str, str], supported: bo
@pytest.mark.asyncio
async def test_provider_count_uses_same_version_and_preserves_native_input(monkeypatch: pytest.MonkeyPatch) -> None:
@pytest.mark.parametrize(
"api_base, count_url",
[
(None, "https://api.anthropic.com/v1/messages/count_tokens"),
("https://gateway.example/v1/messages", "https://gateway.example/v1/messages/count_tokens"),
],
)
async def test_provider_count_uses_same_version_and_preserves_native_input(
monkeypatch: pytest.MonkeyPatch, api_base: str | None, count_url: str
) -> None:
body: Final = _body()
requests: Final[list[httpx.Request]] = []
@ -128,12 +137,12 @@ async def test_provider_count_uses_same_version_and_preserves_native_input(monke
client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider))
monkeypatch.setattr(count_handler, "get_async_httpx_client", lambda **kwargs: client)
try:
assert await count_prompt_tokens(_MODEL, _KEY, body) == 311
assert await count_prompt_tokens(_MODEL, _KEY, body, api_base=api_base) == 311
finally:
await client.client.aclose()
assert len(requests) == 1
assert requests[0].headers["anthropic-version"] == DEFAULT_ANTHROPIC_API_VERSION
assert requests[0].url == "https://api.anthropic.com/v1/messages/count_tokens"
assert requests[0].url == count_url
assert json.loads(requests[0].content) == body

1293
tests/unit/llms/anthropic/test_anthropic_wif.py Normal file
View file

File diff suppressed because it is too large Load diff

10
tests/unit/llms/anthropic/test_cost_calculation_dict_safety.py
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

249
tests/unit/llms/anthropic/test_count_tokens_oauth.py
View file

@ -1,86 +1,271 @@
"""
Tests for Anthropic CountTokens API OAuth token handling.
Tests for the credential every Anthropic count-tokens request carries.
Verifies that get_required_headers() correctly handles OAuth tokens
(sk-ant-oat*) by delegating to optionally_handle_anthropic_oauth().
The count-tokens handler receives the auth header that ``AnthropicModelInfo.get_auth_header``
resolved, so a static key, an OAuth token (sk-ant-oat*), ``ANTHROPIC_AUTH_TOKEN`` and a minted
workload-identity token all reach Anthropic exactly the way chat on the same deployment does.
Regression test for https://github.com/BerriAI/litellm/issues/22040
Regression tests for https://github.com/BerriAI/litellm/issues/22040 and for the
``ANTHROPIC_AUTH_TOKEN`` gap where count-tokens skipped minting but forwarded no credential.
"""
import os
import sys
sys.path.insert(
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
)
import httpx
import pytest
import respx
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")))
import litellm
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
from litellm.llms.anthropic.count_tokens.transformation import (
AnthropicCountTokensConfig,
)
from litellm.types.llms.anthropic import ANTHROPIC_OAUTH_BETA_HEADER
# Fake tokens for testing (not real secrets)
FAKE_OAUTH_TOKEN = "sk-ant-oat01-fake-token-for-testing-123456789abcdef"
FAKE_REGULAR_KEY = "sk-ant-api03-regular-key-for-testing-123456789"
FEDERATED_DEPLOYMENT = {
"litellm_params": {
"model": "anthropic/claude-sonnet-4-5",
"anthropic_federation_rule_id": "fdrl_x",
"anthropic_organization_id": "org-x",
}
}
def count_tokens_headers_for(api_key: str) -> dict[str, str]:
auth_header = AnthropicModelInfo.get_auth_header(api_key=api_key)
assert auth_header is not None
return AnthropicCountTokensConfig().get_count_tokens_headers(auth_header)
@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()
class TestCountTokensOAuthHeaders:
"""Tests that count_tokens headers are correct for both regular and OAuth keys."""
def test_regular_api_key_uses_x_api_key(self):
"""Regular API keys should be sent via x-api-key header."""
config = AnthropicCountTokensConfig()
headers = config.get_required_headers(FAKE_REGULAR_KEY)
headers = count_tokens_headers_for(FAKE_REGULAR_KEY)
assert headers["x-api-key"] == FAKE_REGULAR_KEY
assert "authorization" not in headers
def test_oauth_key_uses_bearer_authorization(self):
"""OAuth tokens (sk-ant-oat*) should be sent via Authorization: Bearer."""
config = AnthropicCountTokensConfig()
headers = config.get_required_headers(FAKE_OAUTH_TOKEN)
headers = count_tokens_headers_for(FAKE_OAUTH_TOKEN)
assert headers.get("authorization") == f"Bearer {FAKE_OAUTH_TOKEN}"
assert "x-api-key" not in headers
def test_oauth_key_sets_oauth_beta_header(self):
"""OAuth tokens should trigger the anthropic-beta oauth header."""
config = AnthropicCountTokensConfig()
headers = config.get_required_headers(FAKE_OAUTH_TOKEN)
headers = count_tokens_headers_for(FAKE_OAUTH_TOKEN)
assert "oauth-2025-04-20" in headers.get("anthropic-beta", "")
assert ANTHROPIC_OAUTH_BETA_HEADER in headers.get("anthropic-beta", "").split(",")
def test_regular_key_preserves_token_counting_beta(self):
"""Regular keys should keep the token-counting beta header."""
config = AnthropicCountTokensConfig()
headers = config.get_required_headers(FAKE_REGULAR_KEY)
headers = count_tokens_headers_for(FAKE_REGULAR_KEY)
assert "token-counting" in headers.get("anthropic-beta", "")
assert headers.get("anthropic-beta") == ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
def test_headers_always_have_content_type(self):
"""Both regular and OAuth paths should have Content-Type."""
config = AnthropicCountTokensConfig()
for key in [FAKE_REGULAR_KEY, FAKE_OAUTH_TOKEN]:
headers = config.get_required_headers(key)
headers = count_tokens_headers_for(key)
assert headers["Content-Type"] == "application/json"
def test_headers_always_have_anthropic_version(self):
"""Both paths should have anthropic-version."""
config = AnthropicCountTokensConfig()
for key in [FAKE_REGULAR_KEY, FAKE_OAUTH_TOKEN]:
headers = config.get_required_headers(key)
headers = count_tokens_headers_for(key)
assert headers["anthropic-version"] == "2023-06-01"
def test_oauth_key_preserves_token_counting_beta(self):
"""OAuth tokens must preserve the token-counting beta alongside the OAuth beta."""
config = AnthropicCountTokensConfig()
headers = config.get_required_headers(FAKE_OAUTH_TOKEN)
headers = count_tokens_headers_for(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}"
betas = headers.get("anthropic-beta", "").split(",")
assert ANTHROPIC_TOKEN_COUNTING_BETA_VERSION in betas, f"token-counting beta missing: {betas}"
assert ANTHROPIC_OAUTH_BETA_HEADER in betas, f"oauth beta missing: {betas}"
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)
monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
minted = "sk-ant-oat01-minted-for-count"
async def fake_mint(_params, _api_base, _model):
return minted
monkeypatch.setattr("litellm.llms.anthropic.common_utils.aget_anthropic_wif_token", fake_mint)
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=FEDERATED_DEPLOYMENT,
request_model="claude-sonnet-4-5",
)
assert result is not None
assert result.total_tokens == 42
assert seen["auth_header"] == {
"authorization": f"Bearer {minted}",
"anthropic-beta": ANTHROPIC_OAUTH_BETA_HEADER,
}
@pytest.mark.asyncio
async def test_an_auth_token_deployment_counts_with_a_bearer_and_never_mints(
self, monkeypatch, httpx_transport_clients
):
"""With only ``ANTHROPIC_AUTH_TOKEN`` set, chat on a federated deployment authenticates with
that token, so count-tokens must send the same Bearer instead of silently returning None."""
from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module
for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"):
monkeypatch.delenv(name, raising=False)
monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", "bearer-token-for-testing")
async def fake_mint(_params, _api_base, _model):
raise AssertionError("an auth-token deployment must never mint a federated token")
monkeypatch.setattr("litellm.llms.anthropic.common_utils.aget_anthropic_wif_token", fake_mint)
with respx.mock(assert_all_called=True) as router:
route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock(
return_value=httpx.Response(200, json={"input_tokens": 11})
)
result = await token_counter_module.AnthropicTokenCounter().count_tokens(
model_to_use="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
contents=None,
deployment=FEDERATED_DEPLOYMENT,
request_model="claude-sonnet-4-5",
)
assert result is not None
assert result.total_tokens == 11
assert result.tokenizer_type == "anthropic_api"
sent = route.calls.last.request.headers
assert sent["authorization"] == "Bearer bearer-token-for-testing"
assert "x-api-key" not in sent
betas = sent["anthropic-beta"].split(",")
assert ANTHROPIC_TOKEN_COUNTING_BETA_VERSION in betas
assert ANTHROPIC_OAUTH_BETA_HEADER not in betas
@pytest.mark.asyncio
async def test_a_failed_mint_degrades_like_an_anthropic_error(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)
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.common_utils.aget_anthropic_wif_token", failing_mint)
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)
async def fake_mint(_params, _api_base, _model):
raise AssertionError("a static key must never mint a federated token")
monkeypatch.setattr("litellm.llms.anthropic.common_utils.aget_anthropic_wif_token", fake_mint)
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=FEDERATED_DEPLOYMENT,
request_model="claude-sonnet-4-5",
)
assert result is not None
assert result.total_tokens == 7
assert seen["auth_header"] == {"x-api-key": vault_key}

65
tests/unit/llms/anthropic/test_message_sanitization.py
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

0
tests/unit/llms/base_llm/auth/__init__.py Normal file
View file

484
tests/unit/llms/base_llm/auth/test_client_credentials.py Normal file
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)

239
tests/unit/llms/base_llm/auth/test_identity_source.py Normal file
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,
}
)
)

188
tests/unit/llms/base_llm/auth/test_internal_issuer.py Normal file
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))

215
tests/unit/llms/base_llm/auth/test_jwt_signing.py Normal file
View file

@ -0,0 +1,215 @@
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, "-I", "-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"})

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