From 6e44b5b6266352a04f5f388653ced66fe189ca94 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 12:05:17 -0700 Subject: [PATCH] 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> --- litellm/batches/batch_utils.py | 8 +- litellm/batches/main.py | 2 + .../litellm_core_utils/get_litellm_params.py | 48 +- litellm/llms/anthropic/batches/handler.py | 18 +- .../llms/anthropic/batches/transformation.py | 28 +- litellm/llms/anthropic/chat/transformation.py | 4 +- litellm/llms/anthropic/common_utils.py | 336 +++- .../llms/anthropic/count_tokens/handler.py | 4 +- .../anthropic/count_tokens/token_counter.py | 31 +- .../anthropic/count_tokens/transformation.py | 12 +- .../messages/transformation.py | 119 +- litellm/llms/anthropic/files/handler.py | 4 +- .../llms/anthropic/files/transformation.py | 62 +- .../llms/anthropic/skills/transformation.py | 56 +- litellm/llms/anthropic/wif.py | 501 +++++ litellm/llms/azure_ai/embed/handler.py | 2 + .../anthropic_messages/transformation.py | 23 + litellm/llms/base_llm/auth/__init__.py | 99 + .../llms/base_llm/auth/client_credentials.py | 227 +++ litellm/llms/base_llm/auth/identity_source.py | 76 + litellm/llms/base_llm/auth/internal_issuer.py | 86 + litellm/llms/base_llm/auth/jwt_signing.py | 115 ++ litellm/llms/base_llm/auth/token_exchange.py | 861 ++++++++ litellm/llms/base_llm/auth/types.py | 99 + litellm/llms/base_llm/base_utils.py | 17 + litellm/llms/custom_httpx/http_handler.py | 2 +- litellm/llms/custom_httpx/llm_http_handler.py | 225 ++- .../llms/openai/chat/gpt_transformation.py | 53 +- litellm/llms/openai/openai.py | 16 +- .../llms/openai/responses/transformation.py | 7 +- litellm/llms/openai/workload_identity.py | 19 +- litellm/main.py | 1 + litellm/models/credentials.py | 6 +- litellm/proxy/auth/auth_utils.py | 15 +- .../common_utils/credential_hydration.py | 158 ++ .../proxy/credential_endpoints/endpoints.py | 215 +- .../health_endpoints/_health_endpoints.py | 5 + .../model_management_endpoints.py | 81 +- .../llm_passthrough_endpoints.py | 68 +- litellm/proxy/proxy_server.py | 16 +- litellm/router.py | 12 +- .../clientside_credential_handler.py | 17 + .../router_utils/fallback_event_handlers.py | 10 +- litellm/secret_managers/main.py | 8 +- litellm/types/llms/anthropic.py | 1 + litellm/types/router.py | 92 +- litellm/types/services.py | 9 + litellm/types/utils.py | 20 +- litellm/utils.py | 5 +- .../endpointaudit/coverage_allowlist.txt | 1 + .../test_litellm/batches/test_batch_utils.py | 57 +- .../integrations/test_prometheus_services.py | 25 + .../test_get_litellm_params.py | 126 ++ .../llms/anthropic/batches/test_handler.py | 146 +- .../anthropic/batches/test_transformation.py | 47 +- .../test_anthropic_guardrail_handler.py | 50 +- .../chat/test_anthropic_chat_handler.py | 423 ++-- .../test_anthropic_chat_transformation.py | 387 +--- ...est_code_interpreter_results_extraction.py | 18 +- ...al_pass_through_adapters_transformation.py | 269 +-- .../test_handler_output_config_passthrough.py | 8 +- .../test_streaming_iterator_combined_chunk.py | 26 +- .../test_streaming_iterator_compaction.py | 34 +- .../test_streaming_iterator_empty_choices.py | 8 +- .../test_streaming_iterator_first_delta.py | 14 +- .../test_streaming_iterator_tool_args.py | 46 +- .../test_clear_tool_uses.py | 8 +- .../context_management/test_compact.py | 73 +- .../context_management/test_dispatcher.py | 4 +- .../messages/test_advisor_integration.py | 27 +- .../test_agentic_streaming_iterator.py | 52 +- ...al_pass_through_messages_transformation.py | 118 ++ .../messages/test_anthropic_messages_speed.py | 12 +- ...t_anthropic_messages_structured_outputs.py | 4 +- .../test_content_after_stop_reason.py | 82 +- .../messages/test_mcp_handler.py | 8 +- .../messages/test_parallel_tool_calls.py | 27 +- .../test_reasoning_auto_summary_messages.py | 38 +- .../test_request_optional_param_utils.py | 33 +- .../messages/test_sse_wrapper.py | 43 +- .../messages/test_streaming_iterator.py | 7 +- .../test_responses_adapters_transformation.py | 8 +- .../test_anthropic_files_transformation.py | 141 +- .../messages/test_advisor_orchestration.py | 24 +- .../anthropic/test_anthropic_common_utils.py | 1570 ++++++++++++++- ...t_anthropic_count_tokens_transformation.py | 74 + .../test_anthropic_files_and_batches.py | 222 ++- .../llms/anthropic/test_anthropic_wif.py | 1222 ++++++++++++ .../anthropic/test_azure_ai_cache_pricing.py | 11 +- .../test_cost_calculation_dict_safety.py | 10 +- .../llms/anthropic/test_count_tokens_oauth.py | 186 +- .../anthropic/test_message_sanitization.py | 65 +- .../llms/base_llm/auth/__init__.py | 0 .../base_llm/auth/test_client_credentials.py | 484 +++++ .../base_llm/auth/test_identity_source.py | 239 +++ .../base_llm/auth/test_internal_issuer.py | 188 ++ .../llms/base_llm/auth/test_jwt_signing.py | 213 ++ .../llms/base_llm/auth/test_token_exchange.py | 1752 +++++++++++++++++ .../custom_httpx/test_llm_http_handler.py | 229 ++- .../openai/test_openai_workload_identity.py | 351 ++++ tests/test_litellm/models/test_models.py | 78 +- .../proxy/auth/test_auth_utils.py | 240 +-- .../common_utils/test_credential_hydration.py | 28 + .../credential_endpoints/test_endpoints.py | 948 ++++++++- .../test_model_management_endpoints.py | 1137 ++++++++--- .../test_llm_pass_through_endpoints.py | 814 ++++---- .../proxy/proxy_server/test_proxy_config.py | 17 + tests/test_litellm/proxy/test_proxy_server.py | 51 + .../test_fallback_event_handlers.py | 56 +- .../test_anthropic_skills_transformation.py | 77 +- tests/test_litellm/test_router.py | 59 + tests/test_litellm/types/test_router.py | 115 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 163 +- 113 files changed, 14209 insertions(+), 2683 deletions(-) create mode 100644 litellm/llms/anthropic/wif.py create mode 100644 litellm/llms/base_llm/auth/__init__.py create mode 100644 litellm/llms/base_llm/auth/client_credentials.py create mode 100644 litellm/llms/base_llm/auth/identity_source.py create mode 100644 litellm/llms/base_llm/auth/internal_issuer.py create mode 100644 litellm/llms/base_llm/auth/jwt_signing.py create mode 100644 litellm/llms/base_llm/auth/token_exchange.py create mode 100644 litellm/llms/base_llm/auth/types.py create mode 100644 litellm/proxy/common_utils/credential_hydration.py create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py create mode 100644 tests/test_litellm/llms/anthropic/test_anthropic_wif.py create mode 100644 tests/test_litellm/llms/base_llm/auth/__init__.py create mode 100644 tests/test_litellm/llms/base_llm/auth/test_client_credentials.py create mode 100644 tests/test_litellm/llms/base_llm/auth/test_identity_source.py create mode 100644 tests/test_litellm/llms/base_llm/auth/test_internal_issuer.py create mode 100644 tests/test_litellm/llms/base_llm/auth/test_jwt_signing.py create mode 100644 tests/test_litellm/llms/base_llm/auth/test_token_exchange.py create mode 100644 tests/test_litellm/proxy/common_utils/test_credential_hydration.py diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 97be5f77d79..52b7c74412b 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -7,7 +7,10 @@ from typing import Any, Final, Literal import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS +from litellm.litellm_core_utils.get_litellm_params import ( + ANTHROPIC_WIF_KWARGS_KEYS, + AWS_CREDENTIAL_KWARGS_KEYS, +) from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.types.llms.openai import Batch from litellm.types.utils import CallTypes, ModelInfo, Usage @@ -507,6 +510,9 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict: "max_retries", "_litellm_internal_model_credentials", *AWS_CREDENTIAL_KWARGS_KEYS, + # A federated deployment holds no api_key, so without these the fetch that reads a + # finished batch's output has nothing to authenticate with and its cost is never billed. + *sorted(ANTHROPIC_WIF_KWARGS_KEYS), ) for key in credential_keys: if key in litellm_params: diff --git a/litellm/batches/main.py b/litellm/batches/main.py index c8360a81c7a..88561ca52ab 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -483,6 +483,7 @@ def _handle_retrieve_batch_providers_without_provider_config( ) api_key = optional_params.api_key or litellm.api_key or litellm.azure_key or get_secret_str("ANTHROPIC_API_KEY") + batch_params: Final = dict(litellm_params) # mutable-ok: handler contract, copied not shared response = anthropic_batches_instance.retrieve_batch( _is_async=_is_async, batch_id=batch_id, @@ -490,6 +491,7 @@ def _handle_retrieve_batch_providers_without_provider_config( api_key=api_key, timeout=timeout, max_retries=optional_params.max_retries, + litellm_params=batch_params, ) else: raise litellm.exceptions.BadRequestError( diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 389e6f7f501..7fbab6e59ee 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -24,10 +24,54 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( # The per-deployment Rust opt-in. RUST_KWARG_KEY: Final = "rust" +# Anthropic workload identity federation config, read from litellm_params by the +# Anthropic auth tier. Registered like `rust`: here so the kwargs funnel carries +# them, and in `all_litellm_params` so they never leak into the provider body. +ANTHROPIC_WIF_KWARGS_KEYS: Final = frozenset( + { + "anthropic_federation_rule_id", + "anthropic_organization_id", + "anthropic_service_account_id", + "anthropic_workspace_id", + "anthropic_identity_token_file", + "anthropic_identity_token", + # Identity-source selection (Phase 1): absent means the legacy + # token_file/env resolver above, byte-identical to today. + "anthropic_identity_source", + # internal_issuer: litellm self-signs the workload assertion. + "anthropic_issuer_url", + "anthropic_issuer_subject", + "anthropic_issuer_audience", + "anthropic_issuer_ttl_seconds", + "anthropic_issuer_signing_key_ref", + # keycloak: litellm fetches the assertion via client_credentials. + "anthropic_keycloak_token_url", + "anthropic_keycloak_client_id", + "anthropic_keycloak_auth_method", + "anthropic_keycloak_client_secret_ref", + "anthropic_keycloak_scope", + # Set server-side when a client redirects api_base, to stop a federated deployment minting + # for a base the caller chose. It has to ride this funnel or it is dropped on the way and + # the deployment federates anyway; being carried here also request-bans it, which is right, + # since a caller must not be able to set it in either direction. + "anthropic_disable_workload_identity_federation", + } +) + +OPENAI_WIF_KWARGS_KEYS: Final = frozenset( + { + "openai_identity_provider_id", + "openai_service_account_id", + "openai_identity_token_file", + } +) + # Keys `completion()` forwards from its own kwargs into `get_litellm_params`, # which are otherwise invisible to it because that call site passes explicit # named arguments rather than `**kwargs`. -FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({RUST_KWARG_KEY}) +FORWARDED_KWARGS_KEYS: Final = ( + AWS_CREDENTIAL_KWARGS_KEYS | ANTHROPIC_WIF_KWARGS_KEYS | OPENAI_WIF_KWARGS_KEYS | frozenset({RUST_KWARG_KEY}) +) # Pre-define optional kwargs keys as frozenset for O(1) lookups # These are extracted from kwargs only if present, avoiding unnecessary .get() calls @@ -65,6 +109,8 @@ OPTIONAL_KWARGS_KEYS: Final = ( } ) | AWS_CREDENTIAL_KWARGS_KEYS + | ANTHROPIC_WIF_KWARGS_KEYS + | OPENAI_WIF_KWARGS_KEYS ) # Backward-compatible alias for existing imports/tests. diff --git a/litellm/llms/anthropic/batches/handler.py b/litellm/llms/anthropic/batches/handler.py index f418e7e08be..3f3716e0de6 100644 --- a/litellm/llms/anthropic/batches/handler.py +++ b/litellm/llms/anthropic/batches/handler.py @@ -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, ) ) diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index 4f4d39f09b0..cefb31bd7c9 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -13,6 +13,8 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest from litellm.types.utils import LiteLLMBatch, LlmProviders, ModelResponse +from ..common_utils import merge_anthropic_beta_headers, without_caller_credential_headers + if TYPE_CHECKING: import tiktoken @@ -70,24 +72,30 @@ class AnthropicBatchesConfig(BaseBatchesConfig): api_base: str | None = None, ) -> dict: """Validate and prepare environment-specific headers and parameters.""" - if api_base is None and isinstance(litellm_params, dict): - api_base = litellm_params.get("api_base") - auth_header: Final = self.anthropic_model_info.get_auth_header(api_key, api_base) + params_mapping: Final = litellm_params if isinstance(litellm_params, dict) else None + if api_base is None and params_mapping is not None: + api_base = params_mapping.get("api_base") + auth_header: Final = self.anthropic_model_info.get_auth_header( + api_key, api_base, litellm_params=params_mapping, allow_workload_identity=True + ) if auth_header is None: raise ValueError( "Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params" ) - _headers: Final = { + merged_beta: Final = merge_anthropic_beta_headers( + merge_anthropic_beta_headers(headers.get("anthropic-beta"), auth_header.get("anthropic-beta")), + "message-batches-2024-09-24", + ) + # The deployment's own credential is applied below, so a caller-supplied one must not + # ride along: without this a minted federation Bearer travels beside the caller's x-api-key. + return { + **without_caller_credential_headers(headers), "accept": "application/json", "anthropic-version": "2023-06-01", "content-type": "application/json", + **auth_header, + "anthropic-beta": merged_beta, } - _headers.update(auth_header) - # Add beta header for message batches - if "anthropic-beta" not in headers: - headers["anthropic-beta"] = "message-batches-2024-09-24" - headers.update(_headers) - return headers def get_complete_batch_url( self, diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 5f7ac73c919..1039569c79a 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -3,7 +3,7 @@ import re import time from collections.abc import Callable, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NoReturn, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, cast import httpx from pydantic import ValidationError @@ -284,6 +284,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): to pass metadata to anthropic, it's {"user_id": "any-relevant-information"} """ + _workload_identity_eligible: ClassVar[bool] = True + max_tokens: int | None = None stop_sequences: list | None = None temperature: int | None = None diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 6079b709bcc..fec230653f6 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -7,10 +7,11 @@ import re from collections.abc import Mapping, MutableMapping, Sequence from datetime import datetime, timezone from types import MappingProxyType -from typing import Any, Final, Literal +from typing import Any, ClassVar, Final, Literal +from urllib.parse import quote import httpx -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError import litellm from litellm.constants import ( @@ -25,8 +26,14 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.prompt_templates.factory import ( THOUGHT_SIGNATURE_SEPARATOR, ) +from litellm.llms.anthropic.wif import ( + aget_anthropic_wif_token, + anthropic_base_without_chat_suffix, + get_anthropic_wif_token, +) from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.proxy._types import SpecialHeaders from litellm.types.llms.anthropic import ( ANTHROPIC_HOSTED_TOOLS, ANTHROPIC_OAUTH_BETA_HEADER, @@ -81,6 +88,29 @@ def _strip_bedrock_id_suffixes(model: str) -> str: ) +_SERVER_OWNED_AUTH_HEADERS: Final = SpecialHeaders.litellm_credential_header_names() +_WIF_ELIGIBILITY_ATTR: Final = "_workload_identity_eligible" + + +def without_caller_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + """``headers`` minus every header that authenticates the caller to litellm. + + The deployment's own credential is applied on top of the result, so a caller-supplied + credential must not survive into the upstream request: without this a minted federation + Bearer travels beside the caller's own ``x-api-key``, and Anthropic sees two credentials. + """ + return MappingProxyType( + {name: value for name, value in headers.items() if name.lower() not in _SERVER_OWNED_AUTH_HEADERS} + ) + + +def config_allows_workload_identity(config: object) -> bool: + """A federation token is an Anthropic-org credential and its exchange POSTs the workload's OIDC + assertion to the deployment's own host, so eligibility is declared per class and read from that + class's own ``__dict__``: a subclass written for another provider inherits nothing.""" + return type(config).__dict__.get(_WIF_ELIGIBILITY_ATTR, False) is True + + def is_anthropic_oauth_key(value: str | None) -> bool: """Check if a value contains an Anthropic OAuth token (sk-ant-oat*).""" if value is None: @@ -90,12 +120,20 @@ def is_anthropic_oauth_key(value: str | None) -> bool: return value.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX) -def _merge_beta_headers(existing: str | None, new_beta: str) -> str: - """Merge a new beta value into an existing comma-separated anthropic-beta header.""" - if not existing: - return new_beta - betas: Final = {b.strip() for b in existing.split(",") if b.strip()} - betas.add(new_beta) +def merge_anthropic_beta_headers(existing: str | Sequence[str] | None, new_beta: str | Sequence[str] | None) -> str: + """Merge anthropic-beta header values, deduplicated and sorted. + + Either side may arrive as a list rather than a comma-separated string: the Skills surface + accepted a list-valued header before it shared this helper, and callers still send one. + """ + values: Final = ( + entry + for side in (existing, new_beta) + if side + for entry in ((side,) if isinstance(side, str) else side) + if isinstance(entry, str) + ) + betas: Final = frozenset(b.strip() for value in values for b in value.split(",") if b.strip()) return ",".join(sorted(betas)) @@ -122,7 +160,9 @@ def optionally_handle_anthropic_oauth(headers: dict, api_key: str | None) -> tup ): headers.pop(name) headers["authorization"] = auth_header - headers["anthropic-beta"] = _merge_beta_headers(headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER) + headers["anthropic-beta"] = merge_anthropic_beta_headers( + headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER + ) headers["anthropic-dangerous-direct-browser-access"] = "true" return headers, api_key # Check api_key directly (standard chat/completion flow) @@ -130,7 +170,9 @@ def optionally_handle_anthropic_oauth(headers: dict, api_key: str | None) -> tup for name in tuple(header_name for header_name in headers if header_name.lower() == "x-api-key"): headers.pop(name) headers["authorization"] = f"Bearer {api_key}" - headers["anthropic-beta"] = _merge_beta_headers(headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER) + headers["anthropic-beta"] = merge_anthropic_beta_headers( + headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER + ) headers["anthropic-dangerous-direct-browser-access"] = "true" return headers, api_key @@ -145,7 +187,79 @@ class AnthropicError(BaseLLMException): super().__init__(status_code=status_code, message=message, headers=headers) +_MODEL_LIST_PAGE_CAP: Final = 20 + + +def _litellm_params_str(litellm_params: Mapping[str, object] | None, key: str) -> str | None: + value: Final = litellm_params.get(key) if litellm_params is not None else None + return value if isinstance(value, str) else None + + +class _AnthropicModelListEntry(BaseModel): + id: str + + +class _AnthropicModelsPage(BaseModel): + data: Sequence[_AnthropicModelListEntry] = Field(default_factory=tuple) + has_more: bool = False + last_id: str | None = None + + +def _sanitized_anthropic_error(response: httpx.Response, detail: str | None = None) -> str: + """A provider error detail built only from structured fields, never ``response.text`` + verbatim: the raw body is untrusted content the caller of ``/v1/models`` did not ask for + and should not have echoed back to it wholesale.""" + if detail is not None: + return f"HTTP {response.status_code}: {detail}" + try: + body: Final = response.json() + except ValueError: + return f"HTTP {response.status_code}" + error: Final = body.get("error") if isinstance(body, dict) else None + message: Final = error.get("message") if isinstance(error, dict) else None + return f"HTTP {response.status_code}: {message}" if isinstance(message, str) else f"HTTP {response.status_code}" + + +def _fetch_anthropic_models_page( + api_base: str, headers: Mapping[str, str], after_id: str | None +) -> _AnthropicModelsPage: + # after_id rides the URL because the client mutates the params mapping it is handed, + # which a read-only one cannot support + query: Final = f"?after_id={quote(after_id)}" if after_id else "" + response: Final = litellm.module_level_client.get( + url=f"{api_base}/v1/models{query}", + headers=headers, + follow_redirects=False, + ) + try: + response.raise_for_status() + except httpx.HTTPStatusError: + raise Exception(f"Failed to fetch models from Anthropic. {_sanitized_anthropic_error(response)}") from None + try: + return _AnthropicModelsPage.model_validate(response.json()) + except ValueError as e: + raise Exception( + f"Failed to fetch models from Anthropic. {_sanitized_anthropic_error(response, detail=str(e))}" + ) from None + + +def _fetch_anthropic_model_ids( + api_base: str, headers: Mapping[str, str], after_id: str | None, pages_left: int +) -> tuple[str, ...]: + collected: tuple[str, ...] = () # rebind-ok: accumulates one page of ids per iteration + cursor: str | None = after_id # rebind-ok: advances to each page's last_id + for _ in range(max(pages_left, 0)): + page = _fetch_anthropic_models_page(api_base, headers, cursor) # rebind-ok: one page per iteration + collected += tuple(entry.id for entry in page.data) + if not page.has_more or page.last_id is None: + return collected + cursor = page.last_id + raise Exception(f"Anthropic /v1/models did not terminate within {_MODEL_LIST_PAGE_CAP} pages.") + + class AnthropicModelInfo(BaseLLMModelInfo): + _workload_identity_eligible: ClassVar[bool] = True + def is_cache_control_set(self, messages: list[AllMessageValues]) -> bool: """ Return if {"cache_control": ..} in message content block @@ -720,7 +834,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): return list(set(betas)) @staticmethod - def _make_api_key_auth_header(api_key: str, api_base: str | None, use_bearer_for_custom_base: bool = False) -> dict: + def _make_api_key_auth_header( + api_key: str, api_base: str | None, use_bearer_for_custom_base: bool = False + ) -> Mapping[str, str]: if use_bearer_for_custom_base and ( api_base and "api.anthropic.com" not in api_base and not api_key.startswith("sk-ant-") ): @@ -728,6 +844,33 @@ class AnthropicModelInfo(BaseLLMModelInfo): return {"authorization": value} return {"x-api-key": api_key} + def _credential_headers( + self, + *, + api_key: str | None, + auth_token: str | None, + api_base: str | None, + use_bearer_for_custom_base: bool, + wif_minted: bool, + betas: set[str], # mutable-ok: the caller's beta accumulator, appended to by the oauth tier + ) -> Mapping[str, str]: + """The credential tier walk: a consumer OAuth token, then ANTHROPIC_AUTH_TOKEN, then an api key. + + A server-minted federation token takes the same Bearer shape as a consumer OAuth token but is + not browser-forwarded, so it does not get the direct-browser-access header. + """ + if api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX): + betas.add(ANTHROPIC_OAUTH_BETA_HEADER) + oauth_headers: Final = {"authorization": f"Bearer {api_key}"} + if wif_minted: + return oauth_headers + return {**oauth_headers, "anthropic-dangerous-direct-browser-access": "true"} + if auth_token and not api_key: + return {"authorization": f"Bearer {auth_token}"} + if api_key: + return self._make_api_key_auth_header(api_key, api_base, use_bearer_for_custom_base) + return {} + def get_anthropic_headers( self, api_key: str | None = None, @@ -749,6 +892,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): container_with_skills_used: bool = False, api_base: str | None = None, use_bearer_for_custom_base: bool = False, + wif_minted: bool = False, ) -> dict: betas: Final = set() # Anthropic no longer requires the prompt-caching beta header @@ -784,20 +928,21 @@ class AnthropicModelInfo(BaseLLMModelInfo): if container_with_skills_used: betas.add("skills-2025-10-02") - _is_oauth: Final = api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX) headers: Final = { "anthropic-version": anthropic_version or "2023-06-01", "accept": "application/json", "content-type": "application/json", } - if _is_oauth: - headers["authorization"] = f"Bearer {api_key}" - headers["anthropic-dangerous-direct-browser-access"] = "true" - betas.add(ANTHROPIC_OAUTH_BETA_HEADER) - elif auth_token and not api_key: - headers["authorization"] = f"Bearer {auth_token}" - elif api_key: - headers.update(self._make_api_key_auth_header(api_key, api_base, use_bearer_for_custom_base)) + headers.update( + self._credential_headers( + api_key=api_key, + auth_token=auth_token, + api_base=api_base, + use_bearer_for_custom_base=use_bearer_for_custom_base, + wif_minted=wif_minted, + betas=betas, + ) + ) if user_anthropic_beta_headers is not None: betas.update(user_anthropic_beta_headers) @@ -824,10 +969,11 @@ class AnthropicModelInfo(BaseLLMModelInfo): api_key: str | None = None, api_base: str | None = None, ) -> dict: - if api_base is None and isinstance(litellm_params, dict): - api_base = litellm_params.get("api_base") + params_mapping: Final = litellm_params if isinstance(litellm_params, dict) else None + if api_base is None and params_mapping is not None: + api_base = params_mapping.get("api_base") use_bearer_for_custom_base: Final[bool] = bool( - isinstance(litellm_params, dict) and litellm_params.get("use_bearer_for_custom_base", False) + params_mapping is not None and params_mapping.get("use_bearer_for_custom_base", False) ) # Check for Anthropic OAuth token in headers headers, api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key) @@ -836,9 +982,23 @@ class AnthropicModelInfo(BaseLLMModelInfo): auth_token: str | None = None if api_key is None: auth_token = AnthropicModelInfo.get_auth_token() - if api_key is None and auth_token is None: + wif_token: Final = ( + get_anthropic_wif_token(params_mapping, api_base, model) + if api_key is None and auth_token is None and config_allows_workload_identity(self) + else None + ) + wif_minted: Final = wif_token is not None + resolved_api_key: Final = wif_token if wif_token is not None else api_key + if resolved_api_key is None and auth_token is None: raise litellm.AuthenticationError( - message="Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` in your environment vars", + message=( + "Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the " + "environment variables or via params. Please set `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` " + "in your environment vars, or configure workload identity federation via " + "`ANTHROPIC_FEDERATION_RULE_ID`, `ANTHROPIC_ORGANIZATION_ID`, " + "`ANTHROPIC_SERVICE_ACCOUNT_ID` and " + "`ANTHROPIC_IDENTITY_TOKEN_FILE` (or `ANTHROPIC_IDENTITY_TOKEN`)" + ), llm_provider="anthropic", model=model, ) @@ -863,7 +1023,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): computer_tool_used=computer_tool_used, prompt_caching_set=prompt_caching_set, pdf_used=pdf_used, - api_key=api_key, + api_key=resolved_api_key, auth_token=auth_token, file_id_used=file_id_used, web_search_tool_used=web_search_tool_used, @@ -878,11 +1038,12 @@ class AnthropicModelInfo(BaseLLMModelInfo): container_with_skills_used=container_with_skills_used, api_base=api_base, use_bearer_for_custom_base=use_bearer_for_custom_base, + wif_minted=wif_minted, ) - headers = {**headers, **anthropic_headers} + caller_headers: Final = without_caller_credential_headers(headers) if wif_minted else headers - return headers + return {**caller_headers, **anthropic_headers} @staticmethod def get_api_base(api_base: str | None = None) -> str | None: @@ -917,52 +1078,121 @@ class AnthropicModelInfo(BaseLLMModelInfo): api_key: str | None = None, api_base: str | None = None, use_bearer_for_custom_base: bool = False, - ) -> dict | None: + litellm_params: Mapping[str, object] | None = None, + allow_workload_identity: bool = False, + ) -> Mapping[str, str] | None: """Resolve Anthropic credentials and return the appropriate auth header dict. Checks ANTHROPIC_API_KEY first (-> x-api-key or Bearer depending on - use_bearer_for_custom_base), then ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer). - Returns None if neither is available. + use_bearer_for_custom_base), then ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer), + then workload identity federation (-> Authorization: Bearer with a minted + sk-ant-oat01 token, honoring anthropic_* litellm_params when provided). Every + Bearer built from an sk-ant-oat token carries the mandatory oauth anthropic-beta. + Returns None if no credential source is available. """ + static_header: Final = AnthropicModelInfo._static_auth_header(api_key, api_base, use_bearer_for_custom_base) + if static_header is not None: + return static_header + if not allow_workload_identity: + return None + wif_token: Final = get_anthropic_wif_token(litellm_params, api_base, "") + if wif_token is not None: + return AnthropicModelInfo._oauth_bearer_header(wif_token) + return None + + @staticmethod + async def aget_auth_header( + api_key: str | None = None, + api_base: str | None = None, + use_bearer_for_custom_base: bool = False, + litellm_params: Mapping[str, object] | None = None, + allow_workload_identity: bool = False, + ) -> Mapping[str, str] | None: + """Async counterpart of get_auth_header: the WIF tier can block on a token + exchange POST, so async callers await it off the event loop.""" + static_header: Final = AnthropicModelInfo._static_auth_header(api_key, api_base, use_bearer_for_custom_base) + if static_header is not None: + return static_header + if not allow_workload_identity: + return None + wif_token: Final = await aget_anthropic_wif_token(litellm_params, api_base, "") + if wif_token is not None: + return AnthropicModelInfo._oauth_bearer_header(wif_token) + return None + + @staticmethod + def _static_auth_header( + api_key: str | None, + api_base: str | None, + use_bearer_for_custom_base: bool, + ) -> Mapping[str, str] | None: resolved_key: Final = AnthropicModelInfo.get_api_key(api_key) if resolved_key is not None: if is_anthropic_oauth_key(resolved_key): - return {"authorization": f"Bearer {resolved_key}"} + return AnthropicModelInfo._oauth_bearer_header(resolved_key) return AnthropicModelInfo._make_api_key_auth_header(resolved_key, api_base, use_bearer_for_custom_base) auth_token: Final = AnthropicModelInfo.get_auth_token() if auth_token is not None: return {"authorization": f"Bearer {auth_token}"} return None + @staticmethod + def _oauth_bearer_header(token: str) -> Mapping[str, str]: + return {"authorization": f"Bearer {token}", "anthropic-beta": ANTHROPIC_OAUTH_BETA_HEADER} + @staticmethod def get_base_model(model: str | None = None) -> str | None: return model.replace("anthropic/", "") if model else None def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]: - api_base = AnthropicModelInfo.get_api_base(api_base) - auth_header: Final = AnthropicModelInfo.get_auth_header(api_key, api_base) - if api_base is None or auth_header is None: - raise ValueError( - "ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN is not set. Please set the environment variable, to query Anthropic's `/models` endpoint." - ) - headers: Final = {"anthropic-version": "2023-06-01"} - headers.update(auth_header) - response: Final = litellm.module_level_client.get( - url=f"{api_base}/v1/models", - headers=headers, + return self._list_models(api_key=api_key, api_base=api_base, litellm_params=None) + + def discover_models( + self, litellm_params: Mapping[str, object] | None = None + ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override + """Live discovery for a configured deployment: unlike ``get_models``, this threads the + full ``litellm_params`` into ``get_auth_header`` so a workload-identity-federation source + configured on the deployment (rather than the environment) is honored, gated the same way + every other Anthropic auth surface is via ``config_allows_workload_identity``.""" + return self._list_models( + api_key=_litellm_params_str(litellm_params, "api_key"), + api_base=_litellm_params_str(litellm_params, "api_base"), + litellm_params=litellm_params, ) - try: - response.raise_for_status() - except httpx.HTTPStatusError: - raise Exception( - f"Failed to fetch models from Anthropic. Status code: {response.status_code}, Response: {response.text}" + def _list_models( + self, + *, + api_key: str | None, + api_base: str | None, + litellm_params: Mapping[str, object] | None, + ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override + resolved_api_base: Final = AnthropicModelInfo.get_api_base(api_base) + auth_header: Final = AnthropicModelInfo.get_auth_header( + api_key, + resolved_api_base, + litellm_params=litellm_params, + allow_workload_identity=config_allows_workload_identity(self), + ) + if resolved_api_base is None or auth_header is None: + raise ValueError( + "ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN (or workload " + "identity federation via ANTHROPIC_FEDERATION_RULE_ID/ANTHROPIC_ORGANIZATION_ID/" + "ANTHROPIC_IDENTITY_TOKEN_FILE) is not set. Please set the environment variable, to query " + "Anthropic's `/models` endpoint." ) - - models: Final[Sequence[Mapping[str, str]]] = response.json()["data"] - - litellm_model_names: Final = ["anthropic/" + model["id"] for model in models] - return litellm_model_names + headers: Final = MappingProxyType({"anthropic-version": "2023-06-01", **auth_header}) + # /v1/models is appended below, so a base the operator already wrote as .../v1 or + # .../v1/messages would otherwise be asked for /v1/v1/models. + model_ids: Final = _fetch_anthropic_model_ids( + anthropic_base_without_chat_suffix(resolved_api_base), + headers, + after_id=None, + pages_left=_MODEL_LIST_PAGE_CAP, + ) + return [ # mutable-ok: matches get_models' list[str] contract shared by every provider override + "anthropic/" + model_id for model_id in model_ids + ] def get_token_counter(self) -> BaseTokenCounter | None: """ diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index 38cd429d99a..39968ac6fb3 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -41,7 +41,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): model: The model identifier (e.g., "claude-3-5-sonnet-20241022") messages: The messages to count tokens for api_key: The Anthropic API key - api_base: Optional custom API base URL + api_base: Optional deployment api_base the count-tokens path is appended to timeout: Optional timeout for the request (defaults to litellm.request_timeout) Returns: @@ -67,7 +67,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): verbose_logger.debug("Transformed request: %s", request_body) # Get endpoint URL - endpoint_url: Final = api_base or self.get_anthropic_count_tokens_endpoint() + endpoint_url: Final = self.get_anthropic_count_tokens_endpoint(api_base) verbose_logger.debug("Making request to: %s", endpoint_url) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 8e8d10c961b..d5b8667a2e1 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -2,10 +2,10 @@ Anthropic Token Counter implementation using the CountTokens API. """ -import os from typing import Any, Final from litellm._logging import verbose_logger +from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.types.utils import LlmProviders, TokenCountResponse @@ -46,28 +46,33 @@ class AnthropicTokenCounter(BaseTokenCounter): Returns: TokenCountResponse with token count, or None if counting fails """ - from litellm.llms.anthropic.common_utils import AnthropicError + from litellm.llms.anthropic.common_utils import AnthropicError, AnthropicModelInfo + from litellm.llms.anthropic.wif import aget_anthropic_wif_token if not messages: return None deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) - - # Get Anthropic API key from deployment config or environment - api_key = litellm_params.get("api_key") - if not api_key: - api_key = os.getenv("ANTHROPIC_API_KEY") - - if not api_key: - verbose_logger.warning("No Anthropic API key found for token counting") - return None + api_base: Final = litellm_params.get("api_base") + static_key: Final = AnthropicModelInfo.get_api_key(litellm_params.get("api_key")) + auth_token_configured: Final = AnthropicModelInfo.get_auth_token() is not None try: + api_key: Final = ( + static_key + if static_key or auth_token_configured + else await aget_anthropic_wif_token(litellm_params, api_base, model_to_use) + ) + if not api_key: + verbose_logger.warning("No Anthropic credential found for token counting") + return None + result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request( model=model_to_use, messages=messages, api_key=api_key, + api_base=api_base, tools=tools, system=system, ) @@ -80,8 +85,8 @@ class AnthropicTokenCounter(BaseTokenCounter): tokenizer_type="anthropic_api", original_response=result, ) - except AnthropicError as e: - verbose_logger.warning("Anthropic CountTokens API error: status=%s, message=%s", e.status_code, e.message) + except (AnthropicError, AuthenticationError) as e: + verbose_logger.warning("Anthropic CountTokens error: status=%s, message=%s", e.status_code, e.message) return TokenCountResponse( total_tokens=0, request_model=request_model, diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 12581b9f658..91f3896f774 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -7,6 +7,7 @@ This module handles the transformation of requests to Anthropic's CountTokens AP from typing import Any, Final from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION +from litellm.llms.anthropic.wif import resolve_anthropic_base class AnthropicCountTokensConfig: @@ -19,14 +20,21 @@ class AnthropicCountTokensConfig: - Response: {"input_tokens": } """ - 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, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 8267da157ad..be44a408f4a 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -1,5 +1,5 @@ from collections.abc import AsyncIterator, Mapping, Sequence -from typing import Any, Final +from typing import Any, ClassVar, Final import httpx @@ -23,12 +23,24 @@ from litellm.types.router import GenericLiteLLMParams from ...common_utils import ( AnthropicError, AnthropicModelInfo, + merge_anthropic_beta_headers, optionally_handle_anthropic_oauth, strip_advisor_blocks_from_messages, ) DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" +_CALLER_CREDENTIAL_HEADERS: Final = frozenset({"x-api-key", "authorization"}) + + +def _carries_caller_credential(headers: Mapping[str, str]) -> bool: + """Whether the caller sent their own Anthropic credential, in which case this passthrough + honors it and never mints. Matched case-insensitively: an SDK caller passing ``X-Api-Key`` + through extra_headers would otherwise slip the check and end up sending their key beside a + minted federation Bearer.""" + return any(name.lower() in _CALLER_CREDENTIAL_HEADERS for name in headers) + + DROP_UNSUPPORTED_ADAPTIVE_EFFORT_WARNING: Final = ( "Dropping adaptive `thinking`/`output_config.effort` for model=%s: the model " "does not support extended thinking, or max_tokens is too small to fit the " @@ -42,6 +54,8 @@ DROP_UNFITTING_REASONING_EFFORT_WARNING: Final = ( class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): + _workload_identity_eligible: ClassVar[bool] = True + @property def custom_llm_provider(self) -> str | None: return "anthropic" @@ -308,32 +322,103 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): # Check for Anthropic OAuth token in Authorization header headers, api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key) - header_names: Final = frozenset(name.lower() for name in headers) - if "x-api-key" not in header_names and "authorization" not in header_names: - auth_header: Final = AnthropicModelInfo.get_auth_header(api_key) - if auth_header is None: - raise AuthenticationError( - message=( - "Missing Anthropic API Key - A call is being made to anthropic but no key is set " - "either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` " - "or `ANTHROPIC_AUTH_TOKEN` in your environment vars" + if not _carries_caller_credential(headers): + self._apply_env_auth_header( + headers, + self._require_auth_header( + AnthropicModelInfo.get_auth_header( + api_key, + api_base=api_base, + litellm_params=litellm_params, + allow_workload_identity=self._allows_workload_identity, ), - llm_provider=self._resolved_provider, model=model, - ) - headers.update(auth_header) + ), + ) + return self._finalize_messages_headers(headers, optional_params), api_base + + async def avalidate_anthropic_messages_environment( + self, + headers: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract + model: str, + messages: list[Any], # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract + optional_params: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract + litellm_params: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract + api_key: str | None = None, + api_base: str | None = None, + ) -> tuple[dict, str | None]: # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract + if type(self).validate_anthropic_messages_environment is not ( + AnthropicMessagesConfig.validate_anthropic_messages_environment + ): + # a subclass sync override must keep winning on the async path + return self.validate_anthropic_messages_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + oauth_headers, oauth_api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key) + + if not _carries_caller_credential(oauth_headers): + self._apply_env_auth_header( + oauth_headers, + self._require_auth_header( + await AnthropicModelInfo.aget_auth_header( + oauth_api_key, + api_base=api_base, + litellm_params=litellm_params, + allow_workload_identity=self._allows_workload_identity, + ), + model=model, + ), + ) + return self._finalize_messages_headers(oauth_headers, optional_params), api_base + + def _require_auth_header(self, auth_header: Mapping[str, str] | None, model: str) -> Mapping[str, str]: + if auth_header is None: + raise AuthenticationError( + message=( + "Missing Anthropic API Key - A call is being made to anthropic but no key is set " + "either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` " + "or `ANTHROPIC_AUTH_TOKEN` in your environment vars" + ), + llm_provider=self._resolved_provider, + model=model, + ) + return auth_header + + @staticmethod + def _apply_env_auth_header(headers: dict, auth_header: Mapping[str, str] | None) -> None: # mutable-ok: out-param + if auth_header is None: + return + merged_beta: Final = merge_anthropic_beta_headers( + headers.get("anthropic-beta"), auth_header.get("anthropic-beta") + ) + headers.update(auth_header) + if merged_beta: + headers["anthropic-beta"] = merged_beta + + @property + def _allows_workload_identity(self) -> bool: + """Subclasses reuse this validate step for their own /v1/messages-compatible providers, so + eligibility is declared per class and never inherited.""" + from litellm.llms.anthropic.common_utils import config_allows_workload_identity + + return config_allows_workload_identity(self) + + def _finalize_messages_headers(self, headers: dict, optional_params: dict) -> dict: # mutable-ok: out-param if "anthropic-version" not in headers: headers["anthropic-version"] = DEFAULT_ANTHROPIC_API_VERSION if "content-type" not in headers: headers["content-type"] = "application/json" - - headers = self._update_headers_with_anthropic_beta( + return self._update_headers_with_anthropic_beta( headers=headers, optional_params=optional_params, ) - return headers, api_base - @staticmethod def _translate_reasoning_effort_to_anthropic( model: str, optional_params: dict, max_tokens: int | None, custom_llm_provider: str diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index dfd62ca575b..2c678310032 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -73,7 +73,9 @@ class AnthropicFilesHandler: # Get Anthropic API credentials api_base = self.anthropic_model_info.get_api_base(api_base) - auth_header: Final = self.anthropic_model_info.get_auth_header(api_key, api_base) + auth_header: Final = await self.anthropic_model_info.aget_auth_header( + api_key, api_base, allow_workload_identity=True + ) if auth_header is None: raise ValueError("Missing Anthropic API Key") diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index 7b5ab78af8d..04d057ed3e9 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -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"] diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index 448e2dc2584..bb38f443a36 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -2,6 +2,7 @@ Anthropic Skills API configuration and transformations """ +from types import MappingProxyType from typing import Final import httpx @@ -35,40 +36,35 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict: """Add Anthropic-specific headers""" - from litellm.llms.anthropic.common_utils import AnthropicModelInfo + from litellm.constants import ANTHROPIC_SKILLS_API_BETA_VERSION + from litellm.llms.anthropic.common_utils import ( + AnthropicModelInfo, + merge_anthropic_beta_headers, + without_caller_credential_headers, + ) - # Get API key from litellm_params if available - api_key = None - api_base = None - if litellm_params is not None: - api_key = litellm_params.api_key - api_base = litellm_params.api_base - - auth_header: Final = AnthropicModelInfo.get_auth_header(api_key, api_base) + auth_header: Final = AnthropicModelInfo.get_auth_header( + api_key=litellm_params.api_key if litellm_params is not None else None, + api_base=litellm_params.api_base if litellm_params is not None else None, + litellm_params=MappingProxyType(dict(litellm_params)) if litellm_params is not None else None, + allow_workload_identity=True, + ) if auth_header is None: raise ValueError("ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN is required for Skills API") - headers.update(auth_header) - headers["anthropic-version"] = "2023-06-01" - - # Add beta header for skills API - from litellm.constants import ANTHROPIC_SKILLS_API_BETA_VERSION - - if "anthropic-beta" not in headers: - headers["anthropic-beta"] = ANTHROPIC_SKILLS_API_BETA_VERSION - elif isinstance(headers["anthropic-beta"], list): - if ANTHROPIC_SKILLS_API_BETA_VERSION not in headers["anthropic-beta"]: - headers["anthropic-beta"].append(ANTHROPIC_SKILLS_API_BETA_VERSION) - elif isinstance(headers["anthropic-beta"], str): - if ANTHROPIC_SKILLS_API_BETA_VERSION not in headers["anthropic-beta"]: - headers["anthropic-beta"] = [ - headers["anthropic-beta"], - ANTHROPIC_SKILLS_API_BETA_VERSION, - ] - - headers["content-type"] = "application/json" - - return headers + merged_beta: Final = merge_anthropic_beta_headers( + merge_anthropic_beta_headers(headers.get("anthropic-beta"), auth_header.get("anthropic-beta")), + ANTHROPIC_SKILLS_API_BETA_VERSION, + ) + # The deployment's own credential is applied here, so a caller-supplied one must not ride + # along upstream beside a minted federation Bearer. + return { # mutable-ok: validate_environment's contract returns a real dict, which httpx then consumes + **without_caller_credential_headers(headers), + **auth_header, + "anthropic-version": "2023-06-01", + "anthropic-beta": merged_beta, + "content-type": "application/json", + } def get_complete_url( self, diff --git a/litellm/llms/anthropic/wif.py b/litellm/llms/anthropic/wif.py new file mode 100644 index 00000000000..2c28aec0413 --- /dev/null +++ b/litellm/llms/anthropic/wif.py @@ -0,0 +1,501 @@ +"""Anthropic workload identity federation: exchanges an external OIDC identity +token for a short-lived ``sk-ant-oat01`` token via the shared RFC 7523 engine.""" + +import os +from collections.abc import Callable, Mapping +from itertools import chain +from types import MappingProxyType +from typing import Final, NoReturn, TypeVar +from urllib.parse import urlsplit, urlunsplit + +from pydantic import BaseModel, ConfigDict, ValidationError +from typing_extensions import assert_never + +import litellm +from litellm.llms.base_llm.auth.client_credentials import keycloak_assertion_source +from litellm.llms.base_llm.auth.identity_source import ( + AnthropicIdentitySourceKind, + InternalIssuerSource, + KeycloakSource, + identity_source_ref, +) +from litellm.llms.base_llm.auth.internal_issuer import internal_issuer_assertion_source +from litellm.llms.base_llm.auth.token_exchange import ( + JwtBearerTokenExchangeEngine, + default_token_exchange_engine, +) +from litellm.llms.base_llm.auth.types import ( + AssertionSourceError, + ExchangeError, + ExchangeResult, + InsecureTokenUrl, + MalformedTokenResponse, + MintedToken, + TokenEndpointError, + TokenExchangeSpec, + TokenTransportError, +) +from litellm.types.llms.anthropic import ANTHROPIC_TOKEN_EXCHANGE_PATH + +_JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer" +_DEFAULT_API_BASE: Final = "https://api.anthropic.com" +_INLINE_ENV_VAR: Final = "ANTHROPIC_IDENTITY_TOKEN" +_DISABLE_WIF_PARAM: Final = "anthropic_disable_workload_identity_federation" +_ACCEPTED_REF_PREFIX: Final = "oidc/" +_CHAT_BASE_SUFFIXES: Final = ("/v1/messages", "/v1") +# Hosts a federated exchange may talk to. api_base decides where the workload's assertion is sent +# AND where the minted org-scoped token is presented, so anyone able to write api_base on a +# federated deployment could otherwise redirect both. Gating each write path does not terminate: +# a deployment, a referenced credential and a future endpoint all reach the same value. This is the +# one place a federated exchange is built, so the trust decision is enforced here instead, and the +# allowlist is server-owned -- read from the environment, never from a model or credential API. +_TRUSTED_EXCHANGE_HOSTS_ENV: Final = "LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS" +_DEFAULT_TRUSTED_EXCHANGE_HOST: Final = "api.anthropic.com" +_REJECTED_REF_PREFIX: Final = "oidc/env_path/" +_IDENTITY_SOURCE_PARAM: Final = "anthropic_identity_source" +_IDENTITY_SOURCE_ENV: Final = "ANTHROPIC_IDENTITY_SOURCE" +_IDENTITY_TOKEN_FILE_PARAM: Final = "anthropic_identity_token_file" +_IDENTITY_TOKEN_PARAM: Final = "anthropic_identity_token" + +# litellm_params key -> InternalIssuerSource/KeycloakSource field name. Every key here must +# also be listed in ANTHROPIC_WIF_KWARGS_KEYS (get_litellm_params.py), which is what makes it +# request-banned and cleared on a client-redirected api_base -- see types/utils.py's +# anthropic_wif_litellm_params, derived from that same set. +_INTERNAL_ISSUER_FIELD_MAP: Final[Mapping[str, str]] = MappingProxyType( + { + "anthropic_issuer_url": "issuer_url", + "anthropic_issuer_subject": "subject", + "anthropic_issuer_audience": "audience", + "anthropic_issuer_ttl_seconds": "ttl_seconds", + "anthropic_issuer_signing_key_ref": "signing_key_ref", + } +) +_KEYCLOAK_FIELD_MAP: Final[Mapping[str, str]] = MappingProxyType( + { + "anthropic_keycloak_token_url": "token_url", + "anthropic_keycloak_client_id": "client_id", + "anthropic_keycloak_auth_method": "auth_method", + "anthropic_keycloak_client_secret_ref": "client_secret_ref", + "anthropic_keycloak_scope": "scope", + } +) +_DENIAL_HINT: Final = ( + " Anthropic answers every denied exchange with the same 401; the reason (for example" + " workspace_id_required or jti_reused) is only shown in the Claude Console under" + " Settings > Workload identity, in the rule's authentication history." +) +_WORKSPACE_HINT: Final = ( + " If the federation rule is enabled in more than one workspace, set anthropic_workspace_id" + " (or ANTHROPIC_WORKSPACE_ID) to the wrkspc_ id of the workspace to mint tokens for, or to 'default'." +) +_SERVICE_ACCOUNT_HINT: Final = ( + " Anthropic's reference lists service_account_id as required: set anthropic_service_account_id" + " (or ANTHROPIC_SERVICE_ACCOUNT_ID) to the svac_ id the federation rule targets." +) +_MISSING_IDS_HINT: Final = ( + " Copy them from the federation rule's detail page under Settings > Workload identity in the" + " Claude Console, or set ANTHROPIC_FEDERATION_RULE_ID and ANTHROPIC_ORGANIZATION_ID." +) +_ALLOWLIST_HINT: Final = ( + " Identity token files must sit under an allowed credential directory" + " (/var/run/secrets or /run/secrets by default);" + " set LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS to extend the allowlist." +) +_EMPTY_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) + +_IdentitySourceVariant = TypeVar("_IdentitySourceVariant", bound="InternalIssuerSource | KeycloakSource") + + +class AnthropicWifParams(BaseModel): + model_config = ConfigDict(frozen=True) + + federation_rule_id: str + organization_id: str + service_account_id: str | None = None + workspace_id: str | None = None + assertion_ref: str + assertion_source: Callable[[], str | None] | None = None + + +def resolve_anthropic_wif_params(litellm_params: Mapping[str, object] | None) -> AnthropicWifParams | None: + if litellm_params is not None and litellm_params.get(_DISABLE_WIF_PARAM) is True: + return None + federation_rule_id: Final = _config_value( + litellm_params, "anthropic_federation_rule_id", "ANTHROPIC_FEDERATION_RULE_ID" + ) + organization_id: Final = _config_value(litellm_params, "anthropic_organization_id", "ANTHROPIC_ORGANIZATION_ID") + if federation_rule_id is None or organization_id is None: + _raise_if_identity_source_configured(litellm_params, federation_rule_id, organization_id) + return None + identity_source: Final = _resolve_identity_source(litellm_params) + if identity_source is None: + return None + assertion_ref, assertion_source = identity_source + return AnthropicWifParams( + federation_rule_id=federation_rule_id, + organization_id=organization_id, + service_account_id=_config_value( + litellm_params, "anthropic_service_account_id", "ANTHROPIC_SERVICE_ACCOUNT_ID" + ), + workspace_id=_config_value(litellm_params, "anthropic_workspace_id", "ANTHROPIC_WORKSPACE_ID"), + assertion_ref=assertion_ref, + assertion_source=assertion_source, + ) + + +def _resolve_identity_source( + litellm_params: Mapping[str, object] | None, +) -> tuple[str, Callable[[], str] | None] | None: + """Dispatches on ``anthropic_identity_source``. Absent (the default) keeps today's + token_file/env resolution byte-identical, with no ``assertion_source`` closure -- the engine + falls back to its own reader exactly as it does today. A recognized kind builds the matching + frozen config, hashes it into the ``oidc//`` cache-key ref (``identity_source_ref``), + and closes the source's fetch/mint function over it. An unset-but-invalid config (unknown + kind, a missing required field, or a field from the other variant) fails closed here rather + than silently falling back to token_file. A deployment whose params carry a legacy token or + token_file ref stays on legacy resolution even when ``ANTHROPIC_IDENTITY_SOURCE`` names a + fleet-wide kind: the env kind only governs deployments that set no identity params of their own.""" + source_kind: Final = _resolve_source_kind(litellm_params) + if source_kind is None: + legacy_ref: Final = _resolve_assertion_ref(litellm_params) + return (legacy_ref, None) if legacy_ref is not None else None + params: Final[Mapping[str, object]] = MappingProxyType( + {key: value for key, value in (litellm_params or _EMPTY_PARAMS).items() if value is not None} + ) + match source_kind: + case AnthropicIdentitySourceKind.internal_issuer.value: + _reject_foreign_variant_fields(params, foreign_field_map=_KEYCLOAK_FIELD_MAP, chosen_kind=source_kind) + issuer_config: Final = _build_variant(InternalIssuerSource, params, _INTERNAL_ISSUER_FIELD_MAP) + return identity_source_ref(issuer_config), internal_issuer_assertion_source(issuer_config) + case AnthropicIdentitySourceKind.keycloak.value: + _reject_foreign_variant_fields( + params, foreign_field_map=_INTERNAL_ISSUER_FIELD_MAP, chosen_kind=source_kind + ) + keycloak_config: Final = _build_variant(KeycloakSource, params, _KEYCLOAK_FIELD_MAP) + return identity_source_ref(keycloak_config), keycloak_assertion_source(keycloak_config) + case _: + _raise_unknown_source_kind(source_kind) + + +def _raise_unknown_source_kind(source_kind: str) -> NoReturn: + raise litellm.AuthenticationError( + message=( + f"{_IDENTITY_SOURCE_PARAM} must be one of " + f"{', '.join(kind.value for kind in AnthropicIdentitySourceKind)}; got {source_kind!r}." + ), + llm_provider="anthropic", + model="", + ) + + +def _raise_if_identity_source_configured( + litellm_params: Mapping[str, object] | None, federation_rule_id: str | None, organization_id: str | None +) -> None: + """A configured identity source is an explicit request to federate, so a missing rule or + organization id fails closed with the ids named, rather than silently skipping federation + and surfacing later as a missing API key.""" + source_kind: Final = _resolve_source_kind(litellm_params) + if source_kind is None: + return + if source_kind not in {kind.value for kind in AnthropicIdentitySourceKind}: + _raise_unknown_source_kind(source_kind) + missing: Final = tuple( + param + for param, value in ( + ("anthropic_federation_rule_id", federation_rule_id), + ("anthropic_organization_id", organization_id), + ) + if value is None + ) + raise litellm.AuthenticationError( + message=( + f"{_IDENTITY_SOURCE_PARAM} is {source_kind!r}, but {' and '.join(missing)} " + f"{'is' if len(missing) == 1 else 'are'} not set.{_MISSING_IDS_HINT}" + ), + llm_provider="anthropic", + model="", + ) + + +def _resolve_source_kind(litellm_params: Mapping[str, object] | None) -> str | None: + param_kind: Final = _param_str(litellm_params, _IDENTITY_SOURCE_PARAM) + if param_kind is not None: + return param_kind + has_param_legacy_ref: Final = any( + _param_str(litellm_params, key) is not None for key in (_IDENTITY_TOKEN_FILE_PARAM, _IDENTITY_TOKEN_PARAM) + ) + return None if has_param_legacy_ref else _env_str(_IDENTITY_SOURCE_ENV) + + +def _reject_foreign_variant_fields( + litellm_params: Mapping[str, object], foreign_field_map: Mapping[str, str], chosen_kind: str +) -> None: + foreign_keys_present: Final = tuple(param for param in foreign_field_map if param in litellm_params) + if foreign_keys_present: + raise litellm.AuthenticationError( + message=( + f"{_IDENTITY_SOURCE_PARAM} is {chosen_kind!r}, but {', '.join(sorted(foreign_keys_present))} " + "belongs to a different identity source and cannot be set alongside it." + ), + llm_provider="anthropic", + model="", + ) + + +def _build_variant( + model: type[_IdentitySourceVariant], + litellm_params: Mapping[str, object], + field_map: Mapping[str, str], +) -> _IdentitySourceVariant: + fields: Final = MappingProxyType( + {field_map[key]: value for key, value in litellm_params.items() if key in field_map} + ) + try: + return model.model_validate(fields) + except ValidationError as e: + # hide_input_in_errors=True on both variant models keeps a secret pasted into the + # wrong field (e.g. a client_secret typed as signing_key_ref) out of str(e). + raise litellm.AuthenticationError( + message=f"Invalid {_IDENTITY_SOURCE_PARAM} configuration: {e}", + llm_provider="anthropic", + model="", + ) from e + + +def build_anthropic_wif_spec(params: AnthropicWifParams, api_base: str) -> TokenExchangeSpec: + return TokenExchangeSpec( + token_url=api_base.rstrip("/") + ANTHROPIC_TOKEN_EXCHANGE_PATH, + assertion_ref=params.assertion_ref, + assertion_field="assertion", + static_body=MappingProxyType( + { + name: value + for name, value in ( + ("grant_type", _JWT_BEARER_GRANT_TYPE), + ("federation_rule_id", params.federation_rule_id), + ("organization_id", params.organization_id), + ("service_account_id", params.service_account_id), + ("workspace_id", params.workspace_id), + ) + if value is not None + } + ), + body_encoding="json", + request_headers=MappingProxyType({}), + assertion_source=params.assertion_source, + cache_key_identity=( + params.federation_rule_id, + params.organization_id, + params.service_account_id or "", + params.workspace_id or "", + ), + ) + + +def get_anthropic_wif_token( + litellm_params: Mapping[str, object] | None, + api_base: str | None, + model: str, + engine: JwtBearerTokenExchangeEngine = default_token_exchange_engine, +) -> str | None: + params: Final = resolve_anthropic_wif_params(litellm_params) + if params is None: + return None + exchange_base: Final = resolve_anthropic_base(api_base) + _raise_if_exchange_host_untrusted(exchange_base, model) + result: Final = engine.get_token(build_anthropic_wif_spec(params, exchange_base)) + return _token_from_result(result, model, params) + + +async def aget_anthropic_wif_token( + litellm_params: Mapping[str, object] | None, + api_base: str | None, + model: str, + engine: JwtBearerTokenExchangeEngine = default_token_exchange_engine, +) -> str | None: + params: Final = resolve_anthropic_wif_params(litellm_params) + if params is None: + return None + exchange_base: Final = resolve_anthropic_base(api_base) + _raise_if_exchange_host_untrusted(exchange_base, model) + result: Final = await engine.aget_token(build_anthropic_wif_spec(params, exchange_base)) + return _token_from_result(result, model, params) + + +def _token_from_result(result: ExchangeResult, model: str, params: AnthropicWifParams) -> str: + match result: + case MintedToken(): + return result.access_token.get_secret_value() + case _: + _raise_anthropic_wif_error( + result, + model=model, + workspace_id_set=params.workspace_id is not None, + service_account_id_set=params.service_account_id is not None, + ) + + +def resolve_anthropic_base(api_base: str | None) -> str: + """The base every Anthropic tier derives its URLs from: the deployment api_base when set, + else ``ANTHROPIC_API_BASE`` / ``ANTHROPIC_BASE_URL``, else Anthropic's host, with trailing + slashes and chat-appended ``/v1/messages`` suffixes stripped, so the token URL, the cache key + and the count-tokens URL all agree for the same deployment.""" + return anthropic_base_without_chat_suffix(api_base or _resolve_default_api_base()) + + +def _trusted_exchange_hosts() -> frozenset[str]: + """Hostnames a federated exchange may reach: Anthropic's own, plus whatever the operator put in + the environment. Comma separated, case folded, entries given as a URL reduced to their host.""" + configured: Final = os.getenv(_TRUSTED_EXCHANGE_HOSTS_ENV) or "" + extra: Final = (entry.strip() for entry in configured.split(",") if entry.strip()) + return frozenset( + chain( + (_DEFAULT_TRUSTED_EXCHANGE_HOST,), + ((urlsplit(entry).hostname or entry.split("/")[0]).lower() for entry in extra), + ) + ) + + +def _raise_if_exchange_host_untrusted(exchange_base: str, model: str) -> None: + """The federated exchange refuses any host the operator has not vouched for, whatever wrote the + deployment's api_base. Exact hostname match, never a substring: ``api.anthropic.com.evil.test`` + contains the real host and must not pass.""" + host: Final = (urlsplit(exchange_base).hostname or "").lower() + if host and host in _trusted_exchange_hosts(): + return + raise litellm.AuthenticationError( + message=( + f"Anthropic workload identity federation refused to use host {host or exchange_base!r}. " + f"A federated exchange sends the workload's identity token to this host and presents the " + f"minted token to it, so only {_DEFAULT_TRUSTED_EXCHANGE_HOST} is trusted by default. To " + f"use a private Anthropic-compatible gateway, add its hostname to the " + f"{_TRUSTED_EXCHANGE_HOSTS_ENV} environment variable (comma separated); that is a " + f"decision to trust it with org-scoped credentials, so it is deliberately server-owned " + f"and cannot be set through the model or credential APIs." + ), + llm_provider="anthropic", + model=model, + ) + + +def _resolve_default_api_base() -> str: + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + return AnthropicModelInfo.get_api_base(None) or _DEFAULT_API_BASE + + +def anthropic_base_without_chat_suffix(base: str) -> str: + """A deployment base with its chat-surface suffix removed, so the token URL and model + discovery both derive from the same value whatever form the operator configured.""" + parts: Final = urlsplit(base) + if not parts.scheme or not parts.netloc: + return base.rstrip("/") + return urlunsplit((parts.scheme, parts.netloc, _strip_path_suffixes(parts.path), "", "")) + + +def _strip_path_suffixes(path: str) -> str: + """Drop the chat-surface suffixes a deployment base may carry, so every tier derives the same + token URL. Each pass removes at most one suffix, so the loop is bounded by the segment count.""" + trimmed = path.rstrip("/") # rebind-ok: fixed-point strip, one suffix per pass + while True: + shortened = next( # rebind-ok: one suffix removed per iteration + (trimmed.removesuffix(suffix) for suffix in _CHAT_BASE_SUFFIXES if trimmed.endswith(suffix)), + trimmed, + ) + if shortened == trimmed: + return trimmed + # Re-strip: a doubled suffix leaves a trailing slash that would stop the next match. + trimmed = shortened.rstrip("/") + + +def _config_value(litellm_params: Mapping[str, object] | None, param_key: str, env_name: str) -> str | None: + return _param_str(litellm_params, param_key) or _env_str(env_name) + + +def _param_str(litellm_params: Mapping[str, object] | None, key: str) -> str | None: + if litellm_params is None: + return None + value: Final = litellm_params.get(key) + return value if isinstance(value, str) and value else None + + +def _env_str(name: str) -> str | None: + from litellm.secret_managers.main import get_secret_str + + value: Final = get_secret_str(name) + return value if isinstance(value, str) and value else None + + +def _resolve_assertion_ref(litellm_params: Mapping[str, object] | None) -> str | None: + file_param: Final = _param_str(litellm_params, _IDENTITY_TOKEN_FILE_PARAM) + if file_param is not None: + return f"oidc/file/{file_param}" + inline_param: Final = _param_str(litellm_params, _IDENTITY_TOKEN_PARAM) + if inline_param is not None: + return _validated_inline_ref(inline_param) + file_env: Final = _env_str("ANTHROPIC_IDENTITY_TOKEN_FILE") + if file_env is not None: + return f"oidc/file/{file_env}" + if _env_str(_INLINE_ENV_VAR) is not None: + return f"oidc/env/{_INLINE_ENV_VAR}" + return None + + +def _validated_inline_ref(value: str) -> str: + if value.startswith(_ACCEPTED_REF_PREFIX) and not value.startswith(_REJECTED_REF_PREFIX): + return value + raise litellm.AuthenticationError( + message=( + "anthropic_identity_token must be an oidc/ secret reference such as oidc/env/VAR_NAME," + " oidc/file//absolute/path, oidc/github/, or oidc/google/." + " Raw identity tokens and oidc/env_path/ references are not accepted;" + " to pass a token directly, export it and reference it as oidc/env/VAR_NAME" + ), + llm_provider="anthropic", + model="", + ) + + +def _raise_anthropic_wif_error( + error: ExchangeError, model: str, workspace_id_set: bool, service_account_id_set: bool +) -> NoReturn: + detail: Final = _error_detail( + error, workspace_id_set=workspace_id_set, service_account_id_set=service_account_id_set + ) + raise litellm.AuthenticationError( + message=f"Anthropic workload identity federation failed. {detail}", + llm_provider="anthropic", + model=model, + ) + + +def _denial_hints(workspace_id_set: bool, service_account_id_set: bool) -> str: + return "".join( + ( + _DENIAL_HINT, + "" if workspace_id_set else _WORKSPACE_HINT, + "" if service_account_id_set else _SERVICE_ACCOUNT_HINT, + ) + ) + + +def _error_detail(error: ExchangeError, workspace_id_set: bool, service_account_id_set: bool) -> str: + match error: + case AssertionSourceError() if error.kind == "disallowed_path": + return f"Could not read the OIDC identity token from {error.source_ref}.{_ALLOWLIST_HINT}" + case AssertionSourceError(): + base: Final = f"Could not obtain the OIDC identity token ({error.kind}) from {error.source_ref}." + return f"{base} {error.detail}" if error.detail else base + case InsecureTokenUrl(): + return f"The token endpoint must use https; refusing to send the identity token to host {error.host!r}." + case TokenEndpointError() if error.status_code == 401: + hints: Final = _denial_hints(workspace_id_set, service_account_id_set) + return f"The token endpoint returned HTTP 401: {error.redacted_body}{hints}" + case TokenEndpointError(): + return f"The token endpoint returned HTTP {error.status_code}: {error.redacted_body}" + case TokenTransportError(): + return f"Could not reach the token endpoint: {error.detail}." + case MalformedTokenResponse(): + return f"The token endpoint returned an unusable response: {error.detail}." + case _: + assert_never(error) diff --git a/litellm/llms/azure_ai/embed/handler.py b/litellm/llms/azure_ai/embed/handler.py index c65edbf56e6..8037be8abeb 100644 --- a/litellm/llms/azure_ai/embed/handler.py +++ b/litellm/llms/azure_ai/embed/handler.py @@ -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 diff --git a/litellm/llms/base_llm/anthropic_messages/transformation.py b/litellm/llms/base_llm/anthropic_messages/transformation.py index 8e7c22930fa..2d2b5afd824 100644 --- a/litellm/llms/base_llm/anthropic_messages/transformation.py +++ b/litellm/llms/base_llm/anthropic_messages/transformation.py @@ -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, diff --git a/litellm/llms/base_llm/auth/__init__.py b/litellm/llms/base_llm/auth/__init__.py new file mode 100644 index 00000000000..291e1492c98 --- /dev/null +++ b/litellm/llms/base_llm/auth/__init__.py @@ -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", +) diff --git a/litellm/llms/base_llm/auth/client_credentials.py b/litellm/llms/base_llm/auth/client_credentials.py new file mode 100644 index 00000000000..b2e683a9ed4 --- /dev/null +++ b/litellm/llms/base_llm/auth/client_credentials.py @@ -0,0 +1,227 @@ +"""Fetches a fresh RFC 6749 client_credentials assertion for Anthropic's ``keycloak`` identity +source: LiteLLM authenticates to Keycloak as its own confidential client and presents the +resulting ``access_token`` as the workload assertion (Phase 1 decision 2). + +The client secret is the operator-supplied pointer at ``KeycloakSource.client_secret_ref``, +resolved the same way every other WIF secret pointer already is (env, a Credential, or whatever +secret manager ``litellm.secret_manager_client`` is globally configured to, Vault included). +Every fetch is a fresh HTTP POST; nothing here caches a fetched token, since the outer +token-exchange engine already caches the Anthropic token it buys with one -- see decision 2's +"no Keycloak-side cache" ruling. +""" + +import base64 +import threading +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias +from urllib.parse import quote, quote_plus, urlencode + +import httpx +from pydantic import BaseModel, SecretStr, ValidationError +from typing_extensions import assert_never + +from litellm.llms.base_llm.auth.identity_source import KeycloakSource, ref_for_error_message +from litellm.llms.base_llm.auth.token_exchange import ( + MAX_RESPONSE_BYTES, + endpoint_url_for_error_message, + redact_oauth_error_body, + require_posted_response, + validate_token_endpoint_url, +) +from litellm.llms.base_llm.auth.types import InsecureTokenUrl, SyncTokenPoster + +if TYPE_CHECKING: + from litellm.llms.custom_httpx.http_handler import HTTPHandler + +SecretReader: TypeAlias = Callable[[str], str | None] # mutable-ok: Callable param-list syntax, not a list + +_GRANT_TYPE: Final = "client_credentials" +_TIMEOUT_SECONDS: Final = 30.0 +_FORM_CONTENT_TYPE: Final = "application/x-www-form-urlencoded" + + +class _ClientCredentialsResponse(BaseModel): + access_token: str + + +def _default_secret_reader(ref: str) -> str | None: + from litellm.secret_managers.main import get_secret_str + + return get_secret_str(ref) + + +def _new_keycloak_handler() -> "HTTPHandler": + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + handler: Final = HTTPHandler(timeout=httpx.Timeout(timeout=30.0, connect=5.0)) + handler.client.follow_redirects = False + return handler + + +class _HttpxSyncKeycloakPoster: + """Dedicated HTTPHandler for the Keycloak token POST: no ``logging_obj`` (so litellm's + request/response logging never sees the client secret or the fetched token), redirects + disabled. A separate instance from the outer engine's own poster, since this is a genuinely + new HTTP call site whose no-logging guarantee must be built here, not assumed inherited.""" + + def __init__(self, handler_factory: Callable[[], "HTTPHandler"] = _new_keycloak_handler) -> None: + self._lock: Final = threading.Lock() + self._handler_factory: Final = handler_factory + self._handler: HTTPHandler | None = None + + def _handler_instance(self) -> "HTTPHandler": + with self._lock: + if self._handler is None: + self._handler = self._handler_factory() + return self._handler + + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + try: + response: Final[httpx.Response | None] = self._handler_instance().post( # pyright: ignore[reportUnknownMemberType] # HTTPHandler.post is legacy-untyped; the result is validated below + url, + content=content, + headers=dict(headers), # mutable-ok: HTTPHandler.post requires a concrete dict + timeout=timeout, + ) + except httpx.HTTPStatusError as e: + return e.response + return require_posted_response(response, "keycloak token endpoint") + + +_DEFAULT_POSTER: Final[SyncTokenPoster] = _HttpxSyncKeycloakPoster() + + +def _form_encode(value: str) -> str: + """RFC 6749 Appendix B before RFC 6749 2.3.1's base64: application/x-www-form-urlencoded + with spaces as ``%20`` rather than ``+``, else a reserved character (":", "+", "%", " ") in + the id or secret corrupts the credential the far side decodes back out of Basic auth.""" + return quote(value, safe="") + + +def _basic_auth_header(client_id: str, client_secret: str) -> str: + encoded_pair: Final = f"{_form_encode(client_id)}:{_form_encode(client_secret)}" + return "Basic " + base64.b64encode(encoded_pair.encode()).decode("ascii") + + +def _prepared_request(config: KeycloakSource, client_secret: str) -> tuple[bytes, Mapping[str, str]]: + scope_field: Final[Mapping[str, str]] = ( + MappingProxyType({"scope": config.scope}) if config.scope else MappingProxyType({}) + ) + match config.auth_method: + case "client_secret_basic": + return ( + urlencode(MappingProxyType({"grant_type": _GRANT_TYPE, **scope_field})).encode(), + MappingProxyType( + { + "content-type": _FORM_CONTENT_TYPE, + "authorization": _basic_auth_header(config.client_id, client_secret), + } + ), + ) + case "client_secret_post": + return ( + urlencode( + MappingProxyType( + { + "grant_type": _GRANT_TYPE, + "client_id": config.client_id, + "client_secret": client_secret, + **scope_field, + } + ) + ).encode(), + MappingProxyType({"content-type": _FORM_CONTENT_TYPE}), + ) + case _: + assert_never(config.auth_method) + + +def _resolve_client_secret(config: KeycloakSource, secret_reader: SecretReader) -> str: + secret: Final = secret_reader(config.client_secret_ref) + if not secret: + raise ValueError(f"keycloak client secret {ref_for_error_message(config.client_secret_ref)} could not be read") + return secret + + +def _wire_forms_of_secret(config: KeycloakSource, client_secret: str) -> tuple[SecretStr, ...]: + """Every shape the secret leaves this process in, so an echo of any of them is caught. + + Neither grant sends the secret verbatim. client_secret_basic base64s ``id:secret``, which + decodes straight back to it, and client_secret_post percent-escapes it. An endpoint echoing + either shape hands over reversible material a raw comparison would miss. + """ + raw: Final = SecretStr(client_secret) + match config.auth_method: + case "client_secret_basic": + encoded_pair: Final = f"{_form_encode(config.client_id)}:{_form_encode(client_secret)}" + return (raw, SecretStr(base64.b64encode(encoded_pair.encode()).decode("ascii"))) + case "client_secret_post": + # urlencode escapes reserved characters and writes a space as "+", so a secret + # containing either leaves in a shape the raw comparison would not recognise coming + # back. quote_plus is what urlencode itself applies. + return (raw, SecretStr(quote_plus(client_secret))) + case _: + assert_never(config.auth_method) + + +def _endpoint_error_message(config: KeycloakSource, response: httpx.Response, client_secret: str) -> str: + endpoint_error: Final = redact_oauth_error_body( + response.status_code, response.text, _wire_forms_of_secret(config, client_secret) + ) + return ( + f"keycloak token endpoint {endpoint_url_for_error_message(config.token_url)} " + f"returned HTTP {endpoint_error.status_code}: {endpoint_error.redacted_body}" + ) + + +def _parse_success_body(response: httpx.Response) -> str: + if len(response.content) > MAX_RESPONSE_BYTES: + raise ValueError("keycloak token response exceeded the size cap") + try: + parsed: Final = _ClientCredentialsResponse.model_validate_json(response.content) + except ValidationError as e: + raise ValueError("keycloak token response failed schema validation") from e + token: Final = parsed.access_token.strip() + if not token: + raise ValueError("keycloak token response carried an empty access_token") + return token + + +def fetch_keycloak_assertion( + config: KeycloakSource, + *, + poster: SyncTokenPoster = _DEFAULT_POSTER, + secret_reader: SecretReader = _default_secret_reader, +) -> str: + """POSTs one fresh client_credentials grant and returns the resulting ``access_token`` as the + workload assertion; the caller must not cache the result -- see the module docstring.""" + match validate_token_endpoint_url(config.token_url): + case InsecureTokenUrl(host=host): + raise ValueError(f"keycloak token_url must use https; refusing to send the client secret to host {host!r}") + case _: + pass + client_secret: Final = _resolve_client_secret(config, secret_reader) + content, headers = _prepared_request(config, client_secret) + try: + response: Final = poster.post(config.token_url, content=content, headers=headers, timeout=_TIMEOUT_SECONDS) + except Exception as e: # noqa: BLE001 # injected posters may raise beyond httpx; every failure becomes a ValueError + raise ValueError( + f"could not reach the keycloak token endpoint {endpoint_url_for_error_message(config.token_url)}: " + f"{type(e).__name__}" + ) from e + if not 200 <= response.status_code < 300: + raise ValueError(_endpoint_error_message(config, response, client_secret)) + return _parse_success_body(response) + + +def keycloak_assertion_source( + config: KeycloakSource, + *, + poster: SyncTokenPoster = _DEFAULT_POSTER, + secret_reader: SecretReader = _default_secret_reader, +) -> Callable[[], str]: + """A zero-arg closure that fetches fresh on every call: the shape an ``oidc/keycloak/...`` + ref dispatches to once wired into ``TokenExchangeSpec.assertion_source`` (Phase 1 decision 7) + -- the caller parses the config and closes this function over it, with no registry involved.""" + return lambda: fetch_keycloak_assertion(config, poster=poster, secret_reader=secret_reader) diff --git a/litellm/llms/base_llm/auth/identity_source.py b/litellm/llms/base_llm/auth/identity_source.py new file mode 100644 index 00000000000..f8f02249b3c --- /dev/null +++ b/litellm/llms/base_llm/auth/identity_source.py @@ -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//`` 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//``: 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 "" diff --git a/litellm/llms/base_llm/auth/internal_issuer.py b/litellm/llms/base_llm/auth/internal_issuer.py new file mode 100644 index 00000000000..444ca7fd5b2 --- /dev/null +++ b/litellm/llms/base_llm/auth/internal_issuer.py @@ -0,0 +1,86 @@ +"""Mints a self-issued workload assertion for Anthropic's ``internal_issuer`` identity source: +LiteLLM signs its own short-lived ES256 JWT instead of reading one from a mounted OIDC file. + +Signing custody is the operator-supplied PEM at ``InternalIssuerSource.signing_key_ref``, +resolved the same way every other WIF secret pointer already is (env, a Credential, or +whatever secret manager ``litellm.secret_manager_client`` is globally configured to, Vault +included) -- see Phase 1 decision 1. Every mint is fresh; nothing here caches a minted JWT, +since the outer token-exchange engine already caches the Anthropic token it buys with one. +""" + +import time +import uuid +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final, TypeAlias + +from litellm.llms.base_llm.auth.identity_source import InternalIssuerSource, ref_for_error_message +from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json, sign_es256_jwt + +SigningKeyReader: TypeAlias = Callable[[str], str | None] # mutable-ok: Callable param-list syntax, not a list + + +def _default_signing_key_reader(ref: str) -> str | None: + from litellm.secret_managers.main import get_secret_str + + return get_secret_str(ref) + + +def _claims(config: InternalIssuerSource, issued_at: int) -> Mapping[str, object]: + return MappingProxyType( + { + key: value + for key, value in ( + ("sub", config.subject), + ("iss", config.issuer_url), + ("aud", config.audience), + ("iat", issued_at), + ("exp", issued_at + config.ttl_seconds), + ("jti", str(uuid.uuid4())), + ) + if value is not None + } + ) + + +def _resolve_signing_key(config: InternalIssuerSource, key_reader: SigningKeyReader) -> str: + pem: Final = key_reader(config.signing_key_ref) + if not pem: + raise ValueError( + f"internal_issuer signing key {ref_for_error_message(config.signing_key_ref)} could not be read" + ) + return pem + + +def mint_internal_issuer_assertion( + config: InternalIssuerSource, + *, + key_reader: SigningKeyReader = _default_signing_key_reader, + clock: Callable[[], float] = time.time, +) -> str: + """Signs one fresh, short-lived assertion; the caller must not cache the result, since a + cached copy would defeat the point of re-minting on every exchange.""" + pem: Final = _resolve_signing_key(config, key_reader) + return sign_es256_jwt(pem, _claims(config, issued_at=int(clock()))) + + +def internal_issuer_assertion_source( + config: InternalIssuerSource, + *, + key_reader: SigningKeyReader = _default_signing_key_reader, + clock: Callable[[], float] = time.time, +) -> Callable[[], str]: + """A zero-arg closure that mints fresh on every call: the shape an ``oidc/internal_issuer/...`` + ref dispatches to once wired into ``TokenExchangeSpec.assertion_source`` (Phase 1 decision 7) + -- the caller parses the config and closes this function over it, with no registry involved.""" + return lambda: mint_internal_issuer_assertion(config, key_reader=key_reader, clock=clock) + + +def internal_issuer_jwks_document( + config: InternalIssuerSource, + *, + key_reader: SigningKeyReader = _default_signing_key_reader, +) -> str: + """The operator-facing JWKS export, resolved from a configured identity source rather than + a raw PEM in hand -- the JSON document to register as Anthropic's inline federation issuer.""" + return jwks_document_json(_resolve_signing_key(config, key_reader)) diff --git a/litellm/llms/base_llm/auth/jwt_signing.py b/litellm/llms/base_llm/auth/jwt_signing.py new file mode 100644 index 00000000000..71aba6fac52 --- /dev/null +++ b/litellm/llms/base_llm/auth/jwt_signing.py @@ -0,0 +1,115 @@ +"""ES256 JWT signing primitives for Anthropic workload identity federation's +``internal_issuer`` identity source (see ``identity_source.InternalIssuerSource``). + +Pure functions over an already-resolved PEM string: no I/O, no secret-manager awareness, no +caching. Given the signing key at, say, $ISSUER_SIGNING_KEY_PEM, an operator publishes the +JWKS document Anthropic's inline federation issuer needs with one line: + + python -c "from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json; \\ + import os; print(jwks_document_json(os.environ['ISSUER_SIGNING_KEY_PEM']))" +""" + +import base64 +import hashlib +import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +if TYPE_CHECKING: + from cryptography.hazmat.primitives.asymmetric import ec + +ALG: Final = "ES256" +MISSING_SIGNING_DEPENDENCIES_MESSAGE: Final = ( + "the internal_issuer identity source needs PyJWT and cryptography, which a base litellm install " + "does not include: pip install 'litellm[proxy]'" +) +_JWK_CURVE_NAME: Final = "P-256" +_JWK_KEY_TYPE: Final = "EC" +_COORDINATE_BYTE_LENGTH: Final = 32 # P-256 field element width, RFC 7518 6.2.1.2/6.2.1.3 + +Jwk: TypeAlias = Mapping[str, str] +Jwks: TypeAlias = Mapping[str, tuple[Jwk, ...]] + + +def load_es256_private_key(pem: str) -> "ec.EllipticCurvePrivateKey": + """Parses an unencrypted PEM EC private key. Never echoes the key material in an error.""" + try: + from cryptography.hazmat.primitives.asymmetric import ec + from cryptography.hazmat.primitives.serialization import load_pem_private_key + except ImportError as e: + raise ImportError(MISSING_SIGNING_DEPENDENCIES_MESSAGE) from e + try: + key: Final = load_pem_private_key(pem.encode(), password=None) + except (ValueError, TypeError) as e: + raise ValueError("internal_issuer signing key is not a valid unencrypted PEM private key") from e + if not isinstance(key, ec.EllipticCurvePrivateKey) or not isinstance(key.curve, ec.SECP256R1): + raise ValueError( # noqa: TRY004 # the reader classifies ValueError into a readable config error; TypeError would not + "internal_issuer signing key must be an EC P-256 (secp256r1) private key for ES256" + ) + return key + + +def _b64url_coordinate(value: int) -> str: + return base64.urlsafe_b64encode(value.to_bytes(_COORDINATE_BYTE_LENGTH, "big")).rstrip(b"=").decode("ascii") + + +def _jwk_thumbprint_members(public_key: "ec.EllipticCurvePublicKey") -> Jwk: + """RFC 7638 3.2's exact EC member set (crv, kty, x, y) and nothing else: an extra member + here would change the thumbprint and desync it from the ``kid`` published in the JWKS.""" + numbers: Final = public_key.public_numbers() + return MappingProxyType( + { + "crv": _JWK_CURVE_NAME, + "kty": _JWK_KEY_TYPE, + "x": _b64url_coordinate(numbers.x), + "y": _b64url_coordinate(numbers.y), + } + ) + + +def rfc7638_thumbprint(public_key: "ec.EllipticCurvePublicKey") -> str: + """RFC 7638: SHA-256 over the lexicographically member-ordered, whitespace-free JSON + rendering of the thumbprint members, base64url-encoded without padding.""" + canonical: Final = json.dumps( + dict(sorted(_jwk_thumbprint_members(public_key).items())), # mutable-ok: json.dumps needs a real dict + separators=(",", ":"), + ) + return base64.urlsafe_b64encode(hashlib.sha256(canonical.encode()).digest()).rstrip(b"=").decode("ascii") + + +def build_jwk(public_key: "ec.EllipticCurvePublicKey", kid: str) -> Jwk: + return MappingProxyType({**_jwk_thumbprint_members(public_key), "use": "sig", "alg": ALG, "kid": kid}) + + +def build_jwks(public_key: "ec.EllipticCurvePublicKey") -> Jwks: + kid: Final = rfc7638_thumbprint(public_key) + return MappingProxyType({"keys": (build_jwk(public_key, kid),)}) + + +def jwks_document_json(pem: str) -> str: + """The operator-facing export: the JSON document to register as Anthropic's inline JWKS. + + ``build_jwks`` returns ``MappingProxyType``/tuple values per this repo's no-mutation + convention; the ``json`` module only knows plain ``dict``/``list``, so those are converted + at this one serialization boundary rather than giving up immutability throughout the module. + """ + key: Final = load_es256_private_key(pem) + jwks: Final = build_jwks(key.public_key()) + return json.dumps( + {"keys": [dict(jwk) for jwk in jwks["keys"]]}, # mutable-ok: json.dumps needs real dicts/lists + indent=2, + ) + + +def sign_es256_jwt(pem: str, claims: Mapping[str, object]) -> str: + """Signs ``claims`` with the PEM key, stamping ``kid`` as its RFC 7638 thumbprint so a + verifier can look the signing key up in the published JWKS by ``kid`` alone.""" + try: + import jwt + except ImportError as e: + raise ImportError(MISSING_SIGNING_DEPENDENCIES_MESSAGE) from e + key: Final = load_es256_private_key(pem) + kid: Final = rfc7638_thumbprint(key.public_key()) + headers: Final = {"kid": kid} # mutable-ok: PyJWT requires a real dict, not a Mapping + return jwt.encode(dict(claims), key, algorithm=ALG, headers=headers) # mutable-ok: PyJWT requires a real dict diff --git a/litellm/llms/base_llm/auth/token_exchange.py b/litellm/llms/base_llm/auth/token_exchange.py new file mode 100644 index 00000000000..be2fb46d1b3 --- /dev/null +++ b/litellm/llms/base_llm/auth/token_exchange.py @@ -0,0 +1,861 @@ +"""RFC 7523 JWT-bearer token exchange engine, shared across providers. + +One sync state machine per process: bounded engine-owned entry map, two-tier +refresh (advisory background refresh + mandatory single-flight), HTTPS pinning, +response caps, and RFC 6749 5.2 redaction. Providers describe a grant profile as +a ``TokenExchangeSpec`` and map the typed ``ExchangeError`` union to their own +public exception contract. +""" + +import asyncio +import hashlib +import json +import re +import threading +import time +from collections.abc import Callable, Coroutine, Mapping, Sequence +from concurrent.futures import Executor, ThreadPoolExecutor +from dataclasses import dataclass +from math import inf +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias +from urllib.parse import unquote, unquote_plus, urlencode, urlsplit, urlunsplit + +import httpx +from pydantic import BaseModel, SecretStr, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm._logging import verbose_logger +from litellm.llms.base_llm.auth.types import ( + AssertionReader, + AssertionSource, + AssertionSourceError, + ExchangeCallType, + ExchangeError, + ExchangeResult, + InsecureTokenUrl, + MalformedTokenResponse, + MintedToken, + SyncTokenPoster, + TokenEndpointError, + TokenExchangeMetricsSink, + TokenExchangeSpec, + TokenTransportError, +) +from litellm.types.services import ServiceTypes + +if TYPE_CHECKING: + from litellm.llms.custom_httpx.http_handler import HTTPHandler + +CALL_TYPE_COLD_MINT: Final[ExchangeCallType] = "cold_mint" +CALL_TYPE_MANDATORY_REFRESH: Final[ExchangeCallType] = "mandatory_refresh" +CALL_TYPE_ADVISORY_REFRESH: Final[ExchangeCallType] = "advisory_refresh" +CALL_TYPE_CACHE_HIT: Final = "cache_hit" + +ADVISORY_REFRESH_SECONDS: Final = 120.0 +MANDATORY_REFRESH_SECONDS: Final = 30.0 +ADVISORY_REFRESH_LIFETIME_FRACTION: Final = 0.5 +MANDATORY_REFRESH_LIFETIME_FRACTION: Final = 0.125 +ADVISORY_REFRESH_BACKOFF_SECONDS: Final = 5.0 +FALLBACK_TOKEN_TTL_SECONDS: Final = 60.0 +# Metrics are best-effort, so the backlog is capped and further events are dropped. Request volume +# must not be able to grow this queue without bound when a telemetry backend stalls. +_METRICS_QUEUE_LIMIT: Final = 1000 +MAX_ASSERTION_BYTES: Final = 16 * 1024 +MAX_RESPONSE_BYTES: Final = 1024 * 1024 + +_REDACTION_CAP: Final = 256 +_FOLLOWER_WAIT_GRACE_SECONDS: Final = 5.0 +_LOCAL_HOSTS: Final = frozenset({"localhost", "127.0.0.1", "::1"}) +_OAUTH_ERROR_FIELDS: Final = ("error", "error_description", "error_uri") +_NESTED_ERROR_FIELDS: Final = ("type", "message") +_CONTENT_TYPES: Final = MappingProxyType({"json": "application/json", "form": "application/x-www-form-urlencoded"}) +_OVERSIZED_BODY_MESSAGE: Final = "oversized error response omitted" +_NON_OBJECT_BODY_MESSAGE: Final = "non-object error response omitted" +_NO_OAUTH_FIELDS_MESSAGE: Final = "error response carried no RFC 6749 fields" +_UNSTRUCTURED_BODY_MESSAGE: Final = "non-JSON error response omitted" +_REFLECTED_VALUE_MESSAGE: Final = "" +# A credential fragment shorter than this is not worth the false positives; longer, and a run +# shared with the assertion is reflection rather than coincidence. +_REFLECTION_MIN_RUN: Final = 8 +# Everything a base64url credential is NOT made of, stripped so a fragment split by delimiters +# still lines up against the assertion. +_CREDENTIAL_CHARS: Final = re.compile(r"[^A-Za-z0-9._~+/=-]") +_SENTINEL_BODY_MESSAGES: Final = frozenset({_OVERSIZED_BODY_MESSAGE, _NON_OBJECT_BODY_MESSAGE}) + + +class _TokenExchangeResponse(BaseModel): + access_token: str + expires_in: int | None = None + token_type: str | None = None + + +_RedactableBody: TypeAlias = Mapping[str, object] | list[object] | str | int | float | bool | None +_REDACTABLE_BODY_ADAPTER: Final = TypeAdapter[_RedactableBody](_RedactableBody) + + +def endpoint_url_for_error_message(url: str) -> str: + """``url`` reduced to scheme, host and path for operator-facing errors. + + A token endpoint is configuration, not a secret, and naming it is what makes these errors + actionable. But nothing stops an operator writing a credential into it, as a query parameter + or as userinfo, and these errors reach model callers, so neither part is echoed. + """ + parsed: Final = urlsplit(url) + host: Final = parsed.hostname or "" + authority: Final = f"{host}:{parsed.port}" if parsed.port is not None else host + return urlunsplit((parsed.scheme, authority, parsed.path, "", "")) + + +def validate_token_endpoint_url(url: str) -> str | InsecureTokenUrl: + parsed: Final = urlsplit(url) + if parsed.scheme == "https": + return url + if parsed.scheme == "http" and (parsed.hostname or "") in _LOCAL_HOSTS: + return url + return InsecureTokenUrl(host=parsed.hostname or "") + + +def redact_oauth_error_body( + status_code: int, + body_text: str, + assertion: SecretStr | Sequence[SecretStr] | None = None, +) -> TokenEndpointError: + """``assertion`` may be every form of the credential that went out on the wire. + + A grant that encodes its credential before sending it (``client_secret_basic`` base64s + ``id:secret``) can have that encoded form echoed back, and it decodes straight to the secret, + so checking only the raw value lets reversible material through. + """ + rendered: Final = _redact_body_text(body_text) + secrets: Final = () if assertion is None else (assertion,) if isinstance(assertion, SecretStr) else tuple(assertion) + redacted: Final = next( + ( + _REFLECTED_VALUE_MESSAGE + for secret in secrets + if _drop_reflected_assertion(rendered, secret) is _REFLECTED_VALUE_MESSAGE + ), + rendered, + ) + return TokenEndpointError(status_code=status_code, redacted_body=redacted) + + +def _drop_reflected_assertion(rendered: str, assertion: SecretStr | None) -> str: + """Catches an endpoint that echoes the submitted credential back, verbatim or in fragments, + however it split or percent-encoded it. + + Both sides are reduced to the characters a credential is made of before comparison. Stripping + only the rendered side would stop matching a secret that carries spaces or punctuation of its + own, which is exactly the hand-set passphrase most at risk of being echoed. + + This stops an accidental or naive echo. It cannot stop an endpoint that deliberately re-encodes + or interleaves the credential, and it is not what keeps the credential from the endpoint, which + already holds it. What it protects is blast radius: keeping the value out of the caller's error + and out of third-party log sinks. + """ + if assertion is None: + return rendered + secret: Final = assertion.get_secret_value() + if not secret: + return rendered + if secret in rendered: + return _REFLECTED_VALUE_MESSAGE + compacted_secret: Final = _CREDENTIAL_CHARS.sub("", secret) + if not compacted_secret: + return rendered + return _REFLECTED_VALUE_MESSAGE if _shares_a_credential_run(rendered, compacted_secret) else rendered + + +def _shares_a_credential_run(rendered: str, compacted_secret: str) -> bool: + """``unquote`` covers a credential sent form-encoded, without every caller enumerating that + shape for itself: percent-escaping is reversible and applies to any field, query string + included. + + A secret shorter than the probe run is compared whole: a window longer than the secret can + never be found inside it, which would leave a short client secret unprotected in every shape + but the verbatim one. + """ + # unquote covers %XX; unquote_plus additionally covers the "+" a form-encoded body uses for a + # space. Both are kept rather than only the wider one, because "+" is a base64 character and + # decoding it away would lose a run that the undecoded candidate still matches on. + run: Final = min(_REFLECTION_MIN_RUN, len(compacted_secret)) + compacted_candidates: Final = tuple( + _CREDENTIAL_CHARS.sub("", candidate) for candidate in (rendered, unquote(rendered), unquote_plus(rendered)) + ) + return any( + compacted[start : start + run] in compacted_secret + for compacted in compacted_candidates + for start in range(len(compacted) - run + 1) + ) + + +def _redact_body_text(body_text: str) -> str: + if body_text in _SENTINEL_BODY_MESSAGES: + return body_text + if len(body_text) > MAX_RESPONSE_BYTES: + return _OVERSIZED_BODY_MESSAGE + try: + parsed: Final = _REDACTABLE_BODY_ADAPTER.validate_json(body_text) + except ValidationError: + return _UNSTRUCTURED_BODY_MESSAGE + match parsed: + case Mapping(): + return _format_oauth_error_fields(parsed) + case _: + return _NON_OBJECT_BODY_MESSAGE + + +def _format_oauth_error_fields(body: Mapping[str, object]) -> str: + fields: Final = tuple( + f"{name}: {_format_oauth_error_value(value)}" + for name in _OAUTH_ERROR_FIELDS + for value in (body.get(name),) + if value is not None + ) + return "; ".join(fields) if fields else _NO_OAUTH_FIELDS_MESSAGE + + +def _format_oauth_error_value(value: object) -> str: + """RFC 6749 types ``error`` as a string, but Anthropic (and other providers) nest their + own ``{"type": ..., "message": ...}`` envelope there; render that rather than a dict repr.""" + if isinstance(value, Mapping): + nested: Final = tuple( + f"{str(part)[:_REDACTION_CAP]}" + for key in _NESTED_ERROR_FIELDS + for part in (value.get(key),) + if part is not None + ) + if nested: + return " - ".join(nested) + return str(value)[:_REDACTION_CAP] + + +def _error_summary(error: ExchangeError) -> str: + match error: + case AssertionSourceError(): + return f"AssertionSourceError: assertion {error.kind} from {error.source_ref}" + case InsecureTokenUrl(): + return f"InsecureTokenUrl: insecure token endpoint host {error.host}" + case TokenEndpointError(): + return f"TokenEndpointError: HTTP {error.status_code}: {error.redacted_body}" + case TokenTransportError(): + return f"TokenTransportError: {error.detail}" + case MalformedTokenResponse(): + return f"MalformedTokenResponse: {error.detail}" + case _: + assert_never(error) + + +class _MetricsFailure(Exception): + """Never raised: typed carriers handed to the service failure hook so the prometheus + ``error_class`` label names the ``ExchangeError`` variant; the message is the redacted + ``_error_summary`` and carries no credential material.""" + + +class TokenExchangeAssertionSourceFailure(_MetricsFailure): ... + + +class TokenExchangeInsecureUrlFailure(_MetricsFailure): ... + + +class TokenExchangeEndpointFailure(_MetricsFailure): ... + + +class TokenExchangeTransportFailure(_MetricsFailure): ... + + +class TokenExchangeMalformedResponseFailure(_MetricsFailure): ... + + +def _failure_exception(error: ExchangeError) -> _MetricsFailure: + summary: Final = _error_summary(error) + match error: + case AssertionSourceError(): + return TokenExchangeAssertionSourceFailure(summary) + case InsecureTokenUrl(): + return TokenExchangeInsecureUrlFailure(summary) + case TokenEndpointError(): + return TokenExchangeEndpointFailure(summary) + case TokenTransportError(): + return TokenExchangeTransportFailure(summary) + case MalformedTokenResponse(): + return TokenExchangeMalformedResponseFailure(summary) + case _: + assert_never(error) + + +def _cache_key(spec: TokenExchangeSpec) -> str: + return hashlib.sha256( + "\x1f".join((spec.token_url, spec.assertion_ref, *spec.cache_key_identity)).encode() + ).hexdigest() + + +def _assertion_fetch(reader: AssertionReader, spec: TokenExchangeSpec) -> AssertionSource: + """``spec.assertion_source`` (an identity source's own fetch/mint closure) takes priority over + the engine-level reader when set; either way, failures are reported against ``spec.assertion_ref``.""" + if spec.assertion_source is not None: + return spec.assertion_source + return lambda: reader(spec.assertion_ref) + + +def _read_assertion(fetch: AssertionSource, ref: str) -> SecretStr | AssertionSourceError: + from litellm.secret_managers.main import OidcPathNotAllowedError + + try: + raw: Final = fetch() + except OidcPathNotAllowedError: + return AssertionSourceError(kind="disallowed_path", source_ref=ref) + except (ValueError, ImportError) as e: + return AssertionSourceError(kind="unreadable", source_ref=ref, detail=str(e)[:_REDACTION_CAP]) + except Exception: # noqa: BLE001 # injected readers (secret managers) raise arbitrarily; all failures become values + return AssertionSourceError(kind="unreadable", source_ref=ref) + if raw is None: + return AssertionSourceError(kind="missing", source_ref=ref) + stripped: Final = raw.strip() + if not stripped: + return AssertionSourceError(kind="empty", source_ref=ref) + if len(stripped.encode("utf-8")) > MAX_ASSERTION_BYTES: + return AssertionSourceError(kind="oversized", source_ref=ref) + return SecretStr(stripped) + + +def _serialize_body(spec: TokenExchangeSpec, assertion: SecretStr) -> bytes: + if spec.body_encoding == "json": + return json.dumps( + { # mutable-ok: transient body dict consumed inline by the serializer + **spec.static_body, + spec.assertion_field: assertion.get_secret_value(), + } + ).encode() + return urlencode( + { # mutable-ok: transient body dict consumed inline by the serializer + **spec.static_body, + spec.assertion_field: assertion.get_secret_value(), + } + ).encode() + + +def _sanitize_expires_in(expires_in: int | None) -> float: + if expires_in is None or expires_in <= 0: + return FALLBACK_TOKEN_TTL_SECONDS + return float(expires_in) + + +@dataclass(frozen=True, slots=True) +class _RefreshWindows: + advisory: float + mandatory: float + + +def _refresh_windows(lifetime_seconds: float | None) -> _RefreshWindows: + """A token whose whole life is shorter than the flat windows sits inside them from the moment it + is minted, so every request would arm another background exchange against the token endpoint. + Scaling each window by a fraction of the observed lifetime makes a 60s token refresh around its + half life instead; at a lifetime of 240s and above both fractions reach the flat windows, so + ordinary long-lived tokens keep exactly the 120s/30s behaviour.""" + if lifetime_seconds is None or lifetime_seconds <= 0.0: + return _RefreshWindows(advisory=ADVISORY_REFRESH_SECONDS, mandatory=MANDATORY_REFRESH_SECONDS) + return _RefreshWindows( + advisory=min(ADVISORY_REFRESH_SECONDS, lifetime_seconds * ADVISORY_REFRESH_LIFETIME_FRACTION), + mandatory=min(MANDATORY_REFRESH_SECONDS, lifetime_seconds * MANDATORY_REFRESH_LIFETIME_FRACTION), + ) + + +def _capped_body_text(response: httpx.Response) -> str: + if len(response.content) > MAX_RESPONSE_BYTES: + return _OVERSIZED_BODY_MESSAGE + return response.text + + +def _default_assertion_reader(ref: str) -> str | None: + from litellm.secret_managers.main import get_secret_str + + return get_secret_str(ref) + + +def _new_exchange_handler() -> "HTTPHandler": + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + handler: Final = HTTPHandler(timeout=httpx.Timeout(timeout=30.0, connect=5.0)) + handler.client.follow_redirects = False + return handler + + +def require_posted_response(response: httpx.Response | None, endpoint_label: str) -> httpx.Response: + """The legacy ``HTTPHandler`` carries no return annotation, so a patched or stubbed client can + hand a poster ``None`` back; a transport error beats dereferencing it.""" + if response is None: + raise httpx.TransportError(f"{endpoint_label} returned no response") + return response + + +class _HttpxSyncTokenPoster: + """Default poster: a dedicated HTTPHandler (no logging_obj, so litellm's + pre/post-call body logging never sees the exchange POST); returns the + response for any status.""" + + def __init__(self, handler_factory: Callable[[], "HTTPHandler"] = _new_exchange_handler) -> None: + self._lock: Final = threading.Lock() + self._handler_factory: Final = handler_factory + self._handler: HTTPHandler | None = None + + def _handler_instance(self) -> "HTTPHandler": + with self._lock: + if self._handler is None: + self._handler = self._handler_factory() + return self._handler + + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + try: + response: Final[httpx.Response | None] = self._handler_instance().post( # pyright: ignore[reportUnknownMemberType] # HTTPHandler.post is legacy-untyped; the result is validated below + url, + content=content, + headers=dict(headers), # mutable-ok: HTTPHandler.post requires a concrete dict + timeout=timeout, + ) + except httpx.HTTPStatusError as e: + return e.response + return require_posted_response(response, "token endpoint") + + +class _ServiceLoggingHooks(Protocol): + """The slice of ``litellm._service_logger.ServiceLogging`` the metrics sink calls; a protocol + so tests inject a recorder instead of monkeypatching.""" + + async def async_service_success_hook(self, service: ServiceTypes, call_type: str, duration: float) -> None: ... + + async def async_service_failure_hook( + self, service: ServiceTypes, duration: float, error: str | Exception, call_type: str + ) -> None: ... + + +_HooksCoroFactory: TypeAlias = Callable[ + [_ServiceLoggingHooks], # mutable-ok: Callable param-list syntax, not a list + Coroutine[object, object, None], +] + + +def _default_service_logging() -> _ServiceLoggingHooks: + from litellm._service_logger import ServiceLogging + + return ServiceLogging() + + +class ServiceLoggingMetricsSink: + """Default sink: bridges engine metrics onto litellm's ServiceTypes pattern + (prometheus ``litellm_anthropic_wif_*`` via ``service_callback``). The engine's entry points + are sync threads with no event loop, and the service hooks are async, so every emission is + fire-and-forget on a dedicated single worker thread that owns its own short-lived loop -- + the mint path only ever pays for an executor queue put.""" + + def __init__( + self, + service_logging_factory: Callable[[], _ServiceLoggingHooks] = _default_service_logging, + executor: Executor | None = None, + ) -> None: + self._lock: Final = threading.Lock() + self._service_logging_factory: Final = service_logging_factory + self._service_logging: _ServiceLoggingHooks | None = None + self._executor: Executor | None = executor + self._queued: int = 0 # rebind-ok: backlog depth, guarded by _lock + + def _service_logging_instance(self) -> _ServiceLoggingHooks: + with self._lock: + if self._service_logging is None: + self._service_logging = self._service_logging_factory() + return self._service_logging + + def _executor_instance(self) -> Executor: + with self._lock: + if self._executor is None: + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="litellm-token-exchange-metrics") + return self._executor + + def _emit(self, coro_factory: _HooksCoroFactory) -> None: + try: + asyncio.run(coro_factory(self._service_logging_instance())) + except Exception as e: # noqa: BLE001 # metrics are best-effort; emission failures must never surface + verbose_logger.debug("token exchange metrics emission failed: %s", e) + + def _submit(self, coro_factory: _HooksCoroFactory) -> None: + """Drop the event rather than queue it once the backlog is full. A stalled telemetry + backend must not let request volume grow an unbounded queue in the proxy: losing a + metric sample is always cheaper than losing the process.""" + with self._lock: + if self._queued >= _METRICS_QUEUE_LIMIT: + verbose_logger.debug("token exchange metrics queue full, dropping event") + return + self._queued += 1 + try: + self._executor_instance().submit(self._emit_and_release, coro_factory) + except Exception as e: # noqa: BLE001 # a rejected submit must not surface to the mint + with self._lock: + self._queued -= 1 + verbose_logger.debug("token exchange metrics submit failed: %s", e) + + def _emit_and_release(self, coro_factory: _HooksCoroFactory) -> None: + try: + self._emit(coro_factory) + finally: + with self._lock: + self._queued -= 1 + + def exchange_success(self, *, call_type: ExchangeCallType, duration_seconds: float) -> None: + def start(hooks: _ServiceLoggingHooks) -> Coroutine[object, object, None]: + return hooks.async_service_success_hook( + service=ServiceTypes.ANTHROPIC_WIF, call_type=call_type, duration=duration_seconds + ) + + self._submit(start) + + def exchange_failure(self, *, call_type: ExchangeCallType, duration_seconds: float, error: ExchangeError) -> None: + failure: Final = _failure_exception(error) + + def start(hooks: _ServiceLoggingHooks) -> Coroutine[object, object, None]: + return hooks.async_service_failure_hook( + service=ServiceTypes.ANTHROPIC_WIF, duration=duration_seconds, error=failure, call_type=call_type + ) + + self._submit(start) + + def cache_hit(self) -> None: + def start(hooks: _ServiceLoggingHooks) -> Coroutine[object, object, None]: + return hooks.async_service_success_hook( + service=ServiceTypes.ANTHROPIC_WIF_CACHE, call_type=CALL_TYPE_CACHE_HIT, duration=0.0 + ) + + self._submit(start) + + +class _Entry: + """Single-flight state for one cache key; mutable by design, confined to the + engine, and only ever mutated under the engine lock.""" + + __slots__ = ("backoff_until", "done", "force_refresh", "in_flight", "last_error", "lifetime_seconds", "token") + + def __init__(self, force_refresh: bool = False) -> None: + self.token: MintedToken | None = None + self.lifetime_seconds: float | None = None + self.in_flight: bool = False + self.done: Final = threading.Event() + self.backoff_until: float = float("-inf") + self.force_refresh: bool = force_refresh + self.last_error: ExchangeError | None = None + + def arm(self) -> None: + self.in_flight = True + self.last_error = None + self.done.clear() + + def _store(self, token: MintedToken, now: float) -> None: + self.token = token + self.lifetime_seconds = None if token.expires_at is None else max(token.expires_at - now, 0.0) + self.last_error = None + + def publish(self, result: ExchangeResult, now: float) -> None: + match result: + case MintedToken(): + self._store(result, now) + case _: + self.last_error = result + self.backoff_until = now + ADVISORY_REFRESH_BACKOFF_SECONDS + self.force_refresh = False + self.in_flight = False + self.done.set() + + def publish_advisory(self, result: ExchangeResult, now: float) -> None: + """A failed advisory refresh records only the backoff, never ``last_error``: a follower whose + cached token expires while this runs must be free to re-lead a fresh mint and recover.""" + match result: + case MintedToken(): + self._store(result, now) + case _: + self.backoff_until = now + ADVISORY_REFRESH_BACKOFF_SECONDS + self.in_flight = False + self.done.set() + + +@dataclass(frozen=True, slots=True) +class _Serve: + token: MintedToken + + +@dataclass(frozen=True, slots=True) +class _ServeAndRefresh: + token: MintedToken + + +@dataclass(frozen=True, slots=True) +class _Lead: + call_type: ExchangeCallType + + +@dataclass(frozen=True, slots=True) +class _Follow: + pass + + +@dataclass(frozen=True, slots=True) +class _Fail: + error: ExchangeError + + +_Decision: TypeAlias = _Serve | _ServeAndRefresh | _Lead | _Follow | _Fail + + +@dataclass(frozen=True, slots=True) +class _Unauthorized: + response: httpx.Response + assertion: SecretStr + + +class JwtBearerTokenExchangeEngine: + def __init__( + self, + poster: SyncTokenPoster | None = None, + assertion_reader: AssertionReader | None = None, + clock: Callable[[], float] = time.monotonic, + refresh_executor: Executor | None = None, + max_entries: int = 64, + metrics_sink: TokenExchangeMetricsSink | None = None, + ) -> None: + self._poster: Final[SyncTokenPoster] = poster if poster is not None else _HttpxSyncTokenPoster() + self._assertion_reader: Final[AssertionReader] = ( + assertion_reader if assertion_reader is not None else _default_assertion_reader + ) + self._clock: Final = clock + self._refresh_executor: Executor | None = refresh_executor + self._max_entries: Final = max_entries + self._metrics_sink: Final[TokenExchangeMetricsSink] = ( + metrics_sink if metrics_sink is not None else ServiceLoggingMetricsSink() + ) + self._lock: Final = threading.Lock() + self._entries: Final[dict[str, _Entry]] = {} # mutable-ok: engine-owned map guarded by _lock + + def get_token(self, spec: TokenExchangeSpec) -> ExchangeResult: + """A follower whose leader published nothing re-classifies rather than recursing, so a + contended entry cannot grow the stack one frame per failed leader.""" + while True: + with self._lock: + entry = self._get_or_create_entry_locked(spec) # rebind-ok: re-read per follower round + decision = self._classify_and_arm_locked(entry) # rebind-ok: re-read per follower round + match decision: + case _Serve(token=token): + self._report_cache_hit() + return token + case _ServeAndRefresh(token=token): + self._report_cache_hit() + self._executor_instance().submit(self._advisory_refresh, spec, entry) + return token + case _Fail(error=error): + return error + case _Lead(call_type=call_type): + return self._lead(spec, entry, call_type) + case _Follow(): + followed = self._await_leader(spec, entry) # rebind-ok: one leader wait per round + if followed is not None: + return followed + case _: + assert_never(decision) + + async def aget_token(self, spec: TokenExchangeSpec) -> ExchangeResult: + return await asyncio.to_thread(self.get_token, spec) + + def invalidate(self, spec: TokenExchangeSpec) -> None: + key: Final = _cache_key(spec) + with self._lock: + if key in self._entries: + self._entries[key] = _Entry(force_refresh=True) + + def _get_or_create_entry_locked(self, spec: TokenExchangeSpec) -> _Entry: + key: Final = _cache_key(spec) + existing: Final = self._entries.get(key) + if existing is not None: + return existing + if len(self._entries) >= self._max_entries: + self._evict_locked() + created: Final = _Entry() + self._entries[key] = created + return created + + def _evict_locked(self) -> None: + now: Final = self._clock() + stale: Final = tuple( + key + for key, entry in self._entries.items() + if not entry.in_flight + and (entry.token is None or (entry.token.expires_at is not None and entry.token.expires_at <= now)) + ) + for key in stale: + del self._entries[key] + if len(self._entries) < self._max_entries: + return + # Evict soonest-to-expire first, and take as many as the overshoot needs rather than one, so a + # burst of distinct identities does not leave the map permanently above max_entries. An entry + # a leader owns or a follower waits on is never a candidate, so a moment where every entry is + # in flight still over-inserts; that residue is bounded by the concurrent mints themselves. + evictable: Final = sorted( + ( + entry.token.expires_at if entry.token is not None and entry.token.expires_at is not None else -inf, + key, + ) + for key, entry in self._entries.items() + if not entry.in_flight + ) + for _, key in evictable[: len(self._entries) - self._max_entries + 1]: + del self._entries[key] + + def _classify_and_arm_locked(self, entry: _Entry) -> _Decision: + token: Final = entry.token + if token is not None and not entry.force_refresh: + if token.expires_at is None: + return _Serve(token=token) + windows: Final = _refresh_windows(entry.lifetime_seconds) + remaining: Final = token.expires_at - self._clock() + if remaining > windows.advisory: + return _Serve(token=token) + if remaining > windows.mandatory: + if entry.in_flight or self._clock() < entry.backoff_until: + return _Serve(token=token) + entry.arm() + return _ServeAndRefresh(token=token) + if entry.in_flight: + return _Follow() + if entry.last_error is not None and self._clock() < entry.backoff_until: + return _Fail(error=entry.last_error) + entry.arm() + return _Lead(call_type=CALL_TYPE_COLD_MINT if token is None else CALL_TYPE_MANDATORY_REFRESH) + + def _executor_instance(self) -> Executor: + with self._lock: + if self._refresh_executor is None: + self._refresh_executor = ThreadPoolExecutor( + max_workers=2, thread_name_prefix="litellm-token-exchange-refresh" + ) + return self._refresh_executor + + def _lead(self, spec: TokenExchangeSpec, entry: _Entry, call_type: ExchangeCallType) -> ExchangeResult: + started: Final = self._clock() + result: Final = self._exchange_never_raises(spec) + duration: Final = self._clock() - started + with self._lock: + entry.publish(result, now=self._clock()) + self._report_exchange(call_type, duration, result) + return result + + def _await_leader(self, spec: TokenExchangeSpec, entry: _Entry) -> "ExchangeResult | None": + """None means the finished round left neither a valid token nor an error + (a failed advisory refresh); the caller re-enters and leads a fresh exchange.""" + leader_finished: Final = entry.done.wait(2 * spec.timeout_seconds + _FOLLOWER_WAIT_GRACE_SECONDS) + with self._lock: + token: Final = entry.token + if token is not None and (token.expires_at is None or token.expires_at > self._clock()): + return token + if entry.last_error is not None: + return entry.last_error + if leader_finished: + return None + return TokenTransportError(detail="timed out waiting for the token exchange leader") + + def _advisory_refresh(self, spec: TokenExchangeSpec, entry: _Entry) -> None: + started: Final = self._clock() + result: Final = self._exchange_never_raises(spec) + duration: Final = self._clock() - started + with self._lock: + now: Final = self._clock() + entry.publish_advisory(result, now=now) + stale_expires_at: Final = entry.token.expires_at if entry.token is not None else None + stale_mandatory: Final = _refresh_windows(entry.lifetime_seconds).mandatory + self._report_exchange(CALL_TYPE_ADVISORY_REFRESH, duration, result) + if isinstance(result, MintedToken): + return + seconds_to_mandatory_wall: Final = ( + max(stale_expires_at - now - stale_mandatory, 0.0) if stale_expires_at is not None else 0.0 + ) + verbose_logger.warning( + "Advisory token refresh against %s failed (%s); serving the cached token for up to " + "%.0fs before the mandatory refresh wall; next attempt after %.0fs backoff", + urlsplit(spec.token_url).hostname or "", + _error_summary(result), + seconds_to_mandatory_wall, + ADVISORY_REFRESH_BACKOFF_SECONDS, + ) + + def _report_exchange(self, call_type: ExchangeCallType, duration_seconds: float, result: ExchangeResult) -> None: + try: + match result: + case MintedToken(): + self._metrics_sink.exchange_success(call_type=call_type, duration_seconds=duration_seconds) + case _: + self._metrics_sink.exchange_failure( + call_type=call_type, duration_seconds=duration_seconds, error=result + ) + except Exception as e: # noqa: BLE001 # metrics are best-effort; a sink failure must never fail a mint + verbose_logger.debug("token exchange metrics emission failed: %s", e) + + def _report_cache_hit(self) -> None: + try: + self._metrics_sink.cache_hit() + except Exception as e: # noqa: BLE001 # metrics are best-effort; a sink failure must never fail a serve + verbose_logger.debug("token exchange cache-hit metric emission failed: %s", e) + + def _exchange_never_raises(self, spec: TokenExchangeSpec) -> ExchangeResult: + """The single-flight leader and the advisory refresher must always publish a result: an + unhandled exception here would leave the entry armed (in_flight, cleared event) forever, so + every subsequent caller for this key would follow a leader that never finishes.""" + try: + return self._exchange(spec) + except Exception as e: # noqa: BLE001 # a leader must resolve its entry; any failure becomes a value + return TokenTransportError(detail=f"{type(e).__name__}: {e}"[:_REDACTION_CAP]) + + def _exchange(self, spec: TokenExchangeSpec) -> ExchangeResult: + first: Final = self._attempt_exchange(spec) + if not isinstance(first, _Unauthorized): + return first + second: Final = self._attempt_exchange(spec) + if isinstance(second, _Unauthorized): + return redact_oauth_error_body( + second.response.status_code, _capped_body_text(second.response), second.assertion + ) + return second + + def _attempt_exchange(self, spec: TokenExchangeSpec) -> "ExchangeResult | _Unauthorized": + assertion: Final = _read_assertion(_assertion_fetch(self._assertion_reader, spec), spec.assertion_ref) + if isinstance(assertion, AssertionSourceError): + return assertion + url_check: Final = validate_token_endpoint_url(spec.token_url) + if isinstance(url_check, InsecureTokenUrl): + return url_check + try: + response: Final = self._poster.post( + spec.token_url, + content=_serialize_body(spec, assertion), + headers=MappingProxyType({"content-type": _CONTENT_TYPES[spec.body_encoding], **spec.request_headers}), + timeout=spec.timeout_seconds, + ) + except Exception as e: # noqa: BLE001 # injected posters may raise beyond httpx; transport failures become values + return TokenTransportError(detail=f"{type(e).__name__}: {e}"[:_REDACTION_CAP]) + if response.status_code == 401: + return _Unauthorized(response=response, assertion=assertion) + return self._parse_response(response, assertion) + + def _parse_response(self, response: httpx.Response, assertion: SecretStr | None = None) -> ExchangeResult: + if not 200 <= response.status_code < 300: + return redact_oauth_error_body(response.status_code, _capped_body_text(response), assertion) + if len(response.content) > MAX_RESPONSE_BYTES: + return MalformedTokenResponse(detail="token response body exceeds the 1 MiB cap") + try: + parsed: Final = _TokenExchangeResponse.model_validate_json(response.content) + except ValidationError: + return MalformedTokenResponse(detail="token response failed RFC 6749 5.1 schema validation") + if parsed.token_type is not None and parsed.token_type.lower() != "bearer": + return MalformedTokenResponse(detail="token response carried a non-bearer token_type") + if not parsed.access_token.strip(): + return MalformedTokenResponse(detail="token response carried an empty access_token") + return MintedToken( + access_token=SecretStr(parsed.access_token), + expires_at=self._clock() + _sanitize_expires_in(parsed.expires_in), + ) + + +default_token_exchange_engine: Final = JwtBearerTokenExchangeEngine() diff --git a/litellm/llms/base_llm/auth/types.py b/litellm/llms/base_llm/auth/types.py new file mode 100644 index 00000000000..9d3a7a5012b --- /dev/null +++ b/litellm/llms/base_llm/auth/types.py @@ -0,0 +1,99 @@ +"""Provider-agnostic types for the RFC 7523 JWT-bearer token exchange engine.""" + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Literal, Protocol, TypeAlias + +import httpx +from pydantic import SecretStr + +BodyEncoding: TypeAlias = Literal["json", "form"] +AssertionReader: TypeAlias = Callable[[str], str | None] # mutable-ok: Callable param-list syntax, not a list +AssertionSource: TypeAlias = Callable[[], str | None] # mutable-ok: Callable param-list syntax, not a list + + +@dataclass(frozen=True, slots=True) +class TokenExchangeSpec: + """One grant profile as pure data: one instance per (provider, deployment, identity). + + ``token_url`` must be derived from deployment config/env only, never per-request caller + input. ``assertion_ref`` is a ``oidc/...`` get_secret ref resolved fresh on every exchange. + + ``assertion_source``, when set, is a zero-arg per-config fetch/mint closure that the engine + prefers over its own engine-level ``AssertionReader`` -- the dispatch mechanism identity + sources beyond token_file/env (e.g. ``internal_issuer``, ``keycloak``) use to plug into the + shared engine without a global registry. ``assertion_ref`` still names the cache-key + discriminator and the ref echoed into operator-facing errors either way. + """ + + token_url: str + assertion_ref: str + assertion_field: str + static_body: Mapping[str, str] + body_encoding: BodyEncoding + request_headers: Mapping[str, str] + cache_key_identity: tuple[str, ...] + timeout_seconds: float = 30.0 + assertion_source: AssertionSource | None = None + + +@dataclass(frozen=True, slots=True) +class MintedToken: + access_token: SecretStr + expires_at: float | None + + +@dataclass(frozen=True, slots=True) +class AssertionSourceError: + kind: Literal["missing", "empty", "oversized", "unreadable", "disallowed_path"] + source_ref: str + detail: str | None = None + + +@dataclass(frozen=True, slots=True) +class InsecureTokenUrl: + host: str + + +@dataclass(frozen=True, slots=True) +class TokenEndpointError: + status_code: int + redacted_body: str + + +@dataclass(frozen=True, slots=True) +class TokenTransportError: + detail: str + + +@dataclass(frozen=True, slots=True) +class MalformedTokenResponse: + detail: str + + +ExchangeError: TypeAlias = ( + AssertionSourceError | InsecureTokenUrl | TokenEndpointError | TokenTransportError | MalformedTokenResponse +) +ExchangeResult: TypeAlias = MintedToken | ExchangeError + +ExchangeCallType: TypeAlias = Literal["cold_mint", "mandatory_refresh", "advisory_refresh"] + + +class TokenExchangeMetricsSink(Protocol): + """Observability seam for the exchange engine. Implementations must be best-effort: never raise + into the mint path, never block the calling thread, and never receive credential material -- + ``ExchangeError`` values are redacted by construction.""" + + def exchange_success(self, *, call_type: ExchangeCallType, duration_seconds: float) -> None: ... + + def exchange_failure( + self, *, call_type: ExchangeCallType, duration_seconds: float, error: ExchangeError + ) -> None: ... + + def cache_hit(self) -> None: ... + + +class SyncTokenPoster(Protocol): + """Returns the response for ANY status; never raises for status.""" + + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: ... diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index c5290b41f7b..bf93308974c 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -5,6 +5,7 @@ Utility functions for base LLM classes. import copy import json from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import Any, Final from openai.lib import _parsing, _pydantic @@ -57,6 +58,22 @@ class BaseLLMModelInfo(ABC): """ return [] + def discover_models( + self, litellm_params: Mapping[str, object] | None = None + ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override + """ + Live model discovery for a configured deployment. Defaults to the api_key/api_base + facade every provider already implements via ``get_models``; a provider whose + discovery needs more of ``litellm_params`` (e.g. Anthropic's workload identity + federation) overrides this instead of widening ``get_models`` for every provider. + """ + api_key: Final = litellm_params.get("api_key") if litellm_params is not None else None + api_base: Final = litellm_params.get("api_base") if litellm_params is not None else None + return self.get_models( + api_key=api_key if isinstance(api_key, str) else None, + api_base=api_base if isinstance(api_base, str) else None, + ) + @staticmethod @abstractmethod def get_api_key(api_key: str | None = None) -> str | None: diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index e1f0fc9e7d3..38b5e9ad673 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1289,7 +1289,7 @@ class HTTPHandler: self, url: str, params: dict | None = None, - headers: dict | None = None, + headers: Mapping[str, Any] | None = None, follow_redirects: bool | None = None, timeout: float | httpx.Timeout | None = None, ): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index f281c249c72..190794cd5e1 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,11 +1,25 @@ import asyncio +import inspect import json import ssl from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager from functools import lru_cache from types import MappingProxyType, ModuleType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, Union, cast, get_type_hints +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + Optional, + Protocol, + TypedDict, + TypeVar, + Union, + cast, + get_type_hints, + runtime_checkable, +) from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import httpx @@ -219,6 +233,55 @@ class _MediaUploadKwargs(TypedDict, total=False): timeout: float | httpx.Timeout +@runtime_checkable +class _AsyncFilesEnvironmentValidator(Protocol): + async def avalidate_environment( + self, + headers: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + model: str, + messages: list, # mutable-ok: mirrors the sync validate_environment contract this overrides + optional_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + litellm_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: ... # mutable-ok: mirrors the sync validate_environment contract this overrides + + +async def _avalidate_files_environment( + provider_config: BaseFilesConfig | BaseBatchesConfig, + *, + headers: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + model: str, + messages: list, # mutable-ok: mirrors the sync validate_environment contract this overrides + optional_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + litellm_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + api_key: str | None, +) -> dict: # mutable-ok: mirrors the sync validate_environment contract this overrides + """Await the provider's async credential hook when it has one (e.g. Anthropic's workload + identity token exchange); otherwise offload the sync hook to a worker thread. Either way + the caller, an async file handler, never blocks the event loop on it.""" + if isinstance(provider_config, _AsyncFilesEnvironmentValidator) and inspect.iscoroutinefunction( + provider_config.avalidate_environment + ): + return await provider_config.avalidate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + ) + return await asyncio.to_thread( + provider_config.validate_environment, + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + ) + + def _google_genai_streaming_hidden_params( *, api_base: str, @@ -2130,7 +2193,7 @@ class BaseLLMHTTPHandler: ( headers, api_base, - ) = anthropic_messages_provider_config.validate_anthropic_messages_environment( + ) = await anthropic_messages_provider_config.avalidate_anthropic_messages_environment( headers=merged_headers or {}, model=model, messages=messages, @@ -3520,6 +3583,19 @@ class BaseLLMHTTPHandler: """ Creates a file using Gemini's two-step upload process """ + if _is_async: + return self._avalidate_and_create_file( + create_file_data=create_file_data, + litellm_params=litellm_params, + provider_config=provider_config, + headers=headers, + api_base=api_base, + api_key=api_key, + logging_obj=logging_obj, + client=client, + timeout=timeout, + ) + # get config from model, custom llm provider headers = provider_config.validate_environment( api_key=api_key, @@ -3549,18 +3625,6 @@ class BaseLLMHTTPHandler: optional_params={}, ) - if _is_async: - return self.async_create_file( - transformed_request=transformed_request, - litellm_params=litellm_params, - provider_config=provider_config, - headers=headers, - api_base=api_base, - logging_obj=logging_obj, - client=client, - timeout=timeout, - ) - if client is None or not isinstance(client, HTTPHandler): sync_httpx_client = _get_httpx_client() else: @@ -3688,6 +3752,52 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params_with_url, ) + async def _avalidate_and_create_file( + self, + *, + create_file_data: CreateFileRequest, + litellm_params: dict, # mutable-ok: mirrors the create_file contract this dispatches for + provider_config: BaseFilesConfig, + headers: dict, # mutable-ok: mirrors the create_file contract this dispatches for + api_base: str | None, + api_key: str | None, + logging_obj: LiteLLMLoggingObj, + client: HTTPHandler | AsyncHTTPHandler | None, + timeout: float | httpx.Timeout | None, + ) -> OpenAIFileObject: + validated_headers: Final = await _avalidate_files_environment( + provider_config, + headers=headers, + model="", + messages=[], # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it + optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it + litellm_params=litellm_params, + api_key=api_key, + ) + complete_api_base: Final = provider_config.get_complete_file_url( + api_base=api_base, + api_key=api_key, + model="", + optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it + litellm_params=litellm_params, + data=create_file_data, + ) + return await self.async_create_file( + transformed_request=provider_config.transform_create_file_request( + model="", + create_file_data=create_file_data, + litellm_params=litellm_params, + optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it + ), + litellm_params=litellm_params, + provider_config=provider_config, + headers=validated_headers, + api_base=complete_api_base, + logging_obj=logging_obj, + client=client, + timeout=timeout, + ) + async def async_create_file( self, transformed_request: Union[bytes, str, dict, "TwoStepFileUploadConfig"], @@ -3938,6 +4048,20 @@ class BaseLLMHTTPHandler: if model is None: raise ValueError("model is required for create_batch") + if _is_async: + return self._avalidate_and_create_batch( + create_batch_data=create_batch_data, + litellm_params=litellm_params, + provider_config=provider_config, + headers=headers, + api_base=api_base, + api_key=api_key, + logging_obj=logging_obj, + client=client, + timeout=timeout, + model=model, + ) + headers = provider_config.validate_environment( api_key=api_key, headers=headers, @@ -3966,19 +4090,6 @@ class BaseLLMHTTPHandler: optional_params={}, ) - if _is_async: - return self.async_create_batch( - transformed_request=transformed_request, - litellm_params=litellm_params, - provider_config=provider_config, - headers=headers, - api_base=api_base, - logging_obj=logging_obj, - client=client, - timeout=timeout, - create_batch_data=create_batch_data, - ) - if client is None or not isinstance(client, HTTPHandler): sync_httpx_client = _get_httpx_client() else: @@ -4115,6 +4226,54 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, ) + async def _avalidate_and_create_batch( + self, + *, + create_batch_data: "CreateBatchRequest", + litellm_params: dict, # mutable-ok: mirrors the create_batch contract this dispatches for + provider_config: "BaseBatchesConfig", + headers: dict, # mutable-ok: mirrors the create_batch contract this dispatches for + api_base: str | None, + api_key: str | None, + logging_obj: "LiteLLMLoggingObj", + client: Union["HTTPHandler", "AsyncHTTPHandler"] | None, + timeout: float | httpx.Timeout | None, + model: str, + ) -> "LiteLLMBatch": + validated_headers: Final = await _avalidate_files_environment( + provider_config, + headers=headers, + model=model, + messages=[], # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it + optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it + litellm_params=litellm_params, + api_key=api_key, + ) + complete_api_base: Final = provider_config.get_complete_batch_url( + api_base=api_base, + api_key=api_key, + model=model, + optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it + litellm_params=litellm_params, + data=create_batch_data, + ) + return await self.async_create_batch( + transformed_request=provider_config.transform_create_batch_request( + model=model, + create_batch_data=create_batch_data, + litellm_params=litellm_params, + optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it + ), + litellm_params=litellm_params, + provider_config=provider_config, + headers=validated_headers, + api_base=complete_api_base, + logging_obj=logging_obj, + client=client, + timeout=timeout, + create_batch_data=create_batch_data, + ) + async def async_create_batch( self, transformed_request: bytes | str | dict, @@ -4712,7 +4871,8 @@ class BaseLLMHTTPHandler: ) # Validate environment and get headers - headers = provider_config.validate_environment( + headers = await _avalidate_files_environment( + provider_config, api_key=litellm_params.get("api_key"), headers=headers, model="", @@ -4836,7 +4996,8 @@ class BaseLLMHTTPHandler: ) # Validate environment and get headers - headers = provider_config.validate_environment( + headers = await _avalidate_files_environment( + provider_config, api_key=litellm_params.get("api_key"), headers=headers, model="", @@ -4960,7 +5121,8 @@ class BaseLLMHTTPHandler: ) # Validate environment and get headers - headers = provider_config.validate_environment( + headers = await _avalidate_files_environment( + provider_config, api_key=litellm_params.get("api_key"), headers=headers, model="", @@ -5091,7 +5253,8 @@ class BaseLLMHTTPHandler: ) # Validate environment and get headers - headers = provider_config.validate_environment( + headers = await _avalidate_files_environment( + provider_config, api_key=litellm_params.get("api_key"), headers=headers, model="", diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 9afc6331d96..d2a1fc98c73 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -54,6 +54,7 @@ from litellm.types.utils import ( from litellm.utils import convert_to_model_response_object from ..common_utils import OpenAIError +from ..workload_identity import get_workload_identity_bearer_token, resolve_openai_workload_identity_config if TYPE_CHECKING: import tiktoken @@ -70,6 +71,11 @@ else: _NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) +def _litellm_params_str(litellm_params: Mapping[str, object] | None, key: str) -> str | None: + value: Final = litellm_params.get(key) if litellm_params is not None else None + return value if isinstance(value, str) else None + + class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): """ Reference: https://platform.openai.com/docs/api-reference/chat/create @@ -747,28 +753,39 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): """ Calls OpenAI's `/v1/models` endpoint and returns the list of models. """ - - if api_base is None: - api_base = "https://api.openai.com" - if api_key is None: - api_key = get_secret_str("OPENAI_API_KEY") - - # Strip api_base to just the base URL (scheme + host + port) - parsed_url: Final = httpx.URL(api_base) - base_url = f"{parsed_url.scheme}://{parsed_url.host}" - if parsed_url.port: - base_url += f":{parsed_url.port}" - - response: Final = litellm.module_level_client.get( - url=f"{base_url}/v1/models", - headers={"Authorization": f"Bearer {api_key}"}, + return self._fetch_model_ids( + api_base=api_base, bearer_token=get_secret_str("OPENAI_API_KEY") if api_key is None else api_key ) + def discover_models( + self, litellm_params: Mapping[str, object] | None = None + ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override + if type(self) is not OpenAIGPTConfig: + return super().discover_models(litellm_params) + api_key: Final = _litellm_params_str(litellm_params, "api_key") + api_base: Final = _litellm_params_str(litellm_params, "api_base") + workload_identity_config: Final = resolve_openai_workload_identity_config( + api_key=api_key, api_base=api_base, litellm_params=litellm_params + ) + if workload_identity_config is None: + return self.get_models(api_key=api_key, api_base=api_base) + return self._fetch_model_ids( + api_base=api_base, bearer_token=get_workload_identity_bearer_token(workload_identity_config) + ) + + @staticmethod + def _fetch_model_ids( + api_base: str | None, bearer_token: str | None + ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override + parsed_url: Final = httpx.URL(api_base or "https://api.openai.com") + port_suffix: Final = f":{parsed_url.port}" if parsed_url.port else "" + response: Final = litellm.module_level_client.get( + url=f"{parsed_url.scheme}://{parsed_url.host}{port_suffix}/v1/models", + headers={"Authorization": f"Bearer {bearer_token}"}, + ) if response.status_code != 200: raise Exception(f"Failed to get models: {response.text}") - - models: Final = response.json()["data"] - return [model["id"] for model in models] + return [model["id"] for model in response.json()["data"]] @staticmethod def get_api_key(api_key: str | None = None) -> str | None: diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index edc8d64d9c2..3ffe4772df1 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -382,8 +382,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): organization: str | None = None, client: OpenAI | AsyncOpenAI | None = None, shared_session: Optional["ClientSession"] = None, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAI | AsyncOpenAI | None: - workload_identity_config: Final = resolve_openai_workload_identity_config(api_key=api_key, api_base=api_base) + workload_identity_config: Final = resolve_openai_workload_identity_config( + api_key=api_key, api_base=api_base, litellm_params=litellm_params + ) client_initialization_params: Final[dict] = locals() if client is None: if not isinstance(max_retries, int): @@ -773,6 +776,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=max_retries, organization=organization, stream_options=stream_options, + litellm_params=litellm_params, ) else: if not isinstance(max_retries, int): @@ -786,6 +790,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=max_retries, organization=organization, client=client, + litellm_params=litellm_params, ) ## LOGGING @@ -927,6 +932,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): organization=organization, client=client, shared_session=shared_session, + litellm_params=litellm_params, ) ## LOGGING @@ -1022,6 +1028,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=None, headers=None, stream_options: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ): data["stream"] = True data.update(self.get_stream_options(stream_options=stream_options, api_base=api_base)) @@ -1035,6 +1042,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=max_retries, organization=organization, client=client, + litellm_params=litellm_params, ) ## LOGGING logging_obj.pre_call( @@ -1107,6 +1115,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): organization=organization, client=client, shared_session=shared_session, + litellm_params=litellm_params, ) ## LOGGING logging_obj.pre_call( @@ -1241,6 +1250,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): client: AsyncOpenAI | None = None, max_retries=None, shared_session: Optional["ClientSession"] = None, + litellm_params: Mapping[str, object] | None = None, ): try: openai_aclient: Final[AsyncOpenAI] = self._get_openai_client( @@ -1251,6 +1261,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=max_retries, client=client, shared_session=shared_session, + litellm_params=litellm_params, ) raw_response: Final = await self.make_openai_embedding_request( openai_aclient=openai_aclient, @@ -1314,6 +1325,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): aembedding=None, max_retries: int | None = None, shared_session: Optional["ClientSession"] = None, + litellm_params: Mapping[str, object] | None = None, ) -> EmbeddingResponse: super().embedding() try: @@ -1340,6 +1352,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): client=client, max_retries=max_retries, shared_session=shared_session, + litellm_params=litellm_params, ) openai_client: Final[OpenAI] = self._get_openai_client( @@ -1349,6 +1362,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): timeout=timeout, max_retries=max_retries, client=client, + litellm_params=litellm_params, ) ## embedding CALL diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index d7c2fcace09..110e359f4af 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -10,6 +10,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.get_litellm_params import OPENAI_WIF_KWARGS_KEYS from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) @@ -494,7 +495,11 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") headers.setdefault("Content-Type", "application/json") workload_identity_config: Final = ( - resolve_openai_workload_identity_config(api_key=api_key, api_base=litellm_params.api_base) + resolve_openai_workload_identity_config( + api_key=api_key, + api_base=litellm_params.api_base, + litellm_params=litellm_params.model_dump(include=set(OPENAI_WIF_KWARGS_KEYS)), + ) if self.custom_llm_provider is LlmProviders.OPENAI else None ) diff --git a/litellm/llms/openai/workload_identity.py b/litellm/llms/openai/workload_identity.py index 283fdfb92c2..e5a4ec3d944 100644 --- a/litellm/llms/openai/workload_identity.py +++ b/litellm/llms/openai/workload_identity.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from functools import lru_cache from typing import TYPE_CHECKING, Final @@ -43,6 +44,7 @@ class OpenAIWorkloadIdentityConfig: def resolve_openai_workload_identity_config( api_key: str | None, api_base: str | None, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAIWorkloadIdentityConfig | None: static_api_key: Final = normalize_nonempty_secret_str(api_key) or normalize_nonempty_secret_str( get_secret_str("OPENAI_API_KEY") @@ -54,10 +56,12 @@ def resolve_openai_workload_identity_config( ) if not _targets_openai_api(effective_api_base): return None - identity_provider_id: Final = get_secret_str("OPENAI_IDENTITY_PROVIDER_ID") - service_account_id: Final = get_secret_str("OPENAI_SERVICE_ACCOUNT_ID") - token_file: Final = get_secret_str("OPENAI_IDENTITY_TOKEN_FILE") - if not identity_provider_id or not service_account_id or not token_file: + identity_provider_id: Final = _config_value( + litellm_params, "openai_identity_provider_id", "OPENAI_IDENTITY_PROVIDER_ID" + ) + service_account_id: Final = _config_value(litellm_params, "openai_service_account_id", "OPENAI_SERVICE_ACCOUNT_ID") + token_file: Final = _config_value(litellm_params, "openai_identity_token_file", "OPENAI_IDENTITY_TOKEN_FILE") + if identity_provider_id is None or service_account_id is None or token_file is None: return None return OpenAIWorkloadIdentityConfig( identity_provider_id=identity_provider_id, @@ -70,6 +74,13 @@ def get_workload_identity_bearer_token(config: OpenAIWorkloadIdentityConfig) -> return _workload_identity_auth(config).get_token() +def _config_value(litellm_params: Mapping[str, object] | None, param_key: str, env_name: str) -> str | None: + param_value: Final = litellm_params.get(param_key) if litellm_params is not None else None + if isinstance(param_value, str) and param_value: + return param_value + return normalize_nonempty_secret_str(get_secret_str(env_name)) + + def _targets_openai_api(api_base: str | None) -> bool: if api_base is None: return True diff --git a/litellm/main.py b/litellm/main.py index 2929790f2bd..e38af6b4a21 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6326,6 +6326,7 @@ def embedding( aembedding=aembedding, max_retries=max_retries, shared_session=shared_session, + litellm_params=litellm_params_dict, ) elif custom_llm_provider == "databricks": api_base = api_base or litellm.api_base or get_secret("DATABRICKS_API_BASE") diff --git a/litellm/models/credentials.py b/litellm/models/credentials.py index 56836234898..a007f2ba5d2 100644 --- a/litellm/models/credentials.py +++ b/litellm/models/credentials.py @@ -5,7 +5,7 @@ These are the canonical credential types for the proxy. They live in the model layer; ``litellm.types.utils`` re-exports them for backwards compatibility. """ -from pydantic import BaseModel, model_validator +from pydantic import BaseModel, Field, model_validator class CredentialBase(BaseModel): @@ -15,6 +15,10 @@ class CredentialBase(BaseModel): class CredentialItem(CredentialBase): credential_values: dict + # PATCH-only instruction naming keys to drop from the stored credential_values. It describes an + # edit rather than the credential, so it stays out of dumps: those feed config loading, the DB + # write, and the in-memory list, none of which have a place for it. + credential_values_to_delete: tuple[str, ...] | None = Field(default=None, exclude=True) class CreateCredentialItem(CredentialBase): diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index d6007a2d56e..accfd6a5f0d 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -33,7 +33,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_ENDPOINT_MARKER, ) from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS -from litellm.types.utils import CustomPricingLiteLLMParams +from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params +from litellm.types.utils import CustomPricingLiteLLMParams, server_owned_wif_litellm_params def is_invalid_virtual_key_error(exception: BaseException | None) -> bool: @@ -226,6 +227,16 @@ def _allow_model_level_clientside_configurable_parameters( # ``extra_body.aws_web_identity_token``) without re-validating, so the # banned-key check has to descend into it the same way it descends into # ``litellm_embedding_config``. +_SERVER_OWNED_WIF_UNCONDITIONAL_BANNED: Final[tuple[str, ...]] = server_owned_wif_litellm_params +# The Bedrock Claude Platform route reads a workspace from workspace_id or aws_workspace_id as +# well, and neither is a federation parameter, so say so rather than leaving that caller stuck. + + +# Re-exported from litellm.types.router, where it lives so the router can call it on a +# post-authentication merge without core importing from the proxy package. +reject_server_owned_wif_params = _reject_server_owned_wif_params + + _NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body") # Metadata containers that carry per-request configuration consumed by the @@ -379,6 +390,7 @@ def _check_banned_params( Shared between the root-level check and the nested-config check so a new banned param only needs to be added in one place. """ + reject_server_owned_wif_params(body) for param in _BANNED_REQUEST_BODY_PARAMS: if param not in body: continue @@ -519,6 +531,7 @@ def is_request_body_safe(request_body: dict, general_settings: dict, llm_router: _reject_url_valued_fallback_target(target) litellm_params: Final = _coerce_metadata_to_dict(request_body.get("litellm_params")) if litellm_params is not None: + reject_server_owned_wif_params(litellm_params) litellm_params_metadata: Final = _coerce_metadata_to_dict(litellm_params.get("metadata")) if litellm_params_metadata is not None: _check_banned_params( diff --git a/litellm/proxy/common_utils/credential_hydration.py b/litellm/proxy/common_utils/credential_hydration.py new file mode 100644 index 00000000000..34cf6b0861c --- /dev/null +++ b/litellm/proxy/common_utils/credential_hydration.py @@ -0,0 +1,158 @@ +"""Shared helper for resolving a named Credential's values server-side. + +Memory first (``litellm.credential_list``, already decrypted -- matching +``CredentialAccessor.get_credential_values``), then a DB decrypt fallback for a pod whose +in-memory list has not yet picked up a credential another pod just wrote or updated. +""" + +import asyncio +from collections.abc import Mapping +from itertools import chain +from types import MappingProxyType +from typing import Final + +import litellm +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.proxy.utils import PrismaClient +from litellm.repositories.credentials_repository import CredentialsRepository +from litellm.types.router import ( + GenericLiteLLMParams, + server_owned_wif_fields_named, + server_owned_wif_fields_present, +) +from litellm.types.utils import CredentialItem, LlmProviders + +_LITELLM_PROVIDER_IDS: Final = frozenset(provider.value for provider in LlmProviders) + + +def stored_credential_provider(credential_provider: object) -> str | None: + """The dashboard stores its display casing (``Anthropic``) on credentials it creates, so the + provider a credential names is the lowercased value when that is a litellm provider id.""" + if not isinstance(credential_provider, str): + return None + lowered: Final = credential_provider.lower() + return lowered if lowered in _LITELLM_PROVIDER_IDS else None + + +def decrypted_or_stored(key: str, value: str) -> str: + """The stored value decrypted, or as stored when it was never encrypted (a config.yaml value).""" + decrypted: Final = decrypt_value_helper(value=value, key=key) + return value if decrypted is None else decrypted + + +def _decrypted(db_credential: CredentialItem) -> CredentialItem: + """The stored credential with every value decrypted, leaving already-plaintext values alone.""" + decrypted_values: Final = MappingProxyType( + {key: decrypted_or_stored(key, value) for key, value in db_credential.credential_values.items()} + ) + return CredentialItem( + credential_name=db_credential.credential_name, + credential_values=decrypted_values, # pyright: ignore[reportArgumentType] # declared dict[str, str], and pydantic copies this mapping into one on validation; LIT002 rules out building that dict here + credential_info=db_credential.credential_info, + ) + + +async def hydrate_named_credential_authoritative( + credential_name: str, + prisma_client: PrismaClient | None, +) -> CredentialItem | None: + """The stored credential, preferring the row over this pod's in-memory copy. + + ``hydrate_named_credential`` reads memory first, which is right when serving a request. A + management operation cannot: on a pod whose in-memory copy predates another pod's update, it + would export the superseded JWKS, or discover models against superseded values. Same reason + ``named_credential_wif_fields`` reads both. + """ + if prisma_client is None: + return await hydrate_named_credential(credential_name, prisma_client) + db_credential: Final = await CredentialsRepository(prisma_client).find_by_name(credential_name) + if db_credential is None: + return await hydrate_named_credential(credential_name, prisma_client) + return _decrypted(db_credential) + + +async def hydrate_named_credential( + credential_name: str, + prisma_client: PrismaClient | None, +) -> CredentialItem | None: + for credential in litellm.credential_list: + if credential.credential_name == credential_name: + return credential + if prisma_client is None: + return None + db_credential: Final = await CredentialsRepository(prisma_client).find_by_name(credential_name) + if db_credential is None: + return None + return _decrypted(db_credential) + + +async def named_credential_wif_fields( + credential_name: str, + prisma_client: PrismaClient | None, +) -> tuple[str, ...]: + """Federation field names a write to ``credential_name`` would touch, from memory AND the row. + + Resolution reads memory first and stops there, which is right when serving a request. An + authorization decision cannot: a pod whose in-memory copy predates an admin adding federation + fields would see none and allow the write. This reads both and returns the union, so the gate + refuses whenever either side says the credential is server-owned. + """ + in_memory: Final = tuple( + name + for credential in litellm.credential_list + if credential.credential_name == credential_name + for name in server_owned_wif_fields_named(credential.credential_values) + ) + if prisma_client is None: + return in_memory + db_credential: Final = await CredentialsRepository(prisma_client).find_by_name(credential_name) + stored: Final = () if db_credential is None else server_owned_wif_fields_named(db_credential.credential_values) + return tuple(dict.fromkeys(in_memory + stored)) + + +async def effective_server_owned_wif_fields( + stored: Mapping[str, object] | None, + incoming: GenericLiteLLMParams | None, + prisma_client: PrismaClient | None, +) -> tuple[str, ...]: + """Federation field names the deployment would carry AFTER this write. + + Authorization has to read the resulting deployment, not the submitted payload. A patch that + names no federation field still lands on a deployment that has them, and a patch that only + attaches ``litellm_credential_name`` inherits whatever that credential holds. + + The two sides are matched differently on purpose. ``stored`` is matched by VALUE, because + ``GenericLiteLLMParams`` declares every federation field, so matching it by key would + report every deployment on the proxy as federated. ``incoming`` is matched by the keys the + write actually set, so an explicit null still counts as touching the field. + """ + from_stored: Final = () if stored is None else server_owned_wif_fields_present(stored) + from_incoming: Final = () if incoming is None else server_owned_wif_fields_named(incoming.model_fields_set) + from_credential: Final = tuple( + chain.from_iterable( + await asyncio.gather( + *( + named_credential_wif_fields(credential_name, prisma_client) + for credential_name in _effective_credential_names(stored, incoming) + ) + ) + ) + ) + return tuple(dict.fromkeys(from_stored + from_incoming + from_credential)) + + +def _effective_credential_names( + stored: Mapping[str, object] | None, + incoming: GenericLiteLLMParams | None, +) -> tuple[str, ...]: + """Both the credential the deployment already carries and the one this write names. + + Taking only the incoming name would let a write clear its way out: detaching a federated + credential, by sending ``litellm_credential_name: null`` alongside an api_key or api_base of + the caller's choosing, would leave nothing federated to find and the write would be allowed. + Detaching an administrator's federated credential is itself an administrator's action, so the + stored name counts whatever the write says. + """ + from_stored: Final = None if stored is None else stored.get("litellm_credential_name") + from_incoming: Final = None if incoming is None else incoming.litellm_credential_name + return tuple(dict.fromkeys(name for name in (from_stored, from_incoming) if isinstance(name, str))) diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 66789748707..20984f3a91c 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -13,16 +13,111 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import _get_masked_values -from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.llms.anthropic.wif import ( + _IDENTITY_SOURCE_PARAM, # pyright: ignore[reportPrivateUsage] # one canonical param name, shared with the litellm_params identity-source resolver + _INTERNAL_ISSUER_FIELD_MAP, # pyright: ignore[reportPrivateUsage] # one canonical field map, shared with the litellm_params identity-source resolver + _build_variant, # pyright: ignore[reportPrivateUsage] # one canonical builder, shared with the litellm_params identity-source resolver +) +from litellm.llms.base_llm.auth.identity_source import ( + AnthropicIdentitySourceKind, + InternalIssuerSource, +) +from litellm.llms.base_llm.auth.internal_issuer import internal_issuer_jwks_document +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.credential_hydration import ( + hydrate_named_credential, + hydrate_named_credential_authoritative, + named_credential_wif_fields, + stored_credential_provider, +) from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.utils import handle_exception_on_proxy, jsonify_object from litellm.repositories.credentials_repository import CredentialsRepository +from litellm.types.router import server_owned_wif_fields_named from litellm.types.utils import CreateCredentialItem, CredentialItem router: Final = APIRouter() +def _reject_non_admin_wif_fields( + wif_fields: tuple[str, ...], + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """A credential referenced by ``litellm_credential_name`` feeds its values into the same + workload identity federation resolution as a deployment's own ``litellm_params``. Only proxy + admins may touch a server-owned WIF field, whether they write it, drop it, or edit a stored + credential that already carries one. + """ + if not wif_fields or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return + raise HTTPException( + status_code=403, + detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict + "error": ( + f"Only proxy admins can change {wif_fields[0]!r}, a server-owned workload identity federation " + "parameter." + ) + }, + ) + + +def _incoming_wif_fields(credential: CredentialItem) -> tuple[str, ...]: + """WIF fields the request payload itself touches: the ones it sets (to any value, ``None`` + included, since the key alone is what the federation resolver reacts to), plus the ones it + names in ``credential_values_to_delete``, since dropping a federation field off the stored + credential breaks every deployment referencing it just as installing one would redirect them. + """ + return server_owned_wif_fields_named(credential.credential_values) + server_owned_wif_fields_named( + credential.credential_values_to_delete or () + ) + + +def _stored_wif_fields(stored_credential: CredentialItem) -> tuple[str, ...]: + return server_owned_wif_fields_named(stored_credential.credential_values) + + +def _reject_overlapping_credential_values(credential: CredentialItem) -> None: + overlap: Final = frozenset(credential.credential_values) & frozenset(credential.credential_values_to_delete or ()) + if overlap: + raise HTTPException( + status_code=400, + detail=f"credential_values_to_delete overlaps credential_values for key(s): {sorted(overlap)}", + ) + + +def _sync_in_memory_credential(credential: CredentialItem, credential_name: str, new_name: str) -> None: + """Mirror a DB credential update into the in-memory ``credential_list`` used by request-time + resolution; a no-op if the credential isn't loaded in memory (e.g. proxy restarted since boot). + """ + existing_in_memory: CredentialItem | None = None + for cred in litellm.credential_list: + if cred.credential_name == credential_name: + existing_in_memory = cred + break + + if existing_in_memory is None: + return + + in_memory_values: Final = dict(existing_in_memory.credential_values or {}) + if credential.credential_values: + in_memory_values.update(credential.credential_values) + for key in credential.credential_values_to_delete or (): + in_memory_values.pop(key, None) + in_memory_info: Final = dict(existing_in_memory.credential_info or {}) + if credential.credential_info: + in_memory_info.update(credential.credential_info) + updated_in_memory: Final = CredentialItem( + credential_name=new_name, + credential_values=in_memory_values, + credential_info=in_memory_info, + ) + # Remove old entry if renamed, then use upsert_credentials to handle duplicates + if new_name != credential_name: + litellm.credential_list = [c for c in litellm.credential_list if c.credential_name != credential_name] + CredentialAccessor.upsert_credentials([updated_in_memory]) + + class CredentialHelperUtils: @staticmethod def encrypt_credential_values(credential: CredentialItem, new_encryption_key: str | None = None) -> CredentialItem: @@ -84,13 +179,19 @@ async def create_credential( status_code=400, detail="Credential values are required. Unable to infer credential values from model ID.", ) + _reject_non_admin_wif_fields(server_owned_wif_fields_named(credential.credential_values), user_api_key_dict) + _reject_non_admin_wif_fields( + await named_credential_wif_fields(credential.credential_name, prisma_client), user_api_key_dict + ) processed_credential: Final = CredentialItem( credential_name=credential.credential_name, credential_values=credential.credential_values, credential_info=credential.credential_info, ) encrypted_credential: Final = CredentialHelperUtils.encrypt_credential_values(processed_credential) - credentials_dict: Final = encrypted_credential.model_dump() + # exclude_none: wif.py rejects foreign-variant fields by presence, so persisting a null + # for every unset variant field would fail the next request against this credential + credentials_dict: Final = encrypted_credential.model_dump(exclude_none=True) credentials_dict_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str "dict[str, object]", jsonify_object(credentials_dict) ) @@ -175,6 +276,75 @@ async def get_credential_by_name( raise handle_exception_on_proxy(e) +@router.get( + "/credentials/{credential_name:path}/jwks", + dependencies=(Depends(user_api_key_auth),), + tags=["credential management"], # mutable-ok: FastAPI's include_router does self.tags.copy(), needs a real list +) +async def get_credential_internal_issuer_jwks( + credential_name: str = Path(..., description="The credential name, percent-decoded; may contain slashes"), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI resolves the dependency from the default +): + """ + Export the public JWKS for an anthropic ``internal_issuer`` credential, so the operator can + register it on the Anthropic federation issuer from the UI. Never touches the private signing + key: only its derived public JWKS leaves this process. 404s for any other credential shape. + """ + from litellm.proxy.proxy_server import prisma_client + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict + "error": "Only proxy admins can export a credential's JWKS." + }, + ) + + try: + credential: Final = await hydrate_named_credential_authoritative(credential_name, prisma_client) + credential_provider: Final = ( + None + if credential is None + else stored_credential_provider(credential.credential_info.get("custom_llm_provider")) + ) + if credential is None or credential_provider != "anthropic": + raise HTTPException( + status_code=404, + detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict + "error": f"No anthropic credential named {credential_name!r}." + }, + ) + configured_source: Final = credential.credential_values.get(_IDENTITY_SOURCE_PARAM) + if configured_source != AnthropicIdentitySourceKind.internal_issuer.value: + raise HTTPException( + status_code=404, + detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict + "error": ( + f"Credential {credential_name!r} is not configured with " + f"{_IDENTITY_SOURCE_PARAM}={AnthropicIdentitySourceKind.internal_issuer.value!r}." + ) + }, + ) + try: + issuer_source: Final = _build_variant( + InternalIssuerSource, credential.credential_values, _INTERNAL_ISSUER_FIELD_MAP + ) + jwks_document: Final = internal_issuer_jwks_document(issuer_source) + except (litellm.AuthenticationError, ValueError) as e: + raise HTTPException( + status_code=400, + detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict + "error": str(e) + }, + ) from e + return Response(content=jwks_document, media_type="application/json") + except HTTPException: + raise + except Exception as e: # noqa: BLE001 # endpoint boundary: every failure becomes the proxy's error contract + verbose_proxy_logger.exception(e) + raise handle_exception_on_proxy(e) + + @router.get( "/credentials/by_model/{model_id}", dependencies=[Depends(user_api_key_auth)], @@ -239,6 +409,9 @@ async def delete_credential( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) + _reject_non_admin_wif_fields( + await named_credential_wif_fields(credential_name, prisma_client), user_api_key_dict + ) deleted: Final = await CredentialsRepository(prisma_client).delete_by_name(credential_name) if deleted is None: raise HTTPException( @@ -249,6 +422,8 @@ async def delete_credential( ## DELETE FROM LITELLM ## litellm.credential_list = [cred for cred in litellm.credential_list if cred.credential_name != credential_name] return {"success": True, "message": "Credential deleted successfully"} + except HTTPException: + raise except Exception as e: raise handle_exception_on_proxy(e) @@ -282,6 +457,9 @@ def update_db_credential( merged_credential.credential_values.update(encrypted_params) + for key in updated_patch.credential_values_to_delete or (): + merged_credential.credential_values.pop(key, None) + # update model info if encrypted_credential.credential_info: """Update credential info""" @@ -310,6 +488,8 @@ async def update_credential( from litellm.proxy.proxy_server import prisma_client try: + _reject_overlapping_credential_values(credential) + _reject_non_admin_wif_fields(_incoming_wif_fields(credential), user_api_key_dict) if prisma_client is None: raise HTTPException( status_code=500, @@ -319,9 +499,14 @@ async def update_credential( db_credential: Final = await credentials_repository.find_by_name(credential_name) if db_credential is None: raise HTTPException(status_code=404, detail="Credential not found in DB.") + _reject_non_admin_wif_fields(_stored_wif_fields(db_credential), user_api_key_dict) + if credential.credential_name != credential_name: + shadowed_credential: Final = await hydrate_named_credential(credential.credential_name, prisma_client) + if shadowed_credential is not None: + _reject_non_admin_wif_fields(_stored_wif_fields(shadowed_credential), user_api_key_dict) merged_credential: Final = update_db_credential(db_credential, credential) credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str - "dict[str, object]", jsonify_object(merged_credential.model_dump()) + "dict[str, object]", jsonify_object(merged_credential.model_dump(exclude_none=True)) ) await credentials_repository.update_by_name( credential_name, @@ -332,29 +517,7 @@ async def update_credential( ) # Sync in-memory credential_list (skip if not in memory - e.g., proxy restarted) - new_name: Final = merged_credential.credential_name - existing_in_memory: CredentialItem | None = None - for cred in litellm.credential_list: - if cred.credential_name == credential_name: - existing_in_memory = cred - break - - if existing_in_memory is not None: - in_memory_values: Final = dict(existing_in_memory.credential_values or {}) - if credential.credential_values: - in_memory_values.update(credential.credential_values) - in_memory_info: Final = dict(existing_in_memory.credential_info or {}) - if credential.credential_info: - in_memory_info.update(credential.credential_info) - updated_in_memory: Final = CredentialItem( - credential_name=new_name, - credential_values=in_memory_values, - credential_info=in_memory_info, - ) - # Remove old entry if renamed, then use upsert_credentials to handle duplicates - if new_name != credential_name: - litellm.credential_list = [c for c in litellm.credential_list if c.credential_name != credential_name] - CredentialAccessor.upsert_credentials([updated_in_memory]) + _sync_in_memory_credential(credential, credential_name, merged_credential.credential_name) return {"success": True, "message": "Credential updated successfully"} except Exception as e: diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 1785a2f0992..daa8ccc7b0e 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -38,6 +38,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_utils import ( _BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check + reject_server_owned_wif_params, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler @@ -2065,6 +2066,7 @@ async def test_model_connection( "Could not find model %s in router: %s. Proceeding with request params only.", model_name, e ) + reject_server_owned_wif_params(request_litellm_params) # Merge: config params (from proxy config) as base, request params override litellm_params = { **_config_base_for_health_check( @@ -2091,6 +2093,9 @@ async def test_model_connection( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + # The Deployment above already carries the caller's merged params, so the effective + # state is model_params itself; there is no separate incoming patch here. + incoming_params=None, ) mode = mode or litellm_params.pop("mode", None) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b77108911aa..399ab8f3706 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -56,6 +56,10 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, publish_config_change, ) +from litellm.proxy.common_utils.credential_hydration import ( + effective_server_owned_wif_fields, + hydrate_named_credential, +) from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -289,6 +293,23 @@ def _raise_on_strategy_router_write_violation( ) +def _reject_non_admin_blocked_flag_on_create( + blocked: bool | None, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """Same proxy-admin-only rule patch_model applies to the blocked flag: a team admin passed + the team-scoped auth check above, but must not be able to create a model already paused + (or explicitly unpaused) out from under the proxy admin. + """ + if blocked is not None and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise ProxyException( + message="Only proxy admins can set a model's blocked flag.", + type=ProxyErrorTypes.auth_error.value, + code=status.HTTP_403_FORBIDDEN, + param="blocked", + ) + + AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301 _CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" _STORED_LITELLM_PARAMS_SQL: Final = ( @@ -831,6 +852,7 @@ async def patch_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + incoming_params=patch_data.litellm_params, ) # Pause/resume (`blocked`) is a proxy-admin-only privilege. Team admins @@ -1141,6 +1163,8 @@ async def _add_model_to_db( } if model_params.model_info.id is not None: _data["model_id"] = model_params.model_info.id + if model_params.blocked is not None: + _data["blocked"] = model_params.blocked _create_data: Final = cast("Mapping[str, object]", _data) # cast-ok: str-keyed json payload built just above if not should_create_model_in_db: return LiteLLM_ProxyModelTable(**_data) @@ -1691,14 +1715,65 @@ class ModelManagementAuthChecks: ) return True + @staticmethod + async def _reject_non_admin_wif_write( + *, + model_params: Deployment, + incoming_params: GenericLiteLLMParams | None, + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, + ) -> None: + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return + stored: Final = model_params.litellm_params.model_dump(exclude_none=True) + wif_fields: Final = await effective_server_owned_wif_fields(stored, incoming_params, prisma_client) + if wif_fields: + # ProxyException rather than HTTPException so the offending field stays a structured + # `param`, which is the contract the narrower gate this replaced already published. + raise ProxyException( + message=( + f"Only proxy admins can modify a deployment configured for workload identity " + f"federation ({wif_fields[0]!r})." + ), + type=ProxyErrorTypes.auth_error.value, + code=status.HTTP_403_FORBIDDEN, + param=wif_fields[0], + ) + # A name the caller expects an admin to create later would resolve to nothing today and + # start federating the moment it exists, so a non-admin may only attach one that is already there. + if incoming_params is not None and "litellm_credential_name" in incoming_params.model_fields_set: + named: Final = incoming_params.litellm_credential_name + if isinstance(named, str) and await hydrate_named_credential(named, prisma_client) is None: + raise ProxyException( + message=f"No credential named {named!r} exists.", + type=ProxyErrorTypes.bad_request_error.value, + code=status.HTTP_400_BAD_REQUEST, + param="litellm_credential_name", + ) + @staticmethod async def can_user_make_model_call( model_params: Deployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, premium_user: bool, + *, + incoming_params: GenericLiteLLMParams | None, allow_missing_team: bool = False, ) -> Literal[True]: + # Federation fields choose which server-side secret is read and where the org-scoped token + # it buys is sent, so only a proxy admin may touch a deployment that has them. Evaluated on + # the RESULTING deployment: a patch naming no federation field still lands on one that has + # them, and a patch attaching a credential by name inherits whatever that credential holds. + # `incoming_params` is keyword-only with no default so a new write path cannot typecheck + # without deciding what it writes. + await ModelManagementAuthChecks._reject_non_admin_wif_write( + model_params=model_params, + incoming_params=incoming_params, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + ## Check team model auth if model_params.model_info is not None and model_params.model_info.team_id is not None: team_obj_row: Final = await _repo_team_table(prisma_client).find_unique( @@ -1793,6 +1868,7 @@ async def delete_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + incoming_params=None, allow_missing_team=True, ) @@ -1914,7 +1990,6 @@ async def delete_team_model_alias( return removed_model_aliases -#### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964 @router.post( "/model/new", description="Allows adding new models to the model list in the config.yaml", @@ -1983,8 +2058,11 @@ async def add_new_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + incoming_params=model_params.litellm_params, ) + _reject_non_admin_blocked_flag_on_create(model_params.blocked, user_api_key_dict) + ModelManagementAuthChecks.can_user_attach_credential( litellm_params=model_params.litellm_params, user_api_key_dict=user_api_key_dict, @@ -2166,6 +2244,7 @@ async def update_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + incoming_params=model_params.litellm_params, ) ModelManagementAuthChecks.can_user_attach_credential( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 32da2658b99..444e8fb79d5 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -32,7 +32,11 @@ from litellm.constants import ( BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES, ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix -from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.anthropic.common_utils import ( + _SERVER_OWNED_AUTH_HEADERS, # pyright: ignore[reportPrivateUsage] # canonical set, must not be duplicated here + AnthropicModelInfo, + merge_anthropic_beta_headers, +) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse @@ -644,6 +648,57 @@ async def is_streaming_request_fn(request: Request) -> bool: return False +def _anthropic_passthrough_headers(auth_header: Mapping[str, str] | None, client_beta: str | None) -> Mapping[str, str]: + """Custom headers take priority over forwarded client headers, so merge the + client's anthropic-beta into the auth header's instead of clobbering it.""" + if auth_header is None: + return MappingProxyType({}) + auth_beta: Final = auth_header.get("anthropic-beta") + if auth_beta is None or client_beta is None: + return auth_header + return MappingProxyType({**auth_header, "anthropic-beta": merge_anthropic_beta_headers(client_beta, auth_beta)}) + + +def _configured_litellm_key_header_name() -> str | None: + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # proxy_server declares it as a bare dict + ) + + configured: Final = general_settings.get( # pyright: ignore[reportUnknownMemberType] # proxy_server general_settings is a bare dict + "litellm_key_header_name" + ) + return configured if isinstance(configured, str) else None + + +def _anthropic_passthrough_header_plan( + request: Request, auth_header: Mapping[str, str] | None, litellm_key_header_name: str | None +) -> tuple[Mapping[str, str], bool]: + """Returns the headers to send upstream plus whether the relay should still forward the + caller's own headers. Once the server owns the Anthropic credential, none of the headers + the proxy accepts a LiteLLM key in (``SpecialHeaders`` plus the configured custom name) + may ride upstream beside it, so the forward merge runs here with those stripped and the + relay is told not to merge again. With no server credential the caller's key is the only + one there is, so forwarding stays on (BYOK).""" + server_headers: Final = _anthropic_passthrough_headers(auth_header, request.headers.get("anthropic-beta")) + if auth_header is None: + return server_headers, True + caller_owned: Final = _SERVER_OWNED_AUTH_HEADERS | frozenset( + (litellm_key_header_name.lower(),) if litellm_key_header_name else () + ) + caller_headers: Final[dict[str, str]] = { # mutable-ok: forward_headers_from_request takes a concrete dict + name: value for name, value in request.headers.items() if name.lower() not in caller_owned + } + merged: Final = cast( # cast-ok: forward_headers_from_request is untyped upstream, its result is a header dict + "dict[str, str]", + HttpPassThroughEndpointHelpers.forward_headers_from_request( # pyright: ignore[reportUnknownMemberType] # untyped upstream + request_headers=caller_headers, + headers=dict(server_headers), # mutable-ok: forward_headers_from_request takes a concrete dict + forward_headers=True, + ), + ) + return MappingProxyType(merged), False + + @router.api_route( "/anthropic/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], @@ -682,12 +737,17 @@ async def anthropic_proxy_route( is_streaming_request: Final = await is_streaming_request_fn(request) ## CREATE PASS-THROUGH - auth_header: Final = AnthropicModelInfo.get_auth_header(anthropic_api_key or None) + auth_header: Final = await AnthropicModelInfo.aget_auth_header( + anthropic_api_key or None, allow_workload_identity=True + ) + upstream_headers, forward_caller_headers = _anthropic_passthrough_header_plan( + request, auth_header, _configured_litellm_key_header_name() + ) endpoint_func: Final = create_pass_through_route( endpoint=endpoint, target=str(updated_url), - custom_headers=auth_header if auth_header is not None else {}, - _forward_headers=True, + custom_headers=upstream_headers, + _forward_headers=forward_caller_headers, is_streaming_request=is_streaming_request, ) # dynamically construct pass-through endpoint based on incoming path received_value: Final = await endpoint_func( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1cd08fe27c0..fe14fa2e23c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -336,6 +336,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber +from litellm.proxy.common_utils.credential_hydration import decrypted_or_stored from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router from litellm.proxy.common_utils.encrypt_decrypt_utils import ( @@ -719,6 +720,7 @@ from litellm.types.router import ( RouterGeneralSettings, RoutingPlugin, SearchToolTypedDict, + holds_secret_pointer, updateDeployment, ) from litellm.types.router import ModelInfo as RouterModelInfo @@ -4679,7 +4681,7 @@ class ProxyConfig: if isinstance(item, dict): item = self._check_for_os_environ_vars(config=item, depth=depth + 1, max_depth=max_depth) # if the value is a string and starts with "os.environ/" - then it's an environment variable - elif isinstance(value, str) and value.startswith("os.environ/"): + elif isinstance(value, str) and value.startswith("os.environ/") and not holds_secret_pointer(key): resolved = get_secret(value) if resolved is None and secret_manager_would_be_consulted(value): verbose_proxy_logger.warning("%s is absent from the configured secret manager", value) @@ -5777,7 +5779,7 @@ class ProxyConfig: for model in model_list: ### LOAD FROM os.environ/ ### for k, v in model["litellm_params"].items(): - if isinstance(v, str) and v.startswith("os.environ/"): + if isinstance(v, str) and v.startswith("os.environ/") and not holds_secret_pointer(k): model["litellm_params"][k] = get_secret(v) validate_deployment_max_agentic_loops(model) validate_deployment_complexity_router_placement(model) @@ -6170,7 +6172,7 @@ class ProxyConfig: for model in model_list: ### LOAD FROM os.environ/ ### for k, v in model["litellm_params"].items(): - if isinstance(v, str) and v.startswith("os.environ/"): + if isinstance(v, str) and v.startswith("os.environ/") and not holds_secret_pointer(k): model["litellm_params"][k] = get_secret(v) ## check if they have model-id's ## @@ -6198,7 +6200,11 @@ class ProxyConfig: return value decrypted_value: Final = decrypt_value_helper(value=value, key=key, return_original_value=True) - if isinstance(decrypted_value, str) and decrypted_value.startswith("os.environ/"): + if ( + isinstance(decrypted_value, str) + and decrypted_value.startswith("os.environ/") + and not holds_secret_pointer(key) + ): return get_secret(decrypted_value) return decrypted_value @@ -7967,7 +7973,7 @@ class ProxyConfig: decrypted_credential_values: Final = {} for k, v in credential_object.credential_values.items(): - decrypted_credential_values[k] = decrypt_value_helper(value=v, key=k) or v + decrypted_credential_values[k] = decrypted_or_stored(k, v) credential_object.credential_values = decrypted_credential_values return credential_object diff --git a/litellm/router.py b/litellm/router.py index 490836f5f0a..cf65a06d604 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -235,6 +235,7 @@ from litellm.types.router import ( RoutingStrategy, SearchToolTypedDict, TaggedPreRoutingStrategy, + holds_secret_pointer, ) from litellm.types.services import ServiceTypes from litellm.types.utils import ( @@ -8650,13 +8651,10 @@ class Router: if ptu_error is not None and is_ptu_cost_attribution_enabled(): raise ValueError(ptu_error) zeroed_pricing: Final = zeroed_ptu_pricing(_model_info, _litellm_params) if config_sourced else None - litellm_params: Final[LiteLLM_Params] = LiteLLM_Params( - **( - _litellm_params - if zeroed_pricing is None - else MappingProxyType({**_litellm_params, **zeroed_pricing}) - ) + merged_params: Final[Mapping[str, Any]] = ( + _litellm_params if zeroed_pricing is None else MappingProxyType({**_litellm_params, **zeroed_pricing}) ) + litellm_params: Final[LiteLLM_Params] = LiteLLM_Params(**merged_params) warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=_litellm_params) deployment = Deployment( **deployment_info, @@ -9226,7 +9224,7 @@ class Router: ## check if litellm params in os.environ if isinstance(_litellm_params, dict): for k, v in _litellm_params.items(): - if isinstance(v, str) and v.startswith("os.environ/"): + if isinstance(v, str) and v.startswith("os.environ/") and not holds_secret_pointer(k): _litellm_params[k] = get_secret(v) _model_info: dict = model.pop("model_info", {}) diff --git a/litellm/router_utils/clientside_credential_handler.py b/litellm/router_utils/clientside_credential_handler.py index 55b246b22a7..186772925a9 100644 --- a/litellm/router_utils/clientside_credential_handler.py +++ b/litellm/router_utils/clientside_credential_handler.py @@ -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 diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 7fda5d96fb0..071752c9f55 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -24,7 +24,7 @@ from litellm.router_utils.cooldown_handlers import ( from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, ) -from litellm.types.router import LiteLLMParamsTypedDict +from litellm.types.router import LiteLLMParamsTypedDict, reject_server_owned_wif_params if TYPE_CHECKING: from litellm.router import Router as _Router @@ -507,6 +507,14 @@ async def run_async_fallback( failed_model_group: Final = get_pre_routing_selection(kwargs) or original_model_group attempted.record(failed_model_group) + # A dict target is merged straight into kwargs below, and kwargs win over the deployment's own + # params, so a stored key/team/global fallback could otherwise set a federation field that the + # request itself is forbidden to carry. Checked here rather than at the merge: inside the loop + # the refusal would be caught as a per-target failure and quietly skipped to the next one. + for target in fallback_model_group: + if isinstance(target, dict): + reject_server_owned_wif_params(target) + for mg in fallback_model_group: if mg == failed_model_group: continue diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index e89fbbdab65..7b5d111c99b 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -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/' (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." ) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index b3462203c4b..0f77ae1d468 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -750,5 +750,6 @@ ANTHROPIC_EFFORT_BETA_HEADER: Final = "effort-2025-11-24" # OAuth constants ANTHROPIC_OAUTH_TOKEN_PREFIX: Final = "sk-ant-oat" ANTHROPIC_OAUTH_BETA_HEADER: Final = "oauth-2025-04-20" +ANTHROPIC_TOKEN_EXCHANGE_PATH: Final = "/v1/oauth/token" ANTHROPIC_PROMPT_CACHING_SCOPE_BETA_HEADER: Final = "prompt-caching-scope-2026-01-05" diff --git a/litellm/types/router.py b/litellm/types/router.py index 728d1037f3d..832399f6dbb 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -4,7 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc import datetime import enum -from collections.abc import Mapping +from collections.abc import Container, Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints @@ -27,6 +27,10 @@ from .utils import ( ModelResponse, StandardLoggingRoutingDecision, ) +from .utils import ( + # private alias: `from .types.router import *` would rebind a public Final in litellm/__init__.py + server_owned_wif_litellm_params as _server_owned_wif_litellm_params, +) class ConfigurableClientsideParamsCustomAuth(TypedDict): @@ -297,6 +301,69 @@ class CredentialLiteLLMParams(BaseModel): ## IBM WATSONX ## watsonx_region_name: str | None = None + ## ANTHROPIC WORKLOAD IDENTITY FEDERATION ## + # Without these, get_deployment_credentials_with_provider silently drops a + # litellm_params-configured WIF setup before files/batches/passthrough callers see + # it, the same #30235-shaped gap azure_ad_token above was added to close. + anthropic_federation_rule_id: str | None = None + anthropic_organization_id: str | None = None + anthropic_service_account_id: str | None = None + anthropic_workspace_id: str | None = None + anthropic_identity_token_file: str | None = None + anthropic_identity_token: str | None = None + anthropic_identity_source: str | None = None + anthropic_issuer_url: str | None = None + anthropic_issuer_subject: str | None = None + anthropic_issuer_audience: str | None = None + anthropic_issuer_ttl_seconds: int | None = None + anthropic_issuer_signing_key_ref: str | None = None + anthropic_keycloak_token_url: str | None = None + anthropic_keycloak_client_id: str | None = None + anthropic_keycloak_auth_method: str | None = None + anthropic_keycloak_client_secret_ref: str | None = None + anthropic_keycloak_scope: str | None = None + # Server-set when a client redirects api_base. Declared so it survives the strict dump the + # other federation fields above are declared for, rather than being rebuilt away in transit. + anthropic_disable_workload_identity_federation: bool | None = None + + ## OPENAI WORKLOAD IDENTITY FEDERATION ## + openai_identity_provider_id: str | None = None + openai_service_account_id: str | None = None + openai_identity_token_file: str | None = None + + +def server_owned_wif_fields_present(fields: Mapping[str, object]) -> tuple[str, ...]: + """Server-owned workload identity federation field names set in ``fields``. + + ``fields`` is a ``litellm_params`` dict (or a credential's ``credential_values`` mapping, + which feeds the same resolution when referenced by name). Derived from + ``server_owned_wif_litellm_params`` rather than hand-copied, so a persistence gate built on + this stays correct when a new WIF field is added there. + """ + return tuple(name for name in _server_owned_wif_litellm_params if fields.get(name) is not None) + + +def server_owned_wif_fields_named(keys: Container[str]) -> tuple[str, ...]: + """Server-owned workload identity federation field names that appear in ``keys``, whatever + value they carry. + + The write gates on credentials need this key-based sibling of ``server_owned_wif_fields_present``: + ``get_litellm_params`` forwards a WIF kwarg on key presence and the federation resolver rejects + a foreign variant's field by key, so a persisted ``{"anthropic_issuer_url": None}`` wedges every + deployment that references the credential even though no value is set. Pass a mapping (its keys + are tested) or a plain collection of key names. + """ + return tuple(name for name in _server_owned_wif_litellm_params if name in keys) + + +_WIF_POINTER_FIELDS: Final = frozenset(name for name in _server_owned_wif_litellm_params if name.endswith("_ref")) + + +def holds_secret_pointer(param_name: str) -> bool: + """A ``*_ref`` federation field is a secret POINTER the identity source dereferences at use + time, so a loader expanding ``os.environ/`` values must leave it as written.""" + return param_name in _WIF_POINTER_FIELDS + _RESERVED_INIT_KEYS: Final = frozenset({"self", "params", "__class__"}) @@ -542,6 +609,9 @@ class Deployment(BaseModel): model_name: str litellm_params: LiteLLM_Params model_info: ModelInfo + # admin-toggled pause flag; mirrors LiteLLM_ProxyModelTable.blocked. None means "don't set it + # on create" -- the Prisma column defaults to False -- rather than "explicitly unblocked". + blocked: bool | None = None model_config = ConfigDict(extra="allow", protected_namespaces=()) @@ -1064,3 +1134,23 @@ class AdaptiveRouterPreferences(BaseModel): quality_tier: int = Field(ge=1, le=3) strengths: list[RequestType] = Field(default_factory=list) + + +_BEDROCK_WORKSPACE_HINT: Final = " On the Bedrock Claude Platform route, pass workspace_id or aws_workspace_id instead." + + +def reject_server_owned_wif_params(body: Mapping[str, object]) -> None: + """Raise ``ValueError`` if a mapping that did not come from deployment config carries a + server-owned workload identity federation field. + + These are never client-settable on any surface, with or without a client-side credential + opt-in. This lives here rather than under ``litellm.proxy`` so the router can call it on a + post-authentication merge without core importing from the proxy package. + """ + for param in _server_owned_wif_litellm_params: + if param in body: + raise ValueError( + f"Rejected Request: {param} is a server-owned workload identity federation parameter " + "and cannot be set in a request body; configure it on the deployment instead." + + (_BEDROCK_WORKSPACE_HINT if param == "anthropic_workspace_id" else "") + ) diff --git a/litellm/types/services.py b/litellm/types/services.py index c558f6fb9d2..b5a3383bc96 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -25,6 +25,8 @@ class ServiceTypes(str, enum.Enum): AUTH = "auth" PROXY_PRE_CALL = "proxy_pre_call" POD_LOCK_MANAGER = "pod_lock_manager" + ANTHROPIC_WIF = "anthropic_wif" + ANTHROPIC_WIF_CACHE = "anthropic_wif_cache" """ Operational metrics for DB Transaction Queues @@ -67,6 +69,13 @@ DEFAULT_SERVICE_CONFIGS: Final = { ServiceTypes.ROUTER.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]}, ServiceTypes.AUTH.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]}, ServiceTypes.PROXY_PRE_CALL.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]}, + ServiceTypes.ANTHROPIC_WIF.value: { # mutable-ok: ServiceConfig mandates the dict-of-list shape + "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] # mutable-ok: ServiceConfig mandates a list + }, + # cache hits are counter-only: no HTTP call happens, so observing a latency would be a lie + ServiceTypes.ANTHROPIC_WIF_CACHE.value: { # mutable-ok: ServiceConfig mandates the dict-of-list shape + "metrics": [ServiceMetrics.COUNTER] # mutable-ok: ServiceConfig mandates a list + }, # Operational metrics for DB Transaction Queues ServiceTypes.POD_LOCK_MANAGER.value: {"metrics": [ServiceMetrics.GAUGE]}, ServiceTypes.IN_MEMORY_DAILY_SPEND_UPDATE_QUEUE.value: {"metrics": [ServiceMetrics.GAUGE]}, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5052cd6ef48..5f5a0646414 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3643,9 +3643,27 @@ bedrock_batch_litellm_params: Final = ( "bedrock_tags", ) +# Anthropic workload identity federation config, read from litellm_params by the +# Anthropic auth tier. Listed for the same reason as the fields above: an +# unrecognized top-level key is swept into extra_body and sent to /v1/messages. +# Derived from get_litellm_params.ANTHROPIC_WIF_KWARGS_KEYS (not hand-typed) so the +# request-body ban list and the clear-on-api_base-override list can never drift from +# the set the kwargs funnel actually forwards. Imported here rather than at module top: +# get_litellm_params.py's own import chain (llms/openai/data_residency -> llms/__init__) +# reaches back into this module for CallTypes, which by this point in the file is +# already bound on the partially-initialized module. +from ..litellm_core_utils.get_litellm_params import ( # noqa: E402 # deferred past CallTypes to break the import cycle + ANTHROPIC_WIF_KWARGS_KEYS, + OPENAI_WIF_KWARGS_KEYS, +) + +anthropic_wif_litellm_params: Final = tuple(sorted(ANTHROPIC_WIF_KWARGS_KEYS)) +openai_wif_litellm_params: Final = tuple(sorted(OPENAI_WIF_KWARGS_KEYS)) +server_owned_wif_litellm_params: Final = anthropic_wif_litellm_params + openai_wif_litellm_params + all_litellm_params = ( agentic_loop_internal_litellm_params - + [TRUSTED_CALLBACK_VARS_FIELD, *bedrock_batch_litellm_params] + + [TRUSTED_CALLBACK_VARS_FIELD, *bedrock_batch_litellm_params, *server_owned_wif_litellm_params] + [ "metadata", "litellm_metadata", diff --git a/litellm/utils.py b/litellm/utils.py index 9d20d32d147..12301681e85 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7361,9 +7361,8 @@ def _get_valid_models_from_provider_api( if cached_result is not None: return cached_result - models: Final = provider_config.get_models( - api_key=litellm_params.api_key if litellm_params is not None else None, - api_base=litellm_params.api_base if litellm_params is not None else None, + models: Final = provider_config.discover_models( + litellm_params=litellm_params.model_dump(exclude_none=True) if litellm_params is not None else None ) _model_cache.set_cached_model_info(custom_llm_provider, litellm_params, models) diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 6bc8947e89f..b9727a8ce70 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -36,6 +36,7 @@ GET /user/spend/report # Admin UI helper endpoints; serve UI forms and caller-scoped views, not desired state GET /budget/settings +GET /credentials/{credential_name}/jwks GET /router/fields GET /guardrails/ui/add_guardrail_settings GET /guardrails/ui/category_yaml/{category_name} diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index c86c7c4df03..0be7a3a345e 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1235,6 +1235,7 @@ async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch): result set - zero cost, zero usage, no models - instead of letting the file fetch raise "Output file id is None" on every aretrieve_batch logging poll. """ + # The output-file fetch must not even be attempted when there is no output file. async def _must_not_fetch(*args, **kwargs): pytest.fail("_fetch_batch_output_file_content should not be called") @@ -1361,7 +1362,10 @@ def test_anthropic_response_body_is_result_message(): def test_anthropic_usage_conversion_includes_cache_tokens(): - body = {"model": "claude-sonnet-4-5-20250929", "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)} + body = { + "model": "claude-sonnet-4-5-20250929", + "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000), + } usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="anthropic") assert usage.prompt_tokens == 11000 assert usage.completion_tokens == 200 @@ -1376,7 +1380,9 @@ def test_bedrock_model_output_line_success_check(): "modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}}, } assert bu._batch_response_was_successful(row, custom_llm_provider="bedrock") is True - assert bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6" + assert ( + bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6" + ) def test_bedrock_cost_uses_deployment_model_name(): @@ -1430,7 +1436,13 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch): rows = [ { "custom_id": "req-1", - "response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}}, + "response": { + "status_code": 200, + "body": { + "model": "gpt-5.2", + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, } ] result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") @@ -1674,7 +1686,10 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> def test_bedrock_converse_shaped_batch_usage_is_parsed(): - body = {"model": "us.amazon.nova-lite-v1:0", "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}} + body = { + "model": "us.amazon.nova-lite-v1:0", + "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}, + } usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock") assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (2202, 540, 2742) @@ -1718,3 +1733,37 @@ def test_unparsable_bedrock_batch_usage_warns(caplog): assert usage.total_tokens == 0 assert "does not understand" in caplog.text assert "inputTextTokenCount" in caplog.text + + +class TestFileAccessCredentialsCarryFederation: + """A federated deployment holds no api_key, so the fetch that reads a finished batch's output + has to inherit the federation fields or it cannot authenticate and the batch is never billed.""" + + def test_federation_fields_survive_extraction(self): + from litellm.batches.batch_utils import _extract_file_access_credentials + + credentials = _extract_file_access_credentials( + { + "model": "anthropic/claude-sonnet-4-5", + "anthropic_federation_rule_id": "fdrl_x", + "anthropic_organization_id": "org-x", + "anthropic_identity_token_file": "/var/run/secrets/anthropic.com/token", + "something_unrelated": "dropped", + } + ) + + assert credentials["anthropic_federation_rule_id"] == "fdrl_x" + assert credentials["anthropic_organization_id"] == "org-x" + assert credentials["anthropic_identity_token_file"] == "/var/run/secrets/anthropic.com/token" + assert "something_unrelated" not in credentials + + def test_every_federation_field_is_carried(self): + """Derived from the kwargs set, so a new federation field is carried without an edit here.""" + from litellm.batches.batch_utils import _extract_file_access_credentials + from litellm.litellm_core_utils.get_litellm_params import ANTHROPIC_WIF_KWARGS_KEYS + + params = {name: f"value-{name}" for name in ANTHROPIC_WIF_KWARGS_KEYS} + + credentials = _extract_file_access_credentials(params) + + assert set(credentials) == set(ANTHROPIC_WIF_KWARGS_KEYS) diff --git a/tests/test_litellm/integrations/test_prometheus_services.py b/tests/test_litellm/integrations/test_prometheus_services.py index 2303061ede8..a5e95de3ea9 100644 --- a/tests/test_litellm/integrations/test_prometheus_services.py +++ b/tests/test_litellm/integrations/test_prometheus_services.py @@ -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", + } diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index 956da571d43..ad4e417a29f 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -244,3 +244,129 @@ class TestRustOptIn: from litellm.types.utils import all_litellm_params assert "rust" in all_litellm_params + + +class TestAnthropicWifKeys: + """The six anthropic_* WIF keys need the same dual registration as `rust`: + carried by the kwargs funnel into litellm_params (where the Anthropic auth + tier reads them) AND listed in all_litellm_params (so the extra_body sweep + never sends them to /v1/messages).""" + + SIX_KEYS = { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_service_account_id": "svcacct_1", + "anthropic_workspace_id": "wrkspc_1", + "anthropic_identity_token_file": "/var/run/secrets/tok", + "anthropic_identity_token": "oidc/env/TOK", + } + + def test_keys_survive_into_litellm_params(self): + params = get_litellm_params(**self.SIX_KEYS) + for key, value in self.SIX_KEYS.items(): + assert params[key] == value + + def test_keys_are_forwarded_from_completion_kwargs(self): + from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS + + assert set(self.SIX_KEYS) <= FORWARDED_KWARGS_KEYS + + def test_keys_stay_out_of_the_provider_body(self): + from litellm.types.utils import all_litellm_params + + for key in self.SIX_KEYS: + assert key in all_litellm_params + + def test_keys_absent_when_not_configured(self): + params = get_litellm_params() + for key in self.SIX_KEYS: + assert key not in params + + +class TestAnthropicWifIdentitySourceKeys: + """Phase 1 adds 11 more anthropic_* WIF keys (the anthropic_identity_source discriminator + plus the internal_issuer/keycloak identity-source fields) that need the same dual + registration as the original six tested above.""" + + NEW_KEYS = { + "anthropic_identity_source": "keycloak", + "anthropic_issuer_url": "https://issuer.example", + "anthropic_issuer_subject": "svc-account", + "anthropic_issuer_audience": "https://api.anthropic.com", + "anthropic_issuer_ttl_seconds": "300", + "anthropic_issuer_signing_key_ref": "oidc/env/ISSUER_KEY", + "anthropic_keycloak_token_url": "https://kc.example/realms/r/protocol/openid-connect/token", + "anthropic_keycloak_client_id": "litellm", + "anthropic_keycloak_auth_method": "client_secret_basic", + "anthropic_keycloak_client_secret_ref": "oidc/env/KC_SECRET", + "anthropic_keycloak_scope": "anthropic-wif", + # Server-set when a client redirects api_base; carried here so it is not dropped in transit + "anthropic_disable_workload_identity_federation": True, + } + + def test_new_keys_are_exactly_the_non_legacy_registered_set(self): + """Fails the moment a key is added to ANTHROPIC_WIF_KWARGS_KEYS without a matching entry + here (or vice versa), catching drift between what wif.py dispatches on and what this + test (and the funnel/provider-body tests below) actually exercises.""" + from litellm.litellm_core_utils.get_litellm_params import ANTHROPIC_WIF_KWARGS_KEYS + + assert set(self.NEW_KEYS) == ANTHROPIC_WIF_KWARGS_KEYS - set(TestAnthropicWifKeys.SIX_KEYS) + + def test_keys_survive_into_litellm_params(self): + params = get_litellm_params(**self.NEW_KEYS) + for key, value in self.NEW_KEYS.items(): + assert params[key] == value + + def test_keys_are_forwarded_from_completion_kwargs(self): + from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS + + assert set(self.NEW_KEYS) <= FORWARDED_KWARGS_KEYS + + def test_keys_stay_out_of_the_provider_body(self): + from litellm.types.utils import all_litellm_params + + for key in self.NEW_KEYS: + assert key in all_litellm_params + + def test_keys_absent_when_not_configured(self): + params = get_litellm_params() + for key in self.NEW_KEYS: + assert key not in params + + +class TestOpenAIWifKeys: + """The three openai_* WIF keys carry a deployment's federation identity through the kwargs + funnel into litellm_params (where the OpenAI client factory reads them) and stay out of the + provider body, exactly like the anthropic_* keys above.""" + + THREE_KEYS = { + "openai_identity_provider_id": "idp_1", + "openai_service_account_id": "user-1", + "openai_identity_token_file": "/var/run/secrets/tokens/openai", + } + + def test_keys_are_exactly_the_registered_set(self): + from litellm.litellm_core_utils.get_litellm_params import OPENAI_WIF_KWARGS_KEYS + + assert set(self.THREE_KEYS) == OPENAI_WIF_KWARGS_KEYS + + def test_keys_survive_into_litellm_params(self): + params = get_litellm_params(**self.THREE_KEYS) + for key, value in self.THREE_KEYS.items(): + assert params[key] == value + + def test_keys_are_forwarded_from_completion_kwargs(self): + from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS + + assert set(self.THREE_KEYS) <= FORWARDED_KWARGS_KEYS + + def test_keys_stay_out_of_the_provider_body(self): + from litellm.types.utils import all_litellm_params + + for key in self.THREE_KEYS: + assert key in all_litellm_params + + def test_keys_absent_when_not_configured(self): + params = get_litellm_params() + for key in self.THREE_KEYS: + assert key not in params diff --git a/tests/test_litellm/llms/anthropic/batches/test_handler.py b/tests/test_litellm/llms/anthropic/batches/test_handler.py index 6fde6350127..398c94a6798 100644 --- a/tests/test_litellm/llms/anthropic/batches/test_handler.py +++ b/tests/test_litellm/llms/anthropic/batches/test_handler.py @@ -14,6 +14,8 @@ asyncio.run) is exercised directly, mirroring the dispatch-contract discipline i tests/test_litellm/batches/test_main.py. """ +import asyncio +import threading from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -35,9 +37,7 @@ def _ok_batch_response(): "ended_at": "2024-09-24T11:00:00Z", "request_counts": {"succeeded": 2, "errored": 0}, }, - request=httpx.Request( - "GET", "https://api.anthropic.com/v1/messages/batches/msgbatch_abc" - ), + request=httpx.Request("GET", "https://api.anthropic.com/v1/messages/batches/msgbatch_abc"), ) @@ -59,9 +59,7 @@ def patched_client(): @pytest.mark.asyncio -async def test_aretrieve_batch_fires_get_with_correct_url_and_headers( - handler, patched_client -): +async def test_aretrieve_batch_fires_get_with_correct_url_and_headers(handler, patched_client): fake_client, factory = patched_client batch = await handler.aretrieve_batch( @@ -76,9 +74,7 @@ async def test_aretrieve_batch_fires_get_with_correct_url_and_headers( fake_client.get.assert_awaited_once() _, call_kwargs = fake_client.get.call_args # Exact URL built by get_retrieve_batch_url. - assert call_kwargs["url"] == ( - "https://api.anthropic.com/v1/messages/batches/msgbatch_abc" - ) + assert call_kwargs["url"] == ("https://api.anthropic.com/v1/messages/batches/msgbatch_abc") # Auth + version + beta headers built by validate_environment. headers = call_kwargs["headers"] assert headers["x-api-key"] == "sk-ant-test" @@ -93,9 +89,7 @@ async def test_aretrieve_batch_fires_get_with_correct_url_and_headers( @pytest.mark.asyncio -async def test_aretrieve_batch_uses_anthropic_provider_for_client( - handler, patched_client -): +async def test_aretrieve_batch_uses_anthropic_provider_for_client(handler, patched_client): from litellm.types.utils import LlmProviders _, factory = patched_client @@ -111,14 +105,10 @@ async def test_aretrieve_batch_uses_anthropic_provider_for_client( @pytest.mark.asyncio -async def test_aretrieve_batch_resolves_api_key_from_model_info( - handler, patched_client -): +async def test_aretrieve_batch_resolves_api_key_from_model_info(handler, patched_client): fake_client, _ = patched_client # api_key=None -> handler falls back to AnthropicModelInfo.get_api_key(). - with patch.object( - handler.anthropic_model_info, "get_api_key", return_value="sk-from-env" - ): + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="sk-from-env"): await handler.aretrieve_batch( batch_id="msgbatch_abc", api_base="https://api.anthropic.com", @@ -134,9 +124,7 @@ async def test_aretrieve_batch_resolves_api_key_from_model_info( async def test_aretrieve_batch_missing_api_key_raises(handler, patched_client): fake_client, _ = patched_client # No api_key and resolver yields None -> hard error before any network call. - with patch.object( - handler.anthropic_model_info, "get_api_key", return_value=None - ): + with patch.object(handler.anthropic_model_info, "get_api_key", return_value=None): with pytest.raises(ValueError, match="Missing Anthropic API Key"): await handler.aretrieve_batch( batch_id="msgbatch_abc", @@ -165,9 +153,7 @@ async def test_aretrieve_batch_resolves_default_api_base(handler, patched_client max_retries=0, ) _, call_kwargs = fake_client.get.call_args - assert call_kwargs["url"] == ( - "https://api.anthropic.com/v1/messages/batches/msgbatch_abc" - ) + assert call_kwargs["url"] == ("https://api.anthropic.com/v1/messages/batches/msgbatch_abc") @pytest.mark.asyncio @@ -176,9 +162,7 @@ async def test_aretrieve_batch_raises_for_status(handler): error_response = httpx.Response( status_code=404, json={"error": "not found"}, - request=httpx.Request( - "GET", "https://api.anthropic.com/v1/messages/batches/missing" - ), + request=httpx.Request("GET", "https://api.anthropic.com/v1/messages/batches/missing"), ) fake_client = MagicMock() fake_client.get = AsyncMock(return_value=error_response) @@ -213,21 +197,15 @@ async def test_aretrieve_batch_invokes_pre_call_logging(handler, patched_client) assert pre_kwargs["input"] == "msgbatch_abc" assert pre_kwargs["api_key"] == "sk-ant-test" # The logged api_base is the full retrieve URL, not the bare base. - assert pre_kwargs["additional_args"]["api_base"] == ( - "https://api.anthropic.com/v1/messages/batches/msgbatch_abc" - ) + assert pre_kwargs["additional_args"]["api_base"] == ("https://api.anthropic.com/v1/messages/batches/msgbatch_abc") @pytest.mark.asyncio -async def test_aretrieve_batch_builds_default_logging_obj_when_absent( - handler, patched_client -): +async def test_aretrieve_batch_builds_default_logging_obj_when_absent(handler, patched_client): # logging_obj=None -> handler constructs a real Logging object; the call # must still complete (no AttributeError on a missing logger). _, _ = patched_client - with patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as logging_cls: + with patch("litellm.litellm_core_utils.litellm_logging.Logging") as logging_cls: logging_cls.return_value = MagicMock() batch = await handler.aretrieve_batch( batch_id="msgbatch_abc", @@ -281,3 +259,99 @@ def test_retrieve_batch_sync_runs_to_result(handler, patched_client): assert isinstance(batch, LiteLLMBatch) assert batch.id == "msgbatch_abc" assert batch.status == "completed" + + +# =========================================================================== # +# aretrieve_batch must not block the event loop on a WIF token exchange +# =========================================================================== # + +_WIF_ENV = { + "ANTHROPIC_FEDERATION_RULE_ID": "fdrl_batches_seam", + "ANTHROPIC_ORGANIZATION_ID": "org-batches-seam", + "ANTHROPIC_IDENTITY_TOKEN": "batches-seam-inline-jwt", +} + + +class _BlockingPoster: + """A token-endpoint poster that blocks until released, so the test can prove + the exchange ran off the event loop's own thread instead of freezing it.""" + + def __init__(self): + self.release = threading.Event() + self.thread_ids = [] + + def post(self, url, *, content, headers, timeout): + self.thread_ids.append(threading.get_ident()) + self.release.wait(timeout=5) + return httpx.Response( + 200, + json={ + "access_token": "sk-ant-oat01-batches-seam", + "token_type": "Bearer", + "expires_in": 3600, + }, + ) + + +@pytest.mark.asyncio +async def test_aretrieve_batch_wif_exchange_does_not_block_event_loop(handler, patched_client, monkeypatch): + """Regression: aretrieve_batch called the synchronous validate_environment + directly, so a cold WIF mint ran inline on the event loop and froze every + other concurrent coroutine until the exchange finished.""" + from litellm.llms.anthropic import common_utils as anthropic_common_utils + from litellm.llms.anthropic.wif import get_anthropic_wif_token + from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + + fake_client, _ = patched_client + for name in ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_API_BASE", + "ANTHROPIC_BASE_URL", + ): + monkeypatch.delenv(name, raising=False) + for name, value in _WIF_ENV.items(): + monkeypatch.setenv(name, value) + + poster = _BlockingPoster() + engine = JwtBearerTokenExchangeEngine(poster=poster) + + def routed_through_injected_engine(litellm_params, api_base, model): + return get_anthropic_wif_token(litellm_params, api_base, model, engine) + + monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", routed_through_injected_engine) + + ticks = [] + + async def ticker(): + for i in range(20): + await asyncio.sleep(0.005) + ticks.append(i) + + ticker_task = asyncio.create_task(ticker()) + await asyncio.sleep(0.02) + + retrieve_task = asyncio.create_task( + handler.aretrieve_batch( + batch_id="msgbatch_abc", + api_base="https://api.anthropic.com", + api_key=None, + timeout=60.0, + max_retries=0, + ) + ) + await asyncio.sleep(0.05) + # The ticker kept advancing while the token exchange was still blocked on + # poster.release, proving the exchange did not run on the event loop. + assert len(ticks) > 0 + assert not retrieve_task.done() + + poster.release.set() + batch = await retrieve_task + await ticker_task + + assert batch.id == "msgbatch_abc" + assert poster.thread_ids + assert poster.thread_ids[0] != threading.get_ident() + sent_headers = fake_client.get.call_args.kwargs["headers"] + assert sent_headers["authorization"] == "Bearer sk-ant-oat01-batches-seam" diff --git a/tests/test_litellm/llms/anthropic/batches/test_transformation.py b/tests/test_litellm/llms/anthropic/batches/test_transformation.py index eacd2c9d03b..dbdb77db093 100644 --- a/tests/test_litellm/llms/anthropic/batches/test_transformation.py +++ b/tests/test_litellm/llms/anthropic/batches/test_transformation.py @@ -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( diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 3044a321aa6..6f4f68c4940 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -48,9 +48,7 @@ class MockDynamicGuardrail(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: Optional[Any] = None, ) -> GenericGuardrailAPIInputs: - self.dynamic_params = self.get_guardrail_dynamic_request_body_params( - request_data - ) + self.dynamic_params = self.get_guardrail_dynamic_request_body_params(request_data) return inputs @@ -197,9 +195,7 @@ class TestAnthropicMessagesHandlerStreamingRequestData: assert guardrail.request_data is not None assert guardrail.request_data["response"] is mock_response - assert ( - guardrail.request_data["litellm_metadata"]["user_api_key_user_id"] == "u-1" - ) + assert guardrail.request_data["litellm_metadata"]["user_api_key_user_id"] == "u-1" @pytest.mark.asyncio async def test_mid_stream_chunk_passes_responses_so_far_and_metadata(self): @@ -211,9 +207,7 @@ class TestAnthropicMessagesHandlerStreamingRequestData: with ( patch.object(handler, "_check_streaming_has_ended", return_value=False), - patch.object( - handler, "get_streaming_string_so_far", return_value="partial text" - ), + patch.object(handler, "get_streaming_string_so_far", return_value="partial text"), ): await handler.process_output_streaming_response( responses_so_far=responses_so_far, @@ -225,9 +219,7 @@ class TestAnthropicMessagesHandlerStreamingRequestData: assert guardrail.request_data is not None assert guardrail.request_data["responses"] is responses_so_far - assert ( - guardrail.request_data["litellm_metadata"]["user_api_key_user_id"] == "u-1" - ) + assert guardrail.request_data["litellm_metadata"]["user_api_key_user_id"] == "u-1" class TestAnthropicMessagesHandlerStreamingOutputProcessing: @@ -276,17 +268,11 @@ class TestAnthropicMessagesHandlerInputProcessing: data = { "model": "claude-3-5-sonnet-20241022", "messages": [{"role": "user", "content": "hello"}], - "litellm_metadata": { - "guardrails": [ - {"cygnal-monitor": {"extra_body": {"policy_id": "policy-123"}}} - ] - }, + "litellm_metadata": {"guardrails": [{"cygnal-monitor": {"extra_body": {"policy_id": "policy-123"}}}]}, } with patch("litellm.proxy.proxy_server.premium_user", True): - await handler.process_input_messages( - data=data, guardrail_to_apply=guardrail - ) + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) assert data.get("litellm_metadata", {}).get("guardrails") assert guardrail.dynamic_params == {"policy_id": "policy-123"} @@ -1183,9 +1169,7 @@ class TestAnthropicMessagesHandlerInputProcessing: # Mock _check_streaming_has_ended to return False (stream not ended) with ( patch.object(handler, "_check_streaming_has_ended", return_value=False), - patch.object( - handler, "get_streaming_string_so_far", return_value="partial text" - ), + patch.object(handler, "get_streaming_string_so_far", return_value="partial text"), ): responses_so_far = [b"data: some chunk"] @@ -1216,9 +1200,7 @@ class TestAnthropicMessagesHandlerInputProcessing: data = { "model": "claude-opus-4-6", - "messages": [ - {"role": "user", "content": "What is the weather in San Francisco?"} - ], + "messages": [{"role": "user", "content": "What is the weather in San Francisco?"}], "tools": [ { "type": "tool_search_tool_regex_20251119", @@ -1381,17 +1363,11 @@ class TestAnthropicMessagesIncrementalScan: ] with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} - await handler.process_input_messages( - data=self._data(turn1, sid), guardrail_to_apply=guardrail - ) + await handler.process_input_messages(data=self._data(turn1, sid), guardrail_to_apply=guardrail) assert mock_api.call_count == 1 - assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ - "What is the capital of France?" - ] + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == ["What is the capital of France?"] mock_api.reset_mock() - await handler.process_input_messages( - data=self._data(turn2, sid), guardrail_to_apply=guardrail - ) + await handler.process_input_messages(data=self._data(turn2, sid), guardrail_to_apply=guardrail) assert mock_api.call_count == 1 assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ "Paris.", @@ -1830,9 +1806,7 @@ class TestAnthropicMessagesScanOnlyToolResults: await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) - assert guardrail.seen_texts == ["fetched POISON page"], ( - "only the tool_result payload may reach the guardrail" - ) + assert guardrail.seen_texts == ["fetched POISON page"], "only the tool_result payload may reach the guardrail" assert guardrail.captured_inputs is not None assert guardrail.captured_inputs.get("tools") is None assert [m["role"] for m in guardrail.captured_inputs["structured_messages"]] == ["tool"] diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index 043537f8c1f..92d0396854f 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -23,9 +23,7 @@ async def test_make_call_passes_logging_obj_to_client_post(): mock_client = AsyncMock() mock_response = MagicMock() mock_response.aiter_lines = MagicMock( - return_value=iter( - [b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n'] - ) + return_value=iter([b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n']) ) mock_client.post.return_value = mock_response @@ -94,9 +92,7 @@ def test_redacted_thinking_content_block_delta(): "data": "EuoBCoYBGAIiQJ/SxkPAgqxhKok29YrpJHRUJ0OT8ahCHKAwyhmRuUhtdmDX9+mn4gDzKNv3fVpQdB01zEPMzNY3QuTCd+1bdtEqQK6JuKHqdndbwpr81oVWb4wxd1GqF/7Jkw74IlQa27oobX+KuRkopr9Dllt/RDe7Se0sI1IkU7tJIAQCoP46OAwSDF51P09q67xhHlQ3ihoM2aOVlkghq/X0w8NlIjBMNvXYNbjhyrOcIg6kPFn2ed/KK7Cm5prYAtXCwkb4Wr5tUSoSHu9T5hKdJRbr6WsqEc7Lle7FULqMLZGkhqXyc3BA", }, } - model_response_iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=False, json_mode=False - ) + model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False) model_response = model_response_iterator.chunk_parser(chunk=chunk) print(f"\n\nmodel_response: {model_response}\n\n") assert model_response.choices[0].delta.thinking_blocks is not None @@ -104,19 +100,14 @@ def test_redacted_thinking_content_block_delta(): print( f"\n\nmodel_response.choices[0].delta.thinking_blocks[0]: {model_response.choices[0].delta.thinking_blocks[0]}\n\n" ) - assert ( - model_response.choices[0].delta.thinking_blocks[0]["type"] - == "redacted_thinking" - ) + assert model_response.choices[0].delta.thinking_blocks[0]["type"] == "redacted_thinking" assert model_response.choices[0].delta.provider_specific_fields is not None assert "thinking_blocks" in model_response.choices[0].delta.provider_specific_fields def test_streaming_thinking_blocks_are_replayable_after_signature_delta(): - model_response_iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) chunks = [ { "type": "content_block_start", @@ -140,17 +131,12 @@ def test_streaming_thinking_blocks_are_replayable_after_signature_delta(): }, ] - parsed_chunks = [ - model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks - ] + parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks] reasoning_content = "".join( - getattr(chunk.choices[0].delta, "reasoning_content", None) or "" - for chunk in parsed_chunks + getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks ) thinking_blocks = tuple( - block - for chunk in parsed_chunks - for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or []) + block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or []) ) expected_delta_blocks = ( {"type": "thinking", "thinking": "Step 1. "}, @@ -164,18 +150,12 @@ def test_streaming_thinking_blocks_are_replayable_after_signature_delta(): assert reasoning_content == "Step 1. Step 2." assert thinking_blocks == (*expected_delta_blocks, expected_thinking_block) - assert parsed_chunks[1].choices[0].delta.provider_specific_fields == { - "thinking_blocks": [expected_delta_blocks[0]] - } - assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == { - "thinking_blocks": [expected_thinking_block] - } + assert parsed_chunks[1].choices[0].delta.provider_specific_fields == {"thinking_blocks": [expected_delta_blocks[0]]} + assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == {"thinking_blocks": [expected_thinking_block]} def test_streaming_unsigned_thinking_deltas_keep_reasoning_content(): - model_response_iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) chunks = [ { "type": "content_block_start", @@ -195,17 +175,12 @@ def test_streaming_unsigned_thinking_deltas_keep_reasoning_content(): {"type": "content_block_stop", "index": 0}, ] - parsed_chunks = [ - model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks - ] + parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks] reasoning_content = "".join( - getattr(chunk.choices[0].delta, "reasoning_content", None) or "" - for chunk in parsed_chunks + getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks ) thinking_blocks = tuple( - block - for chunk in parsed_chunks - for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or []) + block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or []) ) assert reasoning_content == "Step 1. Step 2." @@ -216,9 +191,7 @@ def test_streaming_unsigned_thinking_deltas_keep_reasoning_content(): def test_streaming_truncated_thinking_deltas_keep_reasoning_content(): - model_response_iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) chunks = [ { "type": "content_block_start", @@ -237,17 +210,12 @@ def test_streaming_truncated_thinking_deltas_keep_reasoning_content(): }, ] - parsed_chunks = [ - model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks - ] + parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks] reasoning_content = "".join( - getattr(chunk.choices[0].delta, "reasoning_content", None) or "" - for chunk in parsed_chunks + getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks ) thinking_blocks = tuple( - block - for chunk in parsed_chunks - for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or []) + block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or []) ) assert reasoning_content == "Step 1. Step 2." @@ -258,9 +226,7 @@ def test_streaming_truncated_thinking_deltas_keep_reasoning_content(): def test_handle_json_mode_chunk_response_format_tool(): - model_response_iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=True - ) + model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True) response_format_tool = ChatCompletionToolCallChunk( id="tool_123", type="function", @@ -271,9 +237,7 @@ def test_handle_json_mode_chunk_response_format_tool(): index=0, ) - text, tool_use = model_response_iterator._handle_json_mode_chunk( - "", response_format_tool - ) + text, tool_use = model_response_iterator._handle_json_mode_chunk("", response_format_tool) print(f"\n\nresponse_format_tool text: {text}\n\n") print(f"\n\nresponse_format_tool tool_use: {tool_use}\n\n") @@ -282,15 +246,11 @@ def test_handle_json_mode_chunk_response_format_tool(): def test_handle_json_mode_chunk_regular_tool(): - model_response_iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=True - ) + model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True) regular_tool = ChatCompletionToolCallChunk( id="tool_456", type="function", - function=ChatCompletionToolCallFunctionChunk( - name="get_weather", arguments='{"location": "San Francisco, CA"}' - ), + function=ChatCompletionToolCallFunctionChunk(name="get_weather", arguments='{"location": "San Francisco, CA"}'), index=0, ) @@ -304,17 +264,13 @@ def test_handle_json_mode_chunk_regular_tool(): def test_handle_json_mode_chunk_streaming_response_format_tool(): - model_response_iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=True - ) + model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True) # First chunk: response_format tool with id and name, but no arguments first_chunk = ChatCompletionToolCallChunk( id="tool_123", type="function", - function=ChatCompletionToolCallFunctionChunk( - name=RESPONSE_FORMAT_TOOL_NAME, arguments="" - ), + function=ChatCompletionToolCallFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME, arguments=""), index=0, ) @@ -322,9 +278,7 @@ def test_handle_json_mode_chunk_streaming_response_format_tool(): second_chunk = ChatCompletionToolCallChunk( id=None, type="function", - function=ChatCompletionToolCallFunctionChunk( - name=None, arguments='{"question": "What is the weather?"' - ), + function=ChatCompletionToolCallFunctionChunk(name=None, arguments='{"question": "What is the weather?"'), index=0, ) @@ -332,9 +286,7 @@ def test_handle_json_mode_chunk_streaming_response_format_tool(): third_chunk = ChatCompletionToolCallChunk( id=None, type="function", - function=ChatCompletionToolCallFunctionChunk( - name=None, arguments=', "answer": "It is sunny"}' - ), + function=ChatCompletionToolCallFunctionChunk(name=None, arguments=', "answer": "It is sunny"}'), index=0, ) @@ -365,9 +317,7 @@ def test_handle_json_mode_chunk_streaming_response_format_tool(): def test_handle_json_mode_chunk_streaming_regular_tool(): - model_response_iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=True - ) + model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True) # First chunk: regular tool with id and name, but no arguments first_chunk = ChatCompletionToolCallChunk( @@ -381,9 +331,7 @@ def test_handle_json_mode_chunk_streaming_regular_tool(): second_chunk = ChatCompletionToolCallChunk( id=None, type="function", - function=ChatCompletionToolCallFunctionChunk( - name=None, arguments='{"location": "San Francisco, CA"}' - ), + function=ChatCompletionToolCallFunctionChunk(name=None, arguments='{"location": "San Francisco, CA"}'), index=0, ) @@ -408,27 +356,19 @@ def test_handle_json_mode_chunk_streaming_regular_tool(): def test_response_format_tool_finish_reason(): - model_response_iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=True - ) + model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True) # First chunk: response_format tool response_format_tool = ChatCompletionToolCallChunk( id="tool_123", type="function", - function=ChatCompletionToolCallFunctionChunk( - name=RESPONSE_FORMAT_TOOL_NAME, arguments='{"answer": "test"}' - ), + function=ChatCompletionToolCallFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME, arguments='{"answer": "test"}'), index=0, ) # Process the tool call (should set converted_response_format_tool flag) - text, tool_use = model_response_iterator._handle_json_mode_chunk( - "", response_format_tool - ) - print( - f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n" - ) + text, tool_use = model_response_iterator._handle_json_mode_chunk("", response_format_tool) + print(f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n") # Simulate message_delta chunk with tool_use stop_reason message_delta_chunk = { @@ -447,25 +387,19 @@ def test_response_format_tool_finish_reason(): def test_regular_tool_finish_reason(): - model_response_iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=True - ) + model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True) # First chunk: regular tool (not response_format) regular_tool = ChatCompletionToolCallChunk( id="tool_456", type="function", - function=ChatCompletionToolCallFunctionChunk( - name="get_weather", arguments='{"location": "San Francisco, CA"}' - ), + function=ChatCompletionToolCallFunctionChunk(name="get_weather", arguments='{"location": "San Francisco, CA"}'), index=0, ) # Process the tool call (should NOT set converted_response_format_tool flag) text, tool_use = model_response_iterator._handle_json_mode_chunk("", regular_tool) - print( - f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n" - ) + print(f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n") # Simulate message_delta chunk with tool_use stop_reason message_delta_chunk = { @@ -525,9 +459,7 @@ def test_text_only_streaming_has_index_zero(): for chunk in chunks: parsed = iterator.chunk_parser(chunk) if parsed.choices: - assert ( - parsed.choices[0].index == 0 - ), f"Expected index=0, got {parsed.choices[0].index}" + assert parsed.choices[0].index == 0, f"Expected index=0, got {parsed.choices[0].index}" def test_streaming_thinking_deltas_count_reasoning_tokens_in_usage(): @@ -704,9 +636,7 @@ def test_anthropic_completion_streaming_usage_matches_non_streaming_with_thinkin ] self._write_response( content_type="text/event-stream", - body="".join( - f"data: {json.dumps(event)}\n\n" for event in events - ).encode("utf-8"), + body="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode("utf-8"), ) return @@ -787,13 +717,9 @@ def test_anthropic_completion_streaming_usage_matches_non_streaming_with_thinkin assert content_chunks == [answer_text] assert stream_usage is not None stream_completion_details = stream_usage["completion_tokens_details"] - assert ( - stream_completion_details["reasoning_tokens"] - == non_stream_details.reasoning_tokens - ) + assert stream_completion_details["reasoning_tokens"] == non_stream_details.reasoning_tokens assert stream_completion_details["text_tokens"] == ( - stream_usage["completion_tokens"] - - stream_completion_details["reasoning_tokens"] + stream_usage["completion_tokens"] - stream_completion_details["reasoning_tokens"] ) assert requests_seen == [ { @@ -885,9 +811,9 @@ def test_text_and_tool_streaming_has_index_zero(): for chunk in chunks: parsed = iterator.chunk_parser(chunk) if parsed.choices: - assert ( - parsed.choices[0].index == 0 - ), f"Expected index=0 for chunk type {chunk.get('type')}, got {parsed.choices[0].index}" + assert parsed.choices[0].index == 0, ( + f"Expected index=0 for chunk type {chunk.get('type')}, got {parsed.choices[0].index}" + ) def test_multiple_tools_streaming_has_index_zero(): @@ -940,15 +866,11 @@ def test_multiple_tools_streaming_has_index_zero(): for chunk in chunks: parsed = iterator.chunk_parser(chunk) if parsed.choices: - assert ( - parsed.choices[0].index == 0 - ), f"Expected index=0, got {parsed.choices[0].index}" + assert parsed.choices[0].index == 0, f"Expected index=0, got {parsed.choices[0].index}" def test_streaming_chunks_have_stable_ids(): - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=False, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False) first_chunk = { "type": "content_block_delta", "index": 0, @@ -973,9 +895,7 @@ def test_partial_json_chunk_accumulation(): This tests the fix for https://github.com/BerriAI/litellm/issues/17473 where network fragmentation can cause SSE data to arrive in partial chunks. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) partial_chunk_1 = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hel' partial_chunk_2 = 'lo"}}' @@ -983,31 +903,21 @@ def test_partial_json_chunk_accumulation(): # First partial chunk should return None (still accumulating) result1 = iterator._parse_sse_data(f"data:{partial_chunk_1}") assert result1 is None, "First partial chunk should return None while accumulating" - assert ( - iterator.chunk_type == "accumulated_json" - ), "Should switch to accumulated_json mode" - assert ( - iterator.accumulated_json == partial_chunk_1 - ), "Should have accumulated first part" + assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode" + assert iterator.accumulated_json == partial_chunk_1, "Should have accumulated first part" # Second partial chunk should complete the JSON and return a parsed result result2 = iterator._parse_sse_data(f"data:{partial_chunk_2}") assert result2 is not None, "Second chunk should return parsed result" - assert ( - iterator.accumulated_json == "" - ), "Buffer should be cleared after successful parse" - assert ( - result2.choices[0].delta.content == "Hello" - ), f"Expected 'Hello', got '{result2.choices[0].delta.content}'" + assert iterator.accumulated_json == "", "Buffer should be cleared after successful parse" + assert result2.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result2.choices[0].delta.content}'" def test_complete_json_chunk_no_accumulation(): """ Test that complete JSON chunks are parsed immediately without accumulation. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) complete_chunk = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}' @@ -1015,18 +925,14 @@ def test_complete_json_chunk_no_accumulation(): assert result is not None, "Complete chunk should return parsed result immediately" assert iterator.chunk_type == "valid_json", "Should remain in valid_json mode" assert iterator.accumulated_json == "", "Buffer should remain empty" - assert ( - result.choices[0].delta.content == "Hello" - ), f"Expected 'Hello', got '{result.choices[0].delta.content}'" + assert result.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result.choices[0].delta.content}'" def test_multiple_partial_chunks_accumulation(): """ Test that multiple partial chunks can be accumulated across several iterations. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) # Split a JSON chunk into three parts part1 = '{"type":"content_block_del' @@ -1194,9 +1100,7 @@ def test_web_search_tool_result_no_extra_tool_calls(): The issue was that web_search_tool_result blocks have input_json_delta events with {} that were incorrectly being converted to tool calls. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) # Simulate the streaming sequence: # 1. server_tool_use block starts (web_search) @@ -1271,9 +1175,7 @@ def test_web_search_tool_result_no_extra_tool_calls(): # Should have exactly 2 tool calls: # 1. From content_block_start (server_tool_use) with id and name # 2. From content_block_delta with the actual query - assert ( - len(tool_calls_emitted) == 2 - ), f"Expected 2 tool calls, got {len(tool_calls_emitted)}" + assert len(tool_calls_emitted) == 2, f"Expected 2 tool calls, got {len(tool_calls_emitted)}" # First tool call should have the id and name assert tool_calls_emitted[0]["id"] == "srvtoolu_01ABC123" @@ -1289,9 +1191,7 @@ def test_current_content_block_type_tracking(): """ Test that current_content_block_type is properly tracked and reset. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) # Initially should be None assert iterator.current_content_block_type is None @@ -1344,9 +1244,7 @@ def test_web_search_tool_result_captured_in_provider_specific_fields(): The web_search_tool_result content comes ALL AT ONCE in content_block_start, not in deltas, so we need to capture it there. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) # Simulate the streaming sequence with web_search_tool_result chunks = [ @@ -1417,23 +1315,15 @@ def test_web_search_tool_result_captured_in_provider_specific_fields(): and parsed.choices[0].delta.provider_specific_fields and "web_search_results" in parsed.choices[0].delta.provider_specific_fields ): - web_search_results = parsed.choices[0].delta.provider_specific_fields[ - "web_search_results" - ] + web_search_results = parsed.choices[0].delta.provider_specific_fields["web_search_results"] # Verify web_search_results was captured assert web_search_results is not None, "web_search_results should be captured" assert len(web_search_results) == 1, "Should have 1 web_search_tool_result block" - assert ( - web_search_results[0]["type"] == "web_search_tool_result" - ), "Block type should be web_search_tool_result" - assert ( - web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123" - ), "tool_use_id should match" + assert web_search_results[0]["type"] == "web_search_tool_result", "Block type should be web_search_tool_result" + assert web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123", "tool_use_id should match" assert len(web_search_results[0]["content"]) == 2, "Should have 2 search results" - assert ( - web_search_results[0]["content"][0]["title"] == "Fun Otter Facts" - ), "First result title should match" + assert web_search_results[0]["content"][0]["title"] == "Fun Otter Facts", "First result title should match" def test_web_fetch_tool_result_captured_in_provider_specific_fields(): @@ -1447,9 +1337,7 @@ def test_web_fetch_tool_result_captured_in_provider_specific_fields(): The web_fetch_tool_result content comes ALL AT ONCE in content_block_start, not in deltas, so we need to capture it there. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) # Simulate the streaming sequence with web_fetch_tool_result chunks = [ @@ -1520,25 +1408,15 @@ def test_web_fetch_tool_result_captured_in_provider_specific_fields(): and parsed.choices[0].delta.provider_specific_fields and "web_search_results" in parsed.choices[0].delta.provider_specific_fields ): - web_search_results = parsed.choices[0].delta.provider_specific_fields[ - "web_search_results" - ] + web_search_results = parsed.choices[0].delta.provider_specific_fields["web_search_results"] # Verify web_fetch_tool_result was captured (stored in web_search_results list) assert web_search_results is not None, "web_search_results should be captured" assert len(web_search_results) == 1, "Should have 1 web_fetch_tool_result block" - assert ( - web_search_results[0]["type"] == "web_fetch_tool_result" - ), "Block type should be web_fetch_tool_result" - assert ( - web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123" - ), "tool_use_id should match" - assert ( - web_search_results[0]["content"]["url"] == "https://example.com" - ), "URL should match" - assert ( - web_search_results[0]["content"]["content"]["title"] == "Example Page" - ), "Title should match" + assert web_search_results[0]["type"] == "web_fetch_tool_result", "Block type should be web_fetch_tool_result" + assert web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123", "tool_use_id should match" + assert web_search_results[0]["content"]["url"] == "https://example.com", "URL should match" + assert web_search_results[0]["content"]["content"]["title"] == "Example Page", "Title should match" def test_web_fetch_tool_result_no_extra_tool_calls(): @@ -1551,9 +1429,7 @@ def test_web_fetch_tool_result_no_extra_tool_calls(): The issue was that web_fetch_tool_result blocks have input_json_delta events with {} that were incorrectly being converted to tool calls. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) # to verify it doesn't emit tool calls chunks = [ @@ -1597,9 +1473,9 @@ def test_web_fetch_tool_result_no_extra_tool_calls(): tool_call_count += 1 # Should have 0 tool calls - web_fetch_tool_result should not emit tool calls - assert ( - tool_call_count == 0 - ), f"Expected 0 tool calls, got {tool_call_count}. web_fetch_tool_result should not emit tool calls" + assert tool_call_count == 0, ( + f"Expected 0 tool calls, got {tool_call_count}. web_fetch_tool_result should not emit tool calls" + ) def test_container_in_provider_specific_fields_streaming(): @@ -1609,9 +1485,7 @@ def test_container_in_provider_specific_fields_streaming(): When container with skills is used, the container field should be present in the provider_specific_fields of the message_delta chunk. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) # Simulate streaming chunks chunks = [ @@ -1679,20 +1553,12 @@ def test_container_in_provider_specific_fields_streaming(): and parsed.choices[0].delta.provider_specific_fields and "container" in parsed.choices[0].delta.provider_specific_fields ): - container_field = parsed.choices[0].delta.provider_specific_fields[ - "container" - ] + container_field = parsed.choices[0].delta.provider_specific_fields["container"] # Verify container was captured - assert ( - container_field is not None - ), "container should be captured in provider_specific_fields" - assert ( - container_field["id"] == "container_011CW9hA9zpZ8xD3bjjShy4p" - ), "container id should match" - assert ( - container_field["expires_at"] == "2025-12-16T04:57:16.913181Z" - ), "expires_at should match" + assert container_field is not None, "container should be captured in provider_specific_fields" + assert container_field["id"] == "container_011CW9hA9zpZ8xD3bjjShy4p", "container id should match" + assert container_field["expires_at"] == "2025-12-16T04:57:16.913181Z", "expires_at should match" assert len(container_field["skills"]) == 1, "Should have 1 skill" assert container_field["skills"][0]["skill_id"] == "pptx", "skill_id should be pptx" assert container_field["skills"][0]["version"] == "20251013", "version should match" @@ -1705,9 +1571,7 @@ def test_container_in_provider_specific_fields_non_streaming(): When container with skills is used in non-streaming, the container field should be present in the provider_specific_fields of the response. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=False, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False) # Simulate a message_delta chunk with container (as it would appear in non-streaming) message_delta_chunk = { @@ -1743,21 +1607,13 @@ def test_container_in_provider_specific_fields_non_streaming(): # Verify container is in provider_specific_fields assert model_response.choices[0].delta.provider_specific_fields is not None assert "container" in model_response.choices[0].delta.provider_specific_fields - container_field = model_response.choices[0].delta.provider_specific_fields[ - "container" - ] + container_field = model_response.choices[0].delta.provider_specific_fields["container"] assert container_field["id"] == "container_abc123xyz", "container id should match" - assert ( - container_field["expires_at"] == "2025-12-20T10:30:00.000000Z" - ), "expires_at should match" + assert container_field["expires_at"] == "2025-12-20T10:30:00.000000Z", "expires_at should match" assert len(container_field["skills"]) == 2, "Should have 2 skills" - assert ( - container_field["skills"][0]["skill_id"] == "code_execution" - ), "First skill_id should be code_execution" - assert ( - container_field["skills"][1]["skill_id"] == "pptx" - ), "Second skill_id should be pptx" + assert container_field["skills"][0]["skill_id"] == "code_execution", "First skill_id should be code_execution" + assert container_field["skills"][1]["skill_id"] == "pptx", "Second skill_id should be pptx" def test_container_absent_when_not_provided(): @@ -1766,9 +1622,7 @@ def test_container_absent_when_not_provided(): This ensures we don't add empty or None container fields. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=False, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False) # message_delta without container message_delta_chunk = { @@ -1787,9 +1641,9 @@ def test_container_absent_when_not_provided(): # Verify container is NOT in provider_specific_fields when not provided if model_response.choices[0].delta.provider_specific_fields: - assert ( - "container" not in model_response.choices[0].delta.provider_specific_fields - ), "container should not be present when not provided in delta" + assert "container" not in model_response.choices[0].delta.provider_specific_fields, ( + "container should not be present when not provided in delta" + ) def test_streaming_code_execution_produces_code_interpreter_results(): @@ -1985,8 +1839,7 @@ def test_streaming_multiple_code_executions_no_duplicates(): # Second (final) emission: cumulative list with BOTH results # This is what stream_chunk_builder will pick as "last value wins" assert len(emissions[1]) == 2, ( - f"Expected final emission to have 2 results, got {len(emissions[1])}. " - f"IDs: {[r.id for r in emissions[1]]}" + f"Expected final emission to have 2 results, got {len(emissions[1])}. IDs: {[r.id for r in emissions[1]]}" ) assert emissions[1][0].id == "srvtoolu_01AAA" assert emissions[1][0].code == "echo first" @@ -2150,9 +2003,7 @@ def test_empty_output_produces_null_outputs(): assert code_results is not None, "No code_interpreter_results emitted" assert len(code_results) == 1 assert code_results[0].id == "srvtoolu_01AAA" - assert ( - code_results[0].outputs is None - ), f"Expected outputs=None for empty execution, got {code_results[0].outputs}" + assert code_results[0].outputs is None, f"Expected outputs=None for empty execution, got {code_results[0].outputs}" def test_non_bash_tool_result_skipped(): @@ -2215,12 +2066,10 @@ def test_non_bash_tool_result_skipped(): code_results = psf["code_interpreter_results"] # code_interpreter_results should be emitted but empty (no bash results) - assert ( - code_results is not None - ), "Expected code_interpreter_results key to be emitted" - assert ( - len(code_results) == 0 - ), f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}" + assert code_results is not None, "Expected code_interpreter_results key to be emitted" + assert len(code_results) == 0, ( + f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}" + ) class TestRustChatCompletionsHook: @@ -2257,13 +2106,9 @@ class TestRustChatCompletionsHook: from litellm.rust_bridge import chat_completions as bridge monkeypatch.delenv("LITELLM_RUST", raising=False) - bridge.set_rust_chat_completions( - chat_completions=None, achat_completions=None, decline=None - ) + bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) yield - bridge.set_rust_chat_completions( - chat_completions=None, achat_completions=None, decline=None - ) + bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) @staticmethod def _completion_kwargs(**overrides): @@ -2361,9 +2206,7 @@ class TestRustChatCompletionsHook: from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion seen = self._inject() - AnthropicChatCompletion().completion( - **self._completion_kwargs(optional_params={"max_tokens": 7}) - ) + AnthropicChatCompletion().completion(**self._completion_kwargs(optional_params={"max_tokens": 7})) assert seen["call"][0]["optional_params"]["max_tokens"] == 7 def test_without_the_opt_in_the_core_is_never_consulted(self): @@ -2371,15 +2214,14 @@ class TestRustChatCompletionsHook: from litellm.llms.anthropic.chat.transformation import AnthropicConfig seen = self._inject() - with patch.object( - AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} - ) as transform, patch.object( - AnthropicChatCompletion, "acompletion_function" + with ( + patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ) as transform, + patch.object(AnthropicChatCompletion, "acompletion_function"), ): try: - AnthropicChatCompletion().completion( - **self._completion_kwargs(litellm_params={}) - ) + AnthropicChatCompletion().completion(**self._completion_kwargs(litellm_params={})) except Exception: # The Python path goes on to make an HTTP call; reaching it is # the assertion, so the network failure below is expected. @@ -2393,9 +2235,7 @@ class TestRustChatCompletionsHook: from litellm.llms.anthropic.chat.transformation import AnthropicConfig seen = self._inject(decline_reason="unrecognized request parameter") - with patch.object( - AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} - ): + with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}): try: AnthropicChatCompletion().completion(**self._completion_kwargs()) except Exception: @@ -2408,9 +2248,7 @@ class TestRustChatCompletionsHook: from litellm.llms.anthropic.chat.transformation import AnthropicConfig seen = self._inject() - with patch.object( - AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} - ): + with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}): try: AnthropicChatCompletion().completion( **self._completion_kwargs(optional_params={"max_tokens": 16, "stream": True}) @@ -2424,9 +2262,7 @@ class TestRustChatCompletionsHook: seen = self._inject() logging_obj = MagicMock() - AnthropicChatCompletion().completion( - **self._completion_kwargs(logging_obj=logging_obj) - ) + AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj)) assert logging_obj.pre_call.call_count == 1 assert len(seen["call"]) == 1 @@ -2440,9 +2276,7 @@ class TestRustChatCompletionsHook: self._inject() logging_obj = MagicMock() - AnthropicChatCompletion().completion( - **self._completion_kwargs(logging_obj=logging_obj) - ) + AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj)) assert logging_obj.post_call.call_count == 1 logged = logging_obj.post_call.call_args.kwargs["original_response"] @@ -2467,18 +2301,12 @@ class TestRustChatCompletionsHook: raise _Declined("blank message text") monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) - bridge.set_rust_chat_completions( - decline=lambda **_kwargs: None, chat_completions=declining_native - ) + bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native) logging_obj, calls = self._recording_logging_obj() - with patch.object( - AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} - ): + with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}): try: - AnthropicChatCompletion().completion( - **self._completion_kwargs(logging_obj=logging_obj) - ) + AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj)) except Exception: # The Python path goes on to make an HTTP call; the log count is # the assertion, so a failure past this point is expected. @@ -2503,21 +2331,15 @@ class TestRustChatCompletionsHook: async def declining_native(**_kwargs): raise _Declined("blank message text") - bridge.set_rust_chat_completions( - decline=lambda **_kwargs: None, achat_completions=declining_native - ) + bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=declining_native) sentinel = object() async def python_path(**_kwargs): return sentinel - with patch.object( - AnthropicChatCompletion, "acompletion_function", side_effect=python_path - ) as python_call: - result = await AnthropicChatCompletion().completion( - **self._completion_kwargs(acompletion=True) - ) + with patch.object(AnthropicChatCompletion, "acompletion_function", side_effect=python_path) as python_call: + result = await AnthropicChatCompletion().completion(**self._completion_kwargs(acompletion=True)) assert result is sentinel assert python_call.called, "a failing rust call must re-enter the python path" @@ -2530,20 +2352,15 @@ class TestRustChatCompletionsHook: async def native(**_kwargs): return dict(self.RUST_RESPONSE) - bridge.set_rust_chat_completions( - decline=lambda **_kwargs: None, achat_completions=native - ) + bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=native) with patch.object(AnthropicChatCompletion, "acompletion_function") as python_call: - result = await AnthropicChatCompletion().completion( - **self._completion_kwargs(acompletion=True) - ) + result = await AnthropicChatCompletion().completion(**self._completion_kwargs(acompletion=True)) assert result.choices[0].message.content == "hello from rust" assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} assert not python_call.called - def test_pre_call_logging_fires_once_when_the_sync_rust_call_declines(self, monkeypatch): """One request, one pre_call, on the synchronous path too. Without the suppression the Python path logs a second time for the same attempt.""" @@ -2563,27 +2380,19 @@ class TestRustChatCompletionsHook: def declining_native(**_kwargs): raise _Declined("blank message text") - bridge.set_rust_chat_completions( - decline=lambda **_kwargs: None, chat_completions=declining_native - ) + bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native) logging_obj, calls = self._recording_logging_obj() - with patch.object( - AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} - ): + with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}): try: - AnthropicChatCompletion().completion( - **self._completion_kwargs(logging_obj=logging_obj) - ) + AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj)) except Exception: # The Python path goes on to make an HTTP call; the log count is # the assertion, so a failure past this point is expected. pass assert len(calls["pre_call"]) == 1 - assert calls["pre_call"][0]["additional_args"]["complete_input_dict"]["model"] == ( - "claude-sonnet-4-5" - ) + assert calls["pre_call"][0]["additional_args"]["complete_input_dict"]["model"] == ("claude-sonnet-4-5") def test_pre_call_logging_still_fires_when_rust_is_not_involved(self, monkeypatch): """The suppression must not swallow the log on the ordinary path.""" @@ -2592,9 +2401,7 @@ class TestRustChatCompletionsHook: self._inject() logging_obj, calls = self._recording_logging_obj() - with patch.object( - AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} - ): + with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}): try: AnthropicChatCompletion().completion( **self._completion_kwargs(litellm_params={}, logging_obj=logging_obj) diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 8ea8db5fb65..8031e2b48d6 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1,4 +1,3 @@ - import pytest from unittest.mock import MagicMock, patch @@ -33,13 +32,9 @@ def test_response_format_transformation_unit_test(): "additionalProperties": False, } - result = config._create_json_tool_call_for_response_format( - json_schema=response_format_json_schema - ) + result = config._create_json_tool_call_for_response_format(json_schema=response_format_json_schema) - assert result["input_schema"]["properties"] == { - "agent_doing": {"title": "Agent Doing", "type": "string"} - } + assert result["input_schema"]["properties"] == {"agent_doing": {"title": "Agent Doing", "type": "string"}} print(result) @@ -550,9 +545,7 @@ def test_extract_response_content_with_citations(): }, } - _, citations, _, _, _, _, _, _ = config.extract_response_content( - completion_response - ) + _, citations, _, _, _, _, _, _ = config.extract_response_content(completion_response) assert citations == [ [ { @@ -625,12 +618,8 @@ def test_web_search_tool_transformation(): assert anthropic_web_search_tool["user_location"]["city"] == "San Francisco" -@pytest.mark.parametrize( - "search_context_size, expected_max_uses", [("low", 1), ("medium", 5), ("high", 10)] -) -def test_web_search_tool_transformation_with_search_context_size( - search_context_size, expected_max_uses -): +@pytest.mark.parametrize("search_context_size, expected_max_uses", [("low", 1), ("medium", 5), ("high", 10)]) +def test_web_search_tool_transformation_with_search_context_size(search_context_size, expected_max_uses): from litellm.types.llms.openai import OpenAIWebSearchOptions config = AnthropicConfig() @@ -805,10 +794,7 @@ def test_web_search_tool_result_in_provider_specific_fields(): assert "web_search_results" in provider_fields assert len(provider_fields["web_search_results"]) == 1 assert provider_fields["web_search_results"][0]["type"] == "web_search_tool_result" - assert ( - provider_fields["web_search_results"][0]["tool_use_id"] - == "srvtoolu_provider_test" - ) + assert provider_fields["web_search_results"][0]["tool_use_id"] == "srvtoolu_provider_test" def test_multiple_web_search_tool_results(): @@ -1032,10 +1018,7 @@ def test_transform_response_with_prefix_prompt(): ) assert result is not None - assert ( - result.choices[0].message.content - == "You are a helpful assistant. The grass is green." - ) + assert result.choices[0].message.content == "You are a helpful assistant. The grass is green." def test_get_supported_params_thinking(): @@ -1150,18 +1133,12 @@ def test_anthropic_beta_header_merging_with_output_format(): } } - result_headers = config.update_headers_with_optional_anthropic_beta( - headers, optional_params - ) + result_headers = config.update_headers_with_optional_anthropic_beta(headers, optional_params) # Both beta headers should be present beta_value = result_headers["anthropic-beta"] - assert ( - "context-1m-2025-08-07" in beta_value - ), f"User's context-1m beta header missing from: {beta_value}" - assert ( - "structured-outputs-2025-11-13" in beta_value - ), f"Structured output beta header missing from: {beta_value}" + assert "context-1m-2025-08-07" in beta_value, f"User's context-1m beta header missing from: {beta_value}" + assert "structured-outputs-2025-11-13" in beta_value, f"Structured output beta header missing from: {beta_value}" def test_anthropic_beta_header_merging_with_multiple_features(): @@ -1183,9 +1160,7 @@ def test_anthropic_beta_header_merging_with_multiple_features(): "tools": [{"type": "web_fetch_20250910", "name": "web_fetch"}], } - result_headers = config.update_headers_with_optional_anthropic_beta( - headers, optional_params - ) + result_headers = config.update_headers_with_optional_anthropic_beta(headers, optional_params) beta_value = result_headers["anthropic-beta"] @@ -1228,9 +1203,7 @@ def test_anthropic_structured_output_beta_header(): "strict": True, "schema": { "description": 'Progress report for the thinking process\n\nThis model represents a snapshot of the agent\'s current progress during\nthe thinking process, providing a brief description of the current activity.\n\nAttributes:\n agent_doing: Brief description of what the agent is currently doing.\n Should be kept under 10 words. Example: "Learning about home automation"', - "properties": { - "agent_doing": {"title": "Agent Doing", "type": "string"} - }, + "properties": {"agent_doing": {"title": "Agent Doing", "type": "string"}}, "required": ["agent_doing"], "title": "ThinkingStep", "type": "object", @@ -1244,10 +1217,7 @@ def test_anthropic_structured_output_beta_header(): assert response is not None print(f"response: {response}") print(f"raw_request_headers: {response['raw_request_headers']}") - assert ( - "structured-outputs-2025-11-13" - in response["raw_request_headers"]["anthropic-beta"] - ) + assert "structured-outputs-2025-11-13" in response["raw_request_headers"]["anthropic-beta"] @pytest.mark.parametrize( @@ -1383,9 +1353,7 @@ def test_tool_search_regex_detection(): config = AnthropicModelInfo() # Test with tool search regex tool - tools = [ - {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"} - ] + tools = [{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}] assert config.is_tool_search_used(tools) is True # Test without tool search @@ -1400,9 +1368,7 @@ def test_tool_search_bm25_detection(): config = AnthropicModelInfo() # Test with tool search BM25 tool - tools = [ - {"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"} - ] + tools = [{"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"}] assert config.is_tool_search_used(tools) is True @@ -1594,9 +1560,7 @@ def test_tool_search_complete_response_parsing(): "tool_use_id": "srvtoolu_015i6aVA2niwzv4RG4DtnxDJ", "content": { "type": "tool_search_tool_search_result", - "tool_references": [ - {"type": "tool_reference", "tool_name": "get_weather"} - ], + "tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}], }, }, {"type": "text", "text": "Great! I found a weather tool."}, @@ -1647,9 +1611,7 @@ def test_tool_search_complete_response_parsing(): assert usage.server_tool_use is not None assert usage.server_tool_use.web_search_requests == 0 - assert ( - usage.server_tool_use.tool_search_requests == 1 - ) # Counted from server_tool_use blocks + assert usage.server_tool_use.tool_search_requests == 1 # Counted from server_tool_use blocks def test_allowed_callers_field_preservation(): @@ -1701,9 +1663,7 @@ def test_programmatic_tool_calling_beta_header(): assert is_programmatic is True # Test header generation - headers = model_info.get_anthropic_headers( - api_key="test-key", programmatic_tool_calling_used=True - ) + headers = model_info.get_anthropic_headers(api_key="test-key", programmatic_tool_calling_used=True) assert "anthropic-beta" in headers assert "advanced-tool-use-2025-11-20" in headers["anthropic-beta"] @@ -1847,9 +1807,7 @@ def test_input_examples_beta_header(): assert is_examples_used is True # Test header generation - headers = model_info.get_anthropic_headers( - api_key="test-key", input_examples_used=True - ) + headers = model_info.get_anthropic_headers(api_key="test-key", input_examples_used=True) assert "anthropic-beta" in headers assert "advanced-tool-use-2025-11-20" in headers["anthropic-beta"] @@ -1935,10 +1893,7 @@ def test_input_examples_empty_list_not_added(): transformed_tool, _ = config._map_tool_helper(tool) assert transformed_tool is not None # Empty list should not be added - assert ( - "input_examples" not in transformed_tool - or len(transformed_tool.get("input_examples", [])) == 0 - ) + assert "input_examples" not in transformed_tool or len(transformed_tool.get("input_examples", [])) == 0 # ============ Effort Parameter Tests ============ @@ -1998,9 +1953,7 @@ def test_effort_beta_header_injection(): effort_used = model_info.is_effort_used(optional_params=optional_params, custom_llm_provider="anthropic") assert effort_used is True - headers = model_info.get_anthropic_headers( - api_key="test-key", effort_used=effort_used - ) + headers = model_info.get_anthropic_headers(api_key="test-key", effort_used=effort_used) assert "anthropic-beta" in headers assert "effort-2025-11-24" in headers["anthropic-beta"] @@ -2026,9 +1979,7 @@ def test_effort_validation(): optional_params = {"output_config": {"effort": "invalid"}} - with pytest.raises( - litellm.exceptions.BadRequestError, match="Invalid effort value" - ): + with pytest.raises(litellm.exceptions.BadRequestError, match="Invalid effort value"): config.transform_request( model="claude-opus-4-5-20251101", messages=messages, @@ -2264,16 +2215,8 @@ def test_anthropic_model_supports_speed_param_rejects_non_anthropic_providers( ): """Fast mode is direct-Anthropic-only. Vertex/Azure/Bedrock strip their prefix before the shared transform runs, so the bare Opus id must still be rejected.""" - assert ( - AnthropicConfig._model_supports_speed_param( - "claude-opus-4-8", custom_llm_provider - ) - is False - ) - assert ( - AnthropicConfig._model_supports_speed_param("claude-opus-4-8", "anthropic") - is True - ) + assert AnthropicConfig._model_supports_speed_param("claude-opus-4-8", custom_llm_provider) is False + assert AnthropicConfig._model_supports_speed_param("claude-opus-4-8", "anthropic") is True def test_vertex_anthropic_drops_speed_for_opus_with_drop_params(monkeypatch): @@ -2572,9 +2515,7 @@ def test_supports_effort_level_handles_provider_prefixes(model, level, expected) ("claude-opus-4-5-20251101", None, False), ], ) -def test_validate_effort_for_model_centralises_per_model_gating( - model, effort, expect_error -): +def test_validate_effort_for_model_centralises_per_model_gating(model, effort, expect_error): err = AnthropicConfig._validate_effort_for_model(model, effort, "anthropic") if expect_error: assert err is not None @@ -2623,11 +2564,7 @@ def test_transform_request_injects_dummy_tool_without_tools_param(): litellm.modify_params = prev_modify_params assert "tools" in result - names = [ - t.get("name") - for t in result["tools"] - if isinstance(t, dict) and t.get("name") is not None - ] + names = [t.get("name") for t in result["tools"] if isinstance(t, dict) and t.get("name") is not None] assert "dummy_tool" in names @@ -2716,13 +2653,9 @@ def test_calculate_usage_completion_tokens_details_with_reasoning(): "output_tokens": 500, } # Simulating reasoning content that would count as ~50 tokens - reasoning_content = ( - "Let me think about this step by step. " * 10 - ) # Roughly 50 tokens + reasoning_content = "Let me think about this step by step. " * 10 # Roughly 50 tokens - usage = config.calculate_usage( - usage_object=usage_object, reasoning_content=reasoning_content - ) + usage = config.calculate_usage(usage_object=usage_object, reasoning_content=reasoning_content) # completion_tokens_details should be populated with both reasoning and text tokens assert usage.completion_tokens_details is not None @@ -2773,9 +2706,7 @@ def test_reasoning_effort_maps_to_adaptive_thinking_for_claude_4_6_models(): # reasoning_effort should not be in the result (it's transformed to thinking) assert "reasoning_effort" not in result # Should set output_config with the mapped effort value - assert ( - "output_config" in result - ), f"output_config missing for {model} with effort={effort}" + assert "output_config" in result, f"output_config missing for {model} with effort={effort}" assert result["output_config"]["effort"] == effort_map[effort] @@ -2851,7 +2782,6 @@ def test_raw_adaptive_thinking_untouched_for_46_plus_model(): assert result["thinking"] == {"type": "adaptive"} - @pytest.mark.parametrize( "model, expected", [ @@ -2877,9 +2807,7 @@ def test_raw_adaptive_thinking_untouched_for_46_plus_model(): ("gpt-4o", False), ], ) -def test_is_adaptive_thinking_model_is_sourced_from_cost_map( - local_model_cost_map, model, expected -): +def test_is_adaptive_thinking_model_is_sourced_from_cost_map(local_model_cost_map, model, expected): """Adaptive thinking resolves from the cost map first (an explicit supports_adaptive_thinking entry, or the anthropic-claude fallback rule for unmapped future Claudes), then from a date-safe opus/sonnet/haiku >= 4.6 name version as a @@ -2995,9 +2923,7 @@ def test_reasoning_effort_sets_output_config_for_46_models(): drop_params=False, ) - assert ( - "output_config" in result - ), f"output_config missing for {model} with effort={effort}" + assert "output_config" in result, f"output_config missing for {model} with effort={effort}" assert result["output_config"]["effort"] == effort @@ -3036,9 +2962,7 @@ def test_reasoning_effort_does_not_set_output_config_for_older_models(): drop_params=False, ) - assert ( - "output_config" not in result - ), f"output_config should not be set for {model}" + assert "output_config" not in result, f"output_config should not be set for {model}" @pytest.mark.parametrize( @@ -3078,14 +3002,10 @@ def test_reasoning_effort_accepts_dict_shape_for_adaptive_model(reasoning_effort ) # thinking must be set (adaptive for 4.6+) - assert ( - "thinking" in result - ), f"thinking missing for reasoning_effort={reasoning_effort_value!r}" + assert "thinking" in result, f"thinking missing for reasoning_effort={reasoning_effort_value!r}" assert result["thinking"]["type"] == "adaptive" # output_config must carry the mapped effort - assert ( - "output_config" in result - ), f"output_config missing for reasoning_effort={reasoning_effort_value!r}" + assert "output_config" in result, f"output_config missing for reasoning_effort={reasoning_effort_value!r}" assert result["output_config"]["effort"] == "low" @@ -3114,16 +3034,13 @@ def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model( drop_params=False, ) - assert ( - "thinking" in result - ), f"thinking missing for reasoning_effort={reasoning_effort_value!r}" + assert "thinking" in result, f"thinking missing for reasoning_effort={reasoning_effort_value!r}" assert result["thinking"]["type"] == "enabled" assert "budget_tokens" in result["thinking"] assert result["thinking"]["budget_tokens"] > 0 # Older models must not get adaptive-thinking output_config assert "output_config" not in result, ( - f"output_config should not be set for non-adaptive model " - f"(reasoning_effort={reasoning_effort_value!r})" + f"output_config should not be set for non-adaptive model (reasoning_effort={reasoning_effort_value!r})" ) @@ -3174,12 +3091,8 @@ def test_reasoning_effort_unparseable_dict_is_dropped(bad_value): model="claude-sonnet-4-6-20260219", drop_params=False, ) - assert ( - "thinking" not in result - ), f"thinking should not be set for bad value {bad_value!r}" - assert ( - "output_config" not in result - ), f"output_config should not be set for bad value {bad_value!r}" + assert "thinking" not in result, f"thinking should not be set for bad value {bad_value!r}" + assert "output_config" not in result, f"output_config should not be set for bad value {bad_value!r}" @pytest.mark.parametrize( @@ -3310,9 +3223,7 @@ def test_reasoning_effort_garbage_raises_bad_request(effort): ("max", DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET), ], ) -def test_reasoning_effort_xhigh_max_maps_to_budget_on_budget_model( - effort, expected_budget -): +def test_reasoning_effort_xhigh_max_maps_to_budget_on_budget_model(effort, expected_budget): """``xhigh`` / ``max`` extend the budget_tokens progression on budget-mode models.""" config = AnthropicConfig() @@ -3459,17 +3370,11 @@ def test_code_execution_tool_results_extraction(): # Verify first tool call assert transformed_response.choices[0].message.tool_calls[0].id == "srvtoolu_01ABC" - assert ( - transformed_response.choices[0].message.tool_calls[0].function.name - == "bash_code_execution" - ) + assert transformed_response.choices[0].message.tool_calls[0].function.name == "bash_code_execution" # Verify second tool call assert transformed_response.choices[0].message.tool_calls[1].id == "srvtoolu_01DEF" - assert ( - transformed_response.choices[0].message.tool_calls[1].function.name - == "text_editor_code_execution" - ) + assert transformed_response.choices[0].message.tool_calls[1].function.name == "text_editor_code_execution" # Verify tool results are in provider_specific_fields provider_fields = transformed_response.choices[0].message.provider_specific_fields @@ -3492,10 +3397,7 @@ def test_code_execution_tool_results_extraction(): assert editor_result["content"]["is_file_update"] is False # Verify text content is properly concatenated - assert ( - "I'll calculate that for you." - in transformed_response.choices[0].message.content - ) + assert "I'll calculate that for you." in transformed_response.choices[0].message.content assert "Done!" in transformed_response.choices[0].message.content @@ -3563,10 +3465,7 @@ def test_code_execution_tool_results_in_hidden_params(): assert "provider_specific_fields" in hidden assert "tool_results" in hidden["provider_specific_fields"] assert len(hidden["provider_specific_fields"]["tool_results"]) == 1 - assert ( - hidden["provider_specific_fields"]["tool_results"][0]["content"]["stdout"] - == "hello\n" - ) + assert hidden["provider_specific_fields"]["tool_results"][0]["content"]["stdout"] == "hello\n" def test_tool_search_tool_result_not_in_tool_results(): @@ -3762,10 +3661,7 @@ def test_compaction_block_in_provider_specific_fields(): assert "compaction_blocks" in provider_fields assert len(provider_fields["compaction_blocks"]) == 1 assert provider_fields["compaction_blocks"][0]["type"] == "compaction" - assert ( - "Summary of the conversation" - in provider_fields["compaction_blocks"][0]["content"] - ) + assert "Summary of the conversation" in provider_fields["compaction_blocks"][0]["content"] def test_multiple_compaction_blocks(): @@ -3813,9 +3709,7 @@ def test_compaction_block_request_transformation(): {"role": "user", "content": "What is the weather in San Francisco?"}, { "role": "assistant", - "content": [ - {"type": "text", "text": "I don't have access to real-time data."} - ], + "content": [{"type": "text", "text": "I don't have access to real-time data."}], "provider_specific_fields": { "compaction_blocks": [ { @@ -3828,9 +3722,7 @@ def test_compaction_block_request_transformation(): {"role": "user", "content": "What about New York?"}, ] - result = anthropic_messages_pt( - messages=messages, model="claude-opus-4-6", llm_provider="anthropic" - ) + result = anthropic_messages_pt(messages=messages, model="claude-opus-4-6", llm_provider="anthropic") # Find the assistant message assistant_message = None @@ -3944,9 +3836,7 @@ def test_map_openai_context_management_to_anthropic(): "instructions": "Focus on preserving code snippets", } ] - result = config.map_openai_context_management_to_anthropic( - openai_format_with_instructions - ) + result = config.map_openai_context_management_to_anthropic(openai_format_with_instructions) assert result is not None assert result["edits"][0]["trigger"]["value"] == 150000 @@ -3973,9 +3863,7 @@ def test_map_openai_params_with_context_management(): config = AnthropicConfig() # Test with OpenAI list format - non_default_params = { - "context_management": [{"type": "compaction", "compact_threshold": 200000}] - } + non_default_params = {"context_management": [{"type": "compaction", "compact_threshold": 200000}]} optional_params = {} result = config.map_openai_params( @@ -4012,10 +3900,7 @@ def test_map_openai_params_with_context_management(): ) assert "context_management" in result - assert ( - result["context_management"] - == non_default_params_anthropic["context_management"] - ) + assert result["context_management"] == non_default_params_anthropic["context_management"] def test_cache_control_in_supported_params(): @@ -4126,10 +4011,7 @@ def test_compaction_block_empty_list_not_added(): # Verify compaction_blocks is not in provider_specific_fields when there are none provider_fields = result.choices[0].message.provider_specific_fields if provider_fields: - assert ( - "compaction_blocks" not in provider_fields - or provider_fields.get("compaction_blocks") is None - ) + assert "compaction_blocks" not in provider_fields or provider_fields.get("compaction_blocks") is None def test_fast_mode_beta_header(): @@ -4178,9 +4060,7 @@ def test_fast_mode_usage_calculation(): "output_tokens": 500, } - usage = config.calculate_usage( - usage_object=usage_object, reasoning_content=None, speed="fast" - ) + usage = config.calculate_usage(usage_object=usage_object, reasoning_content=None, speed="fast") assert usage.prompt_tokens == 1000 assert usage.completion_tokens == 500 @@ -4201,9 +4081,7 @@ def test_fast_mode_cost_calculation(): base_completion = 0.025 with ( - patch( - "litellm.llms.anthropic.cost_calculation.generic_cost_per_token" - ) as mock_cost, + patch("litellm.llms.anthropic.cost_calculation.generic_cost_per_token") as mock_cost, patch("litellm.get_model_info") as mock_info, ): mock_cost.return_value = (base_prompt, base_completion) @@ -4243,9 +4121,7 @@ def test_fast_mode_with_inference_geo(): base_completion = 0.025 with ( - patch( - "litellm.llms.anthropic.cost_calculation.generic_cost_per_token" - ) as mock_cost, + patch("litellm.llms.anthropic.cost_calculation.generic_cost_per_token") as mock_cost, patch("litellm.get_model_info") as mock_info, ): mock_cost.return_value = (base_prompt, base_completion) @@ -4436,9 +4312,7 @@ def test_map_tool_helper_enforces_object_type_when_missing(): "name": "search_code", "description": "Search for code patterns", "parameters": { - "properties": { - "query": {"type": "string", "description": "Search query"} - }, + "properties": {"query": {"type": "string", "description": "Search query"}}, "required": ["query"], }, }, @@ -4451,9 +4325,9 @@ def test_map_tool_helper_enforces_object_type_when_missing(): assert "properties" in result["input_schema"] assert "query" in result["input_schema"]["properties"] # Original parameters dict must not be modified in place - assert ( - tool["function"]["parameters"] == original_params - ), "parameters dict was mutated; _map_tool_helper should not modify caller data" + assert tool["function"]["parameters"] == original_params, ( + "parameters dict was mutated; _map_tool_helper should not modify caller data" + ) def test_map_tool_helper_enforces_object_type_when_wrong_type(): @@ -4479,13 +4353,13 @@ def test_map_tool_helper_enforces_object_type_when_wrong_type(): result, _ = config._map_tool_helper(tool) assert result is not None assert result["input_schema"]["type"] == "object" - assert ( - result["input_schema"].get("properties") == {} - ), "properties should be injected as {} when schema has non-object type and no properties key" + assert result["input_schema"].get("properties") == {}, ( + "properties should be injected as {} when schema has non-object type and no properties key" + ) # Original parameters dict must not be modified in place - assert ( - tool["function"]["parameters"] == original_params - ), "parameters dict was mutated; _map_tool_helper should not modify caller data" + assert tool["function"]["parameters"] == original_params, ( + "parameters dict was mutated; _map_tool_helper should not modify caller data" + ) def test_map_tool_helper_preserves_valid_object_schema(): @@ -4552,12 +4426,8 @@ def test_extract_response_content_thinking_block_null_thinking(): {"type": "text", "text": "Hello"}, ] } - text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content( - completion_response_null - ) - assert ( - thinking_blocks is not None - ), "thinking blocks should not be None when thinking=null" + text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(completion_response_null) + assert thinking_blocks is not None, "thinking blocks should not be None when thinking=null" assert len(thinking_blocks) == 1 assert "Hello" in text @@ -4568,12 +4438,8 @@ def test_extract_response_content_thinking_block_null_thinking(): {"type": "text", "text": "World"}, ] } - text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content( - completion_response_missing - ) - assert ( - thinking_blocks is not None - ), "thinking blocks should not be None when thinking key is absent" + text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(completion_response_missing) + assert thinking_blocks is not None, "thinking blocks should not be None when thinking key is absent" assert len(thinking_blocks) == 1 assert "World" in text @@ -4584,9 +4450,7 @@ def test_extract_response_content_thinking_block_null_thinking(): {"type": "text", "text": "Done"}, ] } - text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content( - completion_response_text - ) + text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(completion_response_text) assert thinking_blocks is not None assert len(thinking_blocks) == 1 assert thinking_blocks[0]["thinking"] == "Let me think..." @@ -4645,12 +4509,8 @@ def test_advisor_beta_header_injected(): } ] } - result = config.update_headers_with_optional_anthropic_beta( - headers, optional_params - ) - assert ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value in result.get( - "anthropic-beta", "" - ) + result = config.update_headers_with_optional_anthropic_beta(headers, optional_params) + assert ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value in result.get("anthropic-beta", "") def test_advisor_beta_header_not_injected_without_tool(): @@ -4658,9 +4518,7 @@ def test_advisor_beta_header_not_injected_without_tool(): config = AnthropicConfig() headers: dict = {} optional_params: dict = {"tools": []} - result = config.update_headers_with_optional_anthropic_beta( - headers, optional_params - ) + result = config.update_headers_with_optional_anthropic_beta(headers, optional_params) assert "advisor-tool-2026-03-01" not in result.get("anthropic-beta", "") @@ -4687,9 +4545,7 @@ def test_advisor_tool_result_preserved_in_response(): {"type": "text", "text": "Here is the implementation."}, ] } - text, _, _, _, tool_calls, _, tool_results, _ = config.extract_response_content( - completion_response - ) + text, _, _, _, tool_calls, _, tool_results, _ = config.extract_response_content(completion_response) assert "Consulting advisor." in text assert "Here is the implementation." in text # server_tool_use (advisor) should be a tool_call @@ -4804,9 +4660,7 @@ def test_basic_sanitize_anthropic_tool_name_replaces_invalid_chars(): ) assert ( - _basic_sanitize_anthropic_tool_name( - "github_openapi_mcp-actions/download-job-logs-for-workflow-run" - ) + _basic_sanitize_anthropic_tool_name("github_openapi_mcp-actions/download-job-logs-for-workflow-run") == "github_openapi_mcp-actions_download-job-logs-for-workflow-run" ) # other punctuation @@ -4835,9 +4689,7 @@ def test_build_anthropic_tool_name_maps_no_collisions(): ] ) assert forward == { - "actions/download-job-logs-for-workflow-run": ( - "actions_download-job-logs-for-workflow-run" - ), + "actions/download-job-logs-for-workflow-run": ("actions_download-job-logs-for-workflow-run"), "pulls/list-files": "pulls_list-files", } assert reverse == {v: k for k, v in forward.items()} @@ -4888,9 +4740,7 @@ def test_build_anthropic_tool_name_maps_three_way_collision(): _build_anthropic_tool_name_maps, ) - forward, reverse = _build_anthropic_tool_name_maps( - ["foo_bar", "foo/bar", "foo.bar"] - ) + forward, reverse = _build_anthropic_tool_name_maps(["foo_bar", "foo/bar", "foo.bar"]) assert "foo_bar" not in forward # untouched assert forward["foo/bar"] == "foo_bar_2" assert forward["foo.bar"] == "foo_bar_3" @@ -4963,16 +4813,13 @@ def test_map_openai_params_does_not_pollute_optional_params_with_internal_keys() ) # No internal keys may appear in optional_params for ANY input. for key in optional_params: - assert not key.startswith( - "_anthropic_tool_name" - ), f"optional_params leaked internal key {key!r}: {optional_params}" + assert not key.startswith("_anthropic_tool_name"), ( + f"optional_params leaked internal key {key!r}: {optional_params}" + ) # And no key starting with `_` either; optional_params should only # contain documented Anthropic Messages API parameters. for key in optional_params: - assert not key.startswith("_"), ( - f"optional_params leaked underscore-prefixed key {key!r}: " - f"{optional_params}" - ) + assert not key.startswith("_"), f"optional_params leaked underscore-prefixed key {key!r}: {optional_params}" def test_map_openai_params_no_maps_when_all_names_already_valid(): @@ -5001,11 +4848,7 @@ def test_map_openai_params_no_maps_when_all_names_already_valid(): def test_rewrite_tool_names_in_messages_uses_forward_map(): config = AnthropicConfig() - forward_map = { - "actions/download-job-logs-for-workflow-run": ( - "actions_download-job-logs-for-workflow-run" - ) - } + forward_map = {"actions/download-job-logs-for-workflow-run": ("actions_download-job-logs-for-workflow-run")} messages = [ {"role": "user", "content": "go"}, { @@ -5028,15 +4871,9 @@ def test_rewrite_tool_names_in_messages_uses_forward_map(): out = config._rewrite_tool_names_in_messages(messages, forward_map) # input list must not be mutated - assert ( - messages[1]["tool_calls"][0]["function"]["name"] - == "actions/download-job-logs-for-workflow-run" - ) + assert messages[1]["tool_calls"][0]["function"]["name"] == "actions/download-job-logs-for-workflow-run" # output rewritten according to forward map - assert ( - out[1]["tool_calls"][0]["function"]["name"] - == "actions_download-job-logs-for-workflow-run" - ) + assert out[1]["tool_calls"][0]["function"]["name"] == "actions_download-job-logs-for-workflow-run" # non-tool-call messages pass through unchanged (same object) assert out[0] is messages[0] assert out[2] is messages[2] @@ -5112,9 +4949,7 @@ def test_sanitize_tool_names_in_request_does_not_mutate_caller_tool_dicts(): caller_tools = [caller_tool] optional_params: dict = {"tools": caller_tools} - forward, reverse = config._sanitize_tool_names_in_request( - optional_params=optional_params - ) + forward, reverse = config._sanitize_tool_names_in_request(optional_params=optional_params) assert forward.get(original_name) sanitized = forward[original_name] @@ -5263,10 +5098,7 @@ def test_streaming_iterator_reverse_maps_tool_use_name(): parsed = iterator.chunk_parser(chunk=chunk) tool_calls = parsed.choices[0].delta.tool_calls assert tool_calls is not None and len(tool_calls) == 1 - assert ( - tool_calls[0]["function"]["name"] - == "actions/download-job-logs-for-workflow-run" - ) + assert tool_calls[0]["function"]["name"] == "actions/download-job-logs-for-workflow-run" def test_streaming_iterator_passthrough_when_name_not_in_map(): @@ -5362,9 +5194,9 @@ def test_transform_request_does_not_leak_internal_keys_into_body(): for tool in data.get("tools", []): name = tool.get("name") assert isinstance(name, str) - assert _re.fullmatch( - r"[a-zA-Z0-9_-]{1,128}", name - ), f"sanitized tool name {name!r} still violates Anthropic regex" + assert _re.fullmatch(r"[a-zA-Z0-9_-]{1,128}", name), ( + f"sanitized tool name {name!r} still violates Anthropic regex" + ) # Sent name for the bad tool is the disambiguated form, valid name passes through. sent_names = {t["name"] for t in data["tools"]} @@ -5500,9 +5332,7 @@ def test_transform_request_rewrites_tool_names_in_history(): for block in content: if isinstance(block, dict) and block.get("type") == "tool_use": tool_use_names.append(block.get("name")) - assert ( - tool_use_names - ), "expected at least one tool_use block in transformed messages" + assert tool_use_names, "expected at least one tool_use block in transformed messages" for name in tool_use_names: assert name == "actions_download-job-logs-for-workflow-run", ( f"history tool_use.name {name!r} not rewritten -- Anthropic will " @@ -5526,19 +5356,12 @@ def test_sanitize_tool_names_in_request_skips_hosted_tools(): } forward, reverse = AnthropicConfig._sanitize_tool_names_in_request(optional_params) # Only the custom tool was rewritten. - assert forward == { - "actions/download-job-logs-for-workflow-run": "actions_download-job-logs-for-workflow-run" - } - assert reverse == { - "actions_download-job-logs-for-workflow-run": "actions/download-job-logs-for-workflow-run" - } + assert forward == {"actions/download-job-logs-for-workflow-run": "actions_download-job-logs-for-workflow-run"} + assert reverse == {"actions_download-job-logs-for-workflow-run": "actions/download-job-logs-for-workflow-run"} # Hosted tool's name unchanged. assert optional_params["tools"][0]["name"] == "web_search" # Custom tool's name updated in place. - assert ( - optional_params["tools"][1]["name"] - == "actions_download-job-logs-for-workflow-run" - ) + assert optional_params["tools"][1]["name"] == "actions_download-job-logs-for-workflow-run" def test_sanitize_tool_names_in_request_no_tools_is_noop(): @@ -5772,9 +5595,7 @@ def test_translate_system_message_keeps_billing_header_for_first_party_anthropic assert config.should_strip_billing_metadata() is False result = config.translate_system_message( - messages=_system_with_billing_header( - "You are Claude Code, Anthropic's official CLI for Claude." - ) + messages=_system_with_billing_header("You are Claude Code, Anthropic's official CLI for Claude.") ) texts = [block["text"] for block in result] @@ -5790,9 +5611,7 @@ def test_translate_system_message_strips_billing_header_for_bedrock(): config = BedrockClaudePlatformConfig() assert config.should_strip_billing_metadata() is True - result = config.translate_system_message( - messages=_system_with_billing_header("real system prompt") - ) + result = config.translate_system_message(messages=_system_with_billing_header("real system prompt")) texts = [block["text"] for block in result] assert all(not t.startswith("x-anthropic-billing-header:") for t in texts) @@ -5858,9 +5677,7 @@ def test_translate_system_message_strips_billing_header_for_bedrock_invoke(): config = AmazonAnthropicClaudeConfig() assert config.should_strip_billing_metadata() is True - result = config.translate_system_message( - messages=_system_with_billing_header("real system prompt") - ) + result = config.translate_system_message(messages=_system_with_billing_header("real system prompt")) texts = [block["text"] for block in result] assert all(not t.startswith("x-anthropic-billing-header:") for t in texts) @@ -5914,9 +5731,7 @@ def test_translate_system_message_strips_billing_header_for_bedrock_invoke(): ), ], ) -def test_should_strip_billing_metadata_by_provider( - module_path, class_name, expected_strip -): +def test_should_strip_billing_metadata_by_provider(module_path, class_name, expected_strip): import importlib config_cls = getattr(importlib.import_module(module_path), class_name) @@ -6088,12 +5903,8 @@ def test_sampling_param_gating_driven_by_model_map_flag(monkeypatch): """The drop/raise decision must come from ``supports_sampling_params`` in the model map, not just name matching: a flagged entry gates a model whose name says nothing, and an explicit ``true`` overrides the name fallback.""" - monkeypatch.setitem( - litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False} - ) - monkeypatch.setitem( - litellm.model_cost, "claude-fable-5-test", {"supports_sampling_params": True} - ) + monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False}) + monkeypatch.setitem(litellm.model_cost, "claude-fable-5-test", {"supports_sampling_params": True}) config = AnthropicConfig() flagged_off = config.map_openai_params( @@ -6213,9 +6024,7 @@ def test_is_anthropic_usage_object_rejects_responses_api_usage(): ("claude-sonnet-4-5-20250929", False), ], ) -def test_disabled_thinking_omitted_only_for_always_on_models( - local_model_cost_map, model, expected_dropped -): +def test_disabled_thinking_omitted_only_for_always_on_models(local_model_cost_map, model, expected_dropped): """``thinking={"type": "disabled"}`` is omitted for always-on-thinking models (Fable/Mythos, which 400 on it: the API remedy is to omit the param) and is forwarded verbatim for every model that accepts it.""" diff --git a/tests/test_litellm/llms/anthropic/chat/test_code_interpreter_results_extraction.py b/tests/test_litellm/llms/anthropic/chat/test_code_interpreter_results_extraction.py index 60e45c9b8ce..99266f92c56 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_code_interpreter_results_extraction.py +++ b/tests/test_litellm/llms/anthropic/chat/test_code_interpreter_results_extraction.py @@ -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 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index ea3b19fba2b..92173ca1971 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -6,7 +6,6 @@ import pytest import litellm - from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_PLACEHOLDER, ) @@ -53,9 +52,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block(): tool_calls=[ ChatCompletionDeltaToolCall( id="call_d581d130-e234-4315-94e8-27e7ff7c4e55", - function=Function( - arguments='{"location": "Boston"}', name="get_weather" - ), + function=Function(arguments='{"location": "Boston"}', name="get_weather"), type="function", index=0, ) @@ -69,9 +66,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block(): ( block_type, content_block_start, - ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block( - choices=choices - ) + ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices) print(content_block_start) @@ -101,9 +96,7 @@ def test_translate_streaming_openai_chunk_strips_gemini_thought_from_tool_call_i tool_calls=[ ChatCompletionDeltaToolCall( id=combined, - function=Function( - arguments='{"a": 17, "b": 25}', name="add_numbers" - ), + function=Function(arguments='{"a": 17, "b": 25}', name="add_numbers"), type="function", index=0, ) @@ -117,9 +110,7 @@ def test_translate_streaming_openai_chunk_strips_gemini_thought_from_tool_call_i ( block_type, content_block_start, - ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block( - choices=choices - ) + ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices) assert block_type == "tool_use" assert content_block_start["id"] == base @@ -164,9 +155,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_content_block(): ( block_type, content_block_start, - ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block( - choices=choices - ) + ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices) assert block_type == "thinking" assert content_block_start == { @@ -202,9 +191,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_reasoning_content_only_co ( block_type, content_block_start, - ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block( - choices=choices - ) + ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices) assert block_type == "thinking" assert content_block_start == { @@ -250,9 +237,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_signature_block( ( block_type, content_block_start, - ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block( - choices=choices - ) + ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices) assert block_type == "thinking" assert content_block_start == { @@ -305,9 +290,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block_thinking_an ( block_type, content_block_start, - ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block( - choices=choices - ) + ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices) assert block_type == "thinking" @@ -350,10 +333,7 @@ def test_translate_anthropic_messages_to_openai_thinking_blocks(): assert "thinking_blocks" in result[1] assert len(result[1]["thinking_blocks"]) == 2 assert result[1]["thinking_blocks"][0]["type"] == "thinking" - assert ( - result[1]["thinking_blocks"][0]["thinking"] - == "I will call the get_weather tool." - ) + assert result[1]["thinking_blocks"][0]["thinking"] == "I will call the get_weather tool." assert result[1]["thinking_blocks"][0]["signature"] == "sigsig" assert result[1]["thinking_blocks"][1]["type"] == "redacted_thinking" assert result[1]["thinking_blocks"][1]["data"] == "REDACTED" @@ -456,9 +436,7 @@ def test_translate_anthropic_messages_to_openai_tool_message_placement(): assert tool_message_idx is not None, "Tool message not found" assert user_message_idx is not None, "User message not found" - assert ( - tool_message_idx < user_message_idx - ), "Tool message should be placed before user message" + assert tool_message_idx < user_message_idx, "Tool message should be placed before user message" @pytest.mark.parametrize( @@ -733,9 +711,7 @@ def test_translate_anthropic_to_openai_skips_prompt_cache_key_when_provider_lack def test_translate_anthropic_to_openai_skips_prompt_cache_key_for_chained_litellm_proxy(): - assert "prompt_cache_key" in litellm.get_supported_openai_params( - model="xai", custom_llm_provider="litellm_proxy" - ) + assert "prompt_cache_key" in litellm.get_supported_openai_params(model="xai", custom_llm_provider="litellm_proxy") openai_request = _translate_with_metadata("litellm_proxy/xai", {"user_id": "session-abc"}, "litellm_proxy") assert openai_request["user"] == "session-abc" assert "prompt_cache_key" not in openai_request @@ -780,7 +756,8 @@ def test_translate_openai_content_to_anthropic_empty_function_arguments(): id="call_empty_args", type="function", function=Function( - name="test_function", arguments="" # empty arguments string + name="test_function", + arguments="", # empty arguments string ), ) ], @@ -795,9 +772,7 @@ def test_translate_openai_content_to_anthropic_empty_function_arguments(): assert result[0]["type"] == "tool_use" assert result[0]["id"] == "call_empty_args" assert result[0]["name"] == "test_function" - assert ( - result[0]["input"] == {} - ), "Empty function arguments should result in empty dict" + assert result[0]["input"] == {}, "Empty function arguments should result in empty dict" def test_translate_openai_content_to_anthropic_text_and_tool_calls(): @@ -917,9 +892,7 @@ def test_translate_openai_response_to_anthropic_text_and_tool_calls(): ChatCompletionAssistantToolCall( id="call_tool_combo", type="function", - function=Function( - name="get_weather", arguments='{"location": "Paris"}' - ), + function=Function(name="get_weather", arguments='{"location": "Paris"}'), ) ], ), @@ -929,9 +902,7 @@ def test_translate_openai_response_to_anthropic_text_and_tool_calls(): ) adapter = LiteLLMAnthropicMessagesAdapter() - anthropic_response = adapter.translate_openai_response_to_anthropic( - response=openai_response - ) + anthropic_response = adapter.translate_openai_response_to_anthropic(response=openai_response) anthropic_content = anthropic_response.get("content") assert anthropic_content is not None @@ -972,9 +943,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_with_partial_json(): ( type_of_content, content_block_delta, - ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic( - choices=choices - ) + ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices) print("Type of content:", type_of_content) print("Content block delta:", content_block_delta) @@ -1083,9 +1052,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_delta(): ( type_of_content, content_block_delta, - ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic( - choices=choices - ) + ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices) assert type_of_content == "thinking_delta" assert content_block_delta["type"] == "thinking_delta" @@ -1128,9 +1095,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_with_thinking(): ( type_of_content, content_block_delta, - ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic( - choices=choices - ) + ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices) assert type_of_content == "signature_delta" assert content_block_delta["type"] == "signature_delta" @@ -1194,9 +1159,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_emits_signature_when_thin ( block_type, content_block_start, - ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block( - choices=choices - ) + ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices) assert block_type == "thinking" @@ -1236,9 +1199,7 @@ def test_translate_anthropic_messages_to_openai_user_message_with_base64_image() # Check image content assert result[0]["content"][1]["type"] == "image_url" assert "image_url" in result[0]["content"][1] - assert result[0]["content"][1]["image_url"]["url"].startswith( - "data:image/png;base64," - ) + assert result[0]["content"][1]["image_url"]["url"].startswith("data:image/png;base64,") assert ( "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" in result[0]["content"][1]["image_url"]["url"] @@ -1276,18 +1237,14 @@ def test_translate_anthropic_messages_to_openai_user_message_with_url_image(): # Check image content assert result[0]["content"][1]["type"] == "image_url" assert "image_url" in result[0]["content"][1] - assert ( - result[0]["content"][1]["image_url"]["url"] == "https://example.com/forest.jpg" - ) + assert result[0]["content"][1]["image_url"]["url"] == "https://example.com/forest.jpg" def test_translate_anthropic_messages_to_openai_tool_result_with_base64_image(): """Test that base64 images in tool results are correctly translated to OpenAI format.""" anthropic_messages = [ - AnthropicMessagesUserMessageParam( - role="user", content=[{"type": "text", "text": "Take a screenshot"}] - ), + AnthropicMessagesUserMessageParam(role="user", content=[{"type": "text", "text": "Take a screenshot"}]), AnthopicMessagesAssistantMessageParam( role="assistant", content=[ @@ -1439,9 +1396,7 @@ def test_translate_anthropic_messages_to_openai_mixed_content_with_image(): # Check first image (base64) assert result[0]["content"][1]["type"] == "image_url" - assert result[0]["content"][1]["image_url"]["url"].startswith( - "data:image/png;base64," - ) + assert result[0]["content"][1]["image_url"]["url"].startswith("data:image/png;base64,") # Check middle text assert result[0]["content"][2]["type"] == "text" @@ -1449,9 +1404,7 @@ def test_translate_anthropic_messages_to_openai_mixed_content_with_image(): # Check second image (URL) assert result[0]["content"][3]["type"] == "image_url" - assert ( - result[0]["content"][3]["image_url"]["url"] == "https://example.com/image2.jpg" - ) + assert result[0]["content"][3]["image_url"]["url"] == "https://example.com/image2.jpg" # Check final text assert result[0]["content"][4]["type"] == "text" @@ -1497,10 +1450,7 @@ def test_translate_anthropic_messages_to_openai_tool_use_with_signature(): assert tool_call["id"] == "call_386f67af31f9415781bc35071405" assert "function" in tool_call assert "provider_specific_fields" in tool_call["function"] - assert ( - tool_call["function"]["provider_specific_fields"]["thought_signature"] - == test_signature - ) + assert tool_call["function"]["provider_specific_fields"]["thought_signature"] == test_signature def test_translate_anthropic_messages_to_openai_tool_result_with_multiple_content_items(): @@ -1558,9 +1508,7 @@ def test_translate_anthropic_messages_to_openai_tool_result_with_multiple_conten result = adapter.translate_anthropic_messages_to_openai(messages=anthropic_messages) # Count how many tool messages have the same tool_call_id - tool_messages = [ - msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool" - ] + tool_messages = [msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"] tool_call_ids = [msg.get("tool_call_id") for msg in tool_messages] # The critical assertion: each tool_call_id should appear only ONCE @@ -1576,12 +1524,8 @@ def test_translate_anthropic_messages_to_openai_tool_result_with_multiple_conten # The content should be a list with all items combined tool_message = tool_messages[0] assert tool_message["tool_call_id"] == "toolu_016hYHBkTf4JDF3p22UoYk5C" - assert isinstance( - tool_message["content"], list - ), "Multiple content items should be combined into a list" - assert ( - len(tool_message["content"]) == 3 - ), f"Expected 3 content items, got {len(tool_message['content'])}" + assert isinstance(tool_message["content"], list), "Multiple content items should be combined into a list" + assert len(tool_message["content"]) == 3, f"Expected 3 content items, got {len(tool_message['content'])}" # Verify content types assert tool_message["content"][0]["type"] == "text" @@ -1630,17 +1574,14 @@ def test_translate_anthropic_messages_to_openai_tool_result_single_item_backward adapter = LiteLLMAnthropicMessagesAdapter() result = adapter.translate_anthropic_messages_to_openai(messages=anthropic_messages) - tool_messages = [ - msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool" - ] + tool_messages = [msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"] assert len(tool_messages) == 1 tool_message = tool_messages[0] # Single item should be a string for backward compatibility assert isinstance(tool_message["content"], str), ( - f"Single content item should be a string for backward compatibility, " - f"got {type(tool_message['content'])}" + f"Single content item should be a string for backward compatibility, got {type(tool_message['content'])}" ) assert tool_message["content"] == "72°F and sunny" @@ -1689,9 +1630,7 @@ def test_streaming_chunk_with_both_text_and_tool_calls_issue_18238(): ( block_type, content_block_start, - ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block( - choices=choices - ) + ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices) assert block_type == "tool_use" assert content_block_start["name"] == "Bash" @@ -1735,9 +1674,7 @@ def test_streaming_chunk_with_text_and_empty_tool_calls_returns_text_delta(): ( block_type, content_block_start, - ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block( - choices=choices - ) + ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices) assert block_type == "text" assert content_block_start == {"type": "text", "text": ""} @@ -1748,15 +1685,12 @@ def test_streaming_chunk_with_text_and_empty_tool_calls_returns_text_delta(): # ============================================================================ # Model constant for cache control tests -CACHE_CONTROL_BEDROCK_CONVERSE_MODEL = ( - "bedrock/converse/global.anthropic.claude-opus-4-5-20251101-v1:0" -) +CACHE_CONTROL_BEDROCK_CONVERSE_MODEL = "bedrock/converse/global.anthropic.claude-opus-4-5-20251101-v1:0" CACHE_CONTROL_NON_ANTHROPIC_MODEL = "gpt-4" # Bedrock Application Inference Profile ARN: the string contains neither # "anthropic" nor "claude", so the model can only be recognized via its ARN shape CACHE_CONTROL_BEDROCK_ARN_MODEL = ( - "bedrock/converse/arn:aws:bedrock:us-east-1:123456789012:" - "application-inference-profile/abcdef123456" + "bedrock/converse/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456" ) @@ -1772,9 +1706,7 @@ def test_should_add_cache_control_for_anthropic_model(): "vertex_ai/claude-3-sonnet@20240229", ]: target = {} - adapter._add_cache_control_if_applicable( - {"cache_control": cache_control}, target, model - ) + adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model) assert "cache_control" in target assert target["cache_control"] == cache_control @@ -1790,9 +1722,7 @@ def test_should_not_add_cache_control_for_non_anthropic_model(): "gemini-pro", ]: target = {} - adapter._add_cache_control_if_applicable( - {"cache_control": cache_control}, target, model - ) + adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model) assert "cache_control" not in target @@ -1807,9 +1737,7 @@ def test_should_not_add_cache_control_when_none(): {}, ]: target = {} - adapter._add_cache_control_if_applicable( - source, target, CACHE_CONTROL_BEDROCK_CONVERSE_MODEL - ) + adapter._add_cache_control_if_applicable(source, target, CACHE_CONTROL_BEDROCK_CONVERSE_MODEL) assert "cache_control" not in target @@ -1820,9 +1748,7 @@ def test_should_not_add_cache_control_when_model_none(): for model in [None, ""]: target = {} - adapter._add_cache_control_if_applicable( - {"cache_control": cache_control}, target, model - ) + adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model) assert "cache_control" not in target @@ -1928,12 +1854,7 @@ def test_cache_control_fix_does_not_broaden_claude_detection(): make is_anthropic_claude_model treat ARN profiles as Claude, which would route thinking params through unmodified and break non-Claude Bedrock profiles. """ - assert ( - LiteLLMAnthropicMessagesAdapter.is_anthropic_claude_model( - CACHE_CONTROL_BEDROCK_ARN_MODEL - ) - is False - ) + assert LiteLLMAnthropicMessagesAdapter.is_anthropic_claude_model(CACHE_CONTROL_BEDROCK_ARN_MODEL) is False def test_thinking_preserved_for_bedrock_arn_inference_profile(): @@ -2531,9 +2452,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_reasoning_content_without ( type_of_content, content_block_delta, - ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic( - choices=choices - ) + ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices) assert type_of_content == "thinking_delta" assert content_block_delta["type"] == "thinking_delta" @@ -2565,9 +2484,7 @@ def test_translate_openai_response_to_anthropic_with_reasoning_content_only(): ) adapter = LiteLLMAnthropicMessagesAdapter() - anthropic_response = adapter.translate_openai_response_to_anthropic( - response=openai_response - ) + anthropic_response = adapter.translate_openai_response_to_anthropic(response=openai_response) anthropic_content = anthropic_response.get("content") assert anthropic_content is not None @@ -2580,9 +2497,7 @@ def test_translate_openai_response_to_anthropic_with_reasoning_content_only(): # Second block should be text assert anthropic_content[1]["type"] == "text" - assert ( - anthropic_content[1]["text"] == 'There are **3** "r"s in the word strawberry.' - ) + assert anthropic_content[1]["text"] == 'There are **3** "r"s in the word strawberry.' assert anthropic_response.get("stop_reason") == "end_turn" @@ -2634,9 +2549,7 @@ def test_truncate_tool_name_deterministic(): def test_truncate_tool_name_avoids_collisions(): """Similar long names should produce different truncated names.""" name1 = "process_user_data_with_validation_and_error_handling_for_production_environment" - name2 = ( - "process_user_data_with_validation_and_error_handling_for_staging_environment" - ) + name2 = "process_user_data_with_validation_and_error_handling_for_staging_environment" result1 = truncate_tool_name(name1) result2 = truncate_tool_name(name2) @@ -2656,9 +2569,7 @@ def test_create_tool_name_mapping_no_long_names(): def test_create_tool_name_mapping_with_long_names(): """Mapping should contain entries for truncated names.""" - long_name = ( - "a_very_long_tool_name_that_exceeds_the_64_character_limit_imposed_by_openai" - ) + long_name = "a_very_long_tool_name_that_exceeds_the_64_character_limit_imposed_by_openai" tools = [ {"name": "short_name"}, {"name": long_name}, @@ -2683,9 +2594,7 @@ def test_translate_anthropic_tools_with_long_names(): ] adapter = LiteLLMAnthropicMessagesAdapter() - result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai( - tools=tools, model="gpt-4" - ) + result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(tools=tools, model="gpt-4") assert len(result) == 1 # The tool name should be truncated @@ -2707,9 +2616,7 @@ def test_translate_anthropic_tools_mixed_names(): ] adapter = LiteLLMAnthropicMessagesAdapter() - result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai( - tools=tools, model="gpt-4" - ) + result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(tools=tools, model="gpt-4") assert len(result) == 2 # Short name unchanged @@ -2723,9 +2630,7 @@ def test_translate_anthropic_tools_mixed_names(): def test_translate_openai_response_restores_tool_names(): """Tool names in responses should be restored to original.""" - original_name = ( - "a_very_long_tool_name_that_needs_truncation_for_openai_api_compatibility" - ) + original_name = "a_very_long_tool_name_that_needs_truncation_for_openai_api_compatibility" truncated_name = truncate_tool_name(original_name) tool_name_mapping = {truncated_name: original_name} @@ -2757,9 +2662,7 @@ def test_translate_openai_response_restores_tool_names(): ) adapter = LiteLLMAnthropicMessagesAdapter() - result = adapter.translate_openai_response_to_anthropic( - response=response, tool_name_mapping=tool_name_mapping - ) + result = adapter.translate_openai_response_to_anthropic(response=response, tool_name_mapping=tool_name_mapping) # Find the tool_use block in the response tool_use_blocks = [c for c in result["content"] if c.get("type") == "tool_use"] @@ -2925,9 +2828,7 @@ def test_translate_openai_usage_to_anthropic_cache_tokens_from_dict_details_with "cache_write_tokens": 20.0, } - anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( - usage - ) + anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(usage) assert anthropic_usage["input_tokens"] == 70 assert anthropic_usage["output_tokens"] == 50 @@ -2946,9 +2847,7 @@ def test_translate_openai_usage_to_anthropic_ignores_fractional_cache_tokens(): "cache_creation_tokens": 20.25, } - anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( - usage - ) + anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(usage) assert anthropic_usage["input_tokens"] == 120 assert anthropic_usage["output_tokens"] == 50 @@ -2965,9 +2864,7 @@ def test_translate_openai_usage_to_anthropic_ignores_bool_cache_tokens(): usage.cache_read_input_tokens = True usage.cache_creation_input_tokens = True - anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( - usage - ) + anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(usage) assert anthropic_usage["input_tokens"] == 120 assert anthropic_usage["output_tokens"] == 50 @@ -3186,9 +3083,7 @@ def test_translate_streaming_openai_response_to_anthropic_cache_tokens_with_appl assert message_delta["usage"]["output_tokens"] == 50 assert message_delta["usage"]["cache_read_input_tokens"] == 30 assert message_delta["usage"]["cache_creation_input_tokens"] == 20 - assert message_delta["context_management"]["applied_edits"][0]["type"] == ( - "compact_20260112" - ) + assert message_delta["context_management"]["applied_edits"][0]["type"] == ("compact_20260112") # ===================================================================== @@ -3363,15 +3258,8 @@ class TestTranslateAnthropicOutputFormatToOpenAI: assert schema["required"] == ["user"] assert schema["properties"]["user"]["additionalProperties"] is False assert schema["properties"]["user"]["required"] == ["name", "address"] - assert ( - schema["properties"]["user"]["properties"]["address"][ - "additionalProperties" - ] - is False - ) - assert schema["properties"]["user"]["properties"]["address"]["required"] == [ - "city" - ] + assert schema["properties"]["user"]["properties"]["address"]["additionalProperties"] is False + assert schema["properties"]["user"]["properties"]["address"]["required"] == ["city"] def test_array_items_object_adds_additional_properties_false(self): output_format = { @@ -3446,19 +3334,9 @@ class TestTranslateAnthropicOutputFormatToOpenAI: assert sorted(schema["required"]) == ["age", "email", "name"] def test_invalid_output_format_returns_none(self): - assert ( - self.adapter.translate_anthropic_output_format_to_openai("invalid") is None - ) - assert ( - self.adapter.translate_anthropic_output_format_to_openai({"type": "text"}) - is None - ) - assert ( - self.adapter.translate_anthropic_output_format_to_openai( - {"type": "json_schema"} - ) - is None - ) + assert self.adapter.translate_anthropic_output_format_to_openai("invalid") is None + assert self.adapter.translate_anthropic_output_format_to_openai({"type": "text"}) is None + assert self.adapter.translate_anthropic_output_format_to_openai({"type": "json_schema"}) is None class TestAnthropicStreamWrapperToolArgs: @@ -3662,9 +3540,7 @@ def test_translate_openai_response_to_anthropic_with_polyfill_compaction_block() ) response = _make_simple_openai_response(text="Hello after compaction.") adapter = LiteLLMAnthropicMessagesAdapter() - result = adapter.translate_openai_response_to_anthropic( - response=response, polyfill_result=polyfill - ) + result = adapter.translate_openai_response_to_anthropic(response=response, polyfill_result=polyfill) content = result.get("content") assert content is not None @@ -3696,9 +3572,7 @@ def test_translate_openai_response_to_anthropic_with_polyfill_iterations_usage() ) response = _make_simple_openai_response(prompt_tokens=100, completion_tokens=30) adapter = LiteLLMAnthropicMessagesAdapter() - result = adapter.translate_openai_response_to_anthropic( - response=response, polyfill_result=polyfill - ) + result = adapter.translate_openai_response_to_anthropic(response=response, polyfill_result=polyfill) usage = result.get("usage") assert usage is not None @@ -3753,13 +3627,9 @@ def test_translate_openai_response_to_anthropic_with_polyfill_both_compaction_an {"type": "compaction", "input_tokens": 300, "output_tokens": 75}, ], ) - response = _make_simple_openai_response( - text="After compaction.", prompt_tokens=120, completion_tokens=40 - ) + response = _make_simple_openai_response(text="After compaction.", prompt_tokens=120, completion_tokens=40) adapter = LiteLLMAnthropicMessagesAdapter() - result = adapter.translate_openai_response_to_anthropic( - response=response, polyfill_result=polyfill - ) + result = adapter.translate_openai_response_to_anthropic(response=response, polyfill_result=polyfill) # compaction block must come first content = result.get("content") @@ -3861,7 +3731,9 @@ def test_translate_anthropic_tools_to_openai_omits_unset_strict(): assert function["parameters"]["required"] == ["query"] -TOOL_RESULT_IMAGE_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +TOOL_RESULT_IMAGE_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +) TOOL_RESULT_IMAGE_URL = "https://example.com/screenshot.png" @@ -3869,8 +3741,7 @@ def _anthropic_tool_use_turn(*tool_use_ids): return AnthopicMessagesAssistantMessageParam( role="assistant", content=[ - {"type": "tool_use", "id": tid, "name": "read_file", "input": {"path": "img.png"}} - for tid in tool_use_ids + {"type": "tool_use", "id": tid, "name": "read_file", "input": {"path": "img.png"}} for tid in tool_use_ids ], ) @@ -3990,9 +3861,7 @@ def test_tool_result_parallel_tool_calls_keep_tool_message_adjacency(): result = _run_chat_completions_pipeline( [ _anthropic_tool_use_turn("toolu_01", "toolu_02"), - _anthropic_tool_result_turn( - {"toolu_01": [_base64_image_block()], "toolu_02": [_url_image_block()]} - ), + _anthropic_tool_result_turn({"toolu_01": [_base64_image_block()], "toolu_02": [_url_image_block()]}), ] ) @@ -4156,7 +4025,9 @@ def test_translate_anthropic_to_openai_without_prompt_cache_breakpoint_adds_noth def test_translate_anthropic_messages_to_openai_carries_midturn_system_prompt_cache_breakpoint(): explicit = {"mode": "explicit"} result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( - messages=[{"role": "system", "content": [{"type": "text", "text": "fix", "prompt_cache_breakpoint": explicit}]}], + messages=[ + {"role": "system", "content": [{"type": "text", "text": "fix", "prompt_cache_breakpoint": explicit}]} + ], model="gpt-5.6", ) assert result == [ diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py index a944afc6152..71d7974c2d4 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py @@ -34,9 +34,7 @@ import pytest # Anchor sys.path to this file's location — not the working-directory-relative # pattern Greptile flagged on PR #23706. Resolves correctly regardless of # where pytest is invoked from. -sys.path.insert( - 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../..")) -) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( ANTHROPIC_ONLY_REQUEST_KEYS, @@ -174,9 +172,7 @@ class TestOutputConfigStrippedFromCompletionKwargs: result = _call_prepare( extra_kwargs={ "custom_llm_provider": "azure", - "output_config": { - "format": {"type": "json_schema", "schema": losing_schema} - }, + "output_config": {"format": {"type": "json_schema", "schema": losing_schema}}, }, output_format={"type": "json_schema", "schema": winning_schema}, ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py index 6973340101e..0a4d9c693d6 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -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) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py index 5c53a8fc317..646d63e08bd 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py @@ -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 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py index 3e85872f1e5..ce6d6c0656e 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py @@ -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" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index 17d42f55ae0..3cef89ab255 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -507,11 +507,7 @@ def _thinking_first_chunks() -> List[MagicMock]: def _assert_thinking_first_block_opens_at_index_zero(events: List[dict]) -> None: - starts = [ - (e["index"], e["content_block"]["type"]) - for e in events - if e.get("type") == "content_block_start" - ] + starts = [(e["index"], e["content_block"]["type"]) for e in events if e.get("type") == "content_block_start"] assert starts == [(0, "thinking"), (1, "text")], starts assert "" not in _text_deltas(events) assert _thinking_deltas(events) == ["Let me think", "about it."] @@ -980,9 +976,7 @@ def test_tool_block_start_emitted_without_awaiting_the_next_chunk_sync(): "name": "Write", "input": {}, } - assert stream.pulled == 1, ( - f"content_block_start was withheld until {stream.pulled} upstream chunks had arrived" - ) + assert stream.pulled == 1, f"content_block_start was withheld until {stream.pulled} upstream chunks had arrived" @pytest.mark.asyncio @@ -997,9 +991,7 @@ async def test_tool_block_start_emitted_without_awaiting_the_next_chunk_async(): start = await wrapper.__anext__() assert start["type"] == "content_block_start" assert start["content_block"]["name"] == "Write" - assert stream.pulled == 1, ( - f"content_block_start was withheld until {stream.pulled} upstream chunks had arrived" - ) + assert stream.pulled == 1, f"content_block_start was withheld until {stream.pulled} upstream chunks had arrived" @pytest.mark.parametrize("is_async", [False, True]) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py index 29e9279731d..cf0df4d7ea3 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py @@ -127,21 +127,17 @@ async def test_async_stream_emits_input_json_delta_for_bundled_tool_args(): ): input_json_delta_idx = i - assert ( - tool_start_idx is not None - ), f"Expected content_block_start with type=tool_use; events: {event_types}" - assert ( - input_json_delta_idx is not None - ), f"Expected content_block_delta with input_json_delta; events: {event_types}" - assert ( - input_json_delta_idx == tool_start_idx + 1 - ), "input_json_delta should immediately follow the tool_use content_block_start" + assert tool_start_idx is not None, f"Expected content_block_start with type=tool_use; events: {event_types}" + assert input_json_delta_idx is not None, ( + f"Expected content_block_delta with input_json_delta; events: {event_types}" + ) + assert input_json_delta_idx == tool_start_idx + 1, ( + "input_json_delta should immediately follow the tool_use content_block_start" + ) # Verify the delta carries the tool arguments delta_event = events[input_json_delta_idx] - assert delta_event["delta"][ - "partial_json" - ], "input_json_delta should have non-empty partial_json" + assert delta_event["delta"]["partial_json"], "input_json_delta should have non-empty partial_json" @pytest.mark.asyncio @@ -230,8 +226,7 @@ async def test_async_stream_no_extra_delta_when_tool_args_empty(): and e["delta"].get("type") == "input_json_delta" ] assert len(input_json_deltas) == 1, ( - f"Expected exactly 1 input_json_delta (from the follow-up chunk), " - f"got {len(input_json_deltas)}" + f"Expected exactly 1 input_json_delta (from the follow-up chunk), got {len(input_json_deltas)}" ) assert input_json_deltas[0]["delta"]["partial_json"] == '{"location": "NYC"}' @@ -291,15 +286,13 @@ def test_sync_stream_emits_input_json_delta_for_bundled_tool_args(): ): input_json_delta_idx = i - assert ( - tool_start_idx is not None - ), f"Expected content_block_start with type=tool_use; events: {event_types}" - assert ( - input_json_delta_idx is not None - ), f"Expected content_block_delta with input_json_delta; events: {event_types}" - assert ( - input_json_delta_idx == tool_start_idx + 1 - ), "input_json_delta should immediately follow the tool_use content_block_start" + assert tool_start_idx is not None, f"Expected content_block_start with type=tool_use; events: {event_types}" + assert input_json_delta_idx is not None, ( + f"Expected content_block_delta with input_json_delta; events: {event_types}" + ) + assert input_json_delta_idx == tool_start_idx + 1, ( + "input_json_delta should immediately follow the tool_use content_block_start" + ) assert events[input_json_delta_idx]["delta"]["partial_json"] @@ -343,9 +336,7 @@ def test_sync_stream_no_extra_delta_when_tool_args_empty(): ) wrapper = AnthropicStreamWrapper( - completion_stream=iter( - [text_chunk, tool_name_chunk, tool_args_chunk, finish_chunk] - ), + completion_stream=iter([text_chunk, tool_name_chunk, tool_args_chunk, finish_chunk]), model="test-model", ) @@ -374,7 +365,6 @@ def test_sync_stream_no_extra_delta_when_tool_args_empty(): and e["delta"].get("type") == "input_json_delta" ] assert len(input_json_deltas) == 1, ( - f"Expected exactly 1 input_json_delta (from the follow-up chunk), " - f"got {len(input_json_deltas)}" + f"Expected exactly 1 input_json_delta (from the follow-up chunk), got {len(input_json_deltas)}" ) assert input_json_deltas[0]["delta"]["partial_json"] == '{"location": "NYC"}' diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py index 09ac95ab16e..cc4852f25de 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py @@ -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", diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index 28c82fdf528..da03f3c00ab 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -193,10 +193,7 @@ def test_select_last_user_question_strips_tool_result_from_mixed_turn(): content = selected[0]["content"] assert isinstance(content, list) assert all(b.get("type") != "tool_result" for b in content) - assert any( - b.get("type") == "text" and b.get("text") == "follow-up question" - for b in content - ) + assert any(b.get("type") == "text" and b.get("text") == "follow-up question" for b in content) def test_select_last_user_question_skips_pure_tool_result_turn(): @@ -425,9 +422,7 @@ def test_client_compaction_block_history_without_context_management(): def test_client_compaction_block_history_no_compaction_returns_none(): - result = apply_client_compaction_block_history( - messages=_simple_messages(), system="base" - ) + result = apply_client_compaction_block_history(messages=_simple_messages(), system="base") assert result is None @@ -512,9 +507,7 @@ async def test_slice_only_no_compaction_block_under_threshold(): async def test_full_summary_path(): """Over threshold: summary call fires, compaction_block and iterations_usage returned.""" messages = _simple_messages() - mock_response = _make_mock_response( - "Condensed history", prompt_tokens=200, completion_tokens=50 - ) + mock_response = _make_mock_response("Condensed history", prompt_tokens=200, completion_tokens=50) with ( patch( @@ -1068,13 +1061,11 @@ async def test_summary_call_does_not_emit_consecutive_user_turns(): ) summary_messages = captured_calls[0]["summary_messages"] - user_indices = [ - idx for idx, msg in enumerate(summary_messages) if msg.get("role") == "user" - ] + user_indices = [idx for idx, msg in enumerate(summary_messages) if msg.get("role") == "user"] # No two adjacent indices. - assert all( - b - a > 1 for a, b in zip(user_indices, user_indices[1:]) - ), f"two consecutive user turns produced: {summary_messages}" + assert all(b - a > 1 for a, b in zip(user_indices, user_indices[1:])), ( + f"two consecutive user turns produced: {summary_messages}" + ) async def test_summary_call_sends_default_max_tokens(): @@ -1157,9 +1148,9 @@ def test_summary_max_tokens_setting_falls_back_for_invalid_values(): "litellm.proxy.proxy_server.general_settings", {"context_management_summary_max_tokens": bad}, ): - assert ( - _read_summary_max_tokens_setting() == COMPACT_SUMMARY_MAX_TOKENS - ), f"expected default for invalid override {bad!r}" + assert _read_summary_max_tokens_setting() == COMPACT_SUMMARY_MAX_TOKENS, ( + f"expected default for invalid override {bad!r}" + ) async def test_summary_call_sends_default_timeout(): @@ -1282,9 +1273,7 @@ async def test_summary_model_denied_when_team_not_in_allowlist(): tools=None, system=None, edit_spec=_EDIT_SPEC_DEFAULT, - user_api_key_auth=_fake_user_api_key_auth( - key_models=["all-proxy-models"], team_models=["gpt-4o"] - ), + user_api_key_auth=_fake_user_api_key_auth(key_models=["all-proxy-models"], team_models=["gpt-4o"]), ) mock_call.assert_not_awaited() @@ -1313,9 +1302,7 @@ async def test_summary_model_allowed_when_in_key_allowlist(): tools=None, system=None, edit_spec=_EDIT_SPEC_DEFAULT, - user_api_key_auth=_fake_user_api_key_auth( - key_models=["claude-haiku-4-5", "gpt-4o"] - ), + user_api_key_auth=_fake_user_api_key_auth(key_models=["claude-haiku-4-5", "gpt-4o"]), ) mock_call.assert_awaited_once() @@ -1521,9 +1508,7 @@ async def test_summary_model_denied_when_key_over_model_budget(): limiter = MagicMock() limiter.is_key_within_model_budget = AsyncMock( - side_effect=litellm.BudgetExceededError( - message="over budget", current_cost=10, max_budget=5 - ) + side_effect=litellm.BudgetExceededError(message="over budget", current_cost=10, max_budget=5) ) with ( @@ -1574,9 +1559,7 @@ async def test_summary_model_denied_when_user_over_model_budget(): limiter = MagicMock() limiter.is_user_within_model_budget = AsyncMock( - side_effect=litellm.BudgetExceededError( - message="over budget", current_cost=10, max_budget=5 - ) + side_effect=litellm.BudgetExceededError(message="over budget", current_cost=10, max_budget=5) ) with ( @@ -1617,9 +1600,7 @@ async def test_summary_model_denied_when_user_over_model_budget(): _PROXY_VirtualKeyModelMaxBudgetLimiter, ) - real_params = inspect.signature( - _PROXY_VirtualKeyModelMaxBudgetLimiter.is_user_within_model_budget - ).parameters + real_params = inspect.signature(_PROXY_VirtualKeyModelMaxBudgetLimiter.is_user_within_model_budget).parameters for kwarg in ("user_id", "user_model_max_budget", "model"): assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter no longer accepts" @@ -1641,9 +1622,7 @@ async def test_summary_model_denied_when_end_user_over_model_budget(): limiter = MagicMock() limiter.is_key_within_model_budget = AsyncMock(return_value=True) limiter.is_end_user_within_model_budget = AsyncMock( - side_effect=litellm.BudgetExceededError( - message="over budget", current_cost=10, max_budget=5 - ) + side_effect=litellm.BudgetExceededError(message="over budget", current_cost=10, max_budget=5) ) with ( @@ -1956,9 +1935,7 @@ async def test_model_budget_metadata_propagated_to_summary_call(): parent_litellm_metadata = { "user_api_key": "sk-test", "user_api_key_model_max_budget": {"claude-haiku-4-5": {"budget_limit": 5}}, - "user_api_key_end_user_model_max_budget": { - "claude-haiku-4-5": {"budget_limit": 2} - }, + "user_api_key_end_user_model_max_budget": {"claude-haiku-4-5": {"budget_limit": 2}}, } with ( @@ -1983,12 +1960,8 @@ async def test_model_budget_metadata_propagated_to_summary_call(): ) propagated = mock_call.call_args.kwargs["metadata"] - assert propagated["user_api_key_model_max_budget"] == { - "claude-haiku-4-5": {"budget_limit": 5} - } - assert propagated["user_api_key_end_user_model_max_budget"] == { - "claude-haiku-4-5": {"budget_limit": 2} - } + assert propagated["user_api_key_model_max_budget"] == {"claude-haiku-4-5": {"budget_limit": 5}} + assert propagated["user_api_key_end_user_model_max_budget"] == {"claude-haiku-4-5": {"budget_limit": 2}} async def test_summary_call_propagates_allowed_model_region(): @@ -2460,9 +2433,7 @@ def test_endpoint_returns_anthropic_400_on_context_management_error(): mock_proxy_server.version = "test" with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): - with patch( - "litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_cls: + with patch("litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_cls: mock_instance = MagicMock() mock_instance.base_process_llm_request = AsyncMock( side_effect=AnthropicContextManagementError( @@ -2521,9 +2492,7 @@ def test_endpoint_runs_failure_hook_on_500_context_management_error(): mock_proxy_server.version = "test" with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): - with patch( - "litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_cls: + with patch("litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_cls: mock_instance = MagicMock() mock_instance.base_process_llm_request = AsyncMock( side_effect=AnthropicContextManagementError( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py index 50c72cfe8d0..56b1ce16ebc 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py @@ -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 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py index 414ba8f0f5c..b6df9ecf3bd 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py @@ -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 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py index 015b5754c6e..33f52f90afb 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -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" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py new file mode 100644 index 00000000000..768a2834fbe --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py @@ -0,0 +1,118 @@ +from pathlib import Path +from typing import Final + +import pytest + +import litellm +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, +) +from litellm.llms.anthropic.wif import get_anthropic_wif_token +from litellm.llms.minimax.messages.transformation import MinimaxMessagesConfig +from litellm.llms.tencent.messages.transformation import TencentAnthropicMessagesConfig +from tests.test_litellm.llms.anthropic.test_anthropic_wif import ( + ScriptedPoster, + make_engine, + token_response, + write_token_file, +) + +_WIF_PARAMS: Final[dict] = { + "anthropic_federation_rule_id": "fdrl_abc123", + "anthropic_organization_id": "org-uuid-1", + "anthropic_identity_token_file": "/var/run/secrets/identity-token", +} + + +def test_workload_identity_allowed_for_anthropic() -> None: + assert AnthropicMessagesConfig()._allows_workload_identity is True + + +def test_workload_identity_blocked_for_minimax() -> None: + assert MinimaxMessagesConfig()._allows_workload_identity is False + + +def test_workload_identity_blocked_for_tencent() -> None: + assert TencentAnthropicMessagesConfig()._allows_workload_identity is False + + +def test_minimax_validate_environment_never_attaches_anthropic_wif_credential( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Regression test: before the fix, an Anthropic-WIF-configured proxy would mint a real + Anthropic federation token inside MiniMax's inherited validate_anthropic_messages_environment + and send it as the Authorization header on the MiniMax-routed request. With no MiniMax + credential of its own the deployment must fail closed on the missing key instead.""" + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_prod") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-prod-uuid") + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False) + monkeypatch.delenv("MINIMAX_API_KEY", raising=False) + token_file = write_token_file(tmp_path, "jwt-assertion-value") + litellm_params = {"anthropic_identity_token_file": str(token_file)} + + with pytest.raises(litellm.AuthenticationError, match="Missing Anthropic API Key"): + MinimaxMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model="MiniMax-M2.1", + messages=[], + optional_params={}, + litellm_params=litellm_params, + api_key=None, + api_base="https://api.minimax.io/anthropic", + ) + + +def test_tencent_validate_environment_never_attaches_anthropic_wif_credential( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_prod") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-prod-uuid") + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False) + monkeypatch.delenv("TENCENT_API_KEY", raising=False) + token_file = write_token_file(tmp_path, "jwt-assertion-value") + litellm_params = {"anthropic_identity_token_file": str(token_file)} + + monkeypatch.setattr(litellm, "api_key", None) + with pytest.raises(litellm.AuthenticationError, match="Missing Anthropic API Key"): + TencentAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model="deepseek-v4-pro", + messages=[], + optional_params={}, + litellm_params=litellm_params, + api_key=None, + api_base="https://tokenhub-intl.tencentcloudmaas.com", + ) + + +def test_wif_token_exchange_reaches_only_anthropic_not_minimax_or_tencent( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """get_anthropic_wif_token's engine parameter is the only DI seam in the WIF minting chain; + validate_anthropic_messages_environment always uses the module's default engine, so this + drives that seam directly with the exact litellm_params AnthropicModelInfo.get_auth_header + would receive from each config, proving MiniMax/Tencent never reach the token endpoint even + when a mint would otherwise succeed.""" + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_prod") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-prod-uuid") + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) + token_file = write_token_file(tmp_path, "jwt-assertion-value") + litellm_params = {"anthropic_identity_token_file": str(token_file)} + poster = ScriptedPoster([token_response("sk-ant-oat01-canary")]) + engine = make_engine(poster) + + minted: Final = get_anthropic_wif_token( + litellm_params, + "https://api.anthropic.com", + "claude-sonnet-4-5", + engine, + ) + assert minted == "sk-ant-oat01-canary" + assert len(poster.requests) == 1 + + for config in (MinimaxMessagesConfig(), TencentAnthropicMessagesConfig()): + assert config._allows_workload_identity is False + + assert len(poster.requests) == 1 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py index efd49962ac8..75a254a3291 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py @@ -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 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py index e6d5c6f4ee1..8a0d195107b 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py @@ -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 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py index a0d1f9de6ec..5736de66670 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py @@ -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__": diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py index f8c48e46b2f..e28d3716c1e 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py @@ -59,7 +59,7 @@ def test_anthropic_messages_handler_skips_the_gateway_on_recursion(): "litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp", new=AsyncMock(return_value={"routed": True}), ) as routed: - with pytest.raises(ValueError, match='anthropic_messages_handler is not implemented for sync calls'): + with pytest.raises(ValueError, match="anthropic_messages_handler is not implemented for sync calls"): anthropic_messages_handler( max_tokens=100, messages=[{"role": "user", "content": "hi"}], @@ -78,7 +78,7 @@ def test_anthropic_messages_handler_leaves_native_tools_alone(): "litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp", new=AsyncMock(return_value={"routed": True}), ) as routed: - with pytest.raises(ValueError, match='anthropic_messages_handler is not implemented for sync calls'): + with pytest.raises(ValueError, match="anthropic_messages_handler is not implemented for sync calls"): anthropic_messages_handler( max_tokens=100, messages=[{"role": "user", "content": "hi"}], @@ -115,9 +115,7 @@ def test_build_tool_result_message_uses_anthropic_tool_result_blocks(): message = _build_tool_result_message([{"tool_call_id": "toolu_1", "result": "9 sections", "name": "read_wiki"}]) assert message["role"] == "user" - assert list(message["content"]) == [ - {"type": "tool_result", "tool_use_id": "toolu_1", "content": "9 sections"} - ] + assert list(message["content"]) == [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "9 sections"}] @pytest.mark.asyncio diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py index 137286a18c4..b13dc850f61 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py @@ -1,7 +1,6 @@ from typing import List - from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, ) @@ -53,9 +52,7 @@ def construct_text_chunk(text: str) -> ModelResponseStream: ) -def construct_split_tool_call( - id: str, function_name: str, function_arg_parts: List[str] -) -> List[ModelResponseStream]: +def construct_split_tool_call(id: str, function_name: str, function_arg_parts: List[str]) -> List[ModelResponseStream]: return [ # https://platform.openai.com/docs/guides/function-calling#streaming ModelResponseStream( @@ -144,10 +141,7 @@ def test_anthropic_stream_wrapper_single_tool_call(): get_weather_calls = 0 for chunk in chunks: - if ( - chunk.get("type") == "content_block_start" - and chunk["content_block"]["type"] == "tool_use" - ): + if chunk.get("type") == "content_block_start" and chunk["content_block"]["type"] == "tool_use": if chunk["content_block"]["name"] == "get_weather": get_weather_calls += 1 @@ -203,10 +197,7 @@ def test_anthropic_stream_wrapper_back_to_back_tool_calls(): get_weather_calls = 0 for chunk in chunks: - if ( - chunk.get("type") == "content_block_start" - and chunk["content_block"]["type"] == "tool_use" - ): + if chunk.get("type") == "content_block_start" and chunk["content_block"]["type"] == "tool_use": if chunk["content_block"]["name"] == "get_weather": get_weather_calls += 1 @@ -218,9 +209,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text(): *construct_split_tool_call("tooluse_foo", "get_weather", ['{"city":', '"NY"}']), construct_text_chunk("The weather is nice today."), *construct_split_tool_call("tooluse_bar", "get_weather", ['{"city":', '"SF"}']), - *construct_split_tool_call( - "tooluse_bar", "get_weather", ['{"city":', '"CHI"}'] - ), + *construct_split_tool_call("tooluse_bar", "get_weather", ['{"city":', '"CHI"}']), construct_text_chunk("The weather is not so nice today."), ModelResponseStream( choices=[ @@ -280,8 +269,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text(): text_deltas = [ chunk["delta"]["text"] for chunk in chunks - if chunk.get("type") == "content_block_delta" - and chunk["delta"].get("type") == "text_delta" + if chunk.get("type") == "content_block_delta" and chunk["delta"].get("type") == "text_delta" ] assert text_deltas == [ "The weather is nice today.", @@ -291,10 +279,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text(): get_weather_calls = 0 for chunk in chunks: - if ( - chunk.get("type") == "content_block_start" - and chunk["content_block"]["type"] == "tool_use" - ): + if chunk.get("type") == "content_block_start" and chunk["content_block"]["type"] == "tool_use": if chunk["content_block"]["name"] == "get_weather": get_weather_calls += 1 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py index f478bbb9b50..dacb59c9fc0 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py @@ -29,13 +29,12 @@ def _call_handler_and_capture_optional_params(thinking=None, **extra_kwargs): """ captured = {} - with patch( - "litellm.llms.anthropic.experimental_pass_through.messages.handler." - "base_llm_http_handler" - ) as mock_handler, patch( - "litellm.llms.anthropic.experimental_pass_through.messages.handler." - "ProviderConfigManager" - ) as mock_pcm: + with ( + patch( + "litellm.llms.anthropic.experimental_pass_through.messages.handler.base_llm_http_handler" + ) as mock_handler, + patch("litellm.llms.anthropic.experimental_pass_through.messages.handler.ProviderConfigManager") as mock_pcm, + ): # Make get_provider_anthropic_messages_config return a non-None config # so the handler takes the native Anthropic path mock_pcm.get_provider_anthropic_messages_config.return_value = MagicMock() @@ -71,9 +70,7 @@ class TestReasoningAutoSummaryMessages: def test_adaptive_thinking_gets_display_summarized(self): """reasoning_auto_summary=True + thinking.type='adaptive' -> display='summarized'.""" with patch.object(litellm, "reasoning_auto_summary", True): - params = _call_handler_and_capture_optional_params( - thinking={"type": "adaptive", "budget_tokens": 5000} - ) + params = _call_handler_and_capture_optional_params(thinking={"type": "adaptive", "budget_tokens": 5000}) thinking = params.get("thinking", {}) assert thinking.get("display") == "summarized" assert thinking.get("type") == "adaptive" @@ -82,9 +79,7 @@ class TestReasoningAutoSummaryMessages: def test_enabled_thinking_gets_display_summarized(self): """reasoning_auto_summary=True + thinking.type='enabled' -> display='summarized'.""" with patch.object(litellm, "reasoning_auto_summary", True): - params = _call_handler_and_capture_optional_params( - thinking={"type": "enabled", "budget_tokens": 10000} - ) + params = _call_handler_and_capture_optional_params(thinking={"type": "enabled", "budget_tokens": 10000}) thinking = params.get("thinking", {}) assert thinking.get("display") == "summarized" assert thinking.get("type") == "enabled" @@ -92,18 +87,14 @@ class TestReasoningAutoSummaryMessages: def test_disabled_thinking_no_display(self): """reasoning_auto_summary=True + thinking.type='disabled' -> display NOT set.""" with patch.object(litellm, "reasoning_auto_summary", True): - params = _call_handler_and_capture_optional_params( - thinking={"type": "disabled"} - ) + params = _call_handler_and_capture_optional_params(thinking={"type": "disabled"}) thinking = params.get("thinking", {}) assert "display" not in thinking def test_no_injection_when_flag_false(self): """reasoning_auto_summary=False + active thinking -> display NOT set.""" with patch.object(litellm, "reasoning_auto_summary", False): - params = _call_handler_and_capture_optional_params( - thinking={"type": "enabled", "budget_tokens": 10000} - ) + params = _call_handler_and_capture_optional_params(thinking={"type": "enabled", "budget_tokens": 10000}) thinking = params.get("thinking", {}) assert "display" not in thinking @@ -117,12 +108,11 @@ class TestReasoningAutoSummaryMessages: def test_env_var_enables_auto_summary(self): """LITELLM_REASONING_AUTO_SUMMARY=true env var enables the feature.""" - with patch.object(litellm, "reasoning_auto_summary", False), patch.dict( - os.environ, {"LITELLM_REASONING_AUTO_SUMMARY": "true"} + with ( + patch.object(litellm, "reasoning_auto_summary", False), + patch.dict(os.environ, {"LITELLM_REASONING_AUTO_SUMMARY": "true"}), ): - params = _call_handler_and_capture_optional_params( - thinking={"type": "adaptive", "budget_tokens": 5000} - ) + params = _call_handler_and_capture_optional_params(thinking={"type": "adaptive", "budget_tokens": 5000}) thinking = params.get("thinking", {}) assert thinking.get("display") == "summarized" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py index dc2e107928f..95696a536f1 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py @@ -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 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py index bebdbe9f512..cd5f3f3f327 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py @@ -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 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index be33b2ee3b1..29a892af119 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -165,7 +165,7 @@ async def test_async_sse_wrapper_treats_message_stop_bytes_as_complete(): def test_is_message_stop_chunk(): assert _is_message_stop_chunk({"type": "message_stop"}) is True assert _is_message_stop_chunk({"type": "message_delta"}) is False - assert _is_message_stop_chunk(b'event: message_stop\ndata: {}\n\n') is True + assert _is_message_stop_chunk(b"event: message_stop\ndata: {}\n\n") is True assert _is_message_stop_chunk(b"raw-bytes") is False assert _is_message_stop_chunk("message_stop") is False @@ -177,7 +177,7 @@ def test_is_message_stop_chunk_ignores_substring_in_payload(): not be treated as a terminal stop event. """ delta_frame_with_substring = ( - b'event: content_block_delta\n' + b"event: content_block_delta\n" b'data: {"type": "content_block_delta", "delta": ' b'{"type": "input_json_delta", "partial_json": "\\"message_stop\\""}}\n\n' ) @@ -281,10 +281,11 @@ async def test_async_sse_wrapper_emits_error_when_bytes_stream_only_mentions_mes payload text contains `message_stop` (but never emits the actual `event: message_stop` frame) must still be flagged as incomplete. """ + async def _byte_stream(): yield b'event: message_start\ndata: {"type": "message_start"}\n\n' yield ( - b'event: content_block_delta\n' + b"event: content_block_delta\n" b'data: {"type": "content_block_delta", "delta": ' b'{"type": "input_json_delta", "partial_json": "\\"message_stop\\""}}\n\n' ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 5ecf604f096..5cef7869b42 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -1481,9 +1481,7 @@ class TestToolResultImages: }, { "role": "user", - "content": [ - {"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content} - ], + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content}], }, ] @@ -1857,7 +1855,9 @@ class TestPromptCacheBreakpointToResponses: ] def test_system_without_breakpoint_still_becomes_instructions(self): - request = _make_request(system=[{"type": "text", "text": "Be concise."}, {"type": "text", "text": "Be helpful."}]) + request = _make_request( + system=[{"type": "text", "text": "Be concise."}, {"type": "text", "text": "Be helpful."}] + ) kwargs = _ADAPTER.translate_request(request) assert kwargs["instructions"] == "Be concise.\nBe helpful." assert kwargs["input"] == [ diff --git a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py b/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py index d9763f173a7..f385c2f2211 100644 --- a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py +++ b/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py @@ -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.""" diff --git a/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py b/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py index da5b5ac3867..ab14d914640 100644 --- a/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py +++ b/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py @@ -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" diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 794613942a1..c2dfb91591b 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -13,13 +13,18 @@ Verifies that: import json import os import sys +import threading from types import SimpleNamespace +from typing import Final from unittest.mock import patch +import httpx import pytest sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) +from litellm.proxy._types import SpecialHeaders # noqa: E402 # sys.path must be patched before importing litellm + # 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" @@ -1078,14 +1083,19 @@ class TestGetAuthHeader: assert result is None def test_oauth_token_uses_bearer_not_x_api_key(self): - """OAuth token (sk-ant-oat*) should return Authorization: Bearer, not x-api-key.""" + """OAuth token (sk-ant-oat*) should return Authorization: Bearer with the + mandatory oauth beta, not x-api-key.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo result = AnthropicModelInfo.get_auth_header(api_key=FAKE_OAUTH_TOKEN) - assert result == {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} + assert result == { + "authorization": f"Bearer {FAKE_OAUTH_TOKEN}", + "anthropic-beta": "oauth-2025-04-20", + } def test_oauth_token_from_env_uses_bearer(self): - """OAuth token in ANTHROPIC_API_KEY env var should return Authorization: Bearer.""" + """OAuth token in ANTHROPIC_API_KEY env var should return Authorization: Bearer + with the mandatory oauth beta.""" from unittest.mock import patch as mock_patch from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -1096,7 +1106,10 @@ class TestGetAuthHeader: clear=True, ): result = AnthropicModelInfo.get_auth_header() - assert result == {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} + assert result == { + "authorization": f"Bearer {FAKE_OAUTH_TOKEN}", + "anthropic-beta": "oauth-2025-04-20", + } def test_custom_api_base_get_auth_header_uses_bearer(self): """Non-standard API key and custom api_base returns Bearer when use_bearer_for_custom_base=True.""" @@ -2175,3 +2188,1552 @@ def test_create_anthropic_model_list_response_empty(): assert response["has_more"] is False assert response["first_id"] is None assert response["last_id"] is None + + +# --------------------------------------------------------------------------- # +# Workload identity federation wiring (issue #28607) +# --------------------------------------------------------------------------- # + +FAKE_MINTED_TOKEN = "sk-ant-oat01-wif-minted-token-for-testing-abc123" + +ANTHROPIC_ENV_VARS = ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_API_BASE", + "ANTHROPIC_BASE_URL", + "ANTHROPIC_FEDERATION_RULE_ID", + "ANTHROPIC_ORGANIZATION_ID", + "ANTHROPIC_SERVICE_ACCOUNT_ID", + "ANTHROPIC_WORKSPACE_ID", + "ANTHROPIC_IDENTITY_TOKEN_FILE", + "ANTHROPIC_IDENTITY_TOKEN", + "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", +) + +WIF_ENV = { + "ANTHROPIC_FEDERATION_RULE_ID": "fdrl_wire1", + "ANTHROPIC_ORGANIZATION_ID": "org-wire-1", + "ANTHROPIC_IDENTITY_TOKEN": "inline-wire-jwt", +} + +PROXY_CREDENTIAL_HEADER_NAMES = sorted(SpecialHeaders.litellm_credential_header_names()) + + +class RecordingPoster: + def __init__(self, response): + self.requests = [] + self.thread_ids = [] + self._response = response + + def post(self, url, *, content, headers, timeout): + self.requests.append((url, content, dict(headers))) + self.thread_ids.append(threading.get_ident()) + return self._response + + +@pytest.fixture +def clean_anthropic_env(monkeypatch): + for name in ANTHROPIC_ENV_VARS: + monkeypatch.delenv(name, raising=False) + + +@pytest.fixture +def wif_engine(monkeypatch, clean_anthropic_env): + """Route the wiring's WIF tier through a fresh engine (never the module + singleton, to avoid cross-test cache pollution) and count its consultations.""" + import httpx + + from litellm.llms.anthropic import common_utils as anthropic_common_utils + from litellm.llms.anthropic.wif import get_anthropic_wif_token + from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + + poster = RecordingPoster( + httpx.Response( + 200, + json={"access_token": FAKE_MINTED_TOKEN, "token_type": "Bearer", "expires_in": 3600}, + ) + ) + engine = JwtBearerTokenExchangeEngine(poster=poster) + calls = [] + + def with_injected_engine(litellm_params, api_base, model): + calls.append(model) + return get_anthropic_wif_token(litellm_params, api_base, model, engine) + + monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", with_injected_engine) + return poster, calls + + +@pytest.fixture +def wif_async_engine(monkeypatch, clean_anthropic_env): + """Route both WIF facades through one fresh engine; the poster records the + thread each exchange ran on and sync-facade consultations are counted so + async tests can prove the mint went through the async seam, off the loop.""" + import httpx + + from litellm.llms.anthropic import common_utils as anthropic_common_utils + from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token + from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + + poster = RecordingPoster( + httpx.Response( + 200, + json={"access_token": FAKE_MINTED_TOKEN, "token_type": "Bearer", "expires_in": 3600}, + ) + ) + engine = JwtBearerTokenExchangeEngine(poster=poster) + 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) + return poster, sync_calls + + +def _validate_chat_environment(api_key=None): + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + return AnthropicModelInfo().validate_environment( + headers={}, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=api_key, + api_base=None, + ) + + +class TestWifTierPrecedence: + """WIF is the LOWEST credential tier: any api_key / auth_token source must + win without the engine ever being consulted.""" + + def _set_wif_env(self, monkeypatch): + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + + def test_explicit_api_key_beats_wif(self, monkeypatch, wif_engine): + poster, calls = wif_engine + self._set_wif_env(monkeypatch) + + headers = _validate_chat_environment(api_key=FAKE_REGULAR_KEY) + + assert headers["x-api-key"] == FAKE_REGULAR_KEY + assert calls == [] + assert poster.requests == [] + + def test_api_key_env_beats_wif(self, monkeypatch, wif_engine): + poster, calls = wif_engine + self._set_wif_env(monkeypatch) + monkeypatch.setenv("ANTHROPIC_API_KEY", FAKE_REGULAR_KEY) + + headers = _validate_chat_environment() + + assert headers["x-api-key"] == FAKE_REGULAR_KEY + assert calls == [] + assert poster.requests == [] + + def test_auth_token_env_beats_wif(self, monkeypatch, wif_engine): + poster, calls = wif_engine + self._set_wif_env(monkeypatch) + monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", FAKE_AUTH_TOKEN) + + headers = _validate_chat_environment() + + assert headers["authorization"] == f"Bearer {FAKE_AUTH_TOKEN}" + assert calls == [] + assert poster.requests == [] + + def test_wif_alone_mints_once(self, monkeypatch, wif_engine): + poster, calls = wif_engine + self._set_wif_env(monkeypatch) + + headers = _validate_chat_environment() + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert calls == ["claude-sonnet-4-5"] + assert len(poster.requests) == 1 + assert poster.requests[0][0] == "https://api.anthropic.com/v1/oauth/token" + + +class TestWifZeroBehaviorChange: + def test_unconfigured_raises_same_authentication_error(self, clean_anthropic_env): + """No WIF config and no keys: same AuthenticationError as today (message + extended, type and provider identical).""" + import litellm + + with pytest.raises(litellm.AuthenticationError) as exc_info: + _validate_chat_environment() + + assert exc_info.value.llm_provider == "anthropic" + assert "ANTHROPIC_API_KEY" in exc_info.value.message + assert "ANTHROPIC_AUTH_TOKEN" in exc_info.value.message + assert "ANTHROPIC_FEDERATION_RULE_ID" in exc_info.value.message + assert "ANTHROPIC_ORGANIZATION_ID" in exc_info.value.message + assert "ANTHROPIC_SERVICE_ACCOUNT_ID" in exc_info.value.message + assert "ANTHROPIC_IDENTITY_TOKEN_FILE" in exc_info.value.message + + +class TestWifHeaderContract: + def test_minted_token_headers(self, monkeypatch, wif_engine): + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + + headers = _validate_chat_environment() + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert "oauth-2025-04-20" in headers["anthropic-beta"] + assert "x-api-key" not in headers + assert "anthropic-dangerous-direct-browser-access" not in headers + + def test_consumer_oat_key_keeps_dangerous_header(self, clean_anthropic_env): + """Regression: user-supplied consumer sk-ant-oat keys keep today's behavior.""" + headers = _validate_chat_environment(api_key=FAKE_OAUTH_TOKEN) + + assert headers["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + assert headers["anthropic-dangerous-direct-browser-access"] == "true" + assert "oauth-2025-04-20" in headers["anthropic-beta"] + + +class TestMergeAnthropicBetaHeaders: + """The Skills surface accepted a list-valued anthropic-beta before it shared this helper, + so the helper has to keep taking one: .split() on a list is an AttributeError.""" + + def test_list_valued_existing_header_is_merged(self): + from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers + + assert merge_anthropic_beta_headers(["skills-2025-10-02", "files-api-2025-04-14"], "oauth-2025-04-20") == ( + "files-api-2025-04-14,oauth-2025-04-20,skills-2025-10-02" + ) + + def test_list_and_comma_string_forms_agree(self): + from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers + + as_list = merge_anthropic_beta_headers(["a", "b"], "c") + as_string = merge_anthropic_beta_headers("a,b", "c") + assert as_list == as_string == "a,b,c" + + def test_skills_validate_environment_accepts_a_list_header(self, monkeypatch): + """End of the regression: the Skills surface itself must not raise on the list form.""" + from litellm.llms.anthropic.skills.transformation import AnthropicSkillsConfig + + monkeypatch.setenv("ANTHROPIC_API_KEY", FAKE_REGULAR_KEY) + + headers = AnthropicSkillsConfig().validate_environment( + headers={"anthropic-beta": ["files-api-2025-04-14"]}, + litellm_params=None, + ) + + assert "files-api-2025-04-14" in headers["anthropic-beta"] + assert isinstance(headers["anthropic-beta"], str) + + +class TestWifServerOwnedAuthHeaderStrip: + """A WIF-minted token must never ride alongside a caller-supplied credential + header, but that stripping must fire only when a mint actually happened.""" + + def test_mint_strips_caller_supplied_x_api_key(self, monkeypatch, wif_engine): + """Security regression: without the strip, a caller-forwarded x-api-key + would sit next to the server-minted Authorization on the outgoing request.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + poster, _ = wif_engine + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + caller_key = "sk-ant-CALLER-SUPPLIED" + + headers = AnthropicModelInfo().validate_environment( + headers={"x-api-key": caller_key}, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert "x-api-key" not in headers + assert caller_key not in headers.values() + assert len(poster.requests) == 1 + + def test_server_owned_set_is_every_proxy_credential_header(self): + """The strip list must track the proxy's own key-header list, not a hand-rolled + pair: every header user_api_key_auth accepts a LiteLLM key in must be here.""" + from litellm.llms.anthropic.common_utils import _SERVER_OWNED_AUTH_HEADERS + + assert _SERVER_OWNED_AUTH_HEADERS == SpecialHeaders.litellm_credential_header_names() + assert {"x-litellm-api-key", "api-key", "x-goog-api-key"} < _SERVER_OWNED_AUTH_HEADERS + + @pytest.mark.parametrize("header_name", PROXY_CREDENTIAL_HEADER_NAMES) + def test_mint_strips_every_proxy_credential_header(self, monkeypatch, wif_engine, header_name): + """A LiteLLM virtual key arrives in any of the proxy's accepted key headers; once + a mint happened none of them may reach Anthropic in any header slot.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + caller_key = "sk-litellm-CALLER-VIRTUAL-KEY" + + headers = AnthropicModelInfo().validate_environment( + headers={header_name.title(): caller_key, "user-agent": "caller/1.0"}, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert header_name == "authorization" or header_name not in {name.lower() for name in headers} + assert all(caller_key not in value for value in headers.values()) + assert headers["user-agent"] == "caller/1.0" + + @pytest.mark.parametrize("header_name", PROXY_CREDENTIAL_HEADER_NAMES) + def test_skills_surface_strips_caller_credentials_too(self, monkeypatch, wif_engine, header_name): + """Skills builds its own headers as well; every minting surface needs the same strip.""" + from litellm.llms.anthropic.skills.transformation import AnthropicSkillsConfig + + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + caller_key = "sk-litellm-CALLER-VIRTUAL-KEY" + + headers = AnthropicSkillsConfig().validate_environment( + headers={header_name.title(): caller_key, "user-agent": "caller/1.0"}, + litellm_params=None, + ) + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert header_name == "authorization" or header_name not in {name.lower() for name in headers} + assert all(caller_key not in value for value in headers.values()) + assert headers["user-agent"] == "caller/1.0" + + def test_passthrough_honors_a_case_variant_caller_key_instead_of_minting(self, monkeypatch, wif_engine): + """The passthrough surface hands the caller's own credential upstream rather than minting. + That check was case-sensitive, so X-Api-Key slipped past it and the caller's key would have + travelled beside a minted Bearer.""" + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + poster, _ = wif_engine + caller_key = "sk-ant-CALLER-SUPPLIED" + + headers, _ = AnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={"X-Api-Key": caller_key}, + model="claude-sonnet-4-5", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + assert headers["X-Api-Key"] == caller_key + assert "authorization" not in {name.lower() for name in headers} + assert len(poster.requests) == 0 + + @pytest.mark.parametrize("header_name", PROXY_CREDENTIAL_HEADER_NAMES) + def test_batches_surface_strips_caller_credentials_too(self, monkeypatch, wif_engine, header_name): + """Batches builds its own headers on the create path, so it needs the same strip: the + handler's retrieve path passes none, but this entry point takes the caller's.""" + from litellm.llms.anthropic.batches.transformation import AnthropicBatchesConfig + + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + caller_key = "sk-litellm-CALLER-VIRTUAL-KEY" + + headers = AnthropicBatchesConfig().validate_environment( + headers={header_name.title(): caller_key, "user-agent": "caller/1.0"}, + model="claude-sonnet-4-5", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert header_name == "authorization" or header_name not in {name.lower() for name in headers} + assert all(caller_key not in value for value in headers.values()) + assert headers["user-agent"] == "caller/1.0" + + @pytest.mark.parametrize("header_name", PROXY_CREDENTIAL_HEADER_NAMES) + def test_files_surface_strips_caller_credentials_too(self, monkeypatch, wif_engine, header_name): + """The files surface builds its own headers, so it needs the same strip the chat surface + has: without it a minted federation Bearer travels beside the caller's own credential.""" + from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig + + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + caller_key = "sk-litellm-CALLER-VIRTUAL-KEY" + + headers = AnthropicFilesConfig().validate_environment( + headers={header_name.title(): caller_key, "user-agent": "caller/1.0"}, + model="claude-sonnet-4-5", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert header_name == "authorization" or header_name not in {name.lower() for name in headers} + assert all(caller_key not in value for value in headers.values()) + assert headers["user-agent"] == "caller/1.0" + + def test_no_mint_preserves_caller_supplied_authorization(self, monkeypatch, clean_anthropic_env): + """No-regression: LiteLLM deliberately lets a caller-forwarded credential + header ride alongside a statically configured ANTHROPIC_API_KEY, because the + two occupy different header slots (x-api-key vs authorization) when the key + isn't OAuth-shaped. The strip must stay conditional on an actual WIF mint.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + monkeypatch.setenv("ANTHROPIC_API_KEY", FAKE_REGULAR_KEY) + caller_authorization = "Bearer caller-forwarded-downstream-token" + + headers = AnthropicModelInfo().validate_environment( + headers={"authorization": caller_authorization}, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + assert headers["x-api-key"] == FAKE_REGULAR_KEY + assert headers["authorization"] == caller_authorization + + +class TestWifResolvedApiKeyThreading: + """Regression for the resolved_api_key local (formerly a rebind of the api_key + parameter): a minted token must reach the outgoing headers on both the sync + validate_environment path and the async aget_auth_header path, never a stale + None left over from the original unresolved parameter.""" + + def test_validate_environment_carries_minted_token(self, monkeypatch, wif_engine): + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + + headers = _validate_chat_environment() + + assert "authorization" in headers + assert headers["authorization"] not in (None, "Bearer None") + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + + @pytest.mark.asyncio + async def test_aget_auth_header_carries_minted_token(self, monkeypatch, wif_async_engine): + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + + result = await AnthropicModelInfo.aget_auth_header(allow_workload_identity=True) + + assert result is not None + assert result["authorization"] not in (None, "Bearer None") + assert result["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + + +class TestGetAuthHeaderBetas: + def test_oat_branch_carries_oauth_beta(self, clean_anthropic_env): + """The pre-existing bug: the oat branch returned a bare Bearer without the + mandatory oauth beta.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + result = AnthropicModelInfo.get_auth_header(api_key=FAKE_OAUTH_TOKEN) + + assert result == { + "authorization": f"Bearer {FAKE_OAUTH_TOKEN}", + "anthropic-beta": "oauth-2025-04-20", + } + + def test_wif_fallback_returns_bearer_and_beta(self, monkeypatch, wif_engine): + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + poster, calls = wif_engine + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + + result = AnthropicModelInfo.get_auth_header(allow_workload_identity=True) + + assert result == { + "authorization": f"Bearer {FAKE_MINTED_TOKEN}", + "anthropic-beta": "oauth-2025-04-20", + } + assert len(poster.requests) == 1 + + def test_no_credentials_still_returns_none(self, clean_anthropic_env): + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + assert AnthropicModelInfo.get_auth_header() is None + + +class TestFilesBatchesBetaMerge: + """Regression for the anthropic-beta clobber (files) and drop (batches): + a Bearer oat auth header must keep the oauth beta AND gain the surface beta.""" + + def test_files_merges_oauth_and_files_betas(self, clean_anthropic_env): + from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig + + headers = AnthropicFilesConfig().validate_environment( + headers={}, + model="", + messages=[], + optional_params={}, + litellm_params={}, + api_key=FAKE_OAUTH_TOKEN, + ) + + assert headers["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + betas = set(headers["anthropic-beta"].split(",")) + assert {"oauth-2025-04-20", "files-api-2025-04-14"} <= betas + + def test_batches_merges_oauth_and_batches_betas(self, clean_anthropic_env): + from litellm.llms.anthropic.batches.transformation import AnthropicBatchesConfig + + headers = AnthropicBatchesConfig().validate_environment( + headers={}, + model="", + messages=[], + optional_params={}, + litellm_params={}, + api_key=FAKE_OAUTH_TOKEN, + ) + + assert headers["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + betas = set(headers["anthropic-beta"].split(",")) + assert {"oauth-2025-04-20", "message-batches-2024-09-24"} <= betas + + def test_files_preserves_caller_supplied_beta(self, clean_anthropic_env): + """Regression: files did a two-way merge that dropped the client's own + anthropic-beta; it must three-way merge exactly like batches.""" + from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig + + headers = AnthropicFilesConfig().validate_environment( + headers={"anthropic-beta": "context-1m-2025-08-07"}, + model="", + messages=[], + optional_params={}, + litellm_params={}, + api_key=FAKE_OAUTH_TOKEN, + ) + + betas = set(headers["anthropic-beta"].split(",")) + assert {"context-1m-2025-08-07", "oauth-2025-04-20", "files-api-2025-04-14"} <= betas + + +class TestMessagesEnvAuthBetaMerge: + def test_client_beta_survives_env_auth_injection(self, monkeypatch, clean_anthropic_env): + """Regression: headers.update(auth_header) silently clobbered the client's + anthropic-beta on the native /v1/messages route.""" + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + monkeypatch.setenv("ANTHROPIC_API_KEY", FAKE_OAUTH_TOKEN) + + headers, _ = AnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={"anthropic-beta": "context-1m-2025-08-07"}, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + ) + + assert headers["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + betas = set(headers["anthropic-beta"].split(",")) + assert {"context-1m-2025-08-07", "oauth-2025-04-20"} <= betas + + +WIF_PARAMS_ONLY = { + "anthropic_federation_rule_id": "fdrl_params", + "anthropic_organization_id": "org-params", + "anthropic_identity_token": "oidc/env/WIF_PARAMS_TEST_TOKEN", +} + + +class TestWifLitellmParamsPlumbing: + """Per-deployment anthropic_* litellm_params must reach the WIF tier on every + surface that has them, not only chat.""" + + @pytest.fixture(autouse=True) + def _inline_identity_token(self, monkeypatch): + monkeypatch.setenv("WIF_PARAMS_TEST_TOKEN", "params-jwt") + + def test_files_mints_from_litellm_params(self, wif_engine): + from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig + + poster, _ = wif_engine + headers = AnthropicFilesConfig().validate_environment( + headers={}, + model="", + messages=[], + optional_params={}, + litellm_params=dict(WIF_PARAMS_ONLY), + ) + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert len(poster.requests) == 1 + + def test_batches_mints_from_litellm_params(self, wif_engine): + from litellm.llms.anthropic.batches.transformation import AnthropicBatchesConfig + + poster, _ = wif_engine + headers = AnthropicBatchesConfig().validate_environment( + headers={}, + model="", + messages=[], + optional_params={}, + litellm_params=dict(WIF_PARAMS_ONLY), + ) + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert len(poster.requests) == 1 + + def test_skills_mints_from_litellm_params(self, wif_engine): + from litellm.llms.anthropic.skills.transformation import AnthropicSkillsConfig + from litellm.types.router import GenericLiteLLMParams + + poster, _ = wif_engine + headers = AnthropicSkillsConfig().validate_environment( + headers={}, + litellm_params=GenericLiteLLMParams( + anthropic_federation_rule_id=WIF_PARAMS_ONLY["anthropic_federation_rule_id"], + anthropic_organization_id=WIF_PARAMS_ONLY["anthropic_organization_id"], + anthropic_identity_token=WIF_PARAMS_ONLY["anthropic_identity_token"], + ), + ) + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert len(poster.requests) == 1 + + def test_messages_mints_from_litellm_params(self, wif_engine): + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + poster, _ = wif_engine + headers, _ = AnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params=dict(WIF_PARAMS_ONLY), + ) + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert len(poster.requests) == 1 + + +class TestWifTokenUrlParity: + """Both credential tiers must derive the SAME clean token URL from any form of + the deployment base; a mismatch also duplicates mints because token_url is in + the engine cache key.""" + + @pytest.mark.parametrize( + "configured_base", + [ + "https://gw.example.com", + "https://gw.example.com/", + "https://gw.example.com/v1/messages", + "https://gw.example.com/v1/messages/", + ], + ) + def test_both_tiers_share_one_clean_token_url(self, monkeypatch, wif_engine, configured_base): + # This is about deriving one URL from many spellings of the same base, not about which + # hosts an operator trusts with org-scoped credentials, so the private host is allowlisted. + monkeypatch.setenv("LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", "gw.example.com") + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + poster, _ = wif_engine + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + + AnthropicModelInfo().validate_environment( + headers={}, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"api_base": configured_base}, + api_key=None, + api_base=None, + ) + AnthropicModelInfo.get_auth_header(api_base=configured_base, allow_workload_identity=True) + + assert [url for (url, _, _) in poster.requests] == ["https://gw.example.com/v1/oauth/token"] + + +class TestWifAsyncSeam: + """Async callers must resolve the WIF tier through the async facade so a cold + mint never blocks the event loop.""" + + @pytest.mark.asyncio + async def test_aget_auth_header_runs_exchange_off_event_loop(self, monkeypatch, wif_async_engine): + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + poster, sync_calls = wif_async_engine + for name, value in WIF_ENV.items(): + monkeypatch.setenv(name, value) + + result = await AnthropicModelInfo.aget_auth_header(allow_workload_identity=True) + + assert result == { + "authorization": f"Bearer {FAKE_MINTED_TOKEN}", + "anthropic-beta": "oauth-2025-04-20", + } + assert sync_calls == [] + assert poster.thread_ids == [poster.thread_ids[0]] + assert poster.thread_ids[0] != threading.get_ident() + + @pytest.mark.asyncio + async def test_avalidate_messages_environment_mints_off_loop(self, wif_async_engine, monkeypatch): + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + poster, sync_calls = wif_async_engine + monkeypatch.setenv("WIF_PARAMS_TEST_TOKEN", "params-jwt") + + headers, _ = await AnthropicMessagesConfig().avalidate_anthropic_messages_environment( + headers={}, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params=dict(WIF_PARAMS_ONLY), + ) + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert sync_calls == [] + assert poster.thread_ids[0] != threading.get_ident() + + @pytest.mark.asyncio + async def test_avalidate_delegates_to_subclass_sync_override(self): + """A provider subclass that only overrides the sync method must keep its + behavior when the handler goes through the async variant.""" + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + class MarkerConfig(AnthropicMessagesConfig): + def validate_anthropic_messages_environment( + self, + headers, + model, + messages, + optional_params, + litellm_params, + api_key=None, + api_base=None, + ): + return {"x-marker": "sync"}, api_base + + headers, api_base = await MarkerConfig().avalidate_anthropic_messages_environment( + headers={}, + model="claude-sonnet-4-5", + messages=[], + optional_params={}, + litellm_params={}, + api_base="https://marker.example.com", + ) + + assert headers == {"x-marker": "sync"} + assert api_base == "https://marker.example.com" + + @pytest.mark.asyncio + async def test_base_default_avalidate_delegates_to_sync(self): + from litellm.llms.base_llm.anthropic_messages.transformation import ( + BaseAnthropicMessagesConfig, + ) + + class SyncOnlyConfig(BaseAnthropicMessagesConfig): + def validate_anthropic_messages_environment( + self, + headers, + model, + messages, + optional_params, + litellm_params, + api_key=None, + api_base=None, + ): + return {"x-sync-only": "1"}, api_base + + def get_complete_url(self, api_base, api_key, model, optional_params, litellm_params, stream=None): + return api_base or "" + + def get_supported_anthropic_messages_params(self, model): + return [] + + def transform_anthropic_messages_request( + self, model, messages, anthropic_messages_optional_request_params, litellm_params, headers + ): + return {} + + def transform_anthropic_messages_response(self, model, raw_response, logging_obj): + raise NotImplementedError + + headers, _ = await SyncOnlyConfig().avalidate_anthropic_messages_environment( + headers={}, + model="m", + messages=[], + optional_params={}, + litellm_params={}, + ) + + assert headers == {"x-sync-only": "1"} + + +class TestWifRespxEndToEnd: + def test_completion_mints_and_never_leaks_config(self, monkeypatch, tmp_path, clean_anthropic_env): + """Drives the REAL kwargs funnel through litellm.completion: the mint hits + /v1/oauth/token, the data plane carries the minted Bearer + oauth beta, and + NONE of the six anthropic_* keys leak into the /v1/messages body.""" + import httpx + import respx + + import litellm + from litellm.llms.anthropic import common_utils as anthropic_common_utils + from litellm.llms.anthropic.wif import get_anthropic_wif_token + from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "anthropic_key", None) + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) + token_file = tmp_path / "identity-token" + token_file.write_text("e2e-oidc-assertion", encoding="utf-8") + + engine = JwtBearerTokenExchangeEngine() + monkeypatch.setattr( + anthropic_common_utils, + "get_anthropic_wif_token", + lambda litellm_params, api_base, model: get_anthropic_wif_token(litellm_params, api_base, model, engine), + ) + + wif_kwarg_names: Final = ( + "anthropic_federation_rule_id", + "anthropic_organization_id", + "anthropic_service_account_id", + "anthropic_workspace_id", + "anthropic_identity_token_file", + "anthropic_identity_token", + ) + anthropic_response = { + "id": "msg_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "Hello from WIF"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 5, "output_tokens": 4}, + } + + with respx.mock: + token_route = respx.post("https://api.anthropic.com/v1/oauth/token").mock( + return_value=httpx.Response( + 200, + json={"access_token": FAKE_MINTED_TOKEN, "token_type": "Bearer", "expires_in": 3600}, + ) + ) + messages_route = respx.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response(200, json=anthropic_response) + ) + response = litellm.completion( + model="anthropic/claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + anthropic_federation_rule_id="fdrl_e2e", + anthropic_organization_id="org-e2e", + anthropic_service_account_id="svcacct_e2e", + anthropic_workspace_id="wrkspc_e2e", + anthropic_identity_token_file=str(token_file), + anthropic_identity_token="oidc/env/UNUSED_FALLBACK", + ) + + assert response.choices[0].message.content == "Hello from WIF" + assert token_route.call_count == 1 + exchange_body = json.loads(token_route.calls[0].request.content) + assert exchange_body["assertion"] == "e2e-oidc-assertion" + assert exchange_body["federation_rule_id"] == "fdrl_e2e" + + data_request = messages_route.calls[0].request + assert data_request.headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert "oauth-2025-04-20" in data_request.headers["anthropic-beta"] + assert "x-api-key" not in data_request.headers + assert "anthropic-dangerous-direct-browser-access" not in data_request.headers + data_body = json.loads(data_request.content) + for key in wif_kwarg_names: + assert key not in data_body + + def test_completion_with_trailing_slash_api_base_mints_at_clean_token_url( + self, monkeypatch, tmp_path, clean_anthropic_env + ): + """Regression: a trailing-slash api_base defeated the endswith check in + main.py AND the removesuffix surgery, sending the exchange POST to + .../v1/messages/v1/oauth/token (404).""" + import httpx + import respx + + import litellm + from litellm.llms.anthropic import common_utils as anthropic_common_utils + from litellm.llms.anthropic.wif import get_anthropic_wif_token + from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "anthropic_key", None) + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) + token_file = tmp_path / "identity-token" + token_file.write_text("e2e-oidc-assertion", encoding="utf-8") + + engine = JwtBearerTokenExchangeEngine() + monkeypatch.setattr( + anthropic_common_utils, + "get_anthropic_wif_token", + lambda litellm_params, api_base, model: get_anthropic_wif_token(litellm_params, api_base, model, engine), + ) + + anthropic_response = { + "id": "msg_02", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "Hello from WIF"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 5, "output_tokens": 4}, + } + + with respx.mock: + token_route = respx.post("https://api.anthropic.com/v1/oauth/token").mock( + return_value=httpx.Response( + 200, + json={"access_token": FAKE_MINTED_TOKEN, "token_type": "Bearer", "expires_in": 3600}, + ) + ) + messages_route = respx.post(url__regex=r"https://api\.anthropic\.com/v1/messages.*").mock( + return_value=httpx.Response(200, json=anthropic_response) + ) + response = litellm.completion( + model="anthropic/claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + api_base="https://api.anthropic.com/v1/messages/", + anthropic_federation_rule_id="fdrl_e2e", + anthropic_organization_id="org-e2e", + anthropic_identity_token_file=str(token_file), + ) + + assert response.choices[0].message.content == "Hello from WIF" + assert token_route.call_count == 1 + assert messages_route.calls[0].request.headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + + def test_get_auth_header_with_litellm_params_mints_via_real_engine( + self, monkeypatch, tmp_path, clean_anthropic_env + ): + import httpx + import respx + + from litellm.llms.anthropic import common_utils as anthropic_common_utils + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + from litellm.llms.anthropic.wif import get_anthropic_wif_token + from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) + token_file = tmp_path / "identity-token" + token_file.write_text("e2e-oidc-assertion", encoding="utf-8") + + engine = JwtBearerTokenExchangeEngine() + monkeypatch.setattr( + anthropic_common_utils, + "get_anthropic_wif_token", + lambda litellm_params, api_base, model: get_anthropic_wif_token(litellm_params, api_base, model, engine), + ) + + with respx.mock: + token_route = respx.post("https://api.anthropic.com/v1/oauth/token").mock( + return_value=httpx.Response( + 200, + json={"access_token": FAKE_MINTED_TOKEN, "token_type": "Bearer", "expires_in": 3600}, + ) + ) + result = AnthropicModelInfo.get_auth_header( + allow_workload_identity=True, + litellm_params={ + "anthropic_federation_rule_id": "fdrl_e2e", + "anthropic_organization_id": "org-e2e", + "anthropic_identity_token_file": str(token_file), + }, + ) + + assert result == { + "authorization": f"Bearer {FAKE_MINTED_TOKEN}", + "anthropic-beta": "oauth-2025-04-20", + } + assert token_route.call_count == 1 + exchange_body = json.loads(token_route.calls[0].request.content) + assert exchange_body["federation_rule_id"] == "fdrl_e2e" + + +class TestWifProviderAllowlist: + """A federation token is an Anthropic-org credential, and the exchange POSTs the workload's OIDC + assertion to the deployment's own api_base host. Providers that subclass the Anthropic config for + their own endpoints must therefore never reach the WIF tier, even when it is configured purely + through ANTHROPIC_* environment variables.""" + + @staticmethod + def _env_only_wif(monkeypatch) -> None: # noqa: D401 + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_prod") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-prod-uuid") + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "oidc/env/WIF_TEST_JWT") + monkeypatch.setenv("WIF_TEST_JWT", "jwt-assertion-value") + + def test_vertex_anthropic_never_mints_or_sends_the_assertion(self, monkeypatch, wif_engine): + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( + VertexAIAnthropicConfig, + ) + + import litellm + + poster, calls = wif_engine + self._env_only_wif(monkeypatch) + + with pytest.raises(litellm.AuthenticationError): + VertexAIAnthropicConfig().validate_environment( + headers={}, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + api_key=None, + api_base="https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5", + ) + + assert calls == [] + assert poster.requests == [] + + def test_anthropic_itself_still_mints(self, monkeypatch, wif_engine): + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + poster, calls = wif_engine + self._env_only_wif(monkeypatch) + + headers = AnthropicConfig().validate_environment( + headers={}, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + assert len(poster.requests) == 1 + + def test_auth_header_facade_defaults_to_refusing_to_mint(self, monkeypatch, clean_anthropic_env): + """The facade is reachable from provider code that has nothing to do with Anthropic, so a + caller must state that it authenticates against Anthropic's own API.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + self._env_only_wif(monkeypatch) + + assert AnthropicModelInfo.get_auth_header(None) is None + assert AnthropicModelInfo.get_auth_header(None, allow_workload_identity=False) is None + + def test_eligibility_is_not_inherited_by_a_new_subclass(self): + """A provider added later by subclassing the Anthropic config must not inherit the right to + mint an Anthropic-org credential against its own host.""" + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + from litellm.llms.anthropic.common_utils import config_allows_workload_identity + + class NewCompatibleProvider(AnthropicConfig): + pass + + assert config_allows_workload_identity(AnthropicConfig()) is True + assert config_allows_workload_identity(NewCompatibleProvider()) is False + + def test_model_discovery_gates_on_the_instance(self, monkeypatch, wif_engine): + """get_models is inherited, so it must consult the instance rather than trusting its caller.""" + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( + VertexAIAnthropicConfig, + ) + + poster, calls = wif_engine + self._env_only_wif(monkeypatch) + + with pytest.raises(ValueError, match="ANTHROPIC_API_KEY"): + VertexAIAnthropicConfig().get_models( + api_base="https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" + ) + + assert poster.requests == [] + + +def _models_page_response(page: dict, status_code: int = 200): + import httpx + + return httpx.Response(status_code, json=page, request=httpx.Request("GET", "https://api.anthropic.com/v1/models")) + + +class RecordingModelsClient: + """Records every call and answers with the queued responses in order, cycling the last one + once exhausted so a runaway pagination loop degrades to a repeated page rather than an + IndexError, letting the page-cap test observe the cap firing instead of a test bug.""" + + def __init__(self, pages: list[dict] | None = None, responses=None): + self.calls = [] + self._responses = responses if responses is not None else [_models_page_response(page) for page in pages] + + def get(self, url, headers=None, params=None, follow_redirects=None, timeout=None): + self.calls.append(SimpleNamespace(url=url, headers=headers, params=params, follow_redirects=follow_redirects)) + index = min(len(self.calls) - 1, len(self._responses) - 1) + return self._responses[index] + + +class TestModelDiscovery: + """AnthropicModelInfo.get_models / discover_models: pagination, redirect refusal, and + sanitized errors on the upstream Anthropic /v1/models call itself (issue #28607 gap: a + WIF source configured in litellm_params, rather than the environment, could not + discover).""" + + @pytest.mark.parametrize( + "configured_base", ["https://api.anthropic.com/v1", "https://api.anthropic.com/v1/messages"] + ) + def test_discovery_does_not_double_the_version_segment(self, monkeypatch, clean_anthropic_env, configured_base): + """Regression: /v1/models is appended here, so a base an operator already wrote as + .../v1 (or the chat URL they copied) would be asked for /v1/v1/models and 404.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + monkeypatch.setenv("ANTHROPIC_API_KEY", FAKE_REGULAR_KEY) + client = RecordingModelsClient([{"data": [{"id": "claude-a"}], "has_more": False, "last_id": "claude-a"}]) + monkeypatch.setattr("litellm.module_level_client", client) + + models = AnthropicModelInfo().get_models(api_base=configured_base) + + assert models == ["anthropic/claude-a"] + assert client.calls[0].url == "https://api.anthropic.com/v1/models" + + def test_get_models_paginates_via_has_more_and_last_id(self, monkeypatch, clean_anthropic_env): + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + monkeypatch.setenv("ANTHROPIC_API_KEY", FAKE_REGULAR_KEY) + client = RecordingModelsClient( + [ + {"data": [{"id": "claude-a"}, {"id": "claude-b"}], "has_more": True, "last_id": "claude-b"}, + {"data": [{"id": "claude-c"}], "has_more": False, "last_id": "claude-c"}, + ] + ) + monkeypatch.setattr("litellm.module_level_client", client) + + models = AnthropicModelInfo().get_models(api_base="https://api.anthropic.com") + + assert models == ["anthropic/claude-a", "anthropic/claude-b", "anthropic/claude-c"] + assert len(client.calls) == 2 + assert client.calls[0].url == "https://api.anthropic.com/v1/models" + assert client.calls[1].url == "https://api.anthropic.com/v1/models?after_id=claude-b" + + def test_paginated_fetch_survives_the_real_http_client(self, monkeypatch, clean_anthropic_env): + """Regression: the second page is fetched through the real HTTPHandler, which merges the + URL's query string into the params mapping by mutating it. Handing that client a + read-only mapping raised AttributeError, so discovery blew up for any org holding more + models than one page, while the stubbed client here never exercised the mutation.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + monkeypatch.setenv("ANTHROPIC_API_KEY", FAKE_REGULAR_KEY) + requested: Final = [] # mutable-ok: a test spy recording the URLs the client was asked for + + def respond(request: httpx.Request) -> httpx.Response: + requested.append(str(request.url)) + first: Final = "after_id" not in request.url.params + return httpx.Response( + 200, + json={ + "data": [{"id": "claude-a"}] if first else [{"id": "claude-b"}], + "has_more": first, + "last_id": "claude-a" if first else "claude-b", + }, + ) + + handler: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + monkeypatch.setattr("litellm.module_level_client", handler) + + models = AnthropicModelInfo().get_models(api_base="https://api.anthropic.com") + + assert models == ["anthropic/claude-a", "anthropic/claude-b"] + assert requested == [ + "https://api.anthropic.com/v1/models", + "https://api.anthropic.com/v1/models?after_id=claude-a", + ] + + def test_get_models_refuses_to_follow_redirects(self, monkeypatch, clean_anthropic_env): + """Only the configured api_base is validated, so a redirected /v1/models must not be + allowed to replay the credential to an unvalidated origin -- same rule already applied + to the WIF token exchange itself.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + monkeypatch.setenv("ANTHROPIC_API_KEY", FAKE_REGULAR_KEY) + client = RecordingModelsClient([{"data": [], "has_more": False, "last_id": None}]) + monkeypatch.setattr("litellm.module_level_client", client) + + AnthropicModelInfo().get_models(api_base="https://api.anthropic.com") + + assert client.calls[0].follow_redirects is False + + def test_get_models_page_cap_stops_a_runaway_has_more(self, monkeypatch, clean_anthropic_env): + from litellm.llms.anthropic.common_utils import ( + _MODEL_LIST_PAGE_CAP, + AnthropicModelInfo, + ) + + monkeypatch.setenv("ANTHROPIC_API_KEY", FAKE_REGULAR_KEY) + client = RecordingModelsClient([{"data": [{"id": "claude-loop"}], "has_more": True, "last_id": "claude-loop"}]) + monkeypatch.setattr("litellm.module_level_client", client) + + with pytest.raises(Exception, match="did not terminate"): + AnthropicModelInfo().get_models(api_base="https://api.anthropic.com") + + assert len(client.calls) == _MODEL_LIST_PAGE_CAP + + def test_get_models_error_is_sanitized_not_raw_response_text(self, monkeypatch, clean_anthropic_env): + """A failed discovery call must never echo the raw response body verbatim -- only the + structured error message, so an unrelated/oversized/reflected body is not surfaced.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + monkeypatch.setenv("ANTHROPIC_API_KEY", FAKE_REGULAR_KEY) + reflected_payload = "" * 50 + client = RecordingModelsClient( + responses=[ + _models_page_response( + { + "type": "error", + "error": { + "type": "authentication_error", + "message": "invalid x-api-key", + "reflected": reflected_payload, + }, + }, + status_code=401, + ) + ] + ) + monkeypatch.setattr("litellm.module_level_client", client) + + with pytest.raises(Exception, match="invalid x-api-key") as exc_info: # noqa: B017, PT011 # the callee raises a bare Exception; match pins the sanitized text + AnthropicModelInfo().get_models(api_base="https://api.anthropic.com") + + assert "invalid x-api-key" in str(exc_info.value) + assert reflected_payload not in str(exc_info.value) + + def test_discover_models_threads_litellm_params_into_wif(self, monkeypatch, wif_engine): + """The gap this phase fixes: get_models only ever saw api_key/api_base, so a WIF source + configured in litellm_params (rather than ANTHROPIC_* env vars) could not discover.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + poster, calls = wif_engine + client = RecordingModelsClient([{"data": [{"id": "claude-wif"}], "has_more": False, "last_id": None}]) + monkeypatch.setattr("litellm.module_level_client", client) + monkeypatch.setenv("DISC_JWT", "jwt-assertion-value") + + models = AnthropicModelInfo().discover_models( + litellm_params={ + "anthropic_federation_rule_id": "fdrl_disc", + "anthropic_organization_id": "org-disc", + "anthropic_identity_token": "oidc/env/DISC_JWT", + } + ) + + assert models == ["anthropic/claude-wif"] + assert len(poster.requests) == 1 + assert client.calls[0].headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" + + def test_discover_models_without_litellm_params_behaves_like_get_models(self, monkeypatch, clean_anthropic_env): + """No litellm_params (the wildcard-discovery call shape) must fall back to the + env-only resolution get_models has always used -- zero behavior change for that path.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + monkeypatch.setenv("ANTHROPIC_API_KEY", FAKE_REGULAR_KEY) + client = RecordingModelsClient([{"data": [{"id": "claude-env"}], "has_more": False, "last_id": None}]) + monkeypatch.setattr("litellm.module_level_client", client) + + models = AnthropicModelInfo().discover_models(litellm_params=None) + + assert models == ["anthropic/claude-env"] + assert client.calls[0].headers["x-api-key"] == FAKE_REGULAR_KEY + + def test_discover_models_explicit_api_key_beats_wif(self, monkeypatch, wif_engine): + """Same precedence discover_models must honor as every other Anthropic auth surface: + WIF is the lowest tier.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + poster, calls = wif_engine + client = RecordingModelsClient([{"data": [], "has_more": False, "last_id": None}]) + monkeypatch.setattr("litellm.module_level_client", client) + + AnthropicModelInfo().discover_models( + litellm_params={ + "api_key": FAKE_REGULAR_KEY, + "anthropic_federation_rule_id": "fdrl_disc", + "anthropic_organization_id": "org-disc", + "anthropic_identity_token": "oidc/env/DISC_JWT", + } + ) + + assert client.calls[0].headers["x-api-key"] == FAKE_REGULAR_KEY + assert calls == [] + assert poster.requests == [] + + +class TestWifExchangeTransportHardening: + def test_token_exchange_client_does_not_follow_redirects(self): + """Only the initial token URL is validated, so a 3xx must not be allowed to replay the + assertion to an origin that was never checked.""" + from litellm.llms.base_llm.auth.token_exchange import _HttpxSyncTokenPoster + + handler = _HttpxSyncTokenPoster()._handler_instance() + + assert handler.client.follow_redirects is False + + +class TestWifParamsAreNotClientSettable: + def test_every_minting_param_is_server_owned(self): + """Each of these selects which server-side secret is read, or the scope it is minted for. + The workspace id was once carved out here as inert; it is not. It is the scope of the + minted org credential, and the router merges request kwargs over deployment params, so a + caller who set it picked the scope instead of the administrator.""" + from litellm.proxy.auth.auth_utils import _SERVER_OWNED_WIF_UNCONDITIONAL_BANNED + from litellm.types.utils import anthropic_wif_litellm_params, openai_wif_litellm_params + + assert set(_SERVER_OWNED_WIF_UNCONDITIONAL_BANNED) == set(anthropic_wif_litellm_params) | set( + openai_wif_litellm_params + ) + + +class TestWifServerOwnedParamsAreUnconditional: + """The minting fields choose which server-side secret is read and, with api_base, where it goes, + so no client-side credential opt-in may re-enable them.""" + + @staticmethod + def _body(param: str) -> dict: + return {"model": "claude-sonnet-5", param: "oidc/env/SOME_SERVER_SECRET"} + + @pytest.mark.parametrize( + "param", + [ + "anthropic_identity_token", + "anthropic_identity_token_file", + "anthropic_federation_rule_id", + "anthropic_organization_id", + "anthropic_service_account_id", + # Phase 1 identity-source selection and its two variants' fields: each one + # selects a server-side secret or a destination (a signing key, a client + # secret, a token endpoint), so every one joins the same unconditional ban. + "anthropic_identity_source", + "anthropic_issuer_url", + "anthropic_issuer_subject", + "anthropic_issuer_audience", + "anthropic_issuer_ttl_seconds", + "anthropic_issuer_signing_key_ref", + "anthropic_keycloak_token_url", + "anthropic_keycloak_client_id", + "anthropic_keycloak_auth_method", + "anthropic_keycloak_client_secret_ref", + "anthropic_keycloak_scope", + ], + ) + def test_rejected_even_with_proxy_wide_opt_in(self, param: str): + from litellm.proxy.auth.auth_utils import is_request_body_safe + + with pytest.raises(ValueError, match="server-owned workload identity federation"): + is_request_body_safe( + request_body=self._body(param), + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="claude-sonnet-5", + ) + + def test_rejected_inside_nested_litellm_params(self): + from litellm.proxy.auth.auth_utils import is_request_body_safe + + with pytest.raises(ValueError, match="server-owned workload identity federation"): + is_request_body_safe( + request_body={"model": "claude-sonnet-5", "litellm_params": self._body("anthropic_identity_token")}, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="claude-sonnet-5", + ) + + def test_workspace_id_is_refused_from_a_request_body(self): + """Regression, proven live against Anthropic before this was closed: a caller-supplied + workspace id reached the token endpoint, which answered "workspace_id is not a well-formed + wrkspc_ tagged ID", i.e. the caller's value had become the scope of the minted credential. + router.py merges request kwargs OVER deployment params, so it also beat the configured one.""" + from litellm.proxy.auth.auth_utils import is_request_body_safe + + with pytest.raises(Exception, match="server-owned workload identity federation parameter"): + is_request_body_safe( + request_body={"model": "claude-sonnet-5", "anthropic_workspace_id": "wrkspc_abc"}, + general_settings={}, + llm_router=None, + model="claude-sonnet-5", + ) + + def test_refusal_points_bedrock_callers_at_their_own_spelling(self): + """Banning this spelling must not read as "no workspace selection anywhere": the Bedrock + Claude Platform route takes workspace_id/aws_workspace_id, neither of which is a + federation parameter, so the error names them.""" + from litellm.proxy.auth.auth_utils import is_request_body_safe + + with pytest.raises(Exception, match="workspace_id or aws_workspace_id"): + is_request_body_safe( + request_body={"model": "claude-sonnet-5", "anthropic_workspace_id": "wrkspc_abc"}, + general_settings={}, + llm_router=None, + model="claude-sonnet-5", + ) + + def test_bedrock_workspace_spellings_are_untouched(self): + """The Bedrock route's own spellings stay settable, which is what keeps this ban from + removing a pre-existing capability.""" + from litellm.proxy.auth.auth_utils import is_request_body_safe + + for spelling in ("workspace_id", "aws_workspace_id"): + assert ( + is_request_body_safe( + request_body={"model": "claude-sonnet-5", spelling: "wrkspc_abc"}, + general_settings={}, + llm_router=None, + model="claude-sonnet-5", + ) + is True + ) + + +class TestWifDisabledOnClientRedirectedBase: + def test_the_sentinel_survives_the_kwargs_funnel(self): + """Setting the sentinel is only half of it. get_litellm_params rebuilds litellm_params from + kwargs, so a field it does not carry is dropped on the way and the deployment federates for + the caller-chosen base after all.""" + from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS + from litellm.router_utils.clientside_credential_handler import ( + DISABLE_WORKLOAD_IDENTITY_PARAM, + ) + + assert DISABLE_WORKLOAD_IDENTITY_PARAM in FORWARDED_KWARGS_KEYS + + def test_the_sentinel_is_not_client_settable(self): + """It is server-owned in both directions: a caller must not be able to set it, and must not + be able to clear it either.""" + from litellm.router_utils.clientside_credential_handler import ( + DISABLE_WORKLOAD_IDENTITY_PARAM, + ) + from litellm.types.router import reject_server_owned_wif_params + + with pytest.raises(ValueError, match=DISABLE_WORKLOAD_IDENTITY_PARAM): + reject_server_owned_wif_params({DISABLE_WORKLOAD_IDENTITY_PARAM: False}) + + def test_base_override_clears_wif_and_sets_the_sentinel(self): + """A federation token minted for a client-chosen api_base would send the workload's assertion, + and then the minted bearer, to that host.""" + from litellm.llms.anthropic.wif import resolve_anthropic_wif_params + from litellm.router_utils.clientside_credential_handler import ( + DISABLE_WORKLOAD_IDENTITY_PARAM, + get_dynamic_litellm_params, + ) + + admin_deployment = { + "model": "anthropic/claude-sonnet-5", + "anthropic_federation_rule_id": "fdrl_admin", + "anthropic_organization_id": "org-admin", + "anthropic_identity_token": "oidc/env/WIF_TEST_JWT", + } + + redirected = get_dynamic_litellm_params( + litellm_params=dict(admin_deployment), + request_kwargs={"api_base": "https://not-anthropic.example"}, + ) + + assert redirected[DISABLE_WORKLOAD_IDENTITY_PARAM] is True + assert "anthropic_federation_rule_id" not in redirected + assert resolve_anthropic_wif_params(redirected) is None + + def test_sentinel_blocks_env_var_configured_federation(self, monkeypatch): + """Environment-configured federation cannot be cleared out of a dict, so the sentinel is what + stops it on a redirected deployment.""" + from litellm.llms.anthropic.wif import resolve_anthropic_wif_params + from litellm.router_utils.clientside_credential_handler import DISABLE_WORKLOAD_IDENTITY_PARAM + + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_env") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-env") + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "oidc/env/WIF_TEST_JWT") + monkeypatch.setenv("WIF_TEST_JWT", "jwt-assertion-value") + + assert resolve_anthropic_wif_params({}) is not None + assert resolve_anthropic_wif_params({DISABLE_WORKLOAD_IDENTITY_PARAM: True}) is None + + def test_base_override_clears_internal_issuer_fields(self): + """Same failure mode the legacy-path test above guards against, for the internal_issuer + identity source: a signing_key_ref resolved for a client-chosen api_base would mint an + assertion, and then a bearer token, for that host.""" + from litellm.llms.anthropic.wif import resolve_anthropic_wif_params + from litellm.router_utils.clientside_credential_handler import ( + DISABLE_WORKLOAD_IDENTITY_PARAM, + get_dynamic_litellm_params, + ) + + admin_deployment = { + "model": "anthropic/claude-sonnet-5", + "anthropic_federation_rule_id": "fdrl_admin", + "anthropic_organization_id": "org-admin", + "anthropic_identity_source": "internal_issuer", + "anthropic_issuer_url": "https://issuer.internal.example", + "anthropic_issuer_subject": "workload-a", + "anthropic_issuer_signing_key_ref": "oidc/env/ISSUER_SIGNING_KEY_PEM", + } + + redirected = get_dynamic_litellm_params( + litellm_params=dict(admin_deployment), + request_kwargs={"api_base": "https://not-anthropic.example"}, + ) + + assert redirected[DISABLE_WORKLOAD_IDENTITY_PARAM] is True + assert "anthropic_identity_source" not in redirected + assert "anthropic_issuer_signing_key_ref" not in redirected + assert resolve_anthropic_wif_params(redirected) is None + + def test_base_override_clears_keycloak_fields(self): + """Same as the internal_issuer case above, for the keycloak identity source: a + client_secret_ref resolved for a client-chosen api_base must not follow it there.""" + from litellm.llms.anthropic.wif import resolve_anthropic_wif_params + from litellm.router_utils.clientside_credential_handler import ( + DISABLE_WORKLOAD_IDENTITY_PARAM, + get_dynamic_litellm_params, + ) + + admin_deployment = { + "model": "anthropic/claude-sonnet-5", + "anthropic_federation_rule_id": "fdrl_admin", + "anthropic_organization_id": "org-admin", + "anthropic_identity_source": "keycloak", + "anthropic_keycloak_token_url": "https://keycloak.internal.example/realms/r/protocol/openid-connect/token", + "anthropic_keycloak_client_id": "litellm", + "anthropic_keycloak_client_secret_ref": "oidc/env/KEYCLOAK_CLIENT_SECRET", + } + + redirected = get_dynamic_litellm_params( + litellm_params=dict(admin_deployment), + request_kwargs={"api_base": "https://not-anthropic.example"}, + ) + + assert redirected[DISABLE_WORKLOAD_IDENTITY_PARAM] is True + assert "anthropic_identity_source" not in redirected + assert "anthropic_keycloak_client_secret_ref" not in redirected + assert resolve_anthropic_wif_params(redirected) is None diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py index ddac561f337..5b6b4251e4d 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -1,4 +1,9 @@ +import httpx +import pytest +import respx +import litellm +from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, ) @@ -88,3 +93,72 @@ def test_transform_no_system_no_tools(): assert "system" not in result assert "tools" not in result + + +@pytest.mark.parametrize( + ("api_base", "expected"), + [ + (None, "https://api.anthropic.com/v1/messages/count_tokens"), + ("", "https://api.anthropic.com/v1/messages/count_tokens"), + ("https://gateway.example", "https://gateway.example/v1/messages/count_tokens"), + ("https://gateway.example/", "https://gateway.example/v1/messages/count_tokens"), + ("https://gateway.example/v1", "https://gateway.example/v1/messages/count_tokens"), + ("https://gateway.example/anthropic/v1/messages", "https://gateway.example/anthropic/v1/messages/count_tokens"), + ], +) +def test_endpoint_appends_count_tokens_path_to_deployment_api_base(api_base, expected, monkeypatch): + monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False) + monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) + assert AnthropicCountTokensConfig().get_anthropic_count_tokens_endpoint(api_base) == expected + + +@pytest.mark.parametrize("env_name", ["ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"]) +@pytest.mark.parametrize("api_base", [None, ""]) +def test_endpoint_without_deployment_api_base_follows_env_base(env_name, api_base, monkeypatch): + """Chat and the federated exchange resolve an unset deployment base through the environment, + so an env-only gateway must receive the count too, never Anthropic's public host.""" + monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False) + monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) + monkeypatch.setenv(env_name, "https://env-gateway.example/v1/messages/") + assert ( + AnthropicCountTokensConfig().get_anthropic_count_tokens_endpoint(api_base) + == "https://env-gateway.example/v1/messages/count_tokens" + ) + + +def test_endpoint_prefers_deployment_api_base_over_env_base(monkeypatch): + monkeypatch.setenv("ANTHROPIC_API_BASE", "https://env-gateway.example") + assert ( + AnthropicCountTokensConfig().get_anthropic_count_tokens_endpoint("https://gateway.example/v1") + == "https://gateway.example/v1/messages/count_tokens" + ) + + +@pytest.fixture +def httpx_transport_clients(monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + client_cache = getattr(litellm, "in_memory_llm_clients_cache", None) + if client_cache is not None: + client_cache.flush_cache() + yield + if client_cache is not None: + client_cache.flush_cache() + + +@pytest.mark.asyncio +async def test_handler_posts_to_count_tokens_path_under_deployment_api_base(httpx_transport_clients): + """A deployment api_base names the chat host, so a handler that posts to it verbatim lands on + the host root, gets a 404, and the official count silently degrades to the local tokenizer.""" + with respx.mock: + route = respx.post("https://gateway.example/v1/messages/count_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 7}) + ) + result = await AnthropicCountTokensHandler().handle_count_tokens_request( + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + api_key="sk-ant-api03-test-key", + api_base="https://gateway.example", + ) + + assert route.called + assert result == {"input_tokens": 7} diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py b/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py index 2728ba03ae4..948e30b2c2b 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py @@ -70,9 +70,7 @@ class TestAnthropicFilesHandler: @pytest.fixture def mock_anthropic_batch_results_canceled(self): """Mock Anthropic batch results with canceled status""" - return json.dumps( - {"custom_id": "test-request-3", "result": {"type": "canceled"}} - ).encode("utf-8") + return json.dumps({"custom_id": "test-request-3", "result": {"type": "canceled"}}).encode("utf-8") @pytest.fixture def mock_anthropic_batch_results_mixed(self): @@ -114,9 +112,7 @@ class TestAnthropicFilesHandler: return "\n".join(lines).encode("utf-8") @pytest.mark.asyncio - async def test_afile_content_success( - self, handler, mock_anthropic_batch_results_succeeded - ): + async def test_afile_content_success(self, handler, mock_anthropic_batch_results_succeeded): """Test successful file content retrieval and transformation""" file_content_request: FileContentRequest = { "file_id": "batch_123", @@ -135,16 +131,14 @@ class TestAnthropicFilesHandler: ), ) - with patch( + with patch( # test-quality-ok: the proxy wiring under test is what this patches "litellm.llms.anthropic.files.handler.get_async_httpx_client" - ) as mock_get_client: + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client - with patch.object( - handler.anthropic_model_info, "get_api_key", return_value="test-api-key" - ): + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"): with patch.object( handler.anthropic_model_info, "get_api_base", @@ -161,9 +155,7 @@ class TestAnthropicFilesHandler: # Verify transformation to OpenAI format content = result.response.content.decode("utf-8") - lines = [ - line for line in content.strip().split("\n") if line.strip() - ] + lines = [line for line in content.strip().split("\n") if line.strip()] assert len(lines) == 1 transformed_result = json.loads(lines[0]) @@ -172,18 +164,13 @@ class TestAnthropicFilesHandler: assert "body" in transformed_result["response"] # Verify body has required OpenAI format fields assert "id" in transformed_result["response"]["body"] - assert ( - transformed_result["response"]["body"]["object"] - == "chat.completion" - ) + assert transformed_result["response"]["body"]["object"] == "chat.completion" assert "choices" in transformed_result["response"]["body"] # Verify request_id matches the original message id assert transformed_result["response"]["request_id"] == "msg_123" @pytest.mark.asyncio - async def test_afile_content_with_prefix( - self, handler, mock_anthropic_batch_results_succeeded - ): + async def test_afile_content_with_prefix(self, handler, mock_anthropic_batch_results_succeeded): """Test file content retrieval with anthropic_batch_results: prefix""" file_content_request: FileContentRequest = { "file_id": "anthropic_batch_results:batch_123", @@ -203,14 +190,12 @@ class TestAnthropicFilesHandler: with patch( "litellm.llms.anthropic.files.handler.get_async_httpx_client" - ) as mock_get_client: + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client - with patch.object( - handler.anthropic_model_info, "get_api_key", return_value="test-api-key" - ): + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"): with patch.object( handler.anthropic_model_info, "get_api_base", @@ -228,9 +213,7 @@ class TestAnthropicFilesHandler: assert "batch_123" in call_url @pytest.mark.asyncio - async def test_afile_content_errored_result( - self, handler, mock_anthropic_batch_results_errored - ): + async def test_afile_content_errored_result(self, handler, mock_anthropic_batch_results_errored): """Test transformation of errored batch results""" file_content_request: FileContentRequest = { "file_id": "batch_123", @@ -250,14 +233,12 @@ class TestAnthropicFilesHandler: with patch( "litellm.llms.anthropic.files.handler.get_async_httpx_client" - ) as mock_get_client: + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client - with patch.object( - handler.anthropic_model_info, "get_api_key", return_value="test-api-key" - ): + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"): with patch.object( handler.anthropic_model_info, "get_api_base", @@ -269,29 +250,17 @@ class TestAnthropicFilesHandler: ) content = result.response.content.decode("utf-8") - lines = [ - line for line in content.strip().split("\n") if line.strip() - ] + lines = [line for line in content.strip().split("\n") if line.strip()] assert len(lines) == 1 transformed_result = json.loads(lines[0]) assert transformed_result["custom_id"] == "test-request-2" - assert ( - transformed_result["response"]["status_code"] == 400 - ) # invalid_request_error maps to 400 - assert ( - transformed_result["response"]["body"]["error"]["type"] - == "invalid_request_error" - ) - assert ( - transformed_result["response"]["body"]["error"]["message"] - == "Invalid request" - ) + assert transformed_result["response"]["status_code"] == 400 # invalid_request_error maps to 400 + assert transformed_result["response"]["body"]["error"]["type"] == "invalid_request_error" + assert transformed_result["response"]["body"]["error"]["message"] == "Invalid request" @pytest.mark.asyncio - async def test_afile_content_canceled_result( - self, handler, mock_anthropic_batch_results_canceled - ): + async def test_afile_content_canceled_result(self, handler, mock_anthropic_batch_results_canceled): """Test transformation of canceled batch results""" file_content_request: FileContentRequest = { "file_id": "batch_123", @@ -311,14 +280,12 @@ class TestAnthropicFilesHandler: with patch( "litellm.llms.anthropic.files.handler.get_async_httpx_client" - ) as mock_get_client: + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client - with patch.object( - handler.anthropic_model_info, "get_api_key", return_value="test-api-key" - ): + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"): with patch.object( handler.anthropic_model_info, "get_api_base", @@ -330,23 +297,16 @@ class TestAnthropicFilesHandler: ) content = result.response.content.decode("utf-8") - lines = [ - line for line in content.strip().split("\n") if line.strip() - ] + lines = [line for line in content.strip().split("\n") if line.strip()] assert len(lines) == 1 transformed_result = json.loads(lines[0]) assert transformed_result["custom_id"] == "test-request-3" assert transformed_result["response"]["status_code"] == 400 - assert ( - "Batch request was canceled" - in transformed_result["response"]["body"]["error"]["message"] - ) + assert "Batch request was canceled" in transformed_result["response"]["body"]["error"]["message"] @pytest.mark.asyncio - async def test_afile_content_mixed_results( - self, handler, mock_anthropic_batch_results_mixed - ): + async def test_afile_content_mixed_results(self, handler, mock_anthropic_batch_results_mixed): """Test transformation of mixed batch results (succeeded, errored, expired)""" file_content_request: FileContentRequest = { "file_id": "batch_123", @@ -366,14 +326,12 @@ class TestAnthropicFilesHandler: with patch( "litellm.llms.anthropic.files.handler.get_async_httpx_client" - ) as mock_get_client: + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client - with patch.object( - handler.anthropic_model_info, "get_api_key", return_value="test-api-key" - ): + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"): with patch.object( handler.anthropic_model_info, "get_api_base", @@ -385,9 +343,7 @@ class TestAnthropicFilesHandler: ) content = result.response.content.decode("utf-8") - lines = [ - line for line in content.strip().split("\n") if line.strip() - ] + lines = [line for line in content.strip().split("\n") if line.strip()] assert len(lines) == 3 # Check first result (succeeded) @@ -396,9 +352,7 @@ class TestAnthropicFilesHandler: # Check second result (errored) result2 = json.loads(lines[1]) - assert ( - result2["response"]["status_code"] == 429 - ) # rate_limit_error maps to 429 + assert result2["response"]["status_code"] == 429 # rate_limit_error maps to 429 # Check third result (expired) result3 = json.loads(lines[2]) @@ -415,12 +369,12 @@ class TestAnthropicFilesHandler: } with patch.object( - handler.anthropic_model_info, "get_auth_header", return_value=None + handler.anthropic_model_info, + "aget_auth_header", + new=AsyncMock(return_value=None), ): with pytest.raises(ValueError, match="Missing Anthropic API Key"): - await handler.afile_content( - file_content_request=file_content_request, api_key=None - ) + await handler.afile_content(file_content_request=file_content_request, api_key=None) @pytest.mark.asyncio async def test_afile_content_missing_file_id(self, handler): @@ -432,9 +386,7 @@ class TestAnthropicFilesHandler: } with pytest.raises(ValueError, match="file_id is required"): - await handler.afile_content( - file_content_request=file_content_request, api_key="test-api-key" - ) + await handler.afile_content(file_content_request=file_content_request, api_key="test-api-key") @pytest.mark.asyncio async def test_afile_content_http_error(self, handler): @@ -454,21 +406,17 @@ class TestAnthropicFilesHandler: ), ) mock_response.raise_for_status = MagicMock( - side_effect=httpx.HTTPStatusError( - "Not Found", request=mock_response.request, response=mock_response - ) + side_effect=httpx.HTTPStatusError("Not Found", request=mock_response.request, response=mock_response) ) with patch( "litellm.llms.anthropic.files.handler.get_async_httpx_client" - ) as mock_get_client: + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client - with patch.object( - handler.anthropic_model_info, "get_api_key", return_value="test-api-key" - ): + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"): with patch.object( handler.anthropic_model_info, "get_api_base", @@ -480,6 +428,86 @@ class TestAnthropicFilesHandler: api_key="test-api-key", ) + @pytest.mark.asyncio + async def test_afile_content_resolves_wif_via_async_facade( + self, handler, mock_anthropic_batch_results_succeeded, monkeypatch + ): + """Regression: afile_content ran the blocking WIF mint on the event loop + through the sync get_auth_header; it must go through the async facade.""" + import threading + + from litellm.llms.anthropic import common_utils as anthropic_common_utils + from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token + from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + + for name in ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_API_BASE", + "ANTHROPIC_BASE_URL", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_files") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-files") + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "files-inline-jwt") + + minted = "sk-ant-oat01-files-minted" + thread_ids = [] + + class ThreadRecordingPoster: + def post(self, url, *, content, headers, timeout): + thread_ids.append(threading.get_ident()) + return httpx.Response( + 200, + json={"access_token": minted, "token_type": "Bearer", "expires_in": 3600}, + ) + + engine = JwtBearerTokenExchangeEngine(poster=ThreadRecordingPoster()) + sync_calls = [] + + def sync_shim(litellm_params, api_base, model): + sync_calls.append(model) + return get_anthropic_wif_token(litellm_params, api_base, model, engine) + + async def async_shim(litellm_params, api_base, model): + return await aget_anthropic_wif_token(litellm_params, api_base, model, engine) + + monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim) + monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim) + + mock_response = httpx.Response( + status_code=200, + content=mock_anthropic_batch_results_succeeded, + headers={"content-type": "application/json"}, + request=httpx.Request( + method="GET", + url="https://api.anthropic.com/v1/messages/batches/batch_123/results", + ), + ) + + with patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.llms.anthropic.files.handler.get_async_httpx_client" + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + await handler.afile_content( + file_content_request={ + "file_id": "batch_123", + "extra_headers": None, + "extra_body": None, + }, + api_key=None, + ) + + sent_headers = mock_client.get.call_args.kwargs["headers"] + + assert sent_headers["authorization"] == f"Bearer {minted}" + assert "oauth-2025-04-20" in sent_headers["anthropic-beta"] + assert sync_calls == [] + assert thread_ids and thread_ids[0] != threading.get_ident() + class TestAnthropicBatchesConfig: """Test Anthropic Batches Config for batch retrieval transformation""" @@ -562,15 +590,11 @@ class TestAnthropicBatchesConfig: ) assert url == "https://api.anthropic.com/v1/messages/batches/batch_123" - def test_transform_retrieve_batch_response_in_progress( - self, config, mock_anthropic_batch_response_in_progress - ): + def test_transform_retrieve_batch_response_in_progress(self, config, mock_anthropic_batch_response_in_progress): """Test transformation of in_progress batch response""" mock_response = httpx.Response( status_code=200, - content=json.dumps(mock_anthropic_batch_response_in_progress).encode( - "utf-8" - ), + content=json.dumps(mock_anthropic_batch_response_in_progress).encode("utf-8"), request=httpx.Request( method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_123", @@ -596,9 +620,7 @@ class TestAnthropicBatchesConfig: assert batch.in_progress_at is not None assert batch.completed_at is None - def test_transform_retrieve_batch_response_completed( - self, config, mock_anthropic_batch_response_completed - ): + def test_transform_retrieve_batch_response_completed(self, config, mock_anthropic_batch_response_completed): """Test transformation of completed batch response""" mock_response = httpx.Response( status_code=200, @@ -624,9 +646,7 @@ class TestAnthropicBatchesConfig: assert batch.request_counts.completed == 10 assert batch.request_counts.failed == 0 - def test_transform_retrieve_batch_response_canceling( - self, config, mock_anthropic_batch_response_canceling - ): + def test_transform_retrieve_batch_response_canceling(self, config, mock_anthropic_batch_response_canceling): """Test transformation of canceling batch response""" mock_response = httpx.Response( status_code=200, @@ -663,9 +683,7 @@ class TestAnthropicBatchesConfig: ) logging_obj = MagicMock() - with pytest.raises( - ValueError, match="Failed to parse Anthropic batch response" - ): + with pytest.raises(ValueError, match="Failed to parse Anthropic batch response"): config.transform_retrieve_batch_response( model="claude-3-5-sonnet-20241022", raw_response=mock_response, diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_wif.py b/tests/test_litellm/llms/anthropic/test_anthropic_wif.py new file mode 100644 index 00000000000..95b89e9cc71 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_anthropic_wif.py @@ -0,0 +1,1222 @@ +import concurrent.futures +import json +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import Final + +import httpx +import jwt +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import ec + +import litellm +from litellm.llms.anthropic.wif import ( + AnthropicWifParams, + _raise_anthropic_wif_error, + build_anthropic_wif_spec, + get_anthropic_wif_token, + resolve_anthropic_wif_params, +) +from litellm.llms.base_llm.auth.identity_source import ( + InternalIssuerSource, + KeycloakSource, + identity_source_ref, +) +from litellm.llms.base_llm.auth.jwt_signing import build_jwks, rfc7638_thumbprint +from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine +from litellm.types.router import GenericLiteLLMParams +from litellm.llms.base_llm.auth.types import ( + AssertionSourceError, + ExchangeError, + InsecureTokenUrl, + MalformedTokenResponse, + TokenEndpointError, + TokenTransportError, +) + +WIF_ENV_VARS: Final = ( + "ANTHROPIC_FEDERATION_RULE_ID", + "ANTHROPIC_ORGANIZATION_ID", + "ANTHROPIC_SERVICE_ACCOUNT_ID", + "ANTHROPIC_WORKSPACE_ID", + "ANTHROPIC_IDENTITY_TOKEN_FILE", + "ANTHROPIC_IDENTITY_TOKEN", + "ANTHROPIC_IDENTITY_SOURCE", + "ANTHROPIC_SCOPE", + "ANTHROPIC_API_BASE", + "ANTHROPIC_BASE_URL", + "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", +) + +GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer" + + +@pytest.fixture(autouse=True) +def _clean_wif_env(monkeypatch: pytest.MonkeyPatch) -> None: + for name in WIF_ENV_VARS: + monkeypatch.delenv(name, raising=False) + + +class FakeClock: + def __init__(self, start: float = 1_000.0) -> None: + self.now = start + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +class RecordedRequest: + def __init__(self, url: str, content: bytes, headers: Mapping[str, str], timeout: float) -> None: + self.url = url + self.content = content + self.headers = dict(headers) + self.timeout = timeout + + def json_body(self) -> dict: + return json.loads(self.content) + + +class ScriptedPoster: + 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)) + if len(self._responses) > 1: + return self._responses.pop(0) + return self._responses[0] + + +class ManualExecutor(concurrent.futures.Executor): + def __init__(self) -> None: + self.pending: list[Callable[[], None]] = [] + + def submit(self, fn, /, *args, **kwargs): + future: concurrent.futures.Future = concurrent.futures.Future() + self.pending.append(lambda: fn(*args, **kwargs)) + return future + + +def token_response(token: str = "sk-ant-oat01-minted", expires_in: int | None = 3600) -> httpx.Response: + body: Final[dict[str, str | int]] = { + "access_token": token, + "token_type": "Bearer", + **({} if expires_in is None else {"expires_in": expires_in}), + } + return httpx.Response(200, json=body) + + +def make_engine(poster: ScriptedPoster, clock: FakeClock | None = None) -> JwtBearerTokenExchangeEngine: + return JwtBearerTokenExchangeEngine( + poster=poster, + clock=clock if clock is not None else FakeClock(), + refresh_executor=ManualExecutor(), + ) + + +def write_token_file(directory: Path, content: str, name: str = "identity-token") -> Path: + directory.mkdir(parents=True, exist_ok=True) + token_file = directory / name + token_file.write_text(content, encoding="utf-8") + return token_file + + +class TestWireProtocolExact: + def test_minimal_body_and_headers(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) + monkeypatch.setenv("ANTHROPIC_SCOPE", "user:inference") + token_file = write_token_file(tmp_path, "jwt-assertion-value\n") + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster) + + token = get_anthropic_wif_token( + { + "anthropic_federation_rule_id": "fdrl_abc123", + "anthropic_organization_id": "org-uuid-1", + "anthropic_identity_token_file": str(token_file), + }, + "https://api.anthropic.com", + "claude-sonnet-4-5", + engine, + ) + + assert token == "sk-ant-oat01-minted" + assert len(poster.requests) == 1 + request = poster.requests[0] + assert request.url == "https://api.anthropic.com/v1/oauth/token" + assert "anthropic-beta" not in request.headers + assert request.headers["content-type"] == "application/json" + assert request.json_body() == { + "grant_type": GRANT_TYPE, + "federation_rule_id": "fdrl_abc123", + "organization_id": "org-uuid-1", + "assertion": "jwt-assertion-value", + } + + def test_optional_fields_present_when_set(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) + token_file = write_token_file(tmp_path, "jwt-assertion-value") + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster) + + get_anthropic_wif_token( + { + "anthropic_federation_rule_id": "fdrl_abc123", + "anthropic_organization_id": "org-uuid-1", + "anthropic_service_account_id": "svcacct_1", + "anthropic_workspace_id": "wrkspc_1", + "anthropic_identity_token_file": str(token_file), + }, + "https://api.anthropic.com", + "claude-sonnet-4-5", + engine, + ) + + request = poster.requests[0] + assert "anthropic-beta" not in request.headers + assert request.headers["content-type"] == "application/json" + assert request.json_body() == { + "grant_type": GRANT_TYPE, + "federation_rule_id": "fdrl_abc123", + "organization_id": "org-uuid-1", + "service_account_id": "svcacct_1", + "workspace_id": "wrkspc_1", + "assertion": "jwt-assertion-value", + } + + def test_spec_cache_key_identity(self): + params = AnthropicWifParams( + federation_rule_id="fdrl_1", + organization_id="org-1", + assertion_ref="oidc/env/ANTHROPIC_IDENTITY_TOKEN", + ) + spec = build_anthropic_wif_spec(params, "https://api.anthropic.com") + assert spec.cache_key_identity == ("fdrl_1", "org-1", "", "") + assert spec.body_encoding == "json" + assert spec.assertion_field == "assertion" + + def test_full_params_spec_has_no_request_headers(self): + """The token exchange sends no anthropic-beta header at all (verified against the + live endpoint); this must hold even for a fully populated params set, so a future + edit cannot reintroduce the header gated on service_account_id or workspace_id.""" + params = AnthropicWifParams( + federation_rule_id="fdrl_1", + organization_id="org-1", + service_account_id="svcacct_1", + workspace_id="wrkspc_1", + assertion_ref="oidc/env/ANTHROPIC_IDENTITY_TOKEN", + ) + spec = build_anthropic_wif_spec(params, "https://api.anthropic.com") + assert dict(spec.request_headers) == {} + + +class TestExchangeHostTrust: + """A federated exchange sends the workload's identity token to api_base and presents the minted + org-scoped token to it, so api_base is a trust decision. Anyone able to write api_base, on the + deployment or on a credential it references, could otherwise redirect both, which is why this is + enforced where the exchange is built rather than at each write path.""" + + def _mint(self, api_base: str | None, monkeypatch: pytest.MonkeyPatch) -> str: + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") + poster = ScriptedPoster([token_response()]) + get_anthropic_wif_token( + {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"}, + api_base, + "claude-sonnet-4-5", + make_engine(poster), + ) + return poster.requests[0].url + + def test_anthropic_is_trusted_without_configuration(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", raising=False) + assert self._mint("https://api.anthropic.com", monkeypatch) == "https://api.anthropic.com/v1/oauth/token" + + def test_an_unlisted_host_never_receives_the_identity_token(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", raising=False) + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") + poster = ScriptedPoster([token_response()]) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + get_anthropic_wif_token( + {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"}, + "https://attacker.example", + "claude-sonnet-4-5", + make_engine(poster), + ) + + assert poster.requests == [], "the exchange must be refused before anything is sent" + assert "attacker.example" in str(exc_info.value) + assert "LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS" in str(exc_info.value), ( + "an operator running a private gateway has to be told how to allow it" + ) + + def test_a_lookalike_host_does_not_pass_on_a_substring(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", raising=False) + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") + poster = ScriptedPoster([token_response()]) + + with pytest.raises(litellm.AuthenticationError): + get_anthropic_wif_token( + {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"}, + "https://api.anthropic.com.evil.test", + "claude-sonnet-4-5", + make_engine(poster), + ) + + assert poster.requests == [] + + def test_an_operator_can_allow_a_private_gateway(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", "gateway.internal") + assert self._mint("https://gateway.internal", monkeypatch) == "https://gateway.internal/v1/oauth/token" + + +class TestBaseUrlDerivation: + def _mint(self, api_base: str | None, monkeypatch: pytest.MonkeyPatch) -> str: + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") + # These cases are about how a base is normalised into a token URL, not about which hosts an + # operator trusts, so the private hosts they use are allowlisted explicitly. The trust + # boundary itself is covered by TestExchangeHostTrust. + monkeypatch.setenv( + "LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", + "gw.example.com,env.example.com,base.example.com,model.example.com", + ) + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster) + get_anthropic_wif_token( + {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"}, + api_base, + "claude-sonnet-4-5", + engine, + ) + return poster.requests[0].url + + def test_explicit_api_base_wins(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_API_BASE", "https://env.example.com") + assert self._mint("https://gw.example.com/", monkeypatch) == "https://gw.example.com/v1/oauth/token" + + def test_env_api_base(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_API_BASE", "https://env.example.com") + assert self._mint(None, monkeypatch) == "https://env.example.com/v1/oauth/token" + + def test_empty_api_base_falls_back_like_unset(self, monkeypatch: pytest.MonkeyPatch): + """Chat treats an empty deployment api_base as unset; the exchange must not refuse host ''.""" + monkeypatch.setenv("ANTHROPIC_API_BASE", "https://env.example.com") + assert self._mint("", monkeypatch) == "https://env.example.com/v1/oauth/token" + + def test_env_base_url(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_BASE_URL", "https://base.example.com") + assert self._mint(None, monkeypatch) == "https://base.example.com/v1/oauth/token" + + def test_default_base(self, monkeypatch: pytest.MonkeyPatch): + assert self._mint(None, monkeypatch) == "https://api.anthropic.com/v1/oauth/token" + + @pytest.mark.parametrize( + "api_base", + [ + "https://gw.example.com/v1/messages", + "https://gw.example.com/v1/messages/", + "https://gw.example.com/v1/messages//v1/messages", + ], + ) + def test_chat_appended_bases_normalize_to_clean_token_url(self, api_base: str, monkeypatch: pytest.MonkeyPatch): + """main.py appends /v1/messages before dispatch (twice for trailing-slash + bases); the exchange must still target the deployment base.""" + assert self._mint(api_base, monkeypatch) == "https://gw.example.com/v1/oauth/token" + + def test_trailing_slash_env_base_normalizes(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_BASE_URL", "https://base.example.com/") + assert self._mint(None, monkeypatch) == "https://base.example.com/v1/oauth/token" + + +class TestSecretManagerEnvResolution: + """WIF env vars resolve through get_secret_str so configured secret managers + work, exactly like every sibling Anthropic credential.""" + + def test_values_resolve_through_get_secret_str(self, monkeypatch: pytest.MonkeyPatch): + secrets: Final = { + "ANTHROPIC_FEDERATION_RULE_ID": "fdrl_sm", + "ANTHROPIC_ORGANIZATION_ID": "org-sm", + "ANTHROPIC_IDENTITY_TOKEN": "sm-inline-jwt", + } + monkeypatch.setattr( + "litellm.secret_managers.main.get_secret_str", + lambda secret_name, default_value=None: secrets.get(secret_name, default_value), + ) + + params = resolve_anthropic_wif_params(None) + + assert params == AnthropicWifParams( + federation_rule_id="fdrl_sm", + organization_id="org-sm", + assertion_ref="oidc/env/ANTHROPIC_IDENTITY_TOKEN", + ) + + def test_non_str_secret_value_treated_as_unset(self, monkeypatch: pytest.MonkeyPatch): + secrets: Final = { + "ANTHROPIC_FEDERATION_RULE_ID": {"unexpected": "shape"}, + "ANTHROPIC_ORGANIZATION_ID": "org-sm", + "ANTHROPIC_IDENTITY_TOKEN": "sm-inline-jwt", + } + monkeypatch.setattr( + "litellm.secret_managers.main.get_secret_str", + lambda secret_name, default_value=None: secrets.get(secret_name, default_value), + ) + + assert resolve_anthropic_wif_params(None) is None + + +class TestResolutionMatrix: + def test_params_beat_env_per_field(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_env") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-env") + monkeypatch.setenv("ANTHROPIC_SERVICE_ACCOUNT_ID", "svc-env") + monkeypatch.setenv("ANTHROPIC_WORKSPACE_ID", "wrkspc_env") + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN_FILE", "/var/run/secrets/env-token") + + params = resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_param", + "anthropic_organization_id": "org-param", + "anthropic_service_account_id": "svc-param", + "anthropic_workspace_id": "wrkspc_param", + "anthropic_identity_token_file": "/var/run/secrets/param-token", + } + ) + + assert params == AnthropicWifParams( + federation_rule_id="fdrl_param", + organization_id="org-param", + service_account_id="svc-param", + workspace_id="wrkspc_param", + assertion_ref="oidc/file//var/run/secrets/param-token", + ) + + def test_env_only_config(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_env") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-env") + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "raw-env-jwt") + + params = resolve_anthropic_wif_params(None) + + assert params is not None + assert params.assertion_ref == "oidc/env/ANTHROPIC_IDENTITY_TOKEN" + assert params.service_account_id is None + assert params.workspace_id is None + + def test_file_param_beats_inline_param(self): + params = resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_token_file": "/var/run/secrets/tok", + "anthropic_identity_token": "oidc/env/OTHER", + } + ) + assert params is not None + assert params.assertion_ref == "oidc/file//var/run/secrets/tok" + + def test_inline_param_beats_env_file(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN_FILE", "/var/run/secrets/env-tok") + params = resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_token": "oidc/env/OTHER", + } + ) + assert params is not None + assert params.assertion_ref == "oidc/env/OTHER" + + def test_env_file_beats_env_inline(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN_FILE", "/var/run/secrets/env-tok") + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "raw-env-jwt") + params = resolve_anthropic_wif_params( + {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"} + ) + assert params is not None + assert params.assertion_ref == "oidc/file//var/run/secrets/env-tok" + + def test_param_token_ref_beats_env_identity_source(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_IDENTITY_SOURCE", "internal_issuer") + params = resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_token_file": "/var/run/secrets/dep-tok", + } + ) + assert params is not None + assert params.assertion_ref == "oidc/file//var/run/secrets/dep-tok" + assert params.assertion_source is None + + def test_param_inline_token_beats_env_identity_source(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_IDENTITY_SOURCE", "internal_issuer") + params = resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_token": "oidc/env/OTHER", + } + ) + assert params is not None + assert params.assertion_ref == "oidc/env/OTHER" + assert params.assertion_source is None + + def test_env_identity_source_beats_env_token_refs(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_IDENTITY_SOURCE", "internal_issuer") + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN_FILE", "/var/run/secrets/env-tok") + with pytest.raises(litellm.AuthenticationError): + resolve_anthropic_wif_params( + {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"} + ) + + def test_env_identity_source_dispatches_param_issuer_fields(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_IDENTITY_SOURCE", "internal_issuer") + params = resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_issuer_url": "https://issuer.internal.example", + "anthropic_issuer_subject": "workload-a", + "anthropic_issuer_signing_key_ref": ISSUER_SIGNING_KEY_REF, + } + ) + assert params is not None + assert params.assertion_ref.startswith("oidc/internal_issuer/") + assert params.assertion_source is not None + + def test_empty_workspace_env_coerced_to_none(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_WORKSPACE_ID", "") + params = resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_token": "oidc/env/TOK", + } + ) + assert params is not None + assert params.workspace_id is None + spec = build_anthropic_wif_spec(params, "https://api.anthropic.com") + assert "workspace_id" not in spec.static_body + + @pytest.mark.parametrize( + "litellm_params", + [ + {}, + {"anthropic_federation_rule_id": "fdrl_1"}, + {"anthropic_organization_id": "org-1"}, + {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"}, + {"anthropic_organization_id": "org-1", "anthropic_identity_token": "oidc/env/TOK"}, + {"anthropic_federation_rule_id": "fdrl_1", "anthropic_identity_token": "oidc/env/TOK"}, + ], + ) + def test_gate_unmet_returns_none(self, litellm_params: dict): + assert resolve_anthropic_wif_params(litellm_params) is None + + def test_gate_unmet_facade_returns_none_without_engine_call(self): + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster) + assert get_anthropic_wif_token({}, None, "claude-sonnet-4-5", engine) is None + assert poster.requests == [] + + +class TestServiceAccountIdIsOptional: + """Anthropic's reference docs mark service_account_id required, but a live exchange + against a federation rule targeting a single service account mints successfully + without it; resolution must not gate activation on it, and the wire body must omit + the key entirely rather than send it as null.""" + + def test_activates_and_omits_service_account_id_when_unset(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) + token_file = write_token_file(tmp_path, "jwt-assertion-value") + litellm_params: Final = { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_token_file": str(token_file), + } + + params = resolve_anthropic_wif_params(litellm_params) + assert params is not None + assert params.service_account_id is None + + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster) + token = get_anthropic_wif_token(litellm_params, "https://api.anthropic.com", "claude-sonnet-4-5", engine) + + assert token == "sk-ant-oat01-minted" + assert "service_account_id" not in poster.requests[0].json_body() + + +class TestInlineRefRestrictions: + RAW_JWT: Final = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ3b3JrbG9hZCJ9.c2lnbmF0dXJl" + + @pytest.mark.parametrize("bad_ref", [RAW_JWT, "oidc/env_path/ANTHROPIC_TOKEN_PATH"]) + def test_rejected_inline_refs(self, bad_ref: str): + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + get_anthropic_wif_token( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_token": bad_ref, + }, + None, + "claude-sonnet-4-5", + engine, + ) + + assert "oidc/env/" in exc_info.value.message + assert "oidc/file/" in exc_info.value.message + assert self.RAW_JWT not in exc_info.value.message + assert poster.requests == [] + + +class TestFileAllowlistAndSymlink: + SECRET_CONTENT: Final = "super-secret-jwt-content" + + def _call(self, token_file: Path, poster: ScriptedPoster) -> str | None: + engine = make_engine(poster) + return get_anthropic_wif_token( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_token_file": str(token_file), + }, + "https://api.anthropic.com", + "claude-sonnet-4-5", + engine, + ) + + def test_file_outside_allowlist_rejected(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path / "allowed")) + token_file = write_token_file(tmp_path / "outside", self.SECRET_CONTENT) + poster = ScriptedPoster([token_response()]) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + self._call(token_file, poster) + + assert str(token_file) in exc_info.value.message + assert self.SECRET_CONTENT not in exc_info.value.message + assert poster.requests == [] + + def test_disallowed_path_message_names_allowlist_and_env_var(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + """The disallowed_path error must explain the allowlist and name the env var an + operator would set, not surface as a bare '(disallowed_path)' code dump.""" + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path / "allowed")) + token_file = write_token_file(tmp_path / "outside", self.SECRET_CONTENT) + poster = ScriptedPoster([token_response()]) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + self._call(token_file, poster) + + message = exc_info.value.message + assert "(disallowed_path)" not in message + assert "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS" in message + assert "allowed credential director" in message + + def test_symlink_escape_rejected(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + allowed = tmp_path / "allowed" + allowed.mkdir() + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(allowed)) + outside_file = write_token_file(tmp_path / "outside", self.SECRET_CONTENT) + link = allowed / "identity-token" + link.symlink_to(outside_file) + poster = ScriptedPoster([token_response()]) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + self._call(link, poster) + + assert self.SECRET_CONTENT not in exc_info.value.message + assert poster.requests == [] + + def test_file_inside_allowlist_succeeds(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) + token_file = write_token_file(tmp_path, self.SECRET_CONTENT) + poster = ScriptedPoster([token_response()]) + + assert self._call(token_file, poster) == "sk-ant-oat01-minted" + assert poster.requests[0].json_body()["assertion"] == self.SECRET_CONTENT + + +class TestErrorMappingExhaustive: + @pytest.mark.parametrize( + "error", + [ + AssertionSourceError(kind="missing", source_ref="oidc/env/TOK"), + AssertionSourceError(kind="disallowed_path", source_ref="oidc/file//etc/passwd"), + InsecureTokenUrl(host="token.example"), + TokenEndpointError(status_code=500, redacted_body="error: server_error"), + TokenTransportError(detail="ConnectError: refused"), + MalformedTokenResponse(detail="token response failed RFC 6749 5.1 schema validation"), + ], + ) + def test_every_variant_maps_to_authentication_error(self, error: ExchangeError): + with pytest.raises(litellm.AuthenticationError) as exc_info: + _raise_anthropic_wif_error( + error, model="claude-sonnet-4-5", workspace_id_set=False, service_account_id_set=False + ) + + assert exc_info.value.llm_provider == "anthropic" + assert exc_info.value.model == "claude-sonnet-4-5" + + def test_assertion_source_error_detail_is_rendered_when_present(self): + with pytest.raises(litellm.AuthenticationError) as exc_info: + _raise_anthropic_wif_error( + AssertionSourceError(kind="unreadable", source_ref="oidc/keycloak/abc123", detail="invalid_client"), + model="claude-sonnet-4-5", + workspace_id_set=True, + service_account_id_set=True, + ) + + assert "invalid_client" in exc_info.value.message + + def test_assertion_source_error_without_detail_is_unchanged(self): + """Regression floor: the token_file/env path never populates detail, so its message must stay + byte-identical to before the field existed.""" + with pytest.raises(litellm.AuthenticationError) as exc_info: + _raise_anthropic_wif_error( + AssertionSourceError(kind="unreadable", source_ref="oidc/env/ANTHROPIC_IDENTITY_TOKEN"), + model="claude-sonnet-4-5", + workspace_id_set=True, + service_account_id_set=True, + ) + + assert exc_info.value.message == ( + "litellm.AuthenticationError: Anthropic workload identity federation failed. Could not obtain " + "the OIDC identity token (unreadable) from oidc/env/ANTHROPIC_IDENTITY_TOKEN." + ) + + def test_endpoint_error_raised_through_facade(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") + poster = ScriptedPoster([httpx.Response(500, json={"error": "server_error"})]) + engine = make_engine(poster) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + get_anthropic_wif_token( + {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"}, + None, + "claude-sonnet-4-5", + engine, + ) + + assert exc_info.value.llm_provider == "anthropic" + assert "HTTP 500" in exc_info.value.message + assert "server_error" in exc_info.value.message + + @pytest.mark.parametrize( + "litellm_params,status_code,body", + [ + ( + {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"}, + 401, + {"error": "invalid_grant."}, + ), + ( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_workspace_id": "wrkspc_1", + }, + 500, + {"error": "server_error."}, + ), + ], + ) + def test_token_endpoint_error_message_has_no_doubled_period( + self, litellm_params: dict, status_code: int, body: dict, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") + poster = ScriptedPoster([httpx.Response(status_code, json=body)]) + engine = make_engine(poster) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + get_anthropic_wif_token(litellm_params, None, "claude-sonnet-4-5", engine) + + assert ".." not in exc_info.value.message + + +class TestDenialHints: + """Anthropic answers every denied exchange with an opaque 401 and logs the reason + (workspace_id_required, jti_reused, ...) only in the Console, so the error must say where + to look and name whichever optional id is still unset.""" + + BASE_PARAMS: Final = {"anthropic_federation_rule_id": "fdrl_1", "anthropic_organization_id": "org-1"} + + def _raise(self, litellm_params: dict, status_code: int, monkeypatch: pytest.MonkeyPatch) -> str: + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "inline-jwt") + poster = ScriptedPoster([httpx.Response(status_code, json={"error": "invalid_grant"})]) + engine = make_engine(poster) + with pytest.raises(litellm.AuthenticationError) as exc_info: + get_anthropic_wif_token(litellm_params, None, "claude-sonnet-4-5", engine) + return exc_info.value.message + + def test_401_points_at_console_authentication_history(self, monkeypatch: pytest.MonkeyPatch): + message = self._raise(self.BASE_PARAMS, 401, monkeypatch) + assert "authentication history" in message + assert "workspace_id_required" in message + + def test_500_carries_no_denial_hints(self, monkeypatch: pytest.MonkeyPatch): + message = self._raise(self.BASE_PARAMS, 500, monkeypatch) + assert "authentication history" not in message + assert "ANTHROPIC_WORKSPACE_ID" not in message + assert "ANTHROPIC_SERVICE_ACCOUNT_ID" not in message + + def test_hints_name_both_ids_when_both_unset(self, monkeypatch: pytest.MonkeyPatch): + message = self._raise(self.BASE_PARAMS, 401, monkeypatch) + assert "anthropic_workspace_id" in message + assert "ANTHROPIC_WORKSPACE_ID" in message + assert "anthropic_service_account_id" in message + assert "ANTHROPIC_SERVICE_ACCOUNT_ID" in message + + def test_no_workspace_hint_when_workspace_set(self, monkeypatch: pytest.MonkeyPatch): + message = self._raise({**self.BASE_PARAMS, "anthropic_workspace_id": "wrkspc_1"}, 401, monkeypatch) + assert "ANTHROPIC_WORKSPACE_ID" not in message + assert "ANTHROPIC_SERVICE_ACCOUNT_ID" in message + + def test_no_service_account_hint_when_service_account_set(self, monkeypatch: pytest.MonkeyPatch): + message = self._raise({**self.BASE_PARAMS, "anthropic_service_account_id": "svac_1"}, 401, monkeypatch) + assert "ANTHROPIC_SERVICE_ACCOUNT_ID" not in message + assert "ANTHROPIC_WORKSPACE_ID" in message + + def test_only_console_pointer_when_both_set(self, monkeypatch: pytest.MonkeyPatch): + message = self._raise( + {**self.BASE_PARAMS, "anthropic_workspace_id": "wrkspc_1", "anthropic_service_account_id": "svac_1"}, + 401, + monkeypatch, + ) + assert "authentication history" in message + assert "ANTHROPIC_WORKSPACE_ID" not in message + assert "ANTHROPIC_SERVICE_ACCOUNT_ID" not in message + assert ".." not in message + + +class TestFileRereadOnRefresh: + def test_mandatory_refresh_carries_rotated_assertion(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) + token_file = write_token_file(tmp_path, "first-assertion") + clock = FakeClock(start=1_000.0) + poster = ScriptedPoster( + [token_response("sk-ant-oat01-first", 3600), token_response("sk-ant-oat01-second", 3600)] + ) + engine = make_engine(poster, clock=clock) + litellm_params = { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_token_file": str(token_file), + } + + first = get_anthropic_wif_token(litellm_params, "https://api.anthropic.com", "claude-sonnet-4-5", engine) + token_file.write_text("second-assertion", encoding="utf-8") + clock.advance(3600 - 10) + second = get_anthropic_wif_token(litellm_params, "https://api.anthropic.com", "claude-sonnet-4-5", engine) + + assert first == "sk-ant-oat01-first" + assert second == "sk-ant-oat01-second" + assert len(poster.requests) == 2 + assert poster.requests[1].json_body()["assertion"] == "second-assertion" + + +_ISSUER_PRIVATE_VALUE: Final = 55566677788899900011122233344455566677788899900011122233344455 +ISSUER_SIGNING_KEY_REF: Final = "oidc/env/ISSUER_SIGNING_KEY_PEM" +KEYCLOAK_TOKEN_URL: Final = "https://keycloak.internal.example/realms/litellm/protocol/openid-connect/token" + + +def _issuer_signing_key() -> ec.EllipticCurvePrivateKey: + return ec.derive_private_key(_ISSUER_PRIVATE_VALUE, ec.SECP256R1()) + + +def _issuer_signing_key_pem() -> str: + return ( + _issuer_signing_key() + .private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + .decode() + ) + + +def _get_secret_str_returning(pem: str, ref: str) -> Callable[..., str | None]: + def fake_get_secret_str(secret_name: str, default_value: str | None = None) -> str | None: + return pem if secret_name == ref else default_value + + return fake_get_secret_str + + +class TestIdentitySourceDiscriminatorAbsentIsByteIdenticalToLegacy: + """anthropic_identity_source unset must resolve exactly like today: no new dispatch code + runs, and no assertion_source closure is attached, so the engine falls back to its own + reader precisely as it always has.""" + + def test_file_config_carries_no_assertion_source(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) + token_file = write_token_file(tmp_path, "jwt-assertion-value") + + params = resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_token_file": str(token_file), + } + ) + + assert params == AnthropicWifParams( + federation_rule_id="fdrl_1", + organization_id="org-1", + assertion_ref=f"oidc/file/{token_file}", + ) + assert params.assertion_source is None + + def test_env_config_carries_no_assertion_source(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_env") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-env") + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "raw-env-jwt") + + params = resolve_anthropic_wif_params(None) + + assert params is not None + assert params.assertion_ref == "oidc/env/ANTHROPIC_IDENTITY_TOKEN" + assert params.assertion_source is None + + +class TestInternalIssuerIdentitySourceDispatch: + """A config.yaml-shaped litellm_params block for the internal_issuer identity source.""" + + LITELLM_PARAMS: Final = { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_source": "internal_issuer", + "anthropic_issuer_url": "https://issuer.internal.example", + "anthropic_issuer_subject": "workload-a", + "anthropic_issuer_ttl_seconds": 300, + "anthropic_issuer_signing_key_ref": ISSUER_SIGNING_KEY_REF, + } + + def test_assertion_ref_matches_the_identity_source_hash(self): + params = resolve_anthropic_wif_params(self.LITELLM_PARAMS) + + assert params is not None + expected_config = InternalIssuerSource( + issuer_url="https://issuer.internal.example", + subject="workload-a", + ttl_seconds=300, + signing_key_ref=ISSUER_SIGNING_KEY_REF, + ) + assert params.assertion_ref == identity_source_ref(expected_config) + assert params.assertion_ref.startswith("oidc/internal_issuer/") + + def test_ref_is_stable_and_rolls_on_field_change(self): + first = resolve_anthropic_wif_params(self.LITELLM_PARAMS) + second = resolve_anthropic_wif_params(dict(self.LITELLM_PARAMS)) + changed = resolve_anthropic_wif_params({**self.LITELLM_PARAMS, "anthropic_issuer_subject": "workload-b"}) + + assert first is not None and second is not None and changed is not None + assert first.assertion_ref == second.assertion_ref + assert first.assertion_ref != changed.assertion_ref + + def test_assertion_source_mints_a_verifiable_jwt(self, monkeypatch: pytest.MonkeyPatch): + pem = _issuer_signing_key_pem() + monkeypatch.setattr( + "litellm.secret_managers.main.get_secret_str", + _get_secret_str_returning(pem, ISSUER_SIGNING_KEY_REF), + ) + + params = resolve_anthropic_wif_params(self.LITELLM_PARAMS) + assert params is not None + assert params.assertion_source is not None + + assertion = params.assertion_source() + + assert assertion is not None + public_key = _issuer_signing_key().public_key() + expected_kid = build_jwks(public_key)["keys"][0]["kid"] + assert jwt.get_unverified_header(assertion)["kid"] == expected_kid + assert expected_kid == rfc7638_thumbprint(public_key) + claims = jwt.decode(assertion, public_key, algorithms=["ES256"], options={"verify_aud": False}) + assert claims["sub"] == "workload-a" + assert claims["iss"] == "https://issuer.internal.example" + + def test_full_exchange_sends_the_minted_assertion(self, monkeypatch: pytest.MonkeyPatch): + pem = _issuer_signing_key_pem() + monkeypatch.setattr( + "litellm.secret_managers.main.get_secret_str", + _get_secret_str_returning(pem, ISSUER_SIGNING_KEY_REF), + ) + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster) + + token = get_anthropic_wif_token(self.LITELLM_PARAMS, "https://api.anthropic.com", "claude-sonnet-4-5", engine) + + assert token == "sk-ant-oat01-minted" + sent_assertion = poster.requests[0].json_body()["assertion"] + jwt.decode( + sent_assertion, _issuer_signing_key().public_key(), algorithms=["ES256"], options={"verify_aud": False} + ) + + +class TestKeycloakIdentitySourceDispatch: + """A config.yaml-shaped litellm_params block for the keycloak identity source. The minted + closure's own network behavior is covered by test_client_credentials.py's DI-poster tests; + this only proves wif.py threads the fields into the right config and hash.""" + + LITELLM_PARAMS: Final = { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_source": "keycloak", + "anthropic_keycloak_token_url": KEYCLOAK_TOKEN_URL, + "anthropic_keycloak_client_id": "litellm", + "anthropic_keycloak_client_secret_ref": "oidc/env/KEYCLOAK_CLIENT_SECRET", + } + + def test_assertion_ref_matches_the_identity_source_hash(self): + params = resolve_anthropic_wif_params(self.LITELLM_PARAMS) + + assert params is not None + expected_config = KeycloakSource( + token_url=KEYCLOAK_TOKEN_URL, + client_id="litellm", + client_secret_ref="oidc/env/KEYCLOAK_CLIENT_SECRET", + ) + assert params.assertion_ref == identity_source_ref(expected_config) + assert params.assertion_ref.startswith("oidc/keycloak/") + + def test_assertion_source_is_a_fresh_closure(self): + params = resolve_anthropic_wif_params(self.LITELLM_PARAMS) + + assert params is not None + assert params.assertion_source is not None + assert callable(params.assertion_source) + + def test_auth_method_change_rolls_the_ref(self): + default_method = resolve_anthropic_wif_params(self.LITELLM_PARAMS) + post_method = resolve_anthropic_wif_params( + {**self.LITELLM_PARAMS, "anthropic_keycloak_auth_method": "client_secret_post"} + ) + + assert default_method is not None and post_method is not None + assert default_method.assertion_ref != post_method.assertion_ref + + def test_client_secret_ref_pointer_name_change_rolls_the_ref_without_resolving_it(self): + """The hash covers the pointer NAME, never a resolved secret (decision 7) -- true even + though nothing in this test ever calls get_secret_str.""" + first = resolve_anthropic_wif_params(self.LITELLM_PARAMS) + second = resolve_anthropic_wif_params( + {**self.LITELLM_PARAMS, "anthropic_keycloak_client_secret_ref": "oidc/env/OTHER_SECRET_NAME"} + ) + + assert first is not None and second is not None + assert first.assertion_ref != second.assertion_ref + + + +@pytest.mark.parametrize( + "sparse_params", + [TestInternalIssuerIdentitySourceDispatch.LITELLM_PARAMS, TestKeycloakIdentitySourceDispatch.LITELLM_PARAMS], + ids=["internal_issuer", "keycloak"], +) +def test_dense_router_params_dump_resolves_like_the_sparse_config(sparse_params: Mapping[str, object]): + dense_params = dict(GenericLiteLLMParams(**sparse_params)) + assert any(value is None for value in dense_params.values()) + + dense = resolve_anthropic_wif_params(dense_params) + sparse = resolve_anthropic_wif_params(sparse_params) + + assert dense is not None and sparse is not None + assert dense.assertion_ref == sparse.assertion_ref + + +class TestIdentitySourceValidationFailsClosed: + """Unknown discriminator, a missing required variant field, and a field belonging to the + other variant are all hard config errors at resolution time -- never a silent fallback to + token_file (decision 5).""" + + def test_unknown_discriminator_raises(self): + with pytest.raises(litellm.AuthenticationError, match="anthropic_identity_source"): + resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_source": "bogus", + } + ) + + def test_internal_issuer_missing_required_fields_raises(self): + with pytest.raises(litellm.AuthenticationError): + resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_source": "internal_issuer", + "anthropic_issuer_url": "https://issuer.internal.example", + } + ) + + def test_keycloak_missing_required_fields_raises(self): + with pytest.raises(litellm.AuthenticationError): + resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_source": "keycloak", + "anthropic_keycloak_client_id": "litellm", + } + ) + + def test_mixed_variant_fields_raise(self): + with pytest.raises(litellm.AuthenticationError, match="belongs to a different identity source"): + resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_source": "internal_issuer", + "anthropic_issuer_url": "https://issuer.internal.example", + "anthropic_issuer_subject": "workload-a", + "anthropic_issuer_signing_key_ref": ISSUER_SIGNING_KEY_REF, + "anthropic_keycloak_client_id": "leaked-from-other-variant", + } + ) + + def test_secret_pasted_into_wrong_field_never_appears_in_the_error(self): + secret_value = "super-secret-client-value-xyz" + with pytest.raises(litellm.AuthenticationError) as exc_info: + resolve_anthropic_wif_params( + { + "anthropic_federation_rule_id": "fdrl_1", + "anthropic_organization_id": "org-1", + "anthropic_identity_source": "internal_issuer", + "anthropic_issuer_url": "https://issuer.internal.example", + "anthropic_issuer_subject": "workload-a", + "anthropic_issuer_signing_key_ref": ISSUER_SIGNING_KEY_REF, + "anthropic_issuer_ttl_seconds": secret_value, + } + ) + + assert secret_value not in exc_info.value.message + + +class TestMissingIdsFailClosedWhenIdentitySourceConfigured: + """An explicit identity source is a request to federate. Without the rule or organization id + the exchange cannot even be attempted, so resolution must say which ids are missing instead + of returning None and letting the request die later as a missing API key.""" + + INTERNAL_ISSUER_FIELDS: Final = { + "anthropic_identity_source": "internal_issuer", + "anthropic_issuer_url": "https://issuer.internal.example", + "anthropic_issuer_subject": "workload-a", + "anthropic_issuer_signing_key_ref": ISSUER_SIGNING_KEY_REF, + } + + def test_both_ids_missing_names_both(self): + with pytest.raises(litellm.AuthenticationError) as exc_info: + resolve_anthropic_wif_params(self.INTERNAL_ISSUER_FIELDS) + + message = exc_info.value.message + assert "'internal_issuer'" in message + assert "anthropic_federation_rule_id and anthropic_organization_id are not set" in message + assert "Settings > Workload identity" in message + assert "ANTHROPIC_FEDERATION_RULE_ID" in message + + def test_only_rule_id_missing_names_only_the_rule(self): + with pytest.raises(litellm.AuthenticationError) as exc_info: + resolve_anthropic_wif_params({**self.INTERNAL_ISSUER_FIELDS, "anthropic_organization_id": "org-1"}) + + assert "but anthropic_federation_rule_id is not set" in exc_info.value.message + + def test_only_organization_id_missing_names_only_the_org(self): + with pytest.raises(litellm.AuthenticationError) as exc_info: + resolve_anthropic_wif_params({**self.INTERNAL_ISSUER_FIELDS, "anthropic_federation_rule_id": "fdrl_1"}) + + assert "but anthropic_organization_id is not set" in exc_info.value.message + + def test_keycloak_source_fails_closed_too(self): + with pytest.raises(litellm.AuthenticationError, match="'keycloak', but anthropic_federation_rule_id"): + resolve_anthropic_wif_params( + {"anthropic_identity_source": "keycloak", "anthropic_organization_id": "org-1"} + ) + + def test_env_configured_source_fails_closed(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_IDENTITY_SOURCE", "internal_issuer") + with pytest.raises(litellm.AuthenticationError, match="anthropic_organization_id is not set"): + resolve_anthropic_wif_params({"anthropic_federation_rule_id": "fdrl_1"}) + + def test_env_ids_satisfy_the_gate(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_env") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-env") + params = resolve_anthropic_wif_params(self.INTERNAL_ISSUER_FIELDS) + assert params is not None + assert params.federation_rule_id == "fdrl_env" + + def test_unknown_source_with_missing_ids_reports_the_unknown_source(self): + with pytest.raises(litellm.AuthenticationError, match="must be one of internal_issuer, keycloak"): + resolve_anthropic_wif_params({"anthropic_identity_source": "bogus"}) + + def test_legacy_token_params_without_ids_still_return_none(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_IDENTITY_SOURCE", "internal_issuer") + assert resolve_anthropic_wif_params({"anthropic_identity_token": "oidc/env/TOK"}) is None + + +class TestConfigYamlShapedIdentitySources: + """One litellm_params dict per identity source, shaped exactly like the + model_list[].litellm_params block a proxy config.yaml carries -- proving an operator can + configure each of Phase 1's supported sources.""" + + def test_legacy_token_file_source(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) + token_file = write_token_file(tmp_path, "jwt-assertion-value") + litellm_params = { + "model": "anthropic/claude-sonnet-4-5", + "anthropic_federation_rule_id": "fdrl_prod", + "anthropic_organization_id": "org_prod", + "anthropic_identity_token_file": str(token_file), + } + + params = resolve_anthropic_wif_params(litellm_params) + + assert params is not None + assert params.assertion_ref == f"oidc/file/{token_file}" + assert params.assertion_source is None + + def test_internal_issuer_source(self): + litellm_params = { + "model": "anthropic/claude-sonnet-4-5", + "anthropic_federation_rule_id": "fdrl_prod", + "anthropic_organization_id": "org_prod", + "anthropic_identity_source": "internal_issuer", + "anthropic_issuer_url": "https://litellm.internal.example", + "anthropic_issuer_subject": "litellm-proxy", + "anthropic_issuer_ttl_seconds": 300, + "anthropic_issuer_signing_key_ref": "os.environ/ISSUER_SIGNING_KEY_PEM", + } + + params = resolve_anthropic_wif_params(litellm_params) + + assert params is not None + assert params.assertion_ref.startswith("oidc/internal_issuer/") + assert params.assertion_source is not None + + def test_keycloak_source(self): + litellm_params = { + "model": "anthropic/claude-sonnet-4-5", + "anthropic_federation_rule_id": "fdrl_prod", + "anthropic_organization_id": "org_prod", + "anthropic_identity_source": "keycloak", + "anthropic_keycloak_token_url": KEYCLOAK_TOKEN_URL, + "anthropic_keycloak_client_id": "litellm", + "anthropic_keycloak_auth_method": "client_secret_post", + "anthropic_keycloak_client_secret_ref": "os.environ/KEYCLOAK_CLIENT_SECRET", + "anthropic_keycloak_scope": "anthropic-wif", + } + + params = resolve_anthropic_wif_params(litellm_params) + + assert params is not None + assert params.assertion_ref.startswith("oidc/keycloak/") + assert params.assertion_source is not None diff --git a/tests/test_litellm/llms/anthropic/test_azure_ai_cache_pricing.py b/tests/test_litellm/llms/anthropic/test_azure_ai_cache_pricing.py index 69738118d7a..8b22acc4a23 100644 --- a/tests/test_litellm/llms/anthropic/test_azure_ai_cache_pricing.py +++ b/tests/test_litellm/llms/anthropic/test_azure_ai_cache_pricing.py @@ -3,8 +3,6 @@ Test that Azure AI Anthropic models have cache pricing configured. Verifies the fix for issue #19532. """ - - import litellm from litellm import get_model_info from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map @@ -27,16 +25,11 @@ def reload_model_costs(): ("claude-sonnet-4-5", 3.75e-06, 3e-07), ], ) -def test_azure_ai_claude_cache_pricing( - model, expected_cache_creation_cost, expected_cache_read_cost -): +def test_azure_ai_claude_cache_pricing(model, expected_cache_creation_cost, expected_cache_read_cost): """Test that Azure AI Claude models have correct cache pricing.""" model_info = get_model_info(model=model, custom_llm_provider="azure_ai") assert model_info.get("cache_creation_input_token_cost") is not None assert model_info.get("cache_read_input_token_cost") is not None - assert ( - model_info.get("cache_creation_input_token_cost") - == expected_cache_creation_cost - ) + assert model_info.get("cache_creation_input_token_cost") == expected_cache_creation_cost assert model_info.get("cache_read_input_token_cost") == expected_cache_read_cost diff --git a/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py index 44b8bb3c9a2..58018a665bc 100644 --- a/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py +++ b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py @@ -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 diff --git a/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py b/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py index bcfc56577eb..542f8f55b59 100644 --- a/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py +++ b/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py @@ -10,9 +10,9 @@ Regression test for https://github.com/BerriAI/litellm/issues/22040 import os import sys -sys.path.insert( - 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) -) +import pytest + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, @@ -78,9 +78,177 @@ class TestCountTokensOAuthHeaders: headers = config.get_required_headers(FAKE_OAUTH_TOKEN) beta_value = headers.get("anthropic-beta", "") - assert ( - "token-counting" in beta_value - ), f"token-counting beta missing from OAuth headers: {beta_value}" - assert ( - "oauth-2025-04-20" in beta_value - ), f"oauth beta missing from OAuth headers: {beta_value}" + assert "token-counting" in beta_value, f"token-counting beta missing from OAuth headers: {beta_value}" + assert "oauth-2025-04-20" in beta_value, f"oauth beta missing from OAuth headers: {beta_value}" + + +class TestCountTokensUsesWorkloadIdentity: + """A federated deployment holds no static key. Without minting one, count_tokens returns None + and the caller silently falls back to the local tokenizer, so the number a federated + deployment reports would never come from Anthropic.""" + + @pytest.mark.asyncio + async def test_a_federated_deployment_mints_and_counts(self, monkeypatch): + from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module + + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + minted = "sk-ant-oat01-minted-for-count" + + async def fake_mint(_params, _api_base, _model): + return minted + + monkeypatch.setattr(token_counter_module, "aget_anthropic_wif_token", fake_mint, raising=False) + monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", fake_mint, raising=False) + + seen: dict[str, object] = {} + + async def fake_request(**kwargs): + seen.update(kwargs) + return {"input_tokens": 42} + + monkeypatch.setattr( + token_counter_module.anthropic_count_tokens_handler, + "handle_count_tokens_request", + fake_request, + raising=False, + ) + + result = await token_counter_module.AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={ + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5", + "anthropic_federation_rule_id": "fdrl_x", + "anthropic_organization_id": "org-x", + } + }, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.total_tokens == 42 + assert seen["api_key"] == minted + + @pytest.mark.asyncio + async def test_an_auth_token_deployment_never_mints(self, monkeypatch): + from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module + + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", "bearer-token-for-testing") + mint_calls: list[str] = [] + + async def fake_mint(_params, _api_base, model): + mint_calls.append(model) + return "sk-ant-oat01-should-not-be-minted" + + monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", fake_mint, raising=False) + + result = await token_counter_module.AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={ + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5", + "anthropic_federation_rule_id": "fdrl_x", + "anthropic_organization_id": "org-x", + } + }, + request_model="claude-sonnet-4-5", + ) + + assert result is None + assert mint_calls == [] + + @pytest.mark.asyncio + async def test_a_failed_mint_degrades_like_an_anthropic_error(self, monkeypatch): + import litellm + from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module + + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False) + + async def failing_mint(_params, _api_base, model): + raise litellm.AuthenticationError( + message="federation_rule_id is not a well-formed fdrl_ tagged ID", + llm_provider="anthropic", + model=model, + ) + + monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", failing_mint, raising=False) + + result = await token_counter_module.AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={ + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5", + "anthropic_federation_rule_id": "not-a-rule", + "anthropic_organization_id": "org-x", + } + }, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 401 + assert result.total_tokens == 0 + assert "fdrl_" in (result.error_message or "") + + @pytest.mark.asyncio + async def test_a_vault_backed_static_key_never_mints(self, monkeypatch): + from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module + + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False) + vault_key = "sk-ant-api03-only-in-the-vault" + + def vault_only(secret_name, default_value=None): + return vault_key if secret_name == "ANTHROPIC_API_KEY" else None + + monkeypatch.setattr("litellm.secret_managers.main.get_secret_str", vault_only, raising=False) + + mint_calls: list[str] = [] + + async def fake_mint(_params, _api_base, model): + mint_calls.append(model) + return "sk-ant-oat01-should-not-be-minted" + + monkeypatch.setattr(token_counter_module, "aget_anthropic_wif_token", fake_mint, raising=False) + monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", fake_mint, raising=False) + + seen: dict[str, object] = {} + + async def fake_request(**kwargs): + seen.update(kwargs) + return {"input_tokens": 7} + + monkeypatch.setattr( + token_counter_module.anthropic_count_tokens_handler, + "handle_count_tokens_request", + fake_request, + raising=False, + ) + + result = await token_counter_module.AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={ + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5", + "anthropic_federation_rule_id": "fdrl_x", + "anthropic_organization_id": "org-x", + } + }, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.total_tokens == 7 + assert seen["api_key"] == vault_key + assert mint_calls == [] diff --git a/tests/test_litellm/llms/anthropic/test_message_sanitization.py b/tests/test_litellm/llms/anthropic/test_message_sanitization.py index 79ed321d0ee..7afa60baf7e 100644 --- a/tests/test_litellm/llms/anthropic/test_message_sanitization.py +++ b/tests/test_litellm/llms/anthropic/test_message_sanitization.py @@ -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 diff --git a/tests/test_litellm/llms/base_llm/auth/__init__.py b/tests/test_litellm/llms/base_llm/auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/base_llm/auth/test_client_credentials.py b/tests/test_litellm/llms/base_llm/auth/test_client_credentials.py new file mode 100644 index 00000000000..76c309cae37 --- /dev/null +++ b/tests/test_litellm/llms/base_llm/auth/test_client_credentials.py @@ -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) diff --git a/tests/test_litellm/llms/base_llm/auth/test_identity_source.py b/tests/test_litellm/llms/base_llm/auth/test_identity_source.py new file mode 100644 index 00000000000..bfa8847c69a --- /dev/null +++ b/tests/test_litellm/llms/base_llm/auth/test_identity_source.py @@ -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, + } + ) + ) diff --git a/tests/test_litellm/llms/base_llm/auth/test_internal_issuer.py b/tests/test_litellm/llms/base_llm/auth/test_internal_issuer.py new file mode 100644 index 00000000000..d965a356426 --- /dev/null +++ b/tests/test_litellm/llms/base_llm/auth/test_internal_issuer.py @@ -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)) diff --git a/tests/test_litellm/llms/base_llm/auth/test_jwt_signing.py b/tests/test_litellm/llms/base_llm/auth/test_jwt_signing.py new file mode 100644 index 00000000000..30b400f4adb --- /dev/null +++ b/tests/test_litellm/llms/base_llm/auth/test_jwt_signing.py @@ -0,0 +1,213 @@ +import base64 +import hashlib +import json +import subprocess +import sys +import textwrap +import time +from typing import Final + +import jwt +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import ec, rsa + +from litellm.llms.base_llm.auth.jwt_signing import ( + MISSING_SIGNING_DEPENDENCIES_MESSAGE, + build_jwk, + build_jwks, + jwks_document_json, + load_es256_private_key, + rfc7638_thumbprint, + sign_es256_jwt, +) + +_FIXED_PRIVATE_VALUE: Final = 55090612345678901234567890123456789012345678901234567890123456 +_OTHER_PRIVATE_VALUE: Final = 1 + + +def fixed_private_key(value: int = _FIXED_PRIVATE_VALUE) -> ec.EllipticCurvePrivateKey: + return ec.derive_private_key(value, ec.SECP256R1()) + + +def pem_of(key: ec.EllipticCurvePrivateKey) -> str: + return key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ).decode() + + +def independent_thumbprint(public_key: ec.EllipticCurvePublicKey) -> str: + """Recomputes RFC 7638 by hand, deliberately not sharing a single line of code with + ``jwt_signing.rfc7638_thumbprint`` -- a mutation that broke the real implementation must not + also break this reference, or the two would trivially agree by sharing the bug.""" + numbers: Final = public_key.public_numbers() + x: Final = base64.urlsafe_b64encode(numbers.x.to_bytes(32, "big")).rstrip(b"=").decode() + y: Final = base64.urlsafe_b64encode(numbers.y.to_bytes(32, "big")).rstrip(b"=").decode() + canonical: Final = f'{{"crv":"P-256","kty":"EC","x":"{x}","y":"{y}"}}' + return base64.urlsafe_b64encode(hashlib.sha256(canonical.encode()).digest()).rstrip(b"=").decode() + + +class TestLoadEs256PrivateKey: + def test_valid_ec_p256_pem_loads(self): + key: Final = load_es256_private_key(pem_of(fixed_private_key())) + + assert isinstance(key, ec.EllipticCurvePrivateKey) + assert isinstance(key.curve, ec.SECP256R1) + + def test_garbage_pem_is_rejected(self): + with pytest.raises(ValueError, match="not a valid unencrypted PEM"): + load_es256_private_key("not a pem") + + def test_rsa_key_is_rejected(self): + rsa_pem: Final = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + .decode() + ) + + with pytest.raises(ValueError, match="P-256"): + load_es256_private_key(rsa_pem) + + def test_non_p256_curve_is_rejected(self): + secp384_pem: Final = ( + ec.generate_private_key(ec.SECP384R1()) + .private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + .decode() + ) + + with pytest.raises(ValueError, match="P-256"): + load_es256_private_key(secp384_pem) + + def test_error_never_echoes_key_material(self): + pem: Final = pem_of(fixed_private_key()) + + with pytest.raises(ValueError, match="P-256"): + load_es256_private_key(pem_of(ec.generate_private_key(ec.SECP384R1()))) + with pytest.raises(ValueError, match="not a valid unencrypted PEM") as exc_info: + load_es256_private_key("garbage-not-a-pem") + + assert pem not in str(exc_info.value) + + +class TestRfc7638Thumbprint: + def test_matches_independent_recomputation(self): + public_key: Final = fixed_private_key().public_key() + + assert rfc7638_thumbprint(public_key) == independent_thumbprint(public_key) + + def test_different_keys_have_different_thumbprints(self): + first: Final = fixed_private_key(_FIXED_PRIVATE_VALUE).public_key() + second: Final = fixed_private_key(_OTHER_PRIVATE_VALUE).public_key() + + assert rfc7638_thumbprint(first) != rfc7638_thumbprint(second) + + def test_thumbprint_is_deterministic(self): + public_key: Final = fixed_private_key().public_key() + + assert rfc7638_thumbprint(public_key) == rfc7638_thumbprint(public_key) + + +class TestBuildJwks: + def test_jwks_contains_one_key_matching_the_thumbprint(self): + public_key: Final = fixed_private_key().public_key() + + jwks: Final = build_jwks(public_key) + + assert len(jwks["keys"]) == 1 + assert jwks["keys"][0]["kid"] == rfc7638_thumbprint(public_key) + assert jwks["keys"][0]["kty"] == "EC" + assert jwks["keys"][0]["crv"] == "P-256" + assert jwks["keys"][0]["alg"] == "ES256" + + def test_build_jwk_stamps_the_given_kid_verbatim(self): + jwk: Final = build_jwk(fixed_private_key().public_key(), kid="caller-supplied-kid") + + assert jwk["kid"] == "caller-supplied-kid" + + def test_jwks_document_json_round_trips_through_build_jwks(self): + key: Final = fixed_private_key() + + document: Final = json.loads(jwks_document_json(pem_of(key))) + jwks: Final = build_jwks(key.public_key()) + + assert document == {"keys": [dict(jwk) for jwk in jwks["keys"]]} + + +class TestSignEs256Jwt: + def test_minted_token_verifies_against_the_matching_public_key(self): + key: Final = fixed_private_key() + now: Final = int(time.time()) + claims: Final = {"sub": "workload-a", "iss": "https://issuer.example", "iat": now, "exp": now + 300} + + token: Final = sign_es256_jwt(pem_of(key), claims) + decoded: Final = jwt.decode(token, key.public_key(), algorithms=["ES256"]) + + assert decoded == claims + + def test_header_alg_is_es256(self): + token: Final = sign_es256_jwt(pem_of(fixed_private_key()), {"sub": "x"}) + + assert jwt.get_unverified_header(token)["alg"] == "ES256" + + def test_header_kid_matches_the_published_jwks(self): + key: Final = fixed_private_key() + + token: Final = sign_es256_jwt(pem_of(key), {"sub": "x"}) + + header_kid: Final = jwt.get_unverified_header(token)["kid"] + published_kid: Final = build_jwks(key.public_key())["keys"][0]["kid"] + assert header_kid == published_kid == rfc7638_thumbprint(key.public_key()) + + def test_wrong_key_fails_verification(self): + signing_key: Final = fixed_private_key(_FIXED_PRIVATE_VALUE) + other_key: Final = fixed_private_key(_OTHER_PRIVATE_VALUE) + + token: Final = sign_es256_jwt(pem_of(signing_key), {"sub": "x"}) + + with pytest.raises(jwt.exceptions.InvalidSignatureError): + jwt.decode(token, other_key.public_key(), algorithms=["ES256"]) + + +class TestBaseSdkImport: + """A base ``pip install litellm`` has neither PyJWT nor cryptography (both are proxy extras), + and ``litellm/__init__`` reaches this module through the Anthropic provider, so a + module-level import of either would break ``import litellm`` for every base SDK user.""" + + def test_module_imports_with_pyjwt_and_cryptography_absent(self): + script: Final = textwrap.dedent( + """ + import sys + + class Blocker: + def find_spec(self, name, path=None, target=None): + if name.split(".")[0] in {"jwt", "cryptography"}: + raise ModuleNotFoundError(f"No module named {name!r}") + + sys.meta_path.insert(0, Blocker()) + import litellm + from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json + try: + jwks_document_json("not a key") + except ImportError as e: + print(e) + """ + ) + result: Final = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=False) + assert result.returncode == 0, result.stderr[-2000:] + assert result.stdout.strip() == MISSING_SIGNING_DEPENDENCIES_MESSAGE + + def test_signing_reports_the_missing_extra(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setitem(sys.modules, "jwt", None) + with pytest.raises(ImportError, match="litellm\\[proxy\\]"): + sign_es256_jwt(pem_of(fixed_private_key()), {"sub": "x"}) + diff --git a/tests/test_litellm/llms/base_llm/auth/test_token_exchange.py b/tests/test_litellm/llms/base_llm/auth/test_token_exchange.py new file mode 100644 index 00000000000..129eec88c2b --- /dev/null +++ b/tests/test_litellm/llms/base_llm/auth/test_token_exchange.py @@ -0,0 +1,1752 @@ +import asyncio +import concurrent.futures +import base64 +import json +from urllib.parse import quote, urlencode +import logging +import threading +import time +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final +from urllib.parse import parse_qsl + +import httpx +import pytest +from pydantic import SecretStr + +from litellm.llms.base_llm.auth.token_exchange import ( + _METRICS_QUEUE_LIMIT, + _REDACTION_CAP, + ADVISORY_REFRESH_BACKOFF_SECONDS, + CALL_TYPE_CACHE_HIT, + FALLBACK_TOKEN_TTL_SECONDS, + MAX_ASSERTION_BYTES, + MAX_RESPONSE_BYTES, + JwtBearerTokenExchangeEngine, + ServiceLoggingMetricsSink, + TokenExchangeEndpointFailure, + TokenExchangeTransportFailure, + _default_assertion_reader, + _error_summary, + _HttpxSyncTokenPoster, + _new_exchange_handler, + redact_oauth_error_body, +) +from litellm.llms.base_llm.auth.types import ( + AssertionSource, + AssertionSourceError, + BodyEncoding, + ExchangeError, + ExchangeResult, + InsecureTokenUrl, + MalformedTokenResponse, + MintedToken, + TokenEndpointError, + TokenExchangeSpec, + TokenTransportError, +) +from litellm.secret_managers.main import OidcPathNotAllowedError, _resolve_oidc_file_path +from litellm.types.services import ServiceTypes + +DEFAULT_REF: Final = "oidc/env/TEST_ASSERTION" +DEFAULT_ASSERTION: Final = "test-jwt-assertion" +EXCHANGE_URL: Final = "https://token.example/v1/oauth/token" + + +class FakeClock: + def __init__(self, start: float = 1_000.0) -> None: + self.now = start + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +class RecordedRequest: + def __init__(self, url: str, content: bytes, headers: Mapping[str, str], timeout: float) -> None: + self.url = url + self.content = content + self.headers = dict(headers) + self.timeout = timeout + + def json_body(self) -> dict: + return json.loads(self.content) + + +class ScriptedPoster: + """Returns scripted responses in order (repeating the last one); records requests.""" + + def __init__( + self, + responses: list[httpx.Response], + on_request: Callable[[RecordedRequest], None] | None = None, + ) -> None: + self.requests: list[RecordedRequest] = [] + self._responses = list(responses) + self._on_request = on_request + + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + recorded = RecordedRequest(url, content, headers, timeout) + self.requests.append(recorded) + if self._on_request is not None: + self._on_request(recorded) + if len(self._responses) > 1: + return self._responses.pop(0) + return self._responses[0] + + +class RaisingPoster: + def __init__(self, error: Exception) -> None: + self.calls = 0 + self._error = error + + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + self.calls += 1 + raise self._error + + +class ManualExecutor(concurrent.futures.Executor): + """Records submissions; runs them only when the test says so.""" + + def __init__(self) -> None: + self.pending: list[Callable[[], None]] = [] + + def submit(self, fn, /, *args, **kwargs): + future: concurrent.futures.Future = concurrent.futures.Future() + self.pending.append(lambda: fn(*args, **kwargs)) + return future + + def run_all(self) -> None: + drained = list(self.pending) + self.pending.clear() + for job in drained: + job() + + +class InlineExecutor(concurrent.futures.Executor): + def submit(self, fn, /, *args, **kwargs): + future: concurrent.futures.Future = concurrent.futures.Future() + future.set_result(fn(*args, **kwargs)) + return future + + +class NeverRunsExecutor(concurrent.futures.Executor): + """Accepts work and never runs it, standing in for a telemetry backend that has stalled, so a + test can show the backlog stops growing instead of consuming memory for as long as traffic lasts.""" + + def __init__(self) -> None: + self.submitted = 0 # mutable-ok: a test spy counting accepted work + + def submit(self, fn, /, *args, **kwargs): + self.submitted += 1 + return concurrent.futures.Future() + + +def token_response(token: str = "sk-ant-oat01-minted", expires_in: int | None = 3600) -> httpx.Response: + body: Final[dict[str, str | int]] = { + "access_token": token, + "token_type": "Bearer", + **({} if expires_in is None else {"expires_in": expires_in}), + } + return httpx.Response(200, json=body) + + +def make_spec( + *, + token_url: str = "https://token.example/v1/oauth/token", + assertion_ref: str = DEFAULT_REF, + assertion_field: str = "assertion", + static_body: Mapping[str, str] = MappingProxyType( + { + "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", + "federation_rule_id": "fdrl_1", + "organization_id": "org-1", + } + ), + body_encoding: BodyEncoding = "json", + request_headers: Mapping[str, str] = MappingProxyType( + {"anthropic-beta": "oauth-2025-04-20,oidc-federation-2026-04-01"} + ), + cache_key_identity: tuple[str, ...] = ("fdrl_1", "org-1", "", ""), + timeout_seconds: float = 2.0, + assertion_source: AssertionSource | None = None, +) -> TokenExchangeSpec: + return TokenExchangeSpec( + token_url=token_url, + assertion_ref=assertion_ref, + assertion_field=assertion_field, + static_body=static_body, + body_encoding=body_encoding, + request_headers=request_headers, + cache_key_identity=cache_key_identity, + timeout_seconds=timeout_seconds, + assertion_source=assertion_source, + ) + + +class RecordingMetricsSink: + def __init__(self) -> None: + self.successes: list[tuple[str, float]] = [] + self.failures: list[tuple[str, float, ExchangeError]] = [] + self.cache_hits = 0 + + def exchange_success(self, *, call_type: str, duration_seconds: float) -> None: + self.successes.append((call_type, duration_seconds)) + + def exchange_failure(self, *, call_type: str, duration_seconds: float, error: ExchangeError) -> None: + self.failures.append((call_type, duration_seconds, error)) + + def cache_hit(self) -> None: + self.cache_hits += 1 + + +def make_engine( + poster, + reader: Mapping[str, str] | Callable[[str], str | None] | None = None, + clock: FakeClock | None = None, + executor: concurrent.futures.Executor | None = None, + max_entries: int = 64, + metrics_sink=None, +) -> JwtBearerTokenExchangeEngine: + resolved_reader = reader if callable(reader) else (reader or {DEFAULT_REF: DEFAULT_ASSERTION}).get + return JwtBearerTokenExchangeEngine( + poster=poster, + assertion_reader=resolved_reader, + clock=clock if clock is not None else FakeClock(), + refresh_executor=executor if executor is not None else ManualExecutor(), + max_entries=max_entries, + metrics_sink=metrics_sink if metrics_sink is not None else RecordingMetricsSink(), + ) + + +def mint(engine: JwtBearerTokenExchangeEngine, spec: TokenExchangeSpec) -> MintedToken: + result = engine.get_token(spec) + assert isinstance(result, MintedToken) + return result + + +class TestFreshMintWireExact: + def test_json_body_and_headers(self): + poster = ScriptedPoster([token_response(expires_in=3600)]) + clock = FakeClock(start=1_000.0) + engine = make_engine(poster, clock=clock) + spec = make_spec() + + result = mint(engine, spec) + + assert result.access_token.get_secret_value() == "sk-ant-oat01-minted" + assert result.expires_at == 1_000.0 + 3600 + assert len(poster.requests) == 1 + request = poster.requests[0] + assert request.url == "https://token.example/v1/oauth/token" + assert request.timeout == 2.0 + assert request.headers == { + "content-type": "application/json", + "anthropic-beta": "oauth-2025-04-20,oidc-federation-2026-04-01", + } + assert request.json_body() == { + "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", + "federation_rule_id": "fdrl_1", + "organization_id": "org-1", + "assertion": DEFAULT_ASSERTION, + } + + def test_form_body_and_content_type(self): + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster) + spec = make_spec(body_encoding="form") + + mint(engine, spec) + + request = poster.requests[0] + assert request.headers["content-type"] == "application/x-www-form-urlencoded" + assert dict(parse_qsl(request.content.decode())) == { + "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", + "federation_rule_id": "fdrl_1", + "organization_id": "org-1", + "assertion": DEFAULT_ASSERTION, + } + + +def test_cache_hit_zero_posts(): + poster = ScriptedPoster([token_response(expires_in=3600)]) + clock = FakeClock() + engine = make_engine(poster, clock=clock) + spec = make_spec() + + first = mint(engine, spec) + clock.advance(100.0) + second = mint(engine, spec) + + assert len(poster.requests) == 1 + assert second.access_token.get_secret_value() == first.access_token.get_secret_value() + + +@pytest.mark.parametrize( + "remaining,expect_advisory_submit,expect_new_token", + [ + (121.0, False, False), + (120.0, True, False), + (119.0, True, False), + (31.0, True, False), + (30.0, False, True), + (29.0, False, True), + ], +) +def test_window_boundaries(remaining: float, expect_advisory_submit: bool, expect_new_token: bool): + poster = ScriptedPoster([token_response("old-token", expires_in=3600), token_response("new-token")]) + clock = FakeClock(start=1_000.0) + executor = ManualExecutor() + engine = make_engine(poster, clock=clock, executor=executor) + spec = make_spec() + + mint(engine, spec) + expires_at = 1_000.0 + 3600 + clock.now = expires_at - remaining + result = mint(engine, spec) + + assert len(executor.pending) == (1 if expect_advisory_submit else 0) + expected_token = "new-token" if expect_new_token else "old-token" + assert result.access_token.get_secret_value() == expected_token + assert len(poster.requests) == (2 if expect_new_token else 1) + + +def test_advisory_serve_stale_single_flight_backoff(caplog: pytest.LogCaptureFixture): + poster = ScriptedPoster( + [ + token_response("stale-token", expires_in=3600), + httpx.Response(500, json={"error": "server_error"}), + httpx.Response(500, json={"error": "server_error"}), + ] + ) + clock = FakeClock(start=1_000.0) + executor = ManualExecutor() + engine = make_engine(poster, clock=clock, executor=executor) + spec = make_spec() + + mint(engine, spec) + clock.now = 1_000.0 + 3600 - 100.0 + + first = mint(engine, spec) + second = mint(engine, spec) + assert first.access_token.get_secret_value() == "stale-token" + assert second.access_token.get_secret_value() == "stale-token" + assert len(executor.pending) == 1 + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + executor.run_all() + assert len(poster.requests) == 2 + warning_records = [r for r in caplog.records if r.levelno == logging.WARNING] + assert any("Advisory token refresh" in r.getMessage() for r in warning_records) + assert "server_error" in caplog.text + assert DEFAULT_ASSERTION not in caplog.text + assert "stale-token" not in caplog.text + + within_backoff = mint(engine, spec) + assert within_backoff.access_token.get_secret_value() == "stale-token" + assert len(executor.pending) == 0 + + clock.advance(ADVISORY_REFRESH_BACKOFF_SECONDS) + after_backoff = mint(engine, spec) + assert after_backoff.access_token.get_secret_value() == "stale-token" + assert len(executor.pending) == 1 + executor.run_all() + assert len(poster.requests) == 3 + + +class GatedPoster: + """Blocks the leader inside post() until the test releases it.""" + + def __init__(self, response: httpx.Response) -> None: + self.entered = threading.Event() + self.release = threading.Event() + self.calls = 0 + self._calls_lock = threading.Lock() + self._response = response + + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + with self._calls_lock: + self.calls += 1 + self.entered.set() + assert self.release.wait(timeout=10) + return self._response + + +def _run_concurrent_get_token( + engine: JwtBearerTokenExchangeEngine, spec: TokenExchangeSpec, poster: GatedPoster, thread_count: int +) -> list[ExchangeResult]: + results: list[ExchangeResult] = [] + results_lock = threading.Lock() + start_barrier = threading.Barrier(thread_count) + + def worker() -> None: + start_barrier.wait() + result = engine.get_token(spec) + with results_lock: + results.append(result) + + threads = [threading.Thread(target=worker, daemon=True) for _ in range(thread_count)] + for thread in threads: + thread.start() + assert poster.entered.wait(timeout=10) + time.sleep(0.3) + poster.release.set() + for thread in threads: + thread.join(timeout=10) + assert not thread.is_alive() + return results + + +def test_mandatory_single_leader(): + poster = GatedPoster(token_response("leader-token")) + engine = make_engine(poster) + spec = make_spec() + + results = _run_concurrent_get_token(engine, spec, poster, thread_count=5) + + assert poster.calls == 1 + assert len(results) == 5 + for result in results: + assert isinstance(result, MintedToken) + assert result.access_token.get_secret_value() == "leader-token" + + +def test_mandatory_failure_is_value(): + poster = GatedPoster(httpx.Response(500, json={"error": "server_error"})) + engine = make_engine(poster) + spec = make_spec() + + results = _run_concurrent_get_token(engine, spec, poster, thread_count=3) + + assert len(results) == 3 + for result in results: + assert isinstance(result, TokenEndpointError) + assert result.status_code == 500 + assert "server_error" in result.redacted_body + + +def test_lock_released_around_io(): + inner_spec = make_spec( + token_url="https://inner.example/v1/oauth/token", + cache_key_identity=("fdrl_inner", "org-1", "", ""), + ) + engine_holder: dict[str, JwtBearerTokenExchangeEngine] = {} + inner_results: list[ExchangeResult] = [] + + class ReentrantPoster: + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + if url == "https://token.example/v1/oauth/token": + inner_results.append(engine_holder["engine"].get_token(inner_spec)) + return token_response() + + engine = make_engine(ReentrantPoster()) + engine_holder["engine"] = engine + + outcome: list[ExchangeResult] = [] + thread = threading.Thread(target=lambda: outcome.append(engine.get_token(make_spec())), daemon=True) + thread.start() + thread.join(timeout=10) + + assert not thread.is_alive(), "engine held its lock across poster I/O and deadlocked" + assert len(outcome) == 1 + assert isinstance(outcome[0], MintedToken) + assert len(inner_results) == 1 + assert isinstance(inner_results[0], MintedToken) + + +def test_401_retry_once_with_reread(): + assertions = {DEFAULT_REF: "assertion-v1"} + + def rotate_on_first_request(request: RecordedRequest) -> None: + assertions[DEFAULT_REF] = "assertion-v2" + + poster = ScriptedPoster( + [httpx.Response(401, json={"error": "invalid_grant"}), token_response()], + on_request=rotate_on_first_request, + ) + engine = make_engine(poster, reader=assertions.get) + + result = mint(engine, make_spec()) + + assert result.access_token.get_secret_value() == "sk-ant-oat01-minted" + assert len(poster.requests) == 2 + assert poster.requests[0].json_body()["assertion"] == "assertion-v1" + assert poster.requests[1].json_body()["assertion"] == "assertion-v2" + + +class RotatingAssertionSource: + """A per-call assertion source that mints a fresh value on every read -- the shape + internal_issuer/keycloak identity sources take (a fresh JWT/token minted per call).""" + + def __init__(self, values: list[str]) -> None: + self._values = iter(values) + self.calls = 0 + + def __call__(self) -> str: + self.calls += 1 + return next(self._values) + + +class EchoingUnauthorizedPoster: + """401s every attempt, echoing the submitted assertion back into the error body -- a + token endpoint that reflects the request.""" + + def __init__(self) -> None: + self.requests: list[RecordedRequest] = [] + + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + recorded = RecordedRequest(url, content, headers, timeout) + self.requests.append(recorded) + submitted = recorded.json_body()["assertion"] + return httpx.Response(401, json={"error": "invalid_grant", "error_description": f"bad assertion {submitted}"}) + + +def test_401_retry_redacts_the_assertion_actually_sent_not_a_fresh_reread(): + """Regression: with a rotating identity source, the reflection-drop check must match the + assertion the failing (second) attempt actually sent. Re-reading for the check would mint a + THIRD value that was never sent, so the reflection probe would miss and the actually-sent, + actually-reflected second assertion would leak into the error.""" + poster = EchoingUnauthorizedPoster() + source = RotatingAssertionSource(["assertion-v1", "assertion-v2", "assertion-v3"]) + engine = make_engine(poster) + spec = make_spec(assertion_source=source) + + result = engine.get_token(spec) + + assert isinstance(result, TokenEndpointError) + assert len(poster.requests) == 2 + assert poster.requests[0].json_body()["assertion"] == "assertion-v1" + assert poster.requests[1].json_body()["assertion"] == "assertion-v2" + assert source.calls == 2, "the failing attempt's own assertion must be reused, never re-read a third time" + assert "assertion-v1" not in result.redacted_body + assert "assertion-v2" not in result.redacted_body + assert "assertion-v3" not in result.redacted_body + + +def test_401_twice_is_endpoint_error(): + poster = ScriptedPoster([httpx.Response(401, json={"error": "invalid_grant"})]) + engine = make_engine(poster) + + result = engine.get_token(make_spec()) + + assert isinstance(result, TokenEndpointError) + assert result.status_code == 401 + assert "invalid_grant" in result.redacted_body + assert len(poster.requests) == 2 + + +class TestRedactionAndCaps: + def test_object_body_reduced_to_rfc6749_fields(self): + poster = ScriptedPoster( + [ + httpx.Response( + 400, + json={ + "error": "invalid_grant", + "error_description": "d" * 500, + "error_uri": "https://errors.example/e1", + "assertion_echo": "LEAKED-ASSERTION", + }, + ) + ] + ) + result = make_engine(poster).get_token(make_spec()) + + assert isinstance(result, TokenEndpointError) + assert result.status_code == 400 + assert "invalid_grant" in result.redacted_body + assert "d" * 256 in result.redacted_body + assert "d" * 257 not in result.redacted_body + assert "https://errors.example/e1" in result.redacted_body + assert "LEAKED-ASSERTION" not in result.redacted_body + + def test_nested_error_envelope_renders_readable_text(self): + body = { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": "federation_rule_id is not a well-formed fdrl_ tagged ID", + }, + } + result = redact_oauth_error_body(400, json.dumps(body)) + + assert "invalid_request_error" in result.redacted_body + assert "federation_rule_id is not a well-formed fdrl_ tagged ID" in result.redacted_body + assert "{'" not in result.redacted_body + + def test_flat_rfc6749_shape_still_renders(self): + body = {"error": "invalid_grant", "error_description": "bad request"} + result = redact_oauth_error_body(400, json.dumps(body)) + + assert result.redacted_body == "error: invalid_grant; error_description: bad request" + + def test_nested_error_message_is_capped_at_256_chars(self): + body = {"error": {"type": "invalid_request_error", "message": "m" * 500}} + result = redact_oauth_error_body(400, json.dumps(body)) + + assert "m" * 256 in result.redacted_body + assert "m" * 257 not in result.redacted_body + + def test_json_string_body_is_not_echoed(self): + """A free-text body can carry back whatever was sent, so only structured OAuth fields are + ever rendered into an error an operator or caller will see.""" + result = redact_oauth_error_body(400, json.dumps("s" * 500)) + assert result.redacted_body == "non-object error response omitted" + assert "s" * 32 not in result.redacted_body + + def test_plain_text_body_is_not_echoed(self): + result = redact_oauth_error_body(502, "t" * 500) + assert result.redacted_body == "non-JSON error response omitted" + assert "t" * 32 not in result.redacted_body + + def test_reflected_assertion_is_dropped(self): + """An endpoint that echoes the submitted assertion must not put it in the log or the error.""" + assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9.REFLECTEDPAYLOAD.signature") + body = {"error": "invalid_grant", "error_description": f"bad assertion {assertion.get_secret_value()}"} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert assertion.get_secret_value() not in result.redacted_body + assert "REFLECTEDPAYLOAD" not in result.redacted_body + + def test_assertion_reflected_from_an_offset_is_dropped(self): + """Regression: the probe only looked at the assertion's first 24 characters, so an + endpoint echoing it from any later offset shared no prefix and slipped through.""" + assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "A" * 40 + "PAYLOADMIDDLE" + "B" * 40 + ".signature") + tail = assertion.get_secret_value()[24:] + body = {"error": "invalid_grant", "error_description": tail} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert "PAYLOADMIDDLE" not in result.redacted_body + assert tail[:40] not in result.redacted_body + + def test_a_secret_carrying_spaces_is_dropped_when_echoed_whole(self): + """Regression on the redactor itself: comparing a compacted response against an + uncompacted secret stopped matching hand-set passphrases, which are exactly the secrets + most likely to be echoed and the ones an earlier contiguous match had caught.""" + assertion = SecretStr("correct horse battery staple, 42!") + echoed = assertion.get_secret_value() + body = {"error": "invalid_client", "error_description": f"secret {echoed} rejected"} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert echoed not in result.redacted_body + + def test_a_percent_encoded_secret_is_dropped(self): + """A form-encoded grant puts the secret on the wire percent-escaped, so an echo of that + shape has to be recognised without every caller enumerating it.""" + assertion = SecretStr("sUp3r+S3cret/Value=123") + echoed = quote(assertion.get_secret_value(), safe="") + body = {"error": "invalid_client", "error_description": f"rejected {echoed}"} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert echoed not in result.redacted_body + + def test_a_space_encoded_as_plus_is_dropped(self): + """A form-encoded body writes a space as "+", not %20, so percent-decoding alone does not + recover the secret and a passphrase echoed in its wire shape would travel on.""" + assertion = SecretStr("correct horse battery staple") + echoed = urlencode({"client_secret": assertion.get_secret_value()}).split("=", 1)[1] + body = {"error": "invalid_client", "error_description": f"rejected {echoed}"} + + assert "+" in echoed + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert echoed not in result.redacted_body + + def test_several_wire_forms_are_all_compared(self): + """The caller declares each shape it sent, since an encoding the redactor cannot reverse + (base64 of id:secret) is only knowable there.""" + raw = SecretStr("sUp3rS3cretValue123") + blob = SecretStr(base64.b64encode(b"litellm:sUp3rS3cretValue123").decode()) + body = {"error": "invalid_client", "error_description": f"bad {blob.get_secret_value()}"} + + result = redact_oauth_error_body(400, json.dumps(body), (raw, blob)) + + assert blob.get_secret_value() not in result.redacted_body + + def test_a_fragment_shorter_than_a_long_run_is_dropped(self): + """A slice too short to share a long contiguous run with the assertion is still assertion + material, and repeated errors would hand it over piece by piece.""" + assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "A" * 60 + ".sigsigsig") + fragment = assertion.get_secret_value()[30:48] + body = {"error": "invalid_grant", "error_description": f"rejected near {fragment}"} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert fragment not in result.redacted_body + + def test_a_fragment_broken_up_by_delimiters_is_dropped(self): + """Splitting the echo defeats a contiguous match, so the comparison ignores whatever the + endpoint put between the pieces.""" + assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "A" * 60 + ".sigsigsig") + piece = assertion.get_secret_value()[20:44] + spaced = " ".join(piece[i : i + 6] for i in range(0, 24, 6)) + body = {"error": "invalid_grant", "error_description": spaced} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert spaced not in result.redacted_body + + def test_a_short_secret_is_still_matched_whole(self): + """A Keycloak client secret can be shorter than the probe length; the whole value is + compared in that case rather than a truncated prefix.""" + secret = SecretStr("short-secret") + body = {"error": "invalid_client", "error_description": "rejected short-secret"} + + result = redact_oauth_error_body(400, json.dumps(body), secret) + + assert "short-secret" not in result.redacted_body + + def test_a_short_secret_echoed_in_its_wire_shape_is_dropped(self): + """Regression: the run scan only ever compared eight-character windows, so a secret + with fewer credential characters than that could never match once it came back + percent-encoded rather than verbatim, and the whole-value check needs the raw form.""" + secret = SecretStr("p@ss w0rd!") + echoed = quote(secret.get_secret_value(), safe="") + body = {"error": "invalid_client", "error_description": f"rejected {echoed}"} + + assert secret.get_secret_value() not in echoed + result = redact_oauth_error_body(400, json.dumps(body), secret) + + assert echoed not in result.redacted_body + + def test_an_unrelated_body_is_not_falsely_redacted(self): + """The scan must not fire on a body that merely shares short runs with the assertion.""" + assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "Z" * 60 + ".signature") + body = {"error": "invalid_grant", "error_description": "the federation rule was not found"} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert "the federation rule was not found" in result.redacted_body + + def test_json_array_body_constant_message(self): + result = redact_oauth_error_body(400, json.dumps(["a", "b"])) + assert result.redacted_body == "non-object error response omitted" + + def test_oversized_body_never_parsed(self): + poster = ScriptedPoster([httpx.Response(400, content=b'{"error": "' + b"x" * MAX_RESPONSE_BYTES + b'"}')]) + result = make_engine(poster).get_token(make_spec()) + + assert isinstance(result, TokenEndpointError) + assert result.redacted_body == "oversized error response omitted" + + def test_oversized_success_body_is_malformed(self): + poster = ScriptedPoster( + [httpx.Response(200, content=b'{"access_token": "' + b"x" * MAX_RESPONSE_BYTES + b'"}')] + ) + result = make_engine(poster).get_token(make_spec()) + + assert not isinstance(result, MintedToken) + assert b"x" * 10 not in str(result).encode() + + +@pytest.mark.parametrize("access_token", ["", " "]) +def test_empty_access_token_is_malformed(access_token: str): + poster = ScriptedPoster( + [httpx.Response(200, json={"access_token": access_token, "token_type": "Bearer", "expires_in": 3600})] + ) + result = make_engine(poster).get_token(make_spec()) + + assert isinstance(result, MalformedTokenResponse) + assert "empty access_token" in result.detail + + +def test_sentinel_leak_audit(caplog: pytest.LogCaptureFixture): + jwt_sentinel = "JWT-SENTINEL-2c9f1e7ab4" + token_sentinel = "sk-ant-oat01-TOKEN-SENTINEL-90d4c3aa17" + ref = "oidc/env/SENTINEL_ASSERTION" + + with caplog.at_level(logging.DEBUG): + success_poster = ScriptedPoster([token_response(token_sentinel, expires_in=3600)]) + success_clock = FakeClock() + success_executor = ManualExecutor() + engine = make_engine(success_poster, reader={ref: jwt_sentinel}, clock=success_clock, executor=success_executor) + spec = make_spec(assertion_ref=ref) + minted = mint(engine, spec) + + endpoint_error = make_engine( + ScriptedPoster([httpx.Response(400, json={"error": "invalid_grant"})]), reader={ref: jwt_sentinel} + ).get_token(spec) + transport_error = make_engine(RaisingPoster(RuntimeError("boom")), reader={ref: jwt_sentinel}).get_token(spec) + malformed_error = make_engine( + ScriptedPoster([httpx.Response(200, json={"unexpected": "shape"})]), reader={ref: jwt_sentinel} + ).get_token(spec) + oversized_error = make_engine( + ScriptedPoster([token_response()]), reader={ref: jwt_sentinel + "x" * MAX_ASSERTION_BYTES} + ).get_token(spec) + insecure_error = make_engine(ScriptedPoster([token_response()]), reader={ref: jwt_sentinel}).get_token( + make_spec(assertion_ref=ref, token_url="http://token.example/v1/oauth/token") + ) + + success_poster._responses = [httpx.Response(500, json={"error": "server_error"})] + success_clock.now = success_clock.now + 3600 - 100.0 + stale = engine.get_token(spec) + success_executor.run_all() + + audited_values = [ + str(minted), + repr(minted), + str(minted.access_token), + repr(minted.access_token), + str(endpoint_error), + repr(endpoint_error), + str(transport_error), + repr(transport_error), + str(malformed_error), + repr(malformed_error), + str(oversized_error), + repr(oversized_error), + str(insecure_error), + repr(insecure_error), + str(stale), + repr(stale), + caplog.text, + ] + assert isinstance(oversized_error, AssertionSourceError) + assert oversized_error.kind == "oversized" + for value in audited_values: + assert jwt_sentinel not in value + assert token_sentinel not in value + + +class TestAssertionGuards: + @pytest.mark.parametrize( + "assertion_value,expected_kind", + [ + ("x" * (MAX_ASSERTION_BYTES + 1), "oversized"), + (" \n\t ", "empty"), + (None, "missing"), + ], + ) + def test_bad_assertion_values(self, assertion_value: str | None, expected_kind: str): + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster, reader=lambda ref: assertion_value) + + result = engine.get_token(make_spec()) + + assert isinstance(result, AssertionSourceError) + assert result.kind == expected_kind + assert result.source_ref == DEFAULT_REF + assert len(poster.requests) == 0 + + @pytest.mark.parametrize( + "raised,expected_kind", + [ + (OidcPathNotAllowedError("path outside allowed credential directories"), "disallowed_path"), + (ValueError("Environment variable ANTHROPIC_IDENTITY_TOKEN not found"), "unreadable"), + (ImportError("needs PyJWT and cryptography: pip install 'litellm[proxy]'"), "unreadable"), + (OSError("permission denied"), "unreadable"), + ], + ) + def test_raising_reader(self, raised: Exception, expected_kind: str): + poster = ScriptedPoster([token_response()]) + + def reader(ref: str) -> str | None: + raise raised + + result = make_engine(poster, reader=reader).get_token(make_spec()) + + assert isinstance(result, AssertionSourceError) + assert result.kind == expected_kind + assert len(poster.requests) == 0 + + def test_value_error_message_is_captured_as_detail(self): + poster = ScriptedPoster([token_response()]) + + def reader(ref: str) -> str | None: + raise ValueError("Keycloak token endpoint returned invalid_client") + + result = make_engine(poster, reader=reader).get_token(make_spec()) + + assert isinstance(result, AssertionSourceError) + assert result.detail == "Keycloak token endpoint returned invalid_client" + + def test_import_error_message_is_captured_as_detail(self): + poster = ScriptedPoster([token_response()]) + + def reader(ref: str) -> str | None: + raise ImportError("the internal_issuer identity source needs PyJWT and cryptography: pip install 'litellm[proxy]'") + + result = make_engine(poster, reader=reader).get_token(make_spec()) + + assert isinstance(result, AssertionSourceError) + assert result.detail is not None + assert "litellm[proxy]" in result.detail + + @pytest.mark.parametrize( + "raised", + [OidcPathNotAllowedError("path outside allowed credential directories"), OSError("permission denied")], + ) + def test_non_value_error_never_populates_detail(self, raised: Exception): + """Only the ValueError branch carries operator-diagnosable text; every other reader failure + stays detail=None, matching today's file/env behavior byte-for-byte.""" + poster = ScriptedPoster([token_response()]) + + def reader(ref: str) -> str | None: + raise raised + + result = make_engine(poster, reader=reader).get_token(make_spec()) + + assert isinstance(result, AssertionSourceError) + assert result.detail is None + + def test_value_error_detail_is_capped(self): + poster = ScriptedPoster([token_response()]) + overlong_message = "x" * (_REDACTION_CAP + 100) + + def reader(ref: str) -> str | None: + raise ValueError(overlong_message) + + result = make_engine(poster, reader=reader).get_token(make_spec()) + + assert isinstance(result, AssertionSourceError) + assert result.detail == overlong_message[:_REDACTION_CAP] + + +class TestAssertionSourceOverridesEngineReader: + """``TokenExchangeSpec.assertion_source`` is the dispatch mechanism a per-config identity + source (internal_issuer, keycloak) plugs into the shared engine with -- it must win over the + engine-level reader, and failures must still be reported against ``assertion_ref``.""" + + def test_assertion_source_is_used_instead_of_the_reader(self): + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster, reader=lambda ref: "from-engine-reader") + spec = make_spec(assertion_source=lambda: "from-assertion-source") + + result = mint(engine, spec) + + assert result.access_token.get_secret_value() == "sk-ant-oat01-minted" + assert poster.requests[0].json_body()["assertion"] == "from-assertion-source" + + def test_reader_is_never_called_when_assertion_source_is_set(self): + poster = ScriptedPoster([token_response()]) + calls: list[str] = [] + + def reader(ref: str) -> str | None: + calls.append(ref) + return "from-engine-reader" + + engine = make_engine(poster, reader=reader) + spec = make_spec(assertion_source=lambda: "from-assertion-source") + + mint(engine, spec) + + assert calls == [] + + def test_assertion_source_failure_is_reported_against_assertion_ref(self): + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster, reader=lambda ref: "from-engine-reader") + + def raising_source() -> str | None: + raise ValueError("keycloak token endpoint returned invalid_client") + + spec = make_spec(assertion_source=raising_source, assertion_ref="oidc/keycloak/abc123") + + result = engine.get_token(spec) + + assert isinstance(result, AssertionSourceError) + assert result.source_ref == "oidc/keycloak/abc123" + assert result.detail == "keycloak token endpoint returned invalid_client" + assert len(poster.requests) == 0 + + def test_assertion_source_is_re_invoked_on_401_retry(self): + """The retry's second attempt must also prefer ``assertion_source`` for the assertion it + sends, not silently fall back to the engine reader.""" + values = iter(["assertion-v1", "assertion-v2"]) + poster = ScriptedPoster([httpx.Response(401, json={"error": "invalid_grant"}), token_response()]) + engine = make_engine(poster, reader=lambda ref: "from-engine-reader") + spec = make_spec(assertion_source=lambda: next(values)) + + result = mint(engine, spec) + + assert result.access_token.get_secret_value() == "sk-ant-oat01-minted" + assert poster.requests[0].json_body()["assertion"] == "assertion-v1" + assert poster.requests[1].json_body()["assertion"] == "assertion-v2" + + +class TestOidcFilePathAllowlistRaisesTypedError: + """The engine classifies assertion-source failures by exception type (see + TestAssertionGuards.test_raising_reader); that classification only works if the real + oidc/file allowlist actually raises OidcPathNotAllowedError rather than a bare ValueError.""" + + def test_out_of_allowlist_absolute_path(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", raising=False) + + with pytest.raises(OidcPathNotAllowedError): + _resolve_oidc_file_path("/etc/not-a-credential-dir/token") + + def test_relative_path(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", raising=False) + + with pytest.raises(OidcPathNotAllowedError): + _resolve_oidc_file_path("relative/token/path") + + +class TestHttpsEnforcement: + def test_plain_http_rejected_host_only_zero_posts(self): + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster) + + result = engine.get_token(make_spec(token_url="http://token.example/v1/oauth/token")) + + assert result == InsecureTokenUrl(host="token.example") + assert "/v1/oauth/token" not in str(result) + assert len(poster.requests) == 0 + + @pytest.mark.parametrize( + "url", + [ + "http://localhost:8080/v1/oauth/token", + "http://127.0.0.1/v1/oauth/token", + "http://[::1]/v1/oauth/token", + ], + ) + def test_localhost_http_allowed(self, url: str): + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster) + + result = engine.get_token(make_spec(token_url=url)) + + assert isinstance(result, MintedToken) + assert len(poster.requests) == 1 + + +def test_cache_key_semantics(): + poster = ScriptedPoster([token_response()]) + assertions = {DEFAULT_REF: DEFAULT_ASSERTION, "oidc/env/OTHER": "other-assertion"} + engine = make_engine(poster, reader=assertions.get) + base_spec = make_spec() + + mint(engine, base_spec) + mint(engine, make_spec(cache_key_identity=("fdrl_1", "org-1", "svc-2", ""))) + mint(engine, make_spec(token_url="https://other.example/v1/oauth/token")) + mint(engine, make_spec(assertion_ref="oidc/env/OTHER")) + assert len(poster.requests) == 4 + + assertions[DEFAULT_REF] = "rotated-assertion" + cached = mint(engine, base_spec) + assert len(poster.requests) == 4 + assert cached.access_token.get_secret_value() == "sk-ant-oat01-minted" + + +def test_the_cache_returns_to_its_bound_after_an_all_in_flight_burst(): + """An entry a leader owns is never evictable, so a burst of distinct identities can push the map + past max_entries. It must come back down once those entries are idle, rather than holding the + high-water mark for the life of the process.""" + clock = FakeClock() + engine = make_engine(ScriptedPoster([token_response(expires_in=3600)]), clock=clock, max_entries=4) + + def spec_for(index: int) -> TokenExchangeSpec: + return make_spec(cache_key_identity=("fdrl_1", f"org-{index}", "", "")) + + for index in range(12): + mint(engine, spec_for(index)) + + assert len(engine._entries) <= 4, ( # noqa: SLF001 # the bound under test is internal state + f"the cap is enforced once entries are idle, saw {len(engine._entries)}" + ) + + +def test_bounded_eviction(): + clock = FakeClock() + + class PerCallPoster: + def __init__(self) -> None: + self.calls = 0 + + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + self.calls += 1 + body = json.loads(content) + expires_in = 3600 + int(body["organization_id"].split("-")[1]) + return token_response(f"token-{body['organization_id']}", expires_in=expires_in) + + poster = PerCallPoster() + engine = make_engine(poster, clock=clock, max_entries=64) + + def spec_for(index: int) -> TokenExchangeSpec: + return make_spec( + static_body={ + "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", + "federation_rule_id": "fdrl_1", + "organization_id": f"org-{index}", + }, + cache_key_identity=("fdrl_1", f"org-{index}", "", ""), + ) + + for index in range(65): + mint(engine, spec_for(index)) + assert poster.calls == 65 + + mint(engine, spec_for(0)) + assert poster.calls == 66, "the earliest-expiring entry (index 0) should have been evicted" + + mint(engine, spec_for(2)) + assert poster.calls == 66, "a later-expiring entry should still be cached" + + mint(engine, spec_for(1)) + assert poster.calls == 67, "re-inserting index 0 should have evicted the next earliest-expiring entry" + + +@pytest.mark.parametrize("expires_in", [None, 0, -5]) +def test_missing_or_nonsense_expires_in_gets_fallback_ttl(expires_in: int | None): + poster = ScriptedPoster( + [token_response("short-lived", expires_in=expires_in), token_response("reminted", expires_in=3600)] + ) + clock = FakeClock(start=1_000.0) + engine = make_engine(poster, clock=clock) + spec = make_spec() + + first = mint(engine, spec) + assert first.expires_at == 1_000.0 + FALLBACK_TOKEN_TTL_SECONDS + + clock.advance(FALLBACK_TOKEN_TTL_SECONDS + 1.0) + second = mint(engine, spec) + + assert second.access_token.get_secret_value() == "reminted" + assert len(poster.requests) == 2, "a token without a sane expires_in must never be cached forever" + + +async def test_aget_token_loop_responsive(): + class SleepingPoster: + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + time.sleep(0.3) + return token_response() + + engine = make_engine(SleepingPoster()) + spec = make_spec() + ticks = {"count": 0} + stop = asyncio.Event() + + async def ticker() -> None: + while not stop.is_set(): + ticks["count"] += 1 + await asyncio.sleep(0.01) + + ticker_task = asyncio.create_task(ticker()) + result = await engine.aget_token(spec) + stop.set() + await ticker_task + + assert isinstance(result, MintedToken) + assert result.access_token.get_secret_value() == "sk-ant-oat01-minted" + assert ticks["count"] >= 5, "the event loop was blocked during aget_token" + sync_result = engine.get_token(spec) + assert sync_result == result + + +def test_invalidate_forces_refresh(): + poster = ScriptedPoster([token_response("token-1", expires_in=3600), token_response("token-2", expires_in=3600)]) + engine = make_engine(poster) + spec = make_spec() + + first = mint(engine, spec) + assert first.access_token.get_secret_value() == "token-1" + + engine.invalidate(spec) + second = mint(engine, spec) + assert second.access_token.get_secret_value() == "token-2" + assert len(poster.requests) == 2 + + third = mint(engine, spec) + assert third.access_token.get_secret_value() == "token-2" + assert len(poster.requests) == 2, "force_refresh must be one-shot" + + +def test_invalidate_unknown_spec_is_noop(): + poster = ScriptedPoster([token_response()]) + engine = make_engine(poster) + + engine.invalidate(make_spec()) + + assert len(poster.requests) == 0 + + +def test_advisory_failure_wakes_expired_follower_to_re_lead(): + poster = ScriptedPoster( + [ + token_response("initial-token", expires_in=3600), + httpx.Response(500, json={"error": "server_error"}), + token_response("recovered-token", expires_in=3600), + ] + ) + clock = FakeClock(start=1_000.0) + executor = ManualExecutor() + engine = make_engine(poster, clock=clock, executor=executor) + spec = make_spec() + + mint(engine, spec) + clock.now = 1_000.0 + 3600 - 100.0 + mint(engine, spec) + assert len(executor.pending) == 1 + + clock.advance(200.0) + results: list[ExchangeResult] = [] + follower = threading.Thread(target=lambda: results.append(engine.get_token(spec)), daemon=True) + follower.start() + time.sleep(0.3) + executor.run_all() + follower.join(timeout=10) + + assert not follower.is_alive() + assert len(results) == 1 + result = results[0] + assert isinstance(result, MintedToken), f"follower was handed {result!r} instead of re-leading a fresh mint" + assert result.access_token.get_secret_value() == "recovered-token" + assert len(poster.requests) == 3 + + +class TwoAttemptGatedPoster: + """401 on the first attempt, then blocks the leader's retry until released.""" + + def __init__(self) -> None: + self.entered_second = threading.Event() + self.release = threading.Event() + self.calls = 0 + + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + self.calls += 1 + if self.calls == 1: + return httpx.Response(401, json={"error": "invalid_grant"}) + self.entered_second.set() + assert self.release.wait(timeout=30) + return token_response("slow-leader-token") + + +def test_follower_budget_outlasts_slow_two_attempt_leader(): + poster = TwoAttemptGatedPoster() + engine = make_engine(poster) + spec = make_spec(timeout_seconds=1.0) + + leader_results: list[ExchangeResult] = [] + leader = threading.Thread(target=lambda: leader_results.append(engine.get_token(spec)), daemon=True) + leader.start() + assert poster.entered_second.wait(timeout=10) + + follower_results: list[ExchangeResult] = [] + follower = threading.Thread(target=lambda: follower_results.append(engine.get_token(spec)), daemon=True) + follower.start() + time.sleep(6.5) + poster.release.set() + leader.join(timeout=10) + follower.join(timeout=10) + + assert leader_results and isinstance(leader_results[0], MintedToken) + assert follower_results, "follower never returned" + follower_result = follower_results[0] + assert isinstance(follower_result, MintedToken), ( + f"follower gave up before the leader's two-attempt worst case: {follower_result!r}" + ) + assert follower_result.access_token.get_secret_value() == "slow-leader-token" + + +class FailThenGatePoster: + """500 on the first call, then blocks until released before succeeding.""" + + def __init__(self) -> None: + self.entered_gate = threading.Event() + self.release = threading.Event() + self.calls = 0 + + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + self.calls += 1 + if self.calls == 1: + return httpx.Response(500, json={"error": "server_error"}) + self.entered_gate.set() + assert self.release.wait(timeout=30) + return token_response("round-two-token") + + +def test_new_round_timed_out_follower_never_returns_previous_rounds_error(): + poster = FailThenGatePoster() + clock = FakeClock() + engine = make_engine(poster, clock=clock) + spec = make_spec(timeout_seconds=0.05) + + first = engine.get_token(spec) + assert isinstance(first, TokenEndpointError) + + clock.advance(ADVISORY_REFRESH_BACKOFF_SECONDS + 1.0) + leader = threading.Thread(target=lambda: engine.get_token(spec), daemon=True) + leader.start() + assert poster.entered_gate.wait(timeout=10) + + follower_result = engine.get_token(spec) + + assert isinstance(follower_result, TokenTransportError), ( + f"timed-out follower returned the previous round's error: {follower_result!r}" + ) + assert "timed out" in follower_result.detail + poster.release.set() + leader.join(timeout=10) + + +def test_lead_backoff_fails_fast_within_window_and_expires_after(): + poster = ScriptedPoster([httpx.Response(500, json={"error": "server_error"})]) + clock = FakeClock() + engine = make_engine(poster, clock=clock) + spec = make_spec() + + first = engine.get_token(spec) + assert isinstance(first, TokenEndpointError) + assert len(poster.requests) == 1 + + clock.advance(ADVISORY_REFRESH_BACKOFF_SECONDS - 1.0) + second = engine.get_token(spec) + assert second == first + assert len(poster.requests) == 1, "a request inside the backoff window must make zero POSTs" + + clock.advance(1.0) + third = engine.get_token(spec) + assert isinstance(third, TokenEndpointError) + assert len(poster.requests) == 2 + + +def test_invalidate_bypasses_lead_backoff(): + poster = ScriptedPoster( + [httpx.Response(500, json={"error": "server_error"}), token_response("post-invalidate", expires_in=3600)] + ) + clock = FakeClock() + engine = make_engine(poster, clock=clock) + spec = make_spec() + + first = engine.get_token(spec) + assert isinstance(first, TokenEndpointError) + + engine.invalidate(spec) + second = engine.get_token(spec) + + assert isinstance(second, MintedToken) + assert second.access_token.get_secret_value() == "post-invalidate" + assert len(poster.requests) == 2 + + +class StubExchangeHandler: + """Stands in for the HTTPHandler the default poster builds, so the poster's own contract is + testable without a socket.""" + + def __init__(self, result: httpx.Response | Exception | None) -> None: + self.calls = 0 + self._result = result + + def post(self, url: str, *, content: bytes, headers: dict[str, str], timeout: float) -> httpx.Response | None: + self.calls += 1 + if isinstance(self._result, Exception): + raise self._result + return self._result + + +class TestDefaultTokenPoster: + def test_builds_its_handler_once_and_reuses_it(self): + built: list[StubExchangeHandler] = [] + + def factory() -> StubExchangeHandler: + handler = StubExchangeHandler(httpx.Response(200, json={"access_token": "t"})) + built.append(handler) + return handler + + poster: Final = _HttpxSyncTokenPoster(handler_factory=factory) # pyright: ignore[reportArgumentType] # StubExchangeHandler stands in for the legacy-untyped HTTPHandler + for _ in range(3): + poster.post(EXCHANGE_URL, content=b"", headers={}, timeout=1.0) + + assert len(built) == 1 + assert built[0].calls == 3 + + def test_the_real_handler_refuses_to_follow_redirects(self): + assert _new_exchange_handler().client.follow_redirects is False, ( + "a redirected exchange POST would replay the workload assertion to the redirect target" + ) + + def test_an_http_status_error_becomes_its_response(self): + response: Final = httpx.Response( + 401, json={"error": "invalid_grant"}, request=httpx.Request("POST", EXCHANGE_URL) + ) + poster: Final = _HttpxSyncTokenPoster( + handler_factory=lambda: StubExchangeHandler( # pyright: ignore[reportArgumentType] # StubExchangeHandler stands in for the legacy-untyped HTTPHandler + httpx.HTTPStatusError("boom", request=response.request, response=response) + ) + ) + + assert poster.post(EXCHANGE_URL, content=b"", headers={}, timeout=1.0).status_code == 401 + + def test_a_missing_response_is_a_transport_error(self): + poster: Final = _HttpxSyncTokenPoster(handler_factory=lambda: StubExchangeHandler(None)) # pyright: ignore[reportArgumentType] # StubExchangeHandler stands in for the legacy-untyped HTTPHandler + + with pytest.raises(httpx.TransportError): + poster.post(EXCHANGE_URL, content=b"", headers={}, timeout=1.0) + + +class TestDefaultAssertionReader: + def test_reads_through_litellm_secret_resolution(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("WIF_ASSERTION_FOR_DEFAULT_READER", "header.payload.signature") + + assert _default_assertion_reader("os.environ/WIF_ASSERTION_FOR_DEFAULT_READER") == "header.payload.signature" + + def test_an_unset_reference_reads_as_none(self): + assert _default_assertion_reader("os.environ/DEFINITELY_NOT_SET_WIF_ASSERTION_REF") is None + + +class TestErrorSummary: + def test_every_error_variant_summarises_without_carrying_a_secret(self): + summaries: Final = { + _error_summary(AssertionSourceError(kind="unreadable", source_ref="oidc/file/x")), + _error_summary(InsecureTokenUrl(host="token.internal")), + _error_summary(TokenEndpointError(status_code=401, redacted_body="invalid_grant")), + _error_summary(TokenTransportError(detail="ConnectError: refused")), + _error_summary(MalformedTokenResponse(detail="empty access_token")), + } + + assert {s.split(":")[0] for s in summaries} == { + "AssertionSourceError", + "InsecureTokenUrl", + "TokenEndpointError", + "TokenTransportError", + "MalformedTokenResponse", + }, "each variant names itself so a log line says which stage failed" + + +class TestNonBearerTokenType: + def test_a_non_bearer_token_type_is_refused(self): + poster: Final = ScriptedPoster( + [httpx.Response(200, json={"access_token": "tok", "token_type": "mac", "expires_in": 300})] + ) + engine: Final = JwtBearerTokenExchangeEngine(poster=poster, assertion_reader=lambda _ref: DEFAULT_ASSERTION) + + result: Final = engine.get_token(make_spec()) + + assert isinstance(result, MalformedTokenResponse) + assert "non-bearer" in result.detail + + +class TestShortLivedRefreshWindows: + """A token whose lifetime is at or below the flat 120s advisory window used to be inside that + window from birth, so every request armed another background exchange. The windows now scale + with the observed lifetime; long-lived tokens must keep the flat 120s/30s behaviour.""" + + @staticmethod + def _engine_with( + expires_in: int | None, + ) -> tuple[JwtBearerTokenExchangeEngine, ScriptedPoster, FakeClock, ManualExecutor, TokenExchangeSpec]: + poster = ScriptedPoster([token_response("short-lived", expires_in=expires_in), token_response("reminted")]) + clock = FakeClock(start=1_000.0) + executor = ManualExecutor() + engine = make_engine(poster, clock=clock, executor=executor) + return engine, poster, clock, executor, make_spec() + + def test_fallback_ttl_token_is_served_without_arming_a_refresh(self): + engine, poster, clock, executor, spec = self._engine_with(expires_in=None) + + first = mint(engine, spec) + assert first.expires_at == 1_000.0 + FALLBACK_TOKEN_TTL_SECONDS + + for _ in range(5): + clock.advance(1.0) + assert mint(engine, spec).access_token.get_secret_value() == "short-lived" + + assert executor.pending == [], "a freshly minted fallback-TTL token must not arm a refresh on every request" + assert len(poster.requests) == 1 + + @pytest.mark.parametrize( + "elapsed,expect_advisory_submit", + [(29.0, False), (30.0, True), (52.0, True)], + ) + def test_fallback_ttl_token_refreshes_around_its_half_life(self, elapsed: float, expect_advisory_submit: bool): + engine, poster, clock, executor, spec = self._engine_with(expires_in=None) + + mint(engine, spec) + clock.advance(elapsed) + served = mint(engine, spec) + + assert served.access_token.get_secret_value() == "short-lived" + assert len(executor.pending) == (1 if expect_advisory_submit else 0) + executor.run_all() + assert len(poster.requests) == (2 if expect_advisory_submit else 1) + + @pytest.mark.parametrize( + "elapsed,expect_new_token", + [(52.0, False), (53.0, True)], + ) + def test_fallback_ttl_mandatory_wall_scales_with_the_lifetime(self, elapsed: float, expect_new_token: bool): + engine, poster, clock, executor, spec = self._engine_with(expires_in=None) + + mint(engine, spec) + clock.advance(elapsed) + served = mint(engine, spec) + + assert served.access_token.get_secret_value() == ("reminted" if expect_new_token else "short-lived") + assert len(executor.pending) == (0 if expect_new_token else 1) + + @pytest.mark.parametrize( + "elapsed,expect_advisory_submit", + [(89.0, False), (100.0, True)], + ) + def test_a_200s_token_scales_its_advisory_window_too(self, elapsed: float, expect_advisory_submit: bool): + engine, poster, clock, executor, spec = self._engine_with(expires_in=200) + + mint(engine, spec) + clock.advance(elapsed) + served = mint(engine, spec) + + assert served.access_token.get_secret_value() == "short-lived" + assert len(executor.pending) == (1 if expect_advisory_submit else 0) + + @pytest.mark.parametrize("expires_in", [240, 3600]) + @pytest.mark.parametrize( + "remaining,expect_advisory_submit,expect_new_token", + [ + (121.0, False, False), + (120.0, True, False), + (31.0, True, False), + (30.0, False, True), + ], + ) + def test_long_lived_tokens_keep_the_flat_windows( + self, expires_in: int, remaining: float, expect_advisory_submit: bool, expect_new_token: bool + ): + engine, poster, clock, executor, spec = self._engine_with(expires_in=expires_in) + + mint(engine, spec) + clock.now = 1_000.0 + expires_in - remaining + served = mint(engine, spec) + + assert len(executor.pending) == (1 if expect_advisory_submit else 0) + assert served.access_token.get_secret_value() == ("reminted" if expect_new_token else "short-lived") + assert len(poster.requests) == (2 if expect_new_token else 1) + + +class RaisingMetricsSink: + def exchange_success(self, *, call_type: str, duration_seconds: float) -> None: + raise RuntimeError("metrics sink down") + + def exchange_failure(self, *, call_type: str, duration_seconds: float, error: ExchangeError) -> None: + raise RuntimeError("metrics sink down") + + def cache_hit(self) -> None: + raise RuntimeError("metrics sink down") + + +class TestMetricsEmission: + def test_cold_mint_emits_success_with_duration(self): + clock = FakeClock() + sink = RecordingMetricsSink() + poster = ScriptedPoster([token_response()], on_request=lambda _request: clock.advance(0.25)) + engine = make_engine(poster, clock=clock, metrics_sink=sink) + + mint(engine, make_spec()) + + assert sink.successes == [("cold_mint", 0.25)] + assert sink.failures == [] + assert sink.cache_hits == 0 + + def test_cache_hit_emits_counter_not_a_mint(self): + clock = FakeClock() + sink = RecordingMetricsSink() + engine = make_engine(ScriptedPoster([token_response()]), clock=clock, metrics_sink=sink) + spec = make_spec() + + mint(engine, spec) + clock.advance(100.0) + mint(engine, spec) + + assert sink.cache_hits == 1 + assert len(sink.successes) == 1 + + def test_advisory_refresh_call_type(self): + clock = FakeClock(start=1_000.0) + sink = RecordingMetricsSink() + executor = ManualExecutor() + poster = ScriptedPoster([token_response("old", expires_in=3600), token_response("new")]) + engine = make_engine(poster, clock=clock, executor=executor, metrics_sink=sink) + spec = make_spec() + + mint(engine, spec) + clock.now = 1_000.0 + 3600 - 119.0 + mint(engine, spec) + executor.run_all() + + assert [call_type for call_type, _ in sink.successes] == ["cold_mint", "advisory_refresh"] + assert sink.cache_hits == 1 + + def test_mandatory_refresh_call_type(self): + clock = FakeClock(start=1_000.0) + sink = RecordingMetricsSink() + poster = ScriptedPoster([token_response("old", expires_in=3600), token_response("new")]) + engine = make_engine(poster, clock=clock, metrics_sink=sink) + spec = make_spec() + + mint(engine, spec) + clock.now = 1_000.0 + 3600 - 29.0 + mint(engine, spec) + + assert [call_type for call_type, _ in sink.successes] == ["cold_mint", "mandatory_refresh"] + assert sink.cache_hits == 0 + + def test_failed_exchange_emits_failure_once_and_negative_cache_does_not_reemit(self): + sink = RecordingMetricsSink() + poster = ScriptedPoster([httpx.Response(503, json={"error": "unavailable"})]) + engine = make_engine(poster, metrics_sink=sink) + spec = make_spec() + + first = engine.get_token(spec) + second = engine.get_token(spec) + + assert isinstance(first, TokenEndpointError) + assert isinstance(second, TokenEndpointError) + assert len(sink.failures) == 1 + call_type, _duration, error = sink.failures[0] + assert call_type == "cold_mint" + assert isinstance(error, TokenEndpointError) + assert error.status_code == 503 + assert sink.successes == [] + + def test_failure_payload_carries_no_assertion_material(self): + sink = RecordingMetricsSink() + engine = make_engine(EchoingUnauthorizedPoster(), metrics_sink=sink) + + result = engine.get_token(make_spec()) + + assert isinstance(result, TokenEndpointError) + (failure,) = sink.failures + assert DEFAULT_ASSERTION not in repr(failure) + assert DEFAULT_ASSERTION not in _error_summary(failure[2]) + + def test_raising_sink_never_breaks_mint_serve_or_failure(self): + clock = FakeClock() + engine = make_engine(ScriptedPoster([token_response()]), clock=clock, metrics_sink=RaisingMetricsSink()) + spec = make_spec() + + minted = mint(engine, spec) + clock.advance(100.0) + served = mint(engine, spec) + + assert served.access_token.get_secret_value() == minted.access_token.get_secret_value() + + failing = make_engine(RaisingPoster(httpx.ConnectError("boom")), metrics_sink=RaisingMetricsSink()) + result = failing.get_token(make_spec()) + assert isinstance(result, TokenTransportError) + + +class RecordingServiceHooks: + def __init__(self) -> None: + self.successes: list[tuple[ServiceTypes, str, float]] = [] + self.failures: list[tuple[ServiceTypes, float, str | Exception, str]] = [] + + async def async_service_success_hook(self, service: ServiceTypes, call_type: str, duration: float) -> None: + self.successes.append((service, call_type, duration)) + + async def async_service_failure_hook( + self, service: ServiceTypes, duration: float, error: str | Exception, call_type: str + ) -> None: + self.failures.append((service, duration, error, call_type)) + + +class RaisingServiceHooks: + """Every hook raises, and each call is recorded first so a test can prove the sink kept + calling through rather than bailing after the first failure.""" + + def __init__(self) -> None: + self.attempts: list[str] = [] # mutable-ok: a test spy accumulating calls in order + + async def async_service_success_hook(self, service: ServiceTypes, call_type: str, duration: float) -> None: + self.attempts.append(f"success:{call_type}") + raise RuntimeError("hook down") + + async def async_service_failure_hook( + self, service: ServiceTypes, duration: float, error: str | Exception, call_type: str + ) -> None: + self.attempts.append(f"failure:{call_type}") + raise RuntimeError("hook down") + + +class TestServiceLoggingMetricsSink: + def _sink(self, hooks) -> ServiceLoggingMetricsSink: + return ServiceLoggingMetricsSink(service_logging_factory=lambda: hooks, executor=InlineExecutor()) + + def test_success_maps_to_anthropic_wif_service(self): + hooks = RecordingServiceHooks() + + self._sink(hooks).exchange_success(call_type="cold_mint", duration_seconds=0.2) + + assert hooks.successes == [(ServiceTypes.ANTHROPIC_WIF, "cold_mint", 0.2)] + + def test_a_stalled_backend_stops_accepting_work_instead_of_queueing_without_bound(self): + stalled: Final = NeverRunsExecutor() + sink: Final = ServiceLoggingMetricsSink(service_logging_factory=RecordingServiceHooks, executor=stalled) + + for _ in range(_METRICS_QUEUE_LIMIT + 500): + sink.cache_hit() + + assert stalled.submitted == _METRICS_QUEUE_LIMIT, ( + "once the backlog is full further events are dropped, so request volume cannot grow it" + ) + + def test_a_drained_backlog_accepts_work_again(self): + hooks: Final = RecordingServiceHooks() + sink: Final = ServiceLoggingMetricsSink(service_logging_factory=lambda: hooks, executor=InlineExecutor()) + + for _ in range(_METRICS_QUEUE_LIMIT + 10): + sink.cache_hit() + + assert len(hooks.successes) == _METRICS_QUEUE_LIMIT + 10, ( + "an executor that actually runs releases each slot, so nothing is dropped" + ) + + def test_failure_maps_variant_and_redacted_summary(self): + hooks = RecordingServiceHooks() + error = TokenEndpointError(status_code=503, redacted_body="error: unavailable") + + self._sink(hooks).exchange_failure(call_type="mandatory_refresh", duration_seconds=0.1, error=error) + + ((service, duration, emitted, call_type),) = hooks.failures + assert service is ServiceTypes.ANTHROPIC_WIF + assert duration == 0.1 + assert call_type == "mandatory_refresh" + assert isinstance(emitted, TokenExchangeEndpointFailure) + assert str(emitted) == _error_summary(error) + + def test_transport_failure_gets_its_own_error_class(self): + hooks = RecordingServiceHooks() + + self._sink(hooks).exchange_failure( + call_type="advisory_refresh", duration_seconds=0.05, error=TokenTransportError(detail="ConnectError: boom") + ) + + ((_service, _duration, emitted, _call_type),) = hooks.failures + assert isinstance(emitted, TokenExchangeTransportFailure) + + def test_cache_hit_maps_to_cache_service_with_zero_duration(self): + hooks = RecordingServiceHooks() + + self._sink(hooks).cache_hit() + + assert hooks.successes == [(ServiceTypes.ANTHROPIC_WIF_CACHE, CALL_TYPE_CACHE_HIT, 0.0)] + + def test_end_to_end_reflected_assertion_never_reaches_the_hook(self): + hooks = RecordingServiceHooks() + sink = self._sink(hooks) + engine = make_engine(EchoingUnauthorizedPoster(), metrics_sink=sink) + + result = engine.get_token(make_spec()) + + assert isinstance(result, TokenEndpointError) + ((_service, _duration, emitted, call_type),) = hooks.failures + assert call_type == "cold_mint" + assert DEFAULT_ASSERTION not in str(emitted) + assert DEFAULT_ASSERTION not in repr(emitted) + + def test_raising_hooks_are_swallowed(self): + hooks: Final = RaisingServiceHooks() + sink: Final = self._sink(hooks) + + sink.exchange_success(call_type="cold_mint", duration_seconds=0.2) + sink.cache_hit() + sink.exchange_failure(call_type="cold_mint", duration_seconds=0.1, error=TokenTransportError(detail="boom")) + + assert hooks.attempts == ["success:cold_mint", "success:cache_hit", "failure:cold_mint"], ( + "every event is still handed to the hooks, and one raising hook does not stop the next" + ) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 023e2d8843f..63df501fb3e 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1,6 +1,7 @@ import asyncio import json import logging +import threading import time from unittest.mock import AsyncMock, Mock, patch @@ -17,7 +18,9 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, ) +from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.files.transformation import BaseFilesConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, @@ -455,7 +458,7 @@ async def test_async_anthropic_messages_handler_extra_headers(): # Mock the config mock_config = Mock() - mock_config.validate_anthropic_messages_environment = Mock( + mock_config.avalidate_anthropic_messages_environment = AsyncMock( return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com") ) mock_config.transform_anthropic_messages_request = Mock( @@ -502,7 +505,7 @@ async def test_async_anthropic_messages_handler_extra_headers(): captured_headers.update(kwargs.get("headers", {})) return ({"x-api-key": "test-key"}, "https://api.anthropic.com") - mock_config.validate_anthropic_messages_environment = capture_validate + mock_config.avalidate_anthropic_messages_environment = AsyncMock(side_effect=capture_validate) try: await handler.async_anthropic_messages_handler( @@ -758,7 +761,7 @@ async def test_async_anthropic_messages_handler_passes_litellm_metadata(): handler = BaseLLMHTTPHandler() mock_config = Mock() - mock_config.validate_anthropic_messages_environment = Mock( + mock_config.avalidate_anthropic_messages_environment = AsyncMock( return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com") ) mock_config.transform_anthropic_messages_request = Mock( @@ -837,7 +840,7 @@ async def test_async_anthropic_messages_handler_forwards_router_model_info(): handler = BaseLLMHTTPHandler() mock_config = Mock() - mock_config.validate_anthropic_messages_environment = Mock( + mock_config.avalidate_anthropic_messages_environment = AsyncMock( return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com") ) mock_config.transform_anthropic_messages_request = Mock( @@ -929,7 +932,7 @@ async def test_async_anthropic_messages_handler_header_priority(): captured_headers.update(kwargs.get("headers", {})) return ({"x-api-key": "test-key"}, "https://api.anthropic.com") - mock_config.validate_anthropic_messages_environment = capture_validate + mock_config.avalidate_anthropic_messages_environment = AsyncMock(side_effect=capture_validate) mock_config.transform_anthropic_messages_request = Mock( return_value={"model": "claude-3-opus-20240229", "messages": []} ) @@ -968,7 +971,7 @@ async def test_async_anthropic_messages_handler_drops_top_level_and_nested_param handler = BaseLLMHTTPHandler() mock_config = Mock() - mock_config.validate_anthropic_messages_environment = Mock( + mock_config.avalidate_anthropic_messages_environment = AsyncMock( return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com") ) @@ -1181,9 +1184,7 @@ def test_sync_delete_responses_sets_json_content_type(): ({}, True, None, None), ], ) -def test_resolve_anthropic_messages_timeout( - monkeypatch, litellm_params_kwargs, stream, global_timeout, expected -): +def test_resolve_anthropic_messages_timeout(monkeypatch, litellm_params_kwargs, stream, global_timeout, expected): from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS if global_timeout is None: @@ -1199,9 +1200,7 @@ def test_resolve_anthropic_messages_timeout( ) else: monkeypatch.setattr("litellm.request_timeout", global_timeout, raising=False) - monkeypatch.setattr( - "litellm.request_timeout_explicitly_set", True, raising=False - ) + monkeypatch.setattr("litellm.request_timeout_explicitly_set", True, raising=False) resolved = BaseLLMHTTPHandler._resolve_anthropic_messages_timeout( litellm_params=GenericLiteLLMParams(**litellm_params_kwargs), @@ -1222,13 +1221,11 @@ async def test_async_anthropic_messages_handler_forwards_request_timeout(monkeyp handler = BaseLLMHTTPHandler() mock_config = Mock() - mock_config.validate_anthropic_messages_environment = Mock( + mock_config.avalidate_anthropic_messages_environment = AsyncMock( return_value=({"x-api-key": "k"}, "https://api.anthropic.com") ) mock_config.should_filter_anthropic_beta_headers = Mock(return_value=False) - mock_config.transform_anthropic_messages_request = Mock( - return_value={"model": "claude", "messages": []} - ) + mock_config.transform_anthropic_messages_request = Mock(return_value={"model": "claude", "messages": []}) mock_config.get_complete_url = Mock(return_value="https://api.anthropic.com/v1/messages") mock_config.sign_request = Mock(return_value=({"x-api-key": "k"}, None)) mock_config.max_retry_on_anthropic_messages_http_error = 1 @@ -1270,13 +1267,11 @@ async def test_async_anthropic_messages_handler_forwards_stream_timeout(monkeypa handler = BaseLLMHTTPHandler() mock_config = Mock() - mock_config.validate_anthropic_messages_environment = Mock( + mock_config.avalidate_anthropic_messages_environment = AsyncMock( return_value=({"x-api-key": "k"}, "https://api.anthropic.com") ) mock_config.should_filter_anthropic_beta_headers = Mock(return_value=False) - mock_config.transform_anthropic_messages_request = Mock( - return_value={"model": "claude", "messages": []} - ) + mock_config.transform_anthropic_messages_request = Mock(return_value={"model": "claude", "messages": []}) mock_config.get_complete_url = Mock(return_value="https://api.anthropic.com/v1/messages") mock_config.sign_request = Mock(return_value=({"x-api-key": "k"}, None)) mock_config.max_retry_on_anthropic_messages_http_error = 1 @@ -1678,7 +1673,7 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks( handler = BaseLLMHTTPHandler() mock_config = Mock() - mock_config.validate_anthropic_messages_environment = Mock( + mock_config.avalidate_anthropic_messages_environment = AsyncMock( return_value=({"x-api-key": "sk-test"}, "https://api.anthropic.com") ) mock_config.transform_anthropic_messages_request = Mock( @@ -1686,7 +1681,13 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks( ) mock_config.sign_request = Mock(return_value=({}, None)) - fake_raw_response = {"id": "msg_1", "type": "message", "role": "assistant", "content": [], "stop_reason": "end_turn"} + fake_raw_response = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [], + "stop_reason": "end_turn", + } mock_config.transform_anthropic_messages_response = Mock(return_value=fake_raw_response) mock_logging_obj = Mock() @@ -1706,10 +1707,17 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks( mock_httpx_response.status_code = 200 with ( - patch.object(handler, "_async_post_anthropic_messages_with_http_error_retry", new=AsyncMock(return_value=mock_httpx_response)), + patch.object( + handler, + "_async_post_anthropic_messages_with_http_error_retry", + new=AsyncMock(return_value=mock_httpx_response), + ), patch.object(handler, "_call_agentic_completion_hooks", side_effect=fake_agentic_hooks), patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client"), - patch("litellm.litellm_core_utils.get_provider_specific_headers.ProviderSpecificHeaderUtils.get_provider_specific_headers", return_value=None), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.litellm_core_utils.get_provider_specific_headers.ProviderSpecificHeaderUtils.get_provider_specific_headers", + return_value=None, + ), ): result = await handler.async_anthropic_messages_handler( model="claude-haiku", @@ -1951,7 +1959,9 @@ def test_audio_transcriptions_sends_dict_data_as_json_body(): form-encodes it and silently ignores json=; JSON-body providers (e.g. Google Speech-to-Text) need an application/json body.""" captured = {} - client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_json_transcription_request(captured)))) + client = HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(_capture_json_transcription_request(captured))) + ) response = BaseLLMHTTPHandler().audio_transcriptions( client=client, @@ -2110,6 +2120,105 @@ def test_sync_retrieve_file_content_raises_on_http_error(): assert exc_info.value.status_code == 404 +_FILE_CONTENT_WIF_ENV = { + "ANTHROPIC_FEDERATION_RULE_ID": "fdrl_llm_http_handler_seam", + "ANTHROPIC_ORGANIZATION_ID": "org-llm-http-handler-seam", + "ANTHROPIC_IDENTITY_TOKEN": "llm-http-handler-seam-inline-jwt", +} + + +class _BlockingWifPoster: + """A token-endpoint poster that blocks until released, so the test can prove + the exchange ran off the event loop's own thread instead of freezing it.""" + + def __init__(self): + self.release = threading.Event() + self.thread_ids = [] + + def post(self, url, *, content, headers, timeout): + self.thread_ids.append(threading.get_ident()) + self.release.wait(timeout=5) + return httpx.Response( + 200, + json={ + "access_token": "sk-ant-oat01-llm-http-handler-seam", + "token_type": "Bearer", + "expires_in": 3600, + }, + ) + + +@pytest.mark.asyncio +async def test_async_retrieve_file_content_wif_exchange_does_not_block_event_loop(monkeypatch): + """Regression (Greptile P1): async_retrieve_file_content called the synchronous + validate_environment directly, so a cold WIF mint on this call site froze the + event loop until the exchange finished. It must resolve credentials through the + async facade instead.""" + from litellm.llms.anthropic import common_utils as anthropic_common_utils + from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig + from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token + from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + for name, value in _FILE_CONTENT_WIF_ENV.items(): + monkeypatch.setenv(name, value) + + poster = _BlockingWifPoster() + engine = JwtBearerTokenExchangeEngine(poster=poster) + sync_calls = [] + + def sync_shim(litellm_params, api_base, model): + sync_calls.append(model) + return get_anthropic_wif_token(litellm_params, api_base, model, engine) + + async def async_shim(litellm_params, api_base, model): + return await aget_anthropic_wif_token(litellm_params, api_base, model, engine) + + monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim) + monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim) + + handler = BaseLLMHTTPHandler() + client = Mock(spec=AsyncHTTPHandler) + client.get = AsyncMock(return_value=httpx.Response(status_code=200, content=b"file bytes")) + + ticks = [] + + async def ticker(): + for i in range(20): + await asyncio.sleep(0.005) + ticks.append(i) + + ticker_task = asyncio.create_task(ticker()) + await asyncio.sleep(0.02) + + retrieve_task = asyncio.create_task( + handler.async_retrieve_file_content( + file_content_request={"file_id": "file-abc"}, + provider_config=AnthropicFilesConfig(), + litellm_params={}, + headers={}, + logging_obj=Mock(), + client=client, + ) + ) + await asyncio.sleep(0.05) + # The ticker kept advancing while the token exchange was still blocked on + # poster.release, proving the exchange did not run inline on the event loop. + assert len(ticks) > 0 + assert not retrieve_task.done() + + poster.release.set() + await retrieve_task + await ticker_task + + assert sync_calls == [] + assert poster.thread_ids + assert poster.thread_ids[0] != threading.get_ident() + sent_headers = client.get.call_args.kwargs["headers"] + assert sent_headers["authorization"] == "Bearer sk-ant-oat01-llm-http-handler-seam" + + _UPSTREAM_NOT_FOUND_BODY = { "error": { "message": "Response with id 'resp_abc' not found.", @@ -2257,9 +2366,7 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques ok_response = httpx.Response(200, json={"id": "msg_1"}, request=httpx.Request("POST", request_url)) class FakeAsyncClient: - async def post( - self, url, headers, data, stream=False, logging_obj=None, timeout=None - ): + async def post(self, url, headers, data, stream=False, logging_obj=None, timeout=None): posts.append({"headers": dict(headers), "data": data}) return invalid_signature_response if len(posts) == 1 else ok_response @@ -2576,7 +2683,7 @@ async def test_async_anthropic_messages_handler_carries_deployment_vertex_locati custom_llm_provider="vertex_ai", ) mock_config = Mock() - mock_config.validate_anthropic_messages_environment = Mock( + mock_config.avalidate_anthropic_messages_environment = AsyncMock( return_value=({"authorization": "Bearer t"}, "https://us-east5-aiplatform.googleapis.com") ) mock_config.transform_anthropic_messages_request = Mock( @@ -3100,6 +3207,70 @@ async def test_a_provider_that_keeps_rejecting_is_not_retried_forever_on_the_asy assert len(recorder.bodies) == 2 +def _async_client_returning(response: Mock) -> AsyncMock: + client = AsyncMock(spec=AsyncHTTPHandler) + client.post.return_value = response + return client + + +@pytest.mark.asyncio +async def test_create_file_async_awaits_the_provider_credential_hook_instead_of_blocking(): + provider_config = Mock(spec=BaseFilesConfig) + provider_config.validate_environment.side_effect = AssertionError("sync validate_environment ran on the event loop") + provider_config.avalidate_environment = AsyncMock(return_value={"x-api-key": "federated"}) + provider_config.get_complete_file_url.return_value = "https://files.example/v1/files" + provider_config.transform_create_file_request.return_value = {"file": ("batch.jsonl", b"{}", "application/jsonl")} + file_object = object() + provider_config.transform_create_file_response.return_value = file_object + client = _async_client_returning(Mock(spec=httpx.Response)) + + result = await BaseLLMHTTPHandler().create_file( + create_file_data={"file": b"{}", "purpose": "batch"}, + litellm_params={}, + provider_config=provider_config, + headers={}, + api_base=None, + api_key=None, + logging_obj=Mock(), + _is_async=True, + client=client, + ) + + assert result is file_object + provider_config.validate_environment.assert_not_called() + provider_config.avalidate_environment.assert_awaited_once() + assert client.post.call_args.kwargs["headers"] == {"x-api-key": "federated"} + assert client.post.call_args.kwargs["url"] == "https://files.example/v1/files" + + +@pytest.mark.asyncio +async def test_create_batch_async_validates_credentials_off_the_event_loop(): + provider_config = Mock(spec=BaseBatchesConfig) + provider_config.validate_environment.side_effect = lambda **_: {"x-validated-on": str(threading.get_ident())} + provider_config.get_complete_batch_url.return_value = "https://batches.example/v1/messages/batches" + provider_config.transform_create_batch_request.return_value = {"requests": []} + batch = object() + provider_config.transform_create_batch_response.return_value = batch + client = _async_client_returning(Mock(spec=httpx.Response)) + + result = await BaseLLMHTTPHandler().create_batch( + create_batch_data={"input_file_id": "file_1", "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + litellm_params={}, + provider_config=provider_config, + headers={}, + api_base=None, + api_key=None, + logging_obj=Mock(), + _is_async=True, + client=client, + model="claude-sonnet-4-5", + ) + + assert result is batch + validated_on = client.post.call_args.kwargs["headers"]["x-validated-on"] + assert validated_on != str(threading.get_ident()) + assert client.post.call_args.kwargs["url"] == "https://batches.example/v1/messages/batches" + CONTAINER_NOT_FOUND_BODY = { "error": { "message": "Container with id 'cntr_gone' not found.", diff --git a/tests/test_litellm/llms/openai/test_openai_workload_identity.py b/tests/test_litellm/llms/openai/test_openai_workload_identity.py index db107e00df0..74415d45638 100644 --- a/tests/test_litellm/llms/openai/test_openai_workload_identity.py +++ b/tests/test_litellm/llms/openai/test_openai_workload_identity.py @@ -10,6 +10,7 @@ from openai import AsyncOpenAI, OpenAI import litellm from litellm.llms.litellm_proxy.responses.transformation import LiteLLMProxyResponsesAPIConfig +from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIError from litellm.llms.openai.openai import OpenAIChatCompletion from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig @@ -22,6 +23,17 @@ from litellm.llms.openai.workload_identity import ( from litellm.types.router import GenericLiteLLMParams TOKEN_EXCHANGE_URL: Final = "https://auth.openai.com/oauth/token" +CHAT_COMPLETIONS_URL: Final = "https://api.openai.com/v1/chat/completions" +EMBEDDINGS_URL: Final = "https://api.openai.com/v1/embeddings" +MODELS_URL: Final = "https://api.openai.com/v1/models" +CHAT_COMPLETION_BODY: Final = { + "id": "chatcmpl-wif", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} @pytest.fixture @@ -279,3 +291,342 @@ class TestResponsesValidateEnvironment: headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams() ) assert headers["Authorization"] == "Bearer None" + + +@pytest.fixture +def deployment_wif(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> dict[str, str]: + token_file: Final = tmp_path / "deployment_subject_token.jwt" + token_file.write_text("subject-token-from-deployment-file") + for name in ( + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "OPENAI_IDENTITY_PROVIDER_ID", + "OPENAI_SERVICE_ACCOUNT_ID", + "OPENAI_IDENTITY_TOKEN_FILE", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + _workload_identity_auth.cache_clear() + litellm.in_memory_llm_clients_cache.flush_cache() + return { + "openai_identity_provider_id": "idp_deployment", + "openai_service_account_id": "user-deployment", + "openai_identity_token_file": str(token_file), + } + + +def deployment_config(deployment_wif: dict[str, str]) -> OpenAIWorkloadIdentityConfig: + return OpenAIWorkloadIdentityConfig( + identity_provider_id="idp_deployment", + service_account_id="user-deployment", + token_file=deployment_wif["openai_identity_token_file"], + ) + + +def mock_chat_completions() -> respx.Route: + return respx.post(CHAT_COMPLETIONS_URL).mock(return_value=httpx.Response(200, json=CHAT_COMPLETION_BODY)) + + +def mock_streaming_chat_completions() -> respx.Route: + chunk: Final = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"} + events: Final = ( + {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": None}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + ) + body: Final = "".join(f"data: {json.dumps(event)}\n\n" for event in events) + "data: [DONE]\n\n" + return respx.post(CHAT_COMPLETIONS_URL).mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=body) + ) + + +class TestResolveConfigFromDeployment: + def test_resolves_from_litellm_params_without_env(self, deployment_wif: dict[str, str]) -> None: + assert resolve_openai_workload_identity_config( + api_key=None, api_base=None, litellm_params=deployment_wif + ) == deployment_config(deployment_wif) + + def test_env_alone_disables_nothing_when_params_are_absent(self, deployment_wif: dict[str, str]) -> None: + assert resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=None) is None + + def test_unrelated_litellm_params_do_not_resolve(self, deployment_wif: dict[str, str]) -> None: + assert ( + resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params={"model": "gpt-4o"}) + is None + ) + + def test_litellm_params_beat_env(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + config: Final = resolve_openai_workload_identity_config( + api_key=None, + api_base=None, + litellm_params={ + "openai_identity_provider_id": "idp_deployment", + "openai_service_account_id": "user-deployment", + "openai_identity_token_file": wif_env.token_file, + }, + ) + assert config == OpenAIWorkloadIdentityConfig( + identity_provider_id="idp_deployment", + service_account_id="user-deployment", + token_file=wif_env.token_file, + ) + + def test_partial_litellm_params_fill_from_env_per_field(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + config: Final = resolve_openai_workload_identity_config( + api_key=None, api_base=None, litellm_params={"openai_identity_provider_id": "idp_deployment"} + ) + assert config == OpenAIWorkloadIdentityConfig( + identity_provider_id="idp_deployment", + service_account_id=wif_env.service_account_id, + token_file=wif_env.token_file, + ) + + @pytest.mark.parametrize("blank", ["", None, 7]) + def test_blank_or_non_string_param_falls_back_to_env( + self, wif_env: OpenAIWorkloadIdentityConfig, blank: object + ) -> None: + config: Final = resolve_openai_workload_identity_config( + api_key=None, api_base=None, litellm_params={"openai_identity_provider_id": blank} + ) + assert config == wif_env + + def test_partial_litellm_params_without_env_disable(self, deployment_wif: dict[str, str]) -> None: + partial: Final = {key: value for key, value in deployment_wif.items() if key != "openai_identity_token_file"} + assert resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=partial) is None + + def test_static_api_key_beats_litellm_params(self, deployment_wif: dict[str, str]) -> None: + assert ( + resolve_openai_workload_identity_config(api_key="sk-static", api_base=None, litellm_params=deployment_wif) + is None + ) + + def test_env_openai_api_key_beats_litellm_params( + self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") + assert ( + resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=deployment_wif) is None + ) + + def test_foreign_api_base_disables_deployment_wif(self, deployment_wif: dict[str, str]) -> None: + assert ( + resolve_openai_workload_identity_config( + api_key=None, api_base="https://my-vllm.internal/v1", litellm_params=deployment_wif + ) + is None + ) + + +class TestDeploymentClientConstruction: + def test_sync_client_from_deployment_params(self, deployment_wif: dict[str, str]) -> None: + client: Final = OpenAIChatCompletion()._get_openai_client( + is_async=False, api_key=None, api_base=None, litellm_params=deployment_wif + ) + assert isinstance(client, OpenAI) + assert client.api_key == "workload-identity-auth" + assert client._workload_identity_auth is not None + + def test_async_client_from_deployment_params(self, deployment_wif: dict[str, str]) -> None: + client: Final = OpenAIChatCompletion()._get_openai_client( + is_async=True, api_key=None, api_base=None, litellm_params=deployment_wif + ) + assert isinstance(client, AsyncOpenAI) + assert client._workload_identity_auth is not None + + def test_distinct_deployments_get_distinct_cached_clients(self, deployment_wif: dict[str, str]) -> None: + other_deployment: Final = {**deployment_wif, "openai_service_account_id": "user-other"} + handler: Final = OpenAIChatCompletion() + first: Final = handler._get_openai_client( + is_async=False, api_key=None, api_base=None, litellm_params=deployment_wif + ) + second: Final = handler._get_openai_client( + is_async=False, api_key=None, api_base=None, litellm_params=other_deployment + ) + again: Final = handler._get_openai_client( + is_async=False, api_key=None, api_base=None, litellm_params=dict(deployment_wif) + ) + assert first is not second + assert again is first + + @respx.mock + def test_completion_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("deployment-bearer") + completion_route: Final = mock_chat_completions() + + response: Final = litellm.completion( + model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], **deployment_wif + ) + + assert response.choices[0].message.content == "ok" + request: Final = completion_route.calls.last.request + assert request.headers["Authorization"] == "Bearer deployment-bearer" + assert not any(key.startswith("openai_") for key in json.loads(request.content)) + + @respx.mock + def test_streaming_completion_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("stream-bearer") + stream_route: Final = mock_streaming_chat_completions() + + chunks: Final = tuple( + litellm.completion( + model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **deployment_wif + ) + ) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" + assert stream_route.calls.last.request.headers["Authorization"] == "Bearer stream-bearer" + + @respx.mock + @pytest.mark.asyncio + async def test_async_streaming_completion_kwargs_carry_exchanged_bearer( + self, deployment_wif: dict[str, str] + ) -> None: + mock_token_exchange("async-stream-bearer") + stream_route: Final = mock_streaming_chat_completions() + + stream: Final = await litellm.acompletion( + model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **deployment_wif + ) + chunks: Final = tuple([chunk async for chunk in stream]) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" + assert stream_route.calls.last.request.headers["Authorization"] == "Bearer async-stream-bearer" + + @respx.mock + def test_router_deployment_without_api_key_authenticates_via_token_exchange( + self, deployment_wif: dict[str, str] + ) -> None: + exchange_route: Final = mock_token_exchange("router-bearer") + completion_route: Final = mock_chat_completions() + router: Final = litellm.Router( + model_list=[{"model_name": "wif-gpt", "litellm_params": {"model": "openai/gpt-4o-mini", **deployment_wif}}] + ) + + response: Final = router.completion(model="wif-gpt", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "ok" + assert exchange_route.called + assert completion_route.calls.last.request.headers["Authorization"] == "Bearer router-bearer" + + @respx.mock + def test_embedding_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("embedding-bearer") + embeddings_route: Final = respx.post(EMBEDDINGS_URL).mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + }, + ) + ) + + litellm.embedding(model="openai/text-embedding-3-small", input=["hi"], **deployment_wif) + + assert embeddings_route.calls.last.request.headers["Authorization"] == "Bearer embedding-bearer" + + +class TestResponsesValidateEnvironmentFromDeployment: + @respx.mock + def test_mints_bearer_from_litellm_params(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("responses-bearer") + headers: Final = OpenAIResponsesAPIConfig().validate_environment( + headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams(**deployment_wif) + ) + assert headers["Authorization"] == "Bearer responses-bearer" + + def test_static_key_in_litellm_params_wins(self, deployment_wif: dict[str, str]) -> None: + headers: Final = OpenAIResponsesAPIConfig().validate_environment( + headers={}, + model="gpt-4o-mini", + litellm_params=GenericLiteLLMParams(api_key="sk-responses", **deployment_wif), + ) + assert headers["Authorization"] == "Bearer sk-responses" + + +class TestDiscoverModels: + @staticmethod + def mock_models() -> respx.Route: + return respx.get(MODELS_URL).mock( + return_value=httpx.Response(200, json={"data": [{"id": "gpt-4o-mini"}, {"id": "gpt-4.1"}]}) + ) + + @respx.mock + def test_discovers_with_exchanged_bearer_from_litellm_params(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("discovery-bearer") + models_route: Final = self.mock_models() + + assert OpenAIGPTConfig().discover_models(deployment_wif) == ["gpt-4o-mini", "gpt-4.1"] + assert models_route.calls.last.request.headers["Authorization"] == "Bearer discovery-bearer" + + @respx.mock + def test_discovers_with_env_wif_when_params_carry_no_key(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + mock_token_exchange("env-discovery-bearer") + models_route: Final = self.mock_models() + + OpenAIGPTConfig().discover_models({}) + + assert models_route.calls.last.request.headers["Authorization"] == "Bearer env-discovery-bearer" + + @respx.mock + def test_static_api_key_in_params_skips_token_exchange(self, deployment_wif: dict[str, str]) -> None: + exchange_route: Final = mock_token_exchange() + models_route: Final = self.mock_models() + + OpenAIGPTConfig().discover_models({**deployment_wif, "api_key": "sk-discovery"}) + + assert models_route.calls.last.request.headers["Authorization"] == "Bearer sk-discovery" + assert not exchange_route.called + + @respx.mock + def test_blank_api_base_in_params_discovers_from_openai(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("blank-base-bearer") + models_route: Final = self.mock_models() + + assert OpenAIGPTConfig().discover_models({**deployment_wif, "api_base": ""}) == ["gpt-4o-mini", "gpt-4.1"] + assert models_route.calls.last.request.headers["Authorization"] == "Bearer blank-base-bearer" + + @respx.mock + def test_openai_compatible_subclass_never_mints_wif(self, deployment_wif: dict[str, str]) -> None: + exchange_route: Final = mock_token_exchange() + models_route: Final = self.mock_models() + + class CompatibleConfig(OpenAIGPTConfig): + pass + + CompatibleConfig().discover_models(deployment_wif) + + assert models_route.calls.last.request.headers["Authorization"] == "Bearer None" + assert not exchange_route.called + + + @respx.mock + def test_empty_static_key_never_borrows_the_env_key(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-env-key-that-must-stay-home") + foreign_models: Final = respx.get("https://third-party.example/v1/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "other-model"}]}) + ) + + assert OpenAIGPTConfig().get_models(api_key="", api_base="https://third-party.example") == ["other-model"] + assert foreign_models.calls.last.request.headers["Authorization"] == "Bearer " + + +class TestClientsideBaseOverride: + def test_client_api_base_override_clears_deployment_wif(self, deployment_wif: dict[str, str]) -> None: + from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params + + redirected: Final = get_dynamic_litellm_params( + litellm_params={"model": "openai/gpt-4o-mini", **deployment_wif}, + request_kwargs={"api_base": "https://not-openai.example/v1"}, + ) + + assert not any(key in redirected for key in deployment_wif) + assert ( + resolve_openai_workload_identity_config( + api_key=None, api_base=redirected["api_base"], litellm_params=redirected + ) + is None + ) diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py index 9ae9b732066..6223a28bd6e 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/test_litellm/models/test_models.py @@ -3,6 +3,7 @@ Tests for backend domain models. """ from datetime import datetime +from typing import Final import pytest from pydantic import BaseModel, TypeAdapter @@ -92,9 +93,7 @@ class TestCredentials: assert item.credential_values is None def test_create_credential_item_requires_values_or_model_id(self): - with pytest.raises( - ValueError, match="Either credential_values or model_id must be set" - ): + with pytest.raises(ValueError, match="Either credential_values or model_id must be set"): CreateCredentialItem(credential_name="bad", credential_info={}) @@ -112,12 +111,8 @@ class TestModel: assert model.team_public_model_name == "my-gpt4" def test_is_blocked(self): - model_blocked = LiteLLM_ProxyModelTable( - model_id="m1", model_name="test", litellm_params={}, blocked=True - ) - model_unblocked = LiteLLM_ProxyModelTable( - model_id="m2", model_name="test", litellm_params={}, blocked=False - ) + model_blocked = LiteLLM_ProxyModelTable(model_id="m1", model_name="test", litellm_params={}, blocked=True) + model_unblocked = LiteLLM_ProxyModelTable(model_id="m2", model_name="test", litellm_params={}, blocked=False) assert model_blocked.is_blocked assert not model_unblocked.is_blocked @@ -159,9 +154,7 @@ class TestModel: assert model.blocked is True def test_team_helpers_none_when_no_model_info(self): - model = LiteLLM_ProxyModelTable( - model_id="m1", model_name="gpt-4", litellm_params={}, model_info=None - ) + model = LiteLLM_ProxyModelTable(model_id="m1", model_name="gpt-4", litellm_params={}, model_info=None) assert model.team_id is None assert model.team_public_model_name is None @@ -263,9 +256,7 @@ class TestTeam: assert team.model_max_budget == {"gpt-4": 5.0} def test_cached_team(self): - cached = LiteLLM_TeamTableCachedObj( - team_id="t1", last_refreshed_at=1234567890.0 - ) + cached = LiteLLM_TeamTableCachedObj(team_id="t1", last_refreshed_at=1234567890.0) assert cached.last_refreshed_at == 1234567890.0 def test_deleted_team(self): @@ -316,9 +307,7 @@ class TestUser: assert "password" not in user.model_dump() assert "password" not in user.model_dump_json() - with_keys = LiteLLM_UserTableWithKeyCount( - user_id="u1", user_email="a@b.c", password=secret, key_count=2 - ) + with_keys = LiteLLM_UserTableWithKeyCount(user_id="u1", user_email="a@b.c", password=secret, key_count=2) assert with_keys.password == secret assert "password" not in with_keys.model_dump() assert "password" not in with_keys.model_dump_json() @@ -442,9 +431,7 @@ class TestEndUserTable: class TestBudgetTableFull: def test_full_adds_server_managed_fields(self): now = datetime.now() - budget = LiteLLM_BudgetTableFull( - budget_id="b1", max_budget=10.0, created_at=now, budget_reset_at=now - ) + budget = LiteLLM_BudgetTableFull(budget_id="b1", max_budget=10.0, created_at=now, budget_reset_at=now) assert budget.created_at == now assert budget.budget_reset_at == now assert budget.max_budget == 10.0 @@ -456,9 +443,7 @@ class TestBudgetTableFull: class TestTeamMemberTable: def test_tracks_user_within_team(self): - member = LiteLLM_TeamMemberTable( - user_id="u1", team_id="t1", spend=3.0, budget_id="b1", max_budget=5.0 - ) + member = LiteLLM_TeamMemberTable(user_id="u1", team_id="t1", spend=3.0, budget_id="b1", max_budget=5.0) assert member.user_id == "u1" assert member.team_id == "t1" assert member.spend == 3.0 @@ -548,9 +533,7 @@ class TestSpendLogs: assert log.updated_at == updated_at def test_error_logs_creation(self): - log = LiteLLM_ErrorLogs( - request_id="r1", startTime=None, endTime=None, status_code="500" - ) + log = LiteLLM_ErrorLogs(request_id="r1", startTime=None, endTime=None, status_code="500") assert log.request_id == "r1" assert log.status_code == "500" @@ -568,9 +551,7 @@ class TestManagedTables: def test_managed_object_table_requires_purpose(self): with pytest.raises(ValidationError): - LiteLLM_ManagedObjectTable( - unified_object_id="o1", model_object_id="m1", file_object={} - ) + LiteLLM_ManagedObjectTable(unified_object_id="o1", model_object_id="m1", file_object={}) def test_managed_vector_stores_table(self): table = LiteLLM_ManagedVectorStoresTable( @@ -588,3 +569,40 @@ class TestManagedTables: ) assert table.vector_store_id == "vs1" assert table.custom_llm_provider == "openai" + + +class TestProxyModelTableResponseSerialization: + """FastAPI validates an endpoint's return value against its response model with + ``from_attributes``, so an endpoint that returns an already-built row reaches the + ``mode="before"`` validator as the object itself rather than as a mapping.""" + + def test_validates_from_an_existing_instance(self): + from pydantic import TypeAdapter + + built: Final = LiteLLM_ProxyModelTable( + model_id="m-1", + model_name="claude-sonnet-5-provider", + litellm_params={"model": "anthropic/claude-sonnet-5"}, + blocked=True, + ) + + serialized = TypeAdapter(LiteLLM_ProxyModelTable | None).validate_python(built, from_attributes=True) + + assert serialized is not None + assert serialized.model_id == "m-1" + assert serialized.blocked is True + assert serialized.litellm_params == {"model": "anthropic/claude-sonnet-5"} + + def test_still_parses_json_string_columns(self): + """The DB stores these columns as JSON strings, which is why the validator exists.""" + parsed: Final = LiteLLM_ProxyModelTable.model_validate( + { + "model_id": "m-2", + "model_name": "n", + "litellm_params": '{"model": "anthropic/claude-haiku-4-5"}', + "model_info": '{"id": "m-2"}', + } + ) + + assert parsed.litellm_params == {"model": "anthropic/claude-haiku-4-5"} + assert parsed.model_info == {"id": "m-2"} diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index f513f397b64..f5c7ac713c1 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -29,6 +29,22 @@ from litellm.proxy.auth.auth_utils import ( ) +def test_every_server_owned_wif_kwarg_key_is_request_banned(): + """server_owned_wif_litellm_params (types/utils.py) is derived from ANTHROPIC_WIF_KWARGS_KEYS + and OPENAI_WIF_KWARGS_KEYS (get_litellm_params.py) precisely so a new WIF field can never be + added to the kwargs funnel + without automatically joining the request-body ban list; this guards that invariant itself, + independent of today's field count, so it fails if the derivation is ever reverted to a + hand-typed list that drifts.""" + from litellm.litellm_core_utils.get_litellm_params import ( + ANTHROPIC_WIF_KWARGS_KEYS, + OPENAI_WIF_KWARGS_KEYS, + ) + from litellm.proxy.auth.auth_utils import _SERVER_OWNED_WIF_UNCONDITIONAL_BANNED + + assert ANTHROPIC_WIF_KWARGS_KEYS | OPENAI_WIF_KWARGS_KEYS == set(_SERVER_OWNED_WIF_UNCONDITIONAL_BANNED) + + class TestCustomAuthCommonChecksWarning: """custom_auth_common_checks_warning only warns when custom auth is configured and the common-checks opt-in is off, since that is the only state where @@ -118,9 +134,7 @@ class TestGetKeyModelRpmLimit: """Should fall back to team metadata when key metadata exists but has no model_rpm_limit.""" user_api_key_dict = UserAPIKeyAuth( api_key="sk-123", - metadata={ - "some_other_key": "value" - }, # Has metadata, but not model_rpm_limit + metadata={"some_other_key": "value"}, # Has metadata, but not model_rpm_limit team_metadata={"model_rpm_limit": {"gpt-4": 50}}, ) result = get_key_model_rpm_limit(user_api_key_dict) @@ -202,9 +216,7 @@ class TestGetKeyModelTpmLimit: """Should fall back to team metadata when key metadata exists but has no model_tpm_limit.""" user_api_key_dict = UserAPIKeyAuth( api_key="sk-123", - metadata={ - "some_other_key": "value" - }, # Has metadata, but not model_tpm_limit + metadata={"some_other_key": "value"}, # Has metadata, but not model_tpm_limit team_metadata={"model_tpm_limit": {"gpt-4": 5000}}, ) result = get_key_model_tpm_limit(user_api_key_dict) @@ -315,9 +327,7 @@ class TestGetEndUserIdFromRequestBodyWithStandardHeaders: request_body = {"user": "body-user"} with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers=headers - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers=headers) assert result == "header-customer" def test_should_fall_back_to_body_when_no_standard_header(self): @@ -326,9 +336,7 @@ class TestGetEndUserIdFromRequestBodyWithStandardHeaders: request_body = {"user": "body-user"} with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers=headers - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers=headers) assert result == "body-user" @@ -370,8 +378,7 @@ def test_get_model_from_request_enforces_when_builtin_handler_dispatched(): enforced. Same request path as above, but dispatched to a non-pass-through endpoint: the model must NOT be suppressed.""" - def builtin_chat_completions(): - ... + def builtin_chat_completions(): ... assert ( get_model_from_request( @@ -490,9 +497,7 @@ def test_get_model_from_request_extracts_unified_file_id_models(): "litellm_proxy:application/octet-stream;unified_id,test-id;" "target_model_names,model-a,model-b;llm_output_file_id,file-provider-id" ) - encoded_unified_file_id = ( - base64.urlsafe_b64encode(raw_unified_file_id.encode()).decode().rstrip("=") - ) + encoded_unified_file_id = base64.urlsafe_b64encode(raw_unified_file_id.encode()).decode().rstrip("=") assert get_model_from_request( request_data={"file_id": encoded_unified_file_id}, @@ -552,9 +557,7 @@ def test_get_model_from_request_resolves_video_id_model_with_router(): model_id="veo-3.1-generate-001", ) llm_router = MagicMock() - llm_router.resolve_model_name_from_model_id.return_value = ( - "gcp/google/veo-3.1-generate-001" - ) + llm_router.resolve_model_name_from_model_id.return_value = "gcp/google/veo-3.1-generate-001" assert ( get_model_from_request( @@ -564,9 +567,7 @@ def test_get_model_from_request_resolves_video_id_model_with_router(): ) == "gcp/google/veo-3.1-generate-001" ) - llm_router.resolve_model_name_from_model_id.assert_called_once_with( - "veo-3.1-generate-001" - ) + llm_router.resolve_model_name_from_model_id.assert_called_once_with("veo-3.1-generate-001") _BATCH_DEPLOYMENT_ID = "8d0eaa7e6c6f54a425dfd0062cb6b0dc" @@ -597,9 +598,7 @@ def _encode_managed_id(decoded: str) -> str: return base64.urlsafe_b64encode(decoded.encode()).decode().rstrip("=") -_MANAGED_BATCH_ID = _encode_managed_id( - f"litellm_proxy;model_id:{_BATCH_DEPLOYMENT_ID};llm_batch_id:provider-batch-123" -) +_MANAGED_BATCH_ID = _encode_managed_id(f"litellm_proxy;model_id:{_BATCH_DEPLOYMENT_ID};llm_batch_id:provider-batch-123") _MANAGED_BATCH_OUTPUT_FILE_ID = _encode_managed_id( f"litellm_proxy;model_id:{_BATCH_DEPLOYMENT_ID};llm_batch_id:provider-batch-123;" "llm_output_file_id:provider-file-456" @@ -675,9 +674,7 @@ def test_get_model_from_request_resolves_character_id_model_with_router(): model_id="veo-3.1-generate-001", ) llm_router = MagicMock() - llm_router.resolve_model_name_from_model_id.return_value = ( - "gcp/google/veo-3.1-generate-001" - ) + llm_router.resolve_model_name_from_model_id.return_value = "gcp/google/veo-3.1-generate-001" assert ( get_model_from_request( @@ -687,9 +684,7 @@ def test_get_model_from_request_resolves_character_id_model_with_router(): ) == "gcp/google/veo-3.1-generate-001" ) - llm_router.resolve_model_name_from_model_id.assert_called_once_with( - "veo-3.1-generate-001" - ) + llm_router.resolve_model_name_from_model_id.assert_called_once_with("veo-3.1-generate-001") def test_get_model_from_request_only_runs_media_decoders_for_matching_fields(): @@ -836,9 +831,7 @@ def test_abbreviate_api_key_short_key_is_fully_masked(): def test_get_customer_user_header_returns_none_when_no_customer_role(): from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping - mappings = [ - {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"} - ] + mappings = [{"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}] result = get_customer_user_header_from_mapping(mappings) assert result is None @@ -891,9 +884,7 @@ def test_get_end_user_id_returns_id_from_user_header_mappings(): ), patch("litellm.proxy.proxy_server.general_settings", general_settings), ): - result = get_end_user_id_from_request_body( - request_body={}, request_headers=headers - ) + result = get_end_user_id_from_request_body(request_body={}, request_headers=headers) assert result == "1234" @@ -919,9 +910,7 @@ def test_get_end_user_id_returns_first_customer_header_when_multiple_mappings_ex ), patch("litellm.proxy.proxy_server.general_settings", general_settings), ): - result = get_end_user_id_from_request_body( - request_body={}, request_headers=headers - ) + result = get_end_user_id_from_request_body(request_body={}, request_headers=headers) assert result == "user-456" @@ -942,9 +931,7 @@ def test_get_end_user_id_returns_none_when_no_customer_role_in_mappings(): ), patch("litellm.proxy.proxy_server.general_settings", general_settings), ): - result = get_end_user_id_from_request_body( - request_body={}, request_headers=headers - ) + result = get_end_user_id_from_request_body(request_body={}, request_headers=headers) assert result is None @@ -962,9 +949,7 @@ def test_get_end_user_id_falls_back_to_deprecated_user_header_name(): ), patch("litellm.proxy.proxy_server.general_settings", general_settings), ): - result = get_end_user_id_from_request_body( - request_body={}, request_headers=headers - ) + result = get_end_user_id_from_request_body(request_body={}, request_headers=headers) assert result == "user-legacy" @@ -1108,9 +1093,7 @@ class TestGetEndUserIdDropsMalformedBodyValues: } with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers={} - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers={}) assert result == "alice@example.com" @@ -1120,9 +1103,7 @@ class TestGetEndUserIdDropsMalformedBodyValues: } with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers={} - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers={}) assert result is None @@ -1135,19 +1116,14 @@ class TestGetEndUserIdDropsMalformedBodyValues: """ import litellm - blob = ( - '{"device_id":"d5abe9199ee7759a","account_uuid":"",' - '"session_id":"c284b8cb-a050-4278-8599-cc4e016a10ab"}' - ) + blob = '{"device_id":"d5abe9199ee7759a","account_uuid":"","session_id":"c284b8cb-a050-4278-8599-cc4e016a10ab"}' request_body = {"user": blob} original = litellm.validate_end_user_id_in_db litellm.validate_end_user_id_in_db = False try: with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers={} - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers={}) finally: litellm.validate_end_user_id_in_db = original @@ -1158,8 +1134,7 @@ class TestGetEndUserIdDropsMalformedBodyValues: request_body = { "user": ( - '{"device_id":"d5abe9199ee7759a","account_uuid":"",' - '"session_id":"c284b8cb-a050-4278-8599-cc4e016a10ab"}' + '{"device_id":"d5abe9199ee7759a","account_uuid":"","session_id":"c284b8cb-a050-4278-8599-cc4e016a10ab"}' ), } @@ -1167,9 +1142,7 @@ class TestGetEndUserIdDropsMalformedBodyValues: litellm.validate_end_user_id_in_db = True try: with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers={} - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers={}) finally: litellm.validate_end_user_id_in_db = original @@ -1179,9 +1152,7 @@ class TestGetEndUserIdDropsMalformedBodyValues: request_body = {"user": "alice@example.com"} with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers={} - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers={}) assert result == "alice@example.com" @@ -1193,9 +1164,7 @@ class TestGetEndUserIdDropsMalformedBodyValues: request_body = {"user": codex_id} with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers={} - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers={}) assert result == codex_id @@ -1203,9 +1172,7 @@ class TestGetEndUserIdDropsMalformedBodyValues: request_body = {"user": 12345} with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers={} - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers={}) assert result == "12345" @@ -1216,9 +1183,7 @@ class TestGetEndUserIdDropsMalformedBodyValues: } with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers={} - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers={}) assert result == "alice@example.com" @@ -1228,9 +1193,7 @@ class TestGetEndUserIdDropsMalformedBodyValues: } with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers={} - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers={}) assert result is None @@ -1240,9 +1203,7 @@ class TestGetEndUserIdDropsMalformedBodyValues: } with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers={} - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers={}) assert result is None @@ -1250,9 +1211,7 @@ class TestGetEndUserIdDropsMalformedBodyValues: request_body = {"user": " ", "safety_identifier": "alice@example.com"} with patch("litellm.proxy.proxy_server.general_settings", {}): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers={} - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers={}) assert result == "alice@example.com" @@ -1271,16 +1230,12 @@ class TestGetEndUserIdDropsMalformedBodyValues: ), patch("litellm.proxy.proxy_server.general_settings", general_settings), ): - result = get_end_user_id_from_request_body( - request_body=request_body, request_headers=headers - ) + result = get_end_user_id_from_request_body(request_body=request_body, request_headers=headers) assert result == "alice@example.com" -def _make_deployment_dict( - model_name: str, tpm: Optional[int] = None, rpm: Optional[int] = None -) -> dict: +def _make_deployment_dict(model_name: str, tpm: Optional[int] = None, rpm: Optional[int] = None) -> dict: """Helper to build a minimal deployment dict as returned by router.get_model_list.""" litellm_params: dict = {"model": model_name} if tpm is not None: @@ -1300,9 +1255,7 @@ class TestDeploymentDefaultRpmLimit: """Case 2 from spec: key has no model-specific limits, falls back to deployment default.""" user_api_key_dict = UserAPIKeyAuth(api_key="sk-123") mock_router = MagicMock() - mock_router.get_model_list.return_value = [ - _make_deployment_dict("model1", rpm=200) - ] + mock_router.get_model_list.return_value = [_make_deployment_dict("model1", rpm=200)] with patch(_ROUTER_PATCH, mock_router): result = get_key_model_rpm_limit(user_api_key_dict, model_name="model1") assert result == {"model1": 200} @@ -1314,9 +1267,7 @@ class TestDeploymentDefaultRpmLimit: metadata={"model_rpm_limit": {"model1": 10}}, ) mock_router = MagicMock() - mock_router.get_model_list.return_value = [ - _make_deployment_dict("model1", rpm=200) - ] + mock_router.get_model_list.return_value = [_make_deployment_dict("model1", rpm=200)] with patch(_ROUTER_PATCH, mock_router): result = get_key_model_rpm_limit(user_api_key_dict, model_name="model1") assert result == {"model1": 10} @@ -1336,9 +1287,7 @@ class TestDeploymentDefaultRpmLimit: """No model_name means deployment fallback is skipped.""" user_api_key_dict = UserAPIKeyAuth(api_key="sk-123") mock_router = MagicMock() - mock_router.get_model_list.return_value = [ - _make_deployment_dict("model1", rpm=200) - ] + mock_router.get_model_list.return_value = [_make_deployment_dict("model1", rpm=200)] with patch(_ROUTER_PATCH, mock_router): result = get_key_model_rpm_limit(user_api_key_dict) assert result is None @@ -1399,9 +1348,7 @@ class TestDeploymentDefaultTpmLimit: """Case 2 from spec: key has no model-specific limits, falls back to deployment default.""" user_api_key_dict = UserAPIKeyAuth(api_key="sk-123") mock_router = MagicMock() - mock_router.get_model_list.return_value = [ - _make_deployment_dict("model1", tpm=100) - ] + mock_router.get_model_list.return_value = [_make_deployment_dict("model1", tpm=100)] with patch(_ROUTER_PATCH, mock_router): result = get_key_model_tpm_limit(user_api_key_dict, model_name="model1") assert result == {"model1": 100} @@ -1413,9 +1360,7 @@ class TestDeploymentDefaultTpmLimit: metadata={"model_tpm_limit": {"model1": 20}}, ) mock_router = MagicMock() - mock_router.get_model_list.return_value = [ - _make_deployment_dict("model1", tpm=100) - ] + mock_router.get_model_list.return_value = [_make_deployment_dict("model1", tpm=100)] with patch(_ROUTER_PATCH, mock_router): result = get_key_model_tpm_limit(user_api_key_dict, model_name="model1") assert result == {"model1": 20} @@ -1435,9 +1380,7 @@ class TestDeploymentDefaultTpmLimit: """No model_name means deployment fallback is skipped.""" user_api_key_dict = UserAPIKeyAuth(api_key="sk-123") mock_router = MagicMock() - mock_router.get_model_list.return_value = [ - _make_deployment_dict("model1", tpm=100) - ] + mock_router.get_model_list.return_value = [_make_deployment_dict("model1", tpm=100)] with patch(_ROUTER_PATCH, mock_router): result = get_key_model_tpm_limit(user_api_key_dict) assert result is None @@ -1588,7 +1531,7 @@ class TestCheckCompleteCredentialsBlocksSSRF: "litellm.proxy.auth.auth_utils.validate_url", side_effect=SSRFError(f"blocked: {blocked_url}"), ): - with pytest.raises(ValueError, match='is rejected by the SSRF guard') as exc_info: + with pytest.raises(ValueError, match="is rejected by the SSRF guard") as exc_info: check_complete_credentials( { "model": "gpt-4", @@ -1950,9 +1893,7 @@ class TestIsRequestBodySafeBlocksFallbackSmuggle: is_request_body_safe( request_body={ "model": "gpt-4", - "fallbacks": [ - {"gpt-4": [{"model": "byok", "api_base": "https://my-byok.example"}]} - ], + "fallbacks": [{"gpt-4": [{"model": "byok", "api_base": "https://my-byok.example"}]}], }, general_settings={"allow_client_side_credentials": True}, llm_router=None, @@ -1972,9 +1913,7 @@ class TestIsRequestBodySafeBlocksFallbackSmuggle: "always-fail": [ { "model": "x", - fallback_field: [ - {"x": [{"model": "deepseek-chat", "api_base": "http://attacker"}]} - ], + fallback_field: [{"x": [{"model": "deepseek-chat", "api_base": "http://attacker"}]}], } ] } @@ -2144,7 +2083,7 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: ], ) def test_endpoint_targeting_field_in_request_body_is_rejected(self, field): - with pytest.raises(ValueError, match='Rejected Request') as exc: + with pytest.raises(ValueError, match="Rejected Request") as exc: is_request_body_safe( request_body={"model": "gpt-4", field: "https://attacker.example"}, general_settings={}, @@ -2165,7 +2104,7 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: # on the blocklist into an SSRF / credential-exfil hole. Verify # that supplying an api_key (alongside the banned param) does NOT # bypass the gate — it can only be opened by an admin opt-in. - with pytest.raises(ValueError, match='Rejected Request') as exc: + with pytest.raises(ValueError, match="Rejected Request") as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2590,11 +2529,7 @@ class TestIsRequestBodySafeNestedConfig: when nested.""" with pytest.raises(ValueError, match="langfuse_host"): is_request_body_safe( - request_body={ - "litellm_embedding_config": { - "langfuse_host": "https://attacker.example.com" - } - }, + request_body={"litellm_embedding_config": {"langfuse_host": "https://attacker.example.com"}}, general_settings={}, llm_router=None, model="milvus-store", @@ -2605,11 +2540,7 @@ class TestIsRequestBodySafeNestedConfig: keep the existing escape hatch — same UX as for root-level.""" assert ( is_request_body_safe( - request_body={ - "litellm_embedding_config": { - "api_base": "https://my-azure.example.com" - } - }, + request_body={"litellm_embedding_config": {"api_base": "https://my-azure.example.com"}}, general_settings={"allow_client_side_credentials": True}, llm_router=None, model="milvus-store", @@ -2722,7 +2653,7 @@ class TestObservabilityCallbackBans: ], ) def test_observability_field_in_request_body_root_is_rejected(self, field): - with pytest.raises(ValueError, match='Rejected Request') as exc: + with pytest.raises(ValueError, match="Rejected Request") as exc: is_request_body_safe( request_body={"model": "gpt-4", field: "attacker-value"}, general_settings={}, @@ -2746,13 +2677,11 @@ class TestObservabilityCallbackBans: "user_api_key_auth_metadata", ], ) - def test_observability_field_in_metadata_dict_is_rejected( - self, metadata_key, field - ): + def test_observability_field_in_metadata_dict_is_rejected(self, metadata_key, field): # Verifies the metadata walk: a value smuggled inside ``metadata`` # or ``litellm_metadata`` is just as dangerous as the same field # at the body root, and must hit the same gate. - with pytest.raises(ValueError, match='Rejected Request') as exc: + with pytest.raises(ValueError, match="Rejected Request") as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2787,13 +2716,11 @@ class TestObservabilityCallbackBans: ) def test_observability_field_in_litellm_params_metadata_is_rejected(self): - with pytest.raises(ValueError, match='Rejected Request: turn_off_message_logging is not allowed') as exc: + with pytest.raises(ValueError, match="Rejected Request: turn_off_message_logging is not allowed") as exc: is_request_body_safe( request_body={ "model": "gpt-4", - "litellm_params": { - "metadata": {"turn_off_message_logging": False} - }, + "litellm_params": {"metadata": {"turn_off_message_logging": False}}, }, general_settings={}, llm_router=None, @@ -2805,22 +2732,18 @@ class TestObservabilityCallbackBans: "metadata_key", ["metadata", "litellm_metadata"], ) - def test_observability_field_in_json_string_metadata_is_rejected( - self, metadata_key - ): + def test_observability_field_in_json_string_metadata_is_rejected(self, metadata_key): # Multipart/form-data and ``extra_body`` callers send metadata as a # JSON-encoded string. The bouncer parses it before applying the # banned-params check so the JSON-string path can't smuggle past # the ``isinstance(dict)`` guard. import json - with pytest.raises(ValueError, match='Rejected Request: langfuse_host is not allowed in request') as exc: + with pytest.raises(ValueError, match="Rejected Request: langfuse_host is not allowed in request") as exc: is_request_body_safe( request_body={ "model": "gpt-4", - metadata_key: json.dumps( - {"langfuse_host": "https://attacker.example"} - ), + metadata_key: json.dumps({"langfuse_host": "https://attacker.example"}), }, general_settings={}, llm_router=None, @@ -2887,7 +2810,7 @@ def test_model_level_allow_does_not_skip_subsequent_banned_params(monkeypatch): lambda model, param, request_body_value, llm_router: param == "api_base", ) - with pytest.raises(ValueError, match='Rejected Request: langfuse_host is not allowed in request') as exc: + with pytest.raises(ValueError, match="Rejected Request: langfuse_host is not allowed in request") as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2927,8 +2850,7 @@ def test_observability_ban_covers_canonical_supported_callback_params(): ) for param in _request_blocked_callback_params: assert param in banned, ( - f"{param} is in _request_blocked_callback_params but is not banned " - "at the proxy request-body boundary." + f"{param} is in _request_blocked_callback_params but is not banned at the proxy request-body boundary." ) @@ -2958,7 +2880,7 @@ class TestPricingInjectionBlocked: ], ) def test_pricing_field_rejected_by_default(self, field, value): - with pytest.raises(ValueError, match='Rejected Request') as exc: + with pytest.raises(ValueError, match="Rejected Request") as exc: is_request_body_safe( request_body={"model": "gpt-4", field: value}, general_settings={}, @@ -3026,9 +2948,7 @@ class TestGetRequestRouteTemplate: def test_exception_returns_none(self): req = MagicMock() - type(req).scope = property( - lambda self: (_ for _ in ()).throw(RuntimeError("boom")) - ) + type(req).scope = property(lambda self: (_ for _ in ()).throw(RuntimeError("boom"))) assert get_request_route_template(req) is None @@ -3081,9 +3001,7 @@ class TestGetKeyTagRateLimits: """Tests for get_key_tag_rpm_limit.""" def test_reads_tag_rpm_limit_from_metadata(self): - key = UserAPIKeyAuth( - api_key="sk-123", metadata={"tag_rpm_limit": {"cell-1": 5}} - ) + key = UserAPIKeyAuth(api_key="sk-123", metadata={"tag_rpm_limit": {"cell-1": 5}}) assert get_key_tag_rpm_limit(key) == {"cell-1": 5} def test_returns_none_when_unset(self): @@ -3146,12 +3064,8 @@ class TestIsRequestBodySafeChecksBracketNotationMetadata: def test_bracket_notation_matches_json_encoding_for_deeper_nesting(self): """A value nested below the first level is treated the same either way: the check descends one level into metadata, for both encodings.""" - deep_bracket = { - "litellm_metadata[spend_logs_metadata][langfuse_host]": "https://example.invalid" - } - deep_json = { - "litellm_metadata": {"spend_logs_metadata": {"langfuse_host": "https://example.invalid"}} - } + deep_bracket = {"litellm_metadata[spend_logs_metadata][langfuse_host]": "https://example.invalid"} + deep_json = {"litellm_metadata": {"spend_logs_metadata": {"langfuse_host": "https://example.invalid"}}} kwargs = dict(general_settings={}, llm_router=None, model="gpt-4") assert is_request_body_safe(request_body=deep_bracket, **kwargs) is True assert is_request_body_safe(request_body=deep_json, **kwargs) is True @@ -3200,9 +3114,7 @@ class TestHasUserSetupSso: def test_true_for_saml_metadata_url(self, monkeypatch): from litellm.proxy.auth.auth_utils import has_user_setup_sso - monkeypatch.setenv( - "SAML_IDP_METADATA_URL", "https://idp.example.com/metadata.xml" - ) + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata.xml") assert has_user_setup_sso() is True def test_true_for_saml_metadata_xml(self, monkeypatch): diff --git a/tests/test_litellm/proxy/common_utils/test_credential_hydration.py b/tests/test_litellm/proxy/common_utils/test_credential_hydration.py new file mode 100644 index 00000000000..f40b0114d71 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_credential_hydration.py @@ -0,0 +1,28 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm +from litellm.proxy.common_utils.credential_hydration import hydrate_named_credential_authoritative +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + +@pytest.mark.asyncio +async def test_authoritative_hydrate_returns_an_encrypted_empty_value_as_empty(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-hydration-test-salt") + row = { + "credential_name": "openai-wif", + "credential_values": { + "api_base": encrypt_value_helper(""), + "openai_service_account_id": encrypt_value_helper("user-1"), + }, + "credential_info": {"custom_llm_provider": "openai"}, + } + prisma = MagicMock() + prisma.db.litellm_credentialstable.find_unique = AsyncMock(return_value=row) + + with patch.object(litellm, "credential_list", []): # test-quality-ok: the row under test must win over memory + resolved = await hydrate_named_credential_authoritative("openai-wif", prisma) + + assert resolved is not None + assert resolved.credential_values == {"api_base": "", "openai_service_account_id": "user-1"} diff --git a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py index d67a9afdcc8..da679fbe6e6 100644 --- a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py @@ -1,11 +1,14 @@ """Tests for the credential management endpoints.""" +import json +from contextlib import contextmanager from unittest.mock import AsyncMock, MagicMock, patch import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import ec from fastapi.testclient import TestClient - import litellm from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -19,10 +22,14 @@ def _as_admin(): return UserAPIKeyAuth(api_key="test-key", user_role="proxy_admin") -def _call_as_admin(method: str, path: str, json_body: dict | None = None): +def _as_non_admin(): + return UserAPIKeyAuth(api_key="test-key", user_role="internal_user") + + +def _call_as(method: str, path: str, json_body: dict | None = None, auth=_as_admin): missing = object() previous_override = app.dependency_overrides.get(user_api_key_auth, missing) - app.dependency_overrides[user_api_key_auth] = _as_admin + app.dependency_overrides[user_api_key_auth] = auth try: return client.request(method, path, json=json_body, headers={"Authorization": "Bearer test-key"}) finally: @@ -32,16 +39,26 @@ def _call_as_admin(method: str, path: str, json_body: dict | None = None): app.dependency_overrides[user_api_key_auth] = previous_override -def _patch_credential(name: str, body: dict): - return _call_as_admin("PATCH", f"/credentials/{name}", body) +def _patch_credential(name: str, body: dict, auth=_as_admin): + return _call_as("PATCH", f"/credentials/{name}", body, auth) -def _delete_credential(name: str): - return _call_as_admin("DELETE", f"/credentials/{name}") +def _post_credential(body: dict, auth=_as_admin): + return _call_as("POST", "/credentials", body, auth) + + +def _delete_credential(name: str, auth=_as_admin): + return _call_as("DELETE", f"/credentials/{name}", auth=auth) def _list_credentials(): - return _call_as_admin("GET", "/credentials") + return _call_as("GET", "/credentials") + + +def _prisma_without_credential_rows() -> MagicMock: + prisma_client = MagicMock() + prisma_client.db.litellm_credentialstable.find_unique = AsyncMock(return_value=None) + return prisma_client @pytest.fixture @@ -55,10 +72,11 @@ def credential_store(): in_memory: tuple[object, ...] = (), **repository_calls: AsyncMock, ) -> None: - patch("litellm.proxy.proxy_server.prisma_client", MagicMock() if connected else None).start() + patch("litellm.proxy.proxy_server.prisma_client", _prisma_without_credential_rows() if connected else None).start() patch("litellm.proxy.proxy_server.master_key", "sk-test-master").start() patch.object(litellm, "credential_list", list(in_memory)).start() repository = patch("litellm.proxy.credential_endpoints.endpoints.CredentialsRepository").start() + repository.return_value.find_by_name = AsyncMock(return_value=None) for call_name, result in repository_calls.items(): setattr(repository.return_value, call_name, result) @@ -66,6 +84,52 @@ def credential_store(): patch.stopall() +@contextmanager +def _repository_holding(stored: CredentialItem | None): + """The credentials repository seam, answering ``find_by_name`` with ``stored`` and recording + the writes the handler attempts. Patched at both import sites, since the handlers resolve an + existing credential through ``hydrate_named_credential`` (memory first, then this repository) + and then write through their own ``CredentialsRepository`` binding.""" + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.master_key", "sk-test-master" + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.credential_endpoints.endpoints.CredentialsRepository" + ) as repository, + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.common_utils.credential_hydration.CredentialsRepository", repository + ), + ): + repository.return_value.find_by_name = AsyncMock(return_value=stored) + repository.return_value.create = AsyncMock(return_value=None) + repository.return_value.update_by_name = AsyncMock(return_value=None) + repository.return_value.delete_by_name = AsyncMock(return_value=stored) + yield repository.return_value + + +def test_create_credential_write_omits_the_patch_only_deletion_field(restore_credential_list): + """Regression: CredentialItem.credential_values_to_delete is a PATCH-only field that + defaults to None on every other construction path. A bare .model_dump() (without + exclude_none) on the create path put a `credential_values_to_delete: null` key into the + Prisma write, which litellm_credentialstable has no column for.""" + with _repository_holding(None) as repository: + response = _post_credential( + { + "credential_name": "new-cred", + "credential_values": {"api_key": "sk-new"}, + "credential_info": {"custom_llm_provider": "openai"}, + } + ) + + assert response.status_code == 200, response.text + written_data = repository.create.await_args.kwargs["data"] + assert "credential_values_to_delete" not in written_data + + def test_update_credential_answers_404_when_the_credential_does_not_exist(credential_store): """Regression: the handler used to ``return handle_exception_on_proxy(e)``, which makes the exception the response body and lets FastAPI answer 200, so a write the handler @@ -113,6 +177,872 @@ def test_update_credential_still_answers_200_on_a_successful_write(credential_st assert response.json()["success"] is True +def _get_jwks(name: str): + return _call_as("GET", f"/credentials/{name}/jwks") + + +@pytest.fixture +def restore_credential_list(monkeypatch): + monkeypatch.setattr(litellm, "credential_list", []) + + +def test_update_credential_rejects_overlap_between_update_and_delete(): + """A key in both sets is ambiguous (set to what value, before or after the delete?), so the + endpoint must reject it outright rather than picking a resolution order silently.""" + response = _patch_credential( + "any-name", + { + "credential_name": "any-name", + "credential_values": {"api_key": "sk-new"}, + "credential_values_to_delete": ["api_key"], + "credential_info": {}, + }, + ) + + assert response.status_code == 400, response.text + assert "api_key" in response.json()["error"]["message"] + + +def test_update_credential_deletion_removes_the_key_from_the_db_write(restore_credential_list): + """The bug this closes: switching WIF identity sources (or WIF -> api_key) left the old + variant's fields behind in the DB row, which wif.py then rejects by presence.""" + stored = CredentialItem( + credential_name="wif-cred", + credential_values={"anthropic_identity_source": "keycloak", "anthropic_keycloak_client_id": "old-client"}, + credential_info={"custom_llm_provider": "anthropic"}, + ) + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.credential_endpoints.endpoints.CredentialsRepository" + ) as repository, # test-quality-ok: the proxy wiring under test is what this patches + ): + repository.return_value.find_by_name = AsyncMock(return_value=stored) + update_mock = AsyncMock(return_value=None) + repository.return_value.update_by_name = update_mock + + response = _patch_credential( + "wif-cred", + { + "credential_name": "wif-cred", + "credential_values": {}, + "credential_values_to_delete": ["anthropic_keycloak_client_id"], + "credential_info": {}, + }, + ) + + assert response.status_code == 200, response.text + written_values = json.loads(update_mock.await_args.kwargs["data"]["credential_values"]) + assert "anthropic_keycloak_client_id" not in written_values + assert written_values["anthropic_identity_source"] == "keycloak" + + +def test_update_credential_deletion_updates_in_memory_credential_list(restore_credential_list, monkeypatch): + """The in-memory list is what the request-time auth resolvers read; a deletion that only + landed in the DB would leave the stale field servable until the next process restart.""" + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="wif-cred", + credential_values={ + "anthropic_identity_source": "keycloak", + "anthropic_keycloak_client_id": "old-client", + }, + credential_info={"custom_llm_provider": "anthropic"}, + ) + ], + ) + stored = CredentialItem( + credential_name="wif-cred", + credential_values={"anthropic_identity_source": "keycloak", "anthropic_keycloak_client_id": "old-client"}, + credential_info={"custom_llm_provider": "anthropic"}, + ) + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.credential_endpoints.endpoints.CredentialsRepository" + ) as repository, # test-quality-ok: the proxy wiring under test is what this patches + ): + repository.return_value.find_by_name = AsyncMock(return_value=stored) + repository.return_value.update_by_name = AsyncMock(return_value=None) + + response = _patch_credential( + "wif-cred", + { + "credential_name": "wif-cred", + "credential_values": {}, + "credential_values_to_delete": ["anthropic_keycloak_client_id"], + "credential_info": {}, + }, + ) + + assert response.status_code == 200, response.text + in_memory = next(c for c in litellm.credential_list if c.credential_name == "wif-cred") + assert "anthropic_keycloak_client_id" not in in_memory.credential_values + assert in_memory.credential_values["anthropic_identity_source"] == "keycloak" + + +def test_update_credential_leaves_untouched_fields_alone(): + """Regression for the masked-value hazard: GET /credentials masks values, so a PATCH that + only names the field being changed must not let an untouched field be nulled or overwritten + by anything a round-tripped (masked) form value could contain.""" + stored = CredentialItem( + credential_name="existing", + credential_values={"api_key": "sk-real-value", "api_base": "https://api.anthropic.com"}, + credential_info={"custom_llm_provider": "anthropic"}, + ) + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.master_key", "sk-test-master" + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.credential_endpoints.endpoints.CredentialsRepository" + ) as repository, # test-quality-ok: the proxy wiring under test is what this patches + ): + repository.return_value.find_by_name = AsyncMock(return_value=stored) + update_mock = AsyncMock(return_value=None) + repository.return_value.update_by_name = update_mock + + response = _patch_credential( + "existing", + {"credential_name": "existing", "credential_values": {"api_key": "sk-rotated"}, "credential_info": {}}, + ) + + assert response.status_code == 200, response.text + written_values = json.loads(update_mock.await_args.kwargs["data"]["credential_values"]) + assert written_values["api_base"] == "https://api.anthropic.com" + + +def _generate_es256_pem() -> str: + key = ec.generate_private_key(ec.SECP256R1()) + return key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ).decode() + + +class TestCredentialJwksExport: + def test_jwks_export_succeeds_for_an_internal_issuer_credential(self, restore_credential_list, monkeypatch): + monkeypatch.setenv("JWKS_TEST_SIGNING_KEY", _generate_es256_pem()) + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="anthropic-issuer", + credential_values={ + "anthropic_identity_source": "internal_issuer", + "anthropic_issuer_url": "https://issuer.example.com", + "anthropic_issuer_subject": "my-workload", + "anthropic_issuer_signing_key_ref": "os.environ/JWKS_TEST_SIGNING_KEY", + }, + credential_info={"custom_llm_provider": "anthropic"}, + ) + ], + ) + + response = _get_jwks("anthropic-issuer") + + assert response.status_code == 200, response.text + body = response.json() + assert body["keys"][0]["kty"] == "EC" + assert body["keys"][0]["crv"] == "P-256" + # The private key material must never leave the process via this endpoint. + assert "JWKS_TEST_SIGNING_KEY" not in response.text + assert "PRIVATE KEY" not in response.text + + def test_jwks_export_accepts_the_dashboard_provider_casing(self, restore_credential_list, monkeypatch): + monkeypatch.setenv("JWKS_TEST_SIGNING_KEY", _generate_es256_pem()) + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="anthropic-from-modal", + credential_values={ + "anthropic_identity_source": "internal_issuer", + "anthropic_issuer_url": "https://issuer.example.com", + "anthropic_issuer_subject": "my-workload", + "anthropic_issuer_signing_key_ref": "os.environ/JWKS_TEST_SIGNING_KEY", + }, + credential_info={"custom_llm_provider": "Anthropic"}, + ) + ], + ) + + response = _get_jwks("anthropic-from-modal") + + assert response.status_code == 200, response.text + assert response.json()["keys"][0]["kty"] == "EC" + + def test_jwks_export_404s_for_a_non_anthropic_credential(self, restore_credential_list, monkeypatch): + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="openai-key", + credential_values={"api_key": "sk-x"}, + credential_info={"custom_llm_provider": "openai"}, + ) + ], + ) + + response = _get_jwks("openai-key") + + assert response.status_code == 404, response.text + + def test_jwks_export_404s_for_an_anthropic_credential_without_internal_issuer( + self, restore_credential_list, monkeypatch + ): + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="anthropic-apikey", + credential_values={"api_key": "sk-ant"}, + credential_info={"custom_llm_provider": "anthropic"}, + ) + ], + ) + + response = _get_jwks("anthropic-apikey") + + assert response.status_code == 404, response.text + + def test_jwks_export_404s_for_an_unknown_credential(self, restore_credential_list): + with patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", None + ): # test-quality-ok: the proxy wiring under test is what this patches + response = _get_jwks("does-not-exist") + + assert response.status_code == 404, response.text + + def test_jwks_export_requires_proxy_admin(self, restore_credential_list, monkeypatch): + monkeypatch.setenv("JWKS_TEST_SIGNING_KEY", _generate_es256_pem()) + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="anthropic-issuer", + credential_values={ + "anthropic_identity_source": "internal_issuer", + "anthropic_issuer_url": "https://issuer.example.com", + "anthropic_issuer_subject": "my-workload", + "anthropic_issuer_signing_key_ref": "os.environ/JWKS_TEST_SIGNING_KEY", + }, + credential_info={"custom_llm_provider": "anthropic"}, + ) + ], + ) + + def _as_internal_user(): + return UserAPIKeyAuth(api_key="test-key", user_role="internal_user") + + app.dependency_overrides[user_api_key_auth] = _as_internal_user + try: + response = client.get("/credentials/anthropic-issuer/jwks", headers={"Authorization": "Bearer test-key"}) + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + assert response.status_code == 403, response.text + + +class TestNonAdminCannotPersistWifFieldsOnCredential: + """A credential's ``credential_values`` feeds the same WIF resolution as a deployment's own + ``litellm_params`` when referenced by ``litellm_credential_name``. A non-admin must not be + able to create or update a credential carrying a server-owned WIF field such as + ``anthropic_keycloak_token_url`` (destination) or ``anthropic_keycloak_client_secret_ref`` + (which secret to read and send there).""" + + def test_non_admin_cannot_create_a_credential_with_a_wif_destination(self): + with patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ): # test-quality-ok: the proxy wiring under test is what this patches + response = _post_credential( + { + "credential_name": "attacker-cred", + "credential_values": {"anthropic_keycloak_token_url": "https://evil.example.com/token"}, + "credential_info": {"custom_llm_provider": "anthropic"}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 403, response.text + assert "anthropic_keycloak_token_url" in response.json()["error"]["message"] + + def test_non_admin_cannot_create_a_credential_with_a_wif_secret_ref(self): + with patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ): # test-quality-ok: the proxy wiring under test is what this patches + response = _post_credential( + { + "credential_name": "attacker-cred", + "credential_values": {"anthropic_keycloak_client_secret_ref": "os.environ/LITELLM_MASTER_KEY"}, + "credential_info": {"custom_llm_provider": "anthropic"}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 403, response.text + + def test_non_admin_can_create_a_credential_without_wif_fields(self, restore_credential_list): + with _repository_holding(None) as repository: + response = _post_credential( + { + "credential_name": "ordinary-cred", + "credential_values": {"api_key": "sk-new"}, + "credential_info": {"custom_llm_provider": "openai"}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 200, response.text + repository.create.assert_awaited_once() + + def test_proxy_admin_can_create_a_credential_with_a_wif_destination(self, restore_credential_list): + with _repository_holding(None) as repository: + response = _post_credential( + { + "credential_name": "admin-cred", + "credential_values": {"anthropic_keycloak_token_url": "https://keycloak.internal/token"}, + "credential_info": {"custom_llm_provider": "anthropic"}, + }, + auth=_as_admin, + ) + + assert response.status_code == 200, response.text + repository.create.assert_awaited_once() + + def test_non_admin_cannot_create_a_credential_with_an_openai_token_file(self): + with patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ): + response = _post_credential( + { + "credential_name": "attacker-cred", + "credential_values": {"openai_identity_token_file": "/var/run/secrets/tokens/attacker"}, + "credential_info": {"custom_llm_provider": "openai"}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 403, response.text + assert "openai_identity_token_file" in response.json()["error"]["message"] + + def test_proxy_admin_can_create_a_credential_with_the_openai_identity_trio(self, restore_credential_list): + with _repository_holding(None) as repository: + response = _post_credential( + { + "credential_name": "openai-wif", + "credential_values": { + "openai_identity_provider_id": "idp_1", + "openai_service_account_id": "user-1", + "openai_identity_token_file": "/var/run/secrets/tokens/openai", + }, + "credential_info": {"custom_llm_provider": "openai"}, + }, + auth=_as_admin, + ) + + assert response.status_code == 200, response.text + repository.create.assert_awaited_once() + + def test_non_admin_cannot_update_a_credential_to_add_a_wif_destination(self): + stored = CredentialItem( + credential_name="existing", + credential_values={"api_key": "sk-old"}, + credential_info={"custom_llm_provider": "anthropic"}, + ) + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.credential_endpoints.endpoints.CredentialsRepository" + ) as repository, # test-quality-ok: the proxy wiring under test is what this patches + ): + repository.return_value.find_by_name = AsyncMock(return_value=stored) + update_mock = AsyncMock(return_value=None) + repository.return_value.update_by_name = update_mock + + response = _patch_credential( + "existing", + { + "credential_name": "existing", + "credential_values": {"anthropic_keycloak_token_url": "https://evil.example.com/token"}, + "credential_info": {}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 403, response.text + update_mock.assert_not_awaited() + + def test_proxy_admin_can_update_a_credential_to_add_a_wif_destination(self): + stored = CredentialItem( + credential_name="existing", + credential_values={"api_key": "sk-old"}, + credential_info={"custom_llm_provider": "anthropic"}, + ) + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.master_key", "sk-test-master" + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.credential_endpoints.endpoints.CredentialsRepository" + ) as repository, # test-quality-ok: the proxy wiring under test is what this patches + ): + repository.return_value.find_by_name = AsyncMock(return_value=stored) + update_mock = AsyncMock(return_value=None) + repository.return_value.update_by_name = update_mock + + response = _patch_credential( + "existing", + { + "credential_name": "existing", + "credential_values": {"anthropic_keycloak_token_url": "https://keycloak.internal/token"}, + "credential_info": {}, + }, + auth=_as_admin, + ) + + assert response.status_code == 200, response.text + update_mock.assert_awaited_once() + + +def _wif_credential(name: str = "federated-cred") -> CredentialItem: + return CredentialItem( + credential_name=name, + credential_values={ + "anthropic_keycloak_token_url": "https://keycloak.internal/token", + "api_key": "sk-old", + }, + credential_info={"custom_llm_provider": "anthropic"}, + ) + + +def _plain_credential(name: str = "ordinary-cred") -> CredentialItem: + return CredentialItem( + credential_name=name, + credential_values={"api_key": "sk-old"}, + credential_info={"custom_llm_provider": "openai"}, + ) + + +class TestNonAdminCannotTouchAStoredWifCredential: + """The WIF gate used to read only the incoming ``credential_values``, so a non-admin could + drop a federation field by naming it in ``credential_values_to_delete`` (breaking every + deployment that references the credential), or edit a stored admin-owned WIF credential by + sending a payload carrying no WIF field at all. The gate is evaluated against the effective + surface of the operation: incoming keys (a ``null`` value still persists the key), deleted + keys, and the stored credential, wherever it lives (DB row or config-only ``credential_list`` + entry).""" + + def test_non_admin_cannot_delete_a_wif_field_off_a_credential(self, restore_credential_list): + with _repository_holding(_plain_credential("some-cred")) as repository: + response = _patch_credential( + "some-cred", + { + "credential_name": "some-cred", + "credential_values": {}, + "credential_values_to_delete": ["anthropic_keycloak_token_url"], + "credential_info": {}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 403, response.text + assert "anthropic_keycloak_token_url" in response.text + repository.update_by_name.assert_not_awaited() + + def test_non_admin_cannot_patch_a_stored_wif_credential(self, restore_credential_list): + with _repository_holding(_wif_credential("federated-cred")) as repository: + response = _patch_credential( + "federated-cred", + { + "credential_name": "federated-cred", + "credential_values": {"api_key": "sk-attacker"}, + "credential_info": {}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 403, response.text + assert "anthropic_keycloak_token_url" in response.text + repository.update_by_name.assert_not_awaited() + + def test_proxy_admin_can_delete_a_wif_field_off_a_credential(self, restore_credential_list): + with _repository_holding(_wif_credential("federated-cred")) as repository: + response = _patch_credential( + "federated-cred", + { + "credential_name": "federated-cred", + "credential_values": {}, + "credential_values_to_delete": ["anthropic_keycloak_token_url"], + "credential_info": {}, + }, + auth=_as_admin, + ) + + assert response.status_code == 200, response.text + written_values = json.loads(repository.update_by_name.await_args.kwargs["data"]["credential_values"]) + assert "anthropic_keycloak_token_url" not in written_values + + def test_proxy_admin_can_patch_a_stored_wif_credential(self, restore_credential_list): + with _repository_holding(_wif_credential("federated-cred")) as repository: + response = _patch_credential( + "federated-cred", + { + "credential_name": "federated-cred", + "credential_values": {"api_key": "sk-rotated"}, + "credential_info": {}, + }, + auth=_as_admin, + ) + + assert response.status_code == 200, response.text + written_values = json.loads(repository.update_by_name.await_args.kwargs["data"]["credential_values"]) + assert written_values["anthropic_keycloak_token_url"] is not None + + def test_non_admin_can_still_patch_a_credential_with_no_wif_fields_anywhere(self, restore_credential_list): + with _repository_holding(_plain_credential("ordinary-cred")) as repository: + response = _patch_credential( + "ordinary-cred", + { + "credential_name": "ordinary-cred", + "credential_values": {"api_key": "sk-rotated"}, + "credential_info": {}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 200, response.text + repository.update_by_name.assert_awaited_once() + + def test_non_admin_cannot_delete_a_stored_wif_credential(self, restore_credential_list): + """DELETE takes the whole row, so it drops the admin-owned federation settings as surely + as a targeted key deletion would.""" + with _repository_holding(_wif_credential("federated-cred")) as repository: + response = _delete_credential("federated-cred", auth=_as_non_admin) + + assert response.status_code == 403, response.text + assert "anthropic_keycloak_token_url" in response.text + repository.delete_by_name.assert_not_awaited() + + def test_a_stale_in_memory_copy_does_not_authorize_deleting_a_stored_wif_credential( + self, restore_credential_list, monkeypatch + ): + """Resolution reads memory first and stops, which is right when serving a request. A pod + whose in-memory copy predates an admin adding the federation fields must not read that + stale object and authorize the delete: the gate takes the union of memory and the row.""" + monkeypatch.setattr(litellm, "credential_list", [_plain_credential("federated-cred")]) + + with _repository_holding(_wif_credential("federated-cred")) as repository: + response = _delete_credential("federated-cred", auth=_as_non_admin) + + assert response.status_code == 403, response.text + assert "anthropic_keycloak_token_url" in response.text + repository.delete_by_name.assert_not_awaited() + + def test_proxy_admin_can_delete_a_stored_wif_credential(self, restore_credential_list): + with _repository_holding(_wif_credential("federated-cred")) as repository: + response = _delete_credential("federated-cred", auth=_as_admin) + + assert response.status_code == 200, response.text + repository.delete_by_name.assert_awaited_once_with("federated-cred") + + def test_non_admin_can_still_delete_a_credential_with_no_wif_fields(self, restore_credential_list): + with _repository_holding(_plain_credential("ordinary-cred")) as repository: + response = _delete_credential("ordinary-cred", auth=_as_non_admin) + + assert response.status_code == 200, response.text + repository.delete_by_name.assert_awaited_once_with("ordinary-cred") + + def test_non_admin_cannot_null_out_a_wif_field_on_a_credential(self, restore_credential_list): + """A JSON ``null`` still lands as a key in ``credential_values``. ``get_litellm_params`` + forwards a WIF kwarg on key presence and the federation resolver rejects a foreign + variant's field by key, so a value-based gate let a non-admin persist the key and wedge + every deployment referencing the credential at request time.""" + with _repository_holding(_plain_credential("some-cred")) as repository: + response = _patch_credential( + "some-cred", + { + "credential_name": "some-cred", + "credential_values": {"anthropic_issuer_url": None}, + "credential_info": {}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 403, response.text + assert "anthropic_issuer_url" in response.text + repository.update_by_name.assert_not_awaited() + + def test_non_admin_cannot_patch_a_credential_storing_a_null_wif_field(self, restore_credential_list): + stored = CredentialItem( + credential_name="nulled-cred", + credential_values={"anthropic_issuer_url": None, "api_key": "sk-old"}, + credential_info={"custom_llm_provider": "anthropic"}, + ) + with _repository_holding(stored) as repository: + response = _patch_credential( + "nulled-cred", + { + "credential_name": "nulled-cred", + "credential_values": {"api_key": "sk-attacker"}, + "credential_info": {}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 403, response.text + assert "anthropic_issuer_url" in response.text + repository.update_by_name.assert_not_awaited() + + def test_proxy_admin_can_null_out_a_wif_field_on_a_credential(self, restore_credential_list): + with _repository_holding(_wif_credential("federated-cred")) as repository: + response = _patch_credential( + "federated-cred", + { + "credential_name": "federated-cred", + "credential_values": {"anthropic_keycloak_token_url": None}, + "credential_info": {}, + }, + auth=_as_admin, + ) + + assert response.status_code == 200, response.text + repository.update_by_name.assert_awaited_once() + + def test_non_admin_cannot_delete_a_config_only_wif_credential(self, restore_credential_list, monkeypatch): + """A ``credential_list`` entry from config.yaml has no DB row, so a gate that consulted + only the DB let a non-admin evict the admin-owned federation settings from memory.""" + config_credential = _wif_credential("config-wif") + monkeypatch.setattr(litellm, "credential_list", [config_credential]) + with _repository_holding(None) as repository: + response = _delete_credential("config-wif", auth=_as_non_admin) + + assert response.status_code == 403, response.text + assert "anthropic_keycloak_token_url" in response.text + repository.delete_by_name.assert_not_awaited() + assert litellm.credential_list == [config_credential] + + def test_proxy_admin_can_delete_a_config_only_wif_credential(self, restore_credential_list, monkeypatch): + """The gate lets the admin through to the row delete. The 404 that follows is the rule for + every config-only credential (no row to delete, the entry is back on the next boot), so the + in-memory entry stays put too.""" + config_credential = _wif_credential("config-wif") + monkeypatch.setattr(litellm, "credential_list", [config_credential]) + with _repository_holding(None) as repository: + response = _delete_credential("config-wif", auth=_as_admin) + + assert response.status_code == 404, response.text + repository.delete_by_name.assert_awaited_once_with("config-wif") + assert litellm.credential_list == [config_credential] + + def test_non_admin_cannot_shadow_a_config_only_wif_credential(self, restore_credential_list, monkeypatch): + """POST with the same name carries no WIF field and collides with no DB row, yet + ``CredentialAccessor.upsert_credentials`` would replace the admin entry in memory and + the periodic config sync would then make the takeover permanent.""" + config_credential = _wif_credential("config-wif") + monkeypatch.setattr(litellm, "credential_list", [config_credential]) + with _repository_holding(None) as repository: + response = _post_credential( + { + "credential_name": "config-wif", + "credential_values": {"api_key": "sk-attacker"}, + "credential_info": {"custom_llm_provider": "anthropic"}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 403, response.text + assert "anthropic_keycloak_token_url" in response.text + repository.create.assert_not_awaited() + assert litellm.credential_list == [config_credential] + assert litellm.credential_list[0].credential_values["api_key"] == "sk-old" + + def test_proxy_admin_can_post_over_a_config_only_wif_credential(self, restore_credential_list, monkeypatch): + monkeypatch.setattr(litellm, "credential_list", [_wif_credential("config-wif")]) + with _repository_holding(None) as repository: + response = _post_credential( + { + "credential_name": "config-wif", + "credential_values": {"api_key": "sk-rotated"}, + "credential_info": {"custom_llm_provider": "anthropic"}, + }, + auth=_as_admin, + ) + + assert response.status_code == 200, response.text + repository.create.assert_awaited_once() + assert litellm.credential_list[0].credential_values == {"api_key": "sk-rotated"} + + def test_non_admin_cannot_rename_a_credential_onto_a_config_only_wif_credential( + self, restore_credential_list, monkeypatch + ): + """PATCH is the other way to shadow: renaming an ordinary credential onto the WIF + credential's name makes ``_sync_in_memory_credential`` upsert the attacker's values over + the admin entry, with no WIF field in the payload and no DB row to collide with.""" + config_credential = _wif_credential("config-wif") + monkeypatch.setattr(litellm, "credential_list", [_plain_credential("mine"), config_credential]) + with _repository_holding(_plain_credential("mine")) as repository: + response = _patch_credential( + "mine", + { + "credential_name": "config-wif", + "credential_values": {"api_key": "sk-attacker"}, + "credential_info": {}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 403, response.text + assert "anthropic_keycloak_token_url" in response.text + repository.update_by_name.assert_not_awaited() + assert config_credential in litellm.credential_list + assert litellm.credential_list[1].credential_values["api_key"] == "sk-old" + + def test_proxy_admin_can_rename_a_credential_onto_a_config_only_wif_credential( + self, restore_credential_list, monkeypatch + ): + monkeypatch.setattr(litellm, "credential_list", [_plain_credential("mine"), _wif_credential("config-wif")]) + with _repository_holding(_plain_credential("mine")) as repository: + response = _patch_credential( + "mine", + { + "credential_name": "config-wif", + "credential_values": {"api_key": "sk-rotated"}, + "credential_info": {}, + }, + auth=_as_admin, + ) + + assert response.status_code == 200, response.text + repository.update_by_name.assert_awaited_once() + assert [c.credential_name for c in litellm.credential_list] == ["config-wif"] + + def test_non_admin_cannot_post_a_null_wif_field(self, restore_credential_list): + """Same key-presence rule on the create path: ``{"anthropic_issuer_url": null}`` persists + the key, and the resolver reacts to the key.""" + with _repository_holding(None) as repository: + response = _post_credential( + { + "credential_name": "nulled-cred", + "credential_values": {"anthropic_issuer_url": None, "api_key": "sk-new"}, + "credential_info": {"custom_llm_provider": "anthropic"}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 403, response.text + assert "anthropic_issuer_url" in response.text + repository.create.assert_not_awaited() + assert litellm.credential_list == [] + + def test_non_admin_cannot_shadow_a_db_stored_wif_credential(self, restore_credential_list): + """Same hole for a WIF credential another pod wrote to the DB before this pod's in-memory + list caught up: the existing-credential lookup falls through to the DB.""" + with _repository_holding(_wif_credential("federated-cred")) as repository: + response = _post_credential( + { + "credential_name": "federated-cred", + "credential_values": {"api_key": "sk-attacker"}, + "credential_info": {"custom_llm_provider": "anthropic"}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 403, response.text + repository.create.assert_not_awaited() + + def test_non_admin_can_still_post_a_credential_with_no_wif_fields_anywhere(self, restore_credential_list): + with _repository_holding(None) as repository: + response = _post_credential( + { + "credential_name": "ordinary-cred", + "credential_values": {"api_key": "sk-new"}, + "credential_info": {"custom_llm_provider": "openai"}, + }, + auth=_as_non_admin, + ) + + assert response.status_code == 200, response.text + repository.create.assert_awaited_once() + assert litellm.credential_list[0].credential_name == "ordinary-cred" + + +class TestManagementReadsTheStoredCredential: + """Serving a request reads memory first, which is right. A management operation cannot: on a + pod whose in-memory copy predates another pod's update it would act on superseded values.""" + + @pytest.mark.asyncio + async def test_authoritative_hydrate_prefers_the_row_over_a_stale_memory_copy(self): + import litellm + from litellm.proxy.common_utils.credential_hydration import ( + hydrate_named_credential, + hydrate_named_credential_authoritative, + ) + from litellm.types.utils import CredentialItem + + stale = CredentialItem( + credential_name="anthropic-wif", + credential_values={"anthropic_issuer_url": "https://old.example.com"}, + credential_info={"custom_llm_provider": "anthropic"}, + ) + row = { + "credential_name": "anthropic-wif", + "credential_values": {"anthropic_issuer_url": "https://new.example.com"}, + "credential_info": {"custom_llm_provider": "anthropic"}, + } + + prisma = MagicMock() + prisma.db.litellm_credentialstable.find_unique = AsyncMock(return_value=row) + + with patch.object(litellm, "credential_list", [stale]): # test-quality-ok: the stale copy under test + served = await hydrate_named_credential("anthropic-wif", prisma) + managed = await hydrate_named_credential_authoritative("anthropic-wif", prisma) + + assert served is not None and served.credential_values["anthropic_issuer_url"] == "https://old.example.com" + assert managed is not None and managed.credential_values["anthropic_issuer_url"] == "https://new.example.com" + + @pytest.mark.asyncio + async def test_authoritative_hydrate_falls_back_to_memory_when_the_row_is_absent(self): + import litellm + from litellm.proxy.common_utils.credential_hydration import hydrate_named_credential_authoritative + from litellm.types.utils import CredentialItem + + only_in_memory = CredentialItem( + credential_name="config-yaml-credential", + credential_values={"anthropic_issuer_url": "https://configured.example.com"}, + credential_info={"custom_llm_provider": "anthropic"}, + ) + prisma = MagicMock() + prisma.db.litellm_credentialstable.find_unique = AsyncMock(return_value=None) + + with patch.object(litellm, "credential_list", [only_in_memory]): # test-quality-ok: the config.yaml fallback under test + resolved = await hydrate_named_credential_authoritative("config-yaml-credential", prisma) + + assert resolved is not None + assert resolved.credential_values["anthropic_issuer_url"] == "https://configured.example.com" + + def test_delete_credential_answers_404_when_the_credential_does_not_exist(credential_store): """Regression: prisma's ``delete`` hands back None when the ``where`` clause matched no row instead of raising, and the handler never looked. Deleting a name that was never stored diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 33de2a09626..d1a05cd191b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -7,6 +7,7 @@ from typing import Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException from fastapi.testclient import TestClient from litellm._uuid import uuid @@ -30,7 +31,13 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( ) from litellm.proxy.utils import PrismaClient from litellm.router import Router -from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment +from litellm.types.router import ( + Deployment, + LiteLLM_Params, + ModelInfo, + updateDeployment, + updateLiteLLMParams, +) async def _passthrough_row(update_data): @@ -58,11 +65,7 @@ class MockPrismaClient: return LiteLLM_TeamTable( team_id=where["team_id"], team_alias="test_team", - members_with_roles=[ - Member( - user_id="test_user", role="admin" if self.user_admin else "user" - ) - ], + members_with_roles=[Member(user_id="test_user", role="admin" if self.user_admin else "user")], ) return None @@ -76,10 +79,7 @@ class MockPrismaClient: # Support model_name startswith filter (used by _get_team_deployments) if where and "model_name" in where: model_name_filter = where["model_name"] - if ( - isinstance(model_name_filter, dict) - and "startswith" in model_name_filter - ): + if isinstance(model_name_filter, dict) and "startswith" in model_name_filter: prefix = model_name_filter["startswith"] results = [d for d in results if d.model_name.startswith(prefix)] @@ -124,13 +124,9 @@ class MockProxyConfig: class TestModelManagementAuthChecks: def setup_method(self): """Setup test cases""" - self.admin_user = UserAPIKeyAuth( - user_id="test_admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + self.admin_user = UserAPIKeyAuth(user_id="test_admin", user_role=LitellmUserRoles.PROXY_ADMIN) - self.normal_user = UserAPIKeyAuth( - user_id="test_user", user_role=LitellmUserRoles.INTERNAL_USER - ) + self.normal_user = UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.INTERNAL_USER) self.team_admin_user = UserAPIKeyAuth( user_id="test_user", @@ -149,7 +145,7 @@ class TestModelManagementAuthChecks: @pytest.mark.asyncio async def test_can_user_make_team_model_call_non_premium_fails(self): """Test that non-premium users cannot make team model calls""" - with pytest.raises(Exception, match='You must be a LiteLLM Enterprise user to use this feature\\.') as exc_info: + with pytest.raises(Exception, match="You must be a LiteLLM Enterprise user to use this feature\\.") as exc_info: ModelManagementAuthChecks.can_user_make_team_model_call( team_id="test_team", user_api_key_dict=self.admin_user, @@ -163,9 +159,7 @@ class TestModelManagementAuthChecks: team_obj = LiteLLM_TeamTable( team_id="test_team", team_alias="test_team", - members_with_roles=[ - Member(user_id=self.team_admin_user.user_id, role="admin") - ], + members_with_roles=[Member(user_id=self.team_admin_user.user_id, role="admin")], ) result = ModelManagementAuthChecks.can_user_make_team_model_call( @@ -204,7 +198,7 @@ class TestModelManagementAuthChecks: ) prisma_client = MockPrismaClient(team_exists=True) - with pytest.raises(Exception, match='You must be a LiteLLM Enterprise user to use this feature\\.') as exc_info: + with pytest.raises(Exception, match="You must be a LiteLLM Enterprise user to use this feature\\.") as exc_info: await ModelManagementAuthChecks.allow_team_model_action( model_params=model_params, user_api_key_dict=self.admin_user, @@ -251,6 +245,7 @@ class TestModelManagementAuthChecks: user_api_key_dict=self.admin_user, prisma_client=prisma_client, premium_user=True, + incoming_params=None, ) assert result is True @@ -272,6 +267,7 @@ class TestModelManagementAuthChecks: user_api_key_dict=self.normal_user, prisma_client=prisma_client, premium_user=True, + incoming_params=None, ) assert "403" in str(exc_info.value) @@ -464,29 +460,21 @@ class TestDeleteTeamModelAlias: mock_prisma.db = MockPrismaWrapper(model_aliases_list) # Call the function - await delete_team_model_alias( - public_model_name="public_model_1", prisma_client=mock_prisma - ) + await delete_team_model_alias(public_model_name="public_model_1", prisma_client=mock_prisma) # Verify results mock_db = mock_prisma.db.litellm_modeltable - assert ( - len(mock_db.update_calls) == 2 - ) # Should have 2 update calls since public_model_1 appears twice + assert len(mock_db.update_calls) == 2 # Should have 2 update calls since public_model_1 appears twice # Verify first update first_update = mock_db.update_calls[0] assert first_update["where"] == {"id": 1} - assert json.loads(first_update["data"]["model_aliases"]) == { - "alias2": "public_model_2" - } + assert json.loads(first_update["data"]["model_aliases"]) == {"alias2": "public_model_2"} # Verify second update second_update = mock_db.update_calls[1] assert second_update["where"] == {"id": 2} - assert json.loads(second_update["data"]["model_aliases"]) == { - "alias3": "public_model_3" - } + assert json.loads(second_update["data"]["model_aliases"]) == {"alias3": "public_model_3"} @pytest.mark.asyncio async def test_delete_team_model_alias_no_matches(self): @@ -522,9 +510,7 @@ class TestDeleteTeamModelAlias: mock_prisma.db = MockPrismaWrapper(model_aliases_list) # Call the function with non-existent model - await delete_team_model_alias( - public_model_name="non_existent_model", prisma_client=mock_prisma - ) + await delete_team_model_alias(public_model_name="non_existent_model", prisma_client=mock_prisma) # Verify no updates were made mock_db = mock_prisma.db.litellm_modeltable @@ -1023,18 +1009,12 @@ class TestUpdateModel: updated_row.model_dump_json.return_value = "{}" mock_prisma = MagicMock() - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=existing_row - ) - mock_prisma.db.litellm_proxymodeltable.update = AsyncMock( - return_value=updated_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) mock_router = MagicMock() mock_router.get_model_ids.return_value = [model_id] - admin_user = UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), @@ -1051,9 +1031,7 @@ class TestUpdateModel: ), patch( "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", - new=AsyncMock( - return_value=ReconcileOutcome(still_desired=None, live_after=None) - ), + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ) as mock_clear_cache, ): await update_model( @@ -1098,9 +1076,7 @@ class TestUpdatePublicModelGroups: mock_proxy_config.get_config = mock_get_config mock_proxy_config.save_config = AsyncMock() - admin_user = UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) request = UpdatePublicModelGroupsRequest(model_groups=new_models) @@ -1156,9 +1132,7 @@ class TestUpdatePublicModelGroups: mock_proxy_config.get_config = mock_get_config mock_proxy_config.save_config = AsyncMock() - admin_user = UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) request = UpdateUsefulLinksRequest(useful_links=new_links) @@ -1323,9 +1297,7 @@ class TestTeamModelSiblingRouting: ) # Global deployment should be accessible when team_id is provided - deployments = router._get_all_deployments( - model_name="global-gpt-4o", team_id="teamA" - ) + deployments = router._get_all_deployments(model_name="global-gpt-4o", team_id="teamA") assert len(deployments) == 1 assert deployments[0]["model_name"] == "global-gpt-4o" @@ -1374,9 +1346,9 @@ class TestTeamModelUpdate: patch( "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add" ) as mock_team_model_add, - patch( + patch( # test-quality-ok: the proxy wiring under test is what this patches "litellm.proxy.management_endpoints.model_management_endpoints.update_team" - ) as mock_update_team, + ) as mock_update_team, # test-quality-ok: the proxy wiring under test is what this patches ): result = await _update_team_model_in_db( db_model=db_model, @@ -1407,9 +1379,7 @@ class TestTeamModelUpdate: db_model = Deployment( model_name="model_name_team_123_uuid1", litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"), - model_info=ModelInfo( - team_id="team_123", team_public_model_name="old-public-name" - ), + model_info=ModelInfo(team_id="team_123", team_public_model_name="old-public-name"), ) # Create a sibling deployment that still uses the old public name @@ -1420,9 +1390,7 @@ class TestTeamModelUpdate: "team_public_model_name": "old-public-name", } - prisma_client = MockPrismaClient( - team_exists=True, sibling_deployments=[sibling_deployment] - ) + prisma_client = MockPrismaClient(team_exists=True, sibling_deployments=[sibling_deployment]) patch_data = updateDeployment( model_name="new-public-name", @@ -1437,10 +1405,10 @@ class TestTeamModelUpdate: with ( patch( "litellm.proxy.management_endpoints.model_management_endpoints.team_model_delete" - ) as mock_delete, + ) as mock_delete, # test-quality-ok: the proxy wiring under test is what this patches patch( "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add" - ) as mock_add, + ) as mock_add, # test-quality-ok: the proxy wiring under test is what this patches ): await _update_existing_team_model_assignment( team_id="team_123", @@ -1482,10 +1450,10 @@ class TestTeamModelUpdate: with ( patch( "litellm.proxy.management_endpoints.model_management_endpoints.team_model_delete" - ) as mock_delete, + ) as mock_delete, # test-quality-ok: the proxy wiring under test is what this patches patch( "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add" - ) as mock_add, + ) as mock_add, # test-quality-ok: the proxy wiring under test is what this patches ): await _update_existing_team_model_assignment( team_id="team_123", @@ -1503,7 +1471,6 @@ class TestTeamModelUpdate: """The team's model list autocommits, so it is written only after the row write succeeded: a refused write (the heuristic_v2 slot 403, a DB error) must not leave the team listing a name whose row never changed.""" - from fastapi import HTTPException from litellm.proxy.management_endpoints.model_management_endpoints import ( _update_team_model_in_db, @@ -1579,20 +1546,14 @@ class TestTeamModelUpdate: db_model = Deployment( model_name="model_name_team_123_uuid1", litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"), - model_info=ModelInfo( - team_id="team_123", team_public_model_name="old-public-name" - ), + model_info=ModelInfo(team_id="team_123", team_public_model_name="old-public-name"), ) sibling_deployment = MagicMock() sibling_deployment.model_name = "model_name_team_123_uuid2" - sibling_deployment.model_info = ( - '{"team_id":"team_123","team_public_model_name":"old-public-name"}' - ) + sibling_deployment.model_info = '{"team_id":"team_123","team_public_model_name":"old-public-name"}' - prisma_client = MockPrismaClient( - team_exists=True, sibling_deployments=[sibling_deployment] - ) + prisma_client = MockPrismaClient(team_exists=True, sibling_deployments=[sibling_deployment]) patch_data = updateDeployment( model_name="new-public-name", @@ -1607,10 +1568,10 @@ class TestTeamModelUpdate: with ( patch( "litellm.proxy.management_endpoints.model_management_endpoints.team_model_delete" - ) as mock_delete, + ) as mock_delete, # test-quality-ok: the proxy wiring under test is what this patches patch( "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add" - ) as mock_add, + ) as mock_add, # test-quality-ok: the proxy wiring under test is what this patches ): await _update_existing_team_model_assignment( team_id="team_123", @@ -1687,10 +1648,7 @@ class TestTeamModelUpdate: ), ) - assert ( - _get_public_model_name(patch_data=patch_data, db_model=db_model) - == "gpt-5.2-low-rpm-testing" - ) + assert _get_public_model_name(patch_data=patch_data, db_model=db_model) == "gpt-5.2-low-rpm-testing" def test_get_public_model_name_preserves_db_public_name_when_internal_name_unchanged( self, @@ -1717,10 +1675,7 @@ class TestTeamModelUpdate: model_info=ModelInfo(team_id="test-team"), ) - assert ( - _get_public_model_name(patch_data=patch_data, db_model=db_model) - == "gpt-5.2-low-rpm-testing" - ) + assert _get_public_model_name(patch_data=patch_data, db_model=db_model) == "gpt-5.2-low-rpm-testing" def test_get_public_model_name_allows_top_level_rename(self): """A genuine rename via the top-level model_name field (no @@ -1745,10 +1700,7 @@ class TestTeamModelUpdate: model_info=ModelInfo(team_id="test-team"), ) - assert ( - _get_public_model_name(patch_data=patch_data, db_model=db_model) - == "new-public-name" - ) + assert _get_public_model_name(patch_data=patch_data, db_model=db_model) == "new-public-name" def test_get_public_model_name_top_level_rename_wins_over_stale_model_info(self): """Regression (codex review): on a dashboard rename the UI sends the new @@ -1765,9 +1717,7 @@ class TestTeamModelUpdate: db_model = Deployment( model_name="model_name_team-a_abc123", litellm_params=LiteLLM_Params(model="azure/gpt-4.1"), - model_info=ModelInfo( - team_id="team-a", team_public_model_name="old-public-name" - ), + model_info=ModelInfo(team_id="team-a", team_public_model_name="old-public-name"), ) patch_data = updateDeployment( model_name="new-public-name", @@ -1777,10 +1727,7 @@ class TestTeamModelUpdate: ), ) - assert ( - _get_public_model_name(patch_data=patch_data, db_model=db_model) - == "new-public-name" - ) + assert _get_public_model_name(patch_data=patch_data, db_model=db_model) == "new-public-name" def test_get_public_model_name_falls_back_to_db_public_name(self): """When patch_data carries no name hints at all (neither model_name @@ -1803,10 +1750,7 @@ class TestTeamModelUpdate: model_info=ModelInfo(team_id="test-team"), ) - assert ( - _get_public_model_name(patch_data=patch_data, db_model=db_model) - == "gpt-5.2-low-rpm-testing" - ) + assert _get_public_model_name(patch_data=patch_data, db_model=db_model) == "gpt-5.2-low-rpm-testing" def test_get_public_model_name_last_resort_returns_db_model_name(self): """Legacy rows may have no team_public_model_name anywhere; the @@ -1826,10 +1770,7 @@ class TestTeamModelUpdate: model_info=ModelInfo(team_id="test-team"), ) - assert ( - _get_public_model_name(patch_data=patch_data, db_model=db_model) - == "legacy-model" - ) + assert _get_public_model_name(patch_data=patch_data, db_model=db_model) == "legacy-model" def test_get_public_model_name_ignores_different_internal_shape_name(self): """A stale client may PATCH an internal-shaped model_name that does not @@ -1853,10 +1794,7 @@ class TestTeamModelUpdate: model_info=ModelInfo(team_id="test-team"), ) - assert ( - _get_public_model_name(patch_data=patch_data, db_model=db_model) - == "gpt-5.2-low-rpm-testing" - ) + assert _get_public_model_name(patch_data=patch_data, db_model=db_model) == "gpt-5.2-low-rpm-testing" def test_get_public_model_name_ignores_internal_shape_patch_public(self): """If a corrupted row round-trips an internal-shaped value in @@ -1882,10 +1820,7 @@ class TestTeamModelUpdate: ), ) - assert ( - _get_public_model_name(patch_data=patch_data, db_model=db_model) - == "gpt-5.2-low-rpm-testing" - ) + assert _get_public_model_name(patch_data=patch_data, db_model=db_model) == "gpt-5.2-low-rpm-testing" @pytest.mark.asyncio async def test_dashboard_edit_preserves_public_name_and_acl(self): @@ -1953,9 +1888,7 @@ class TestTeamModelUpdate: # the merged model_info written to the DB must keep the public name model_info_json = result.get("model_info", "") parsed_model_info = json.loads(model_info_json) - assert ( - parsed_model_info.get("team_public_model_name") == "gpt-5.2-low-rpm-testing" - ) + assert parsed_model_info.get("team_public_model_name") == "gpt-5.2-low-rpm-testing" # the internal model_name must not have been overwritten (caller # intentionally clears patch_data.model_name so the DB row's name @@ -1997,9 +1930,7 @@ class TestModelInfoEndpoint: model_info=ModelInfo(id="gpt-4"), ) - result = await model_info( - model_id="gpt-4", user_api_key_dict=user_api_key_dict - ) + result = await model_info(model_id="gpt-4", user_api_key_dict=user_api_key_dict) assert result["id"] == "gpt-4" assert result["object"] == "model" @@ -2009,7 +1940,6 @@ class TestModelInfoEndpoint: @pytest.mark.asyncio async def test_model_info_inaccessible_model_returns_404(self): """Test model_info returns 404 for inaccessible models""" - from fastapi import HTTPException from litellm.proxy.proxy_server import model_info @@ -2074,9 +2004,7 @@ class TestModelInfoEndpoint: model_info=ModelInfo(id="team-model-1"), ) - result = await model_info( - model_id="team-model-1", user_api_key_dict=user_api_key_dict - ) + result = await model_info(model_id="team-model-1", user_api_key_dict=user_api_key_dict) assert result["id"] == "team-model-1" assert result["object"] == "model" @@ -2108,9 +2036,7 @@ class TestAddAndDeleteModelLifecycle: ) model_id = "lifecycle-test-model-123" - admin_user = UserAPIKeyAuth( - user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_user = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) # Build a real LiteLLM_ProxyModelTable for the DB mock to return db_row = LiteLLM_ProxyModelTable( @@ -2127,9 +2053,7 @@ class TestAddAndDeleteModelLifecycle: mock_prisma.db.litellm_proxymodeltable = AsyncMock() mock_prisma.db.query_raw = AsyncMock(return_value=[]) mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=db_row) - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=db_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) mock_proxy_config = MagicMock() @@ -2153,14 +2077,11 @@ class TestAddAndDeleteModelLifecycle: patch(f"{_PS}.llm_router", mock_router), patch(_ENCRYPT, side_effect=lambda value, **kwargs: value), ): - # --- ADD --- add_result = await add_new_model( model_params=Deployment( model_name="lifecycle-model", - litellm_params=LiteLLM_Params( - model="openai/gpt-4.1-nano", api_key="fake-key" - ), + litellm_params=LiteLLM_Params(model="openai/gpt-4.1-nano", api_key="fake-key"), model_info={"id": model_id}, ), user_api_key_dict=admin_user, @@ -2175,9 +2096,7 @@ class TestAddAndDeleteModelLifecycle: assert "deleted successfully" in delete_result["message"] # --- DELETE again should fail (model not found) --- - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None) from litellm.proxy.proxy_server import ProxyException with pytest.raises(ProxyException) as exc_info: @@ -2239,24 +2158,18 @@ class TestDeleteTeamBYOKModelGhost: mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() mock_prisma.db.query_raw = AsyncMock(return_value=[]) - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=db_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) # After the row delete no team deployment remains -> nothing backs the public name. mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_teamtable = AsyncMock() mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=updated_team_row - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated_team_row) # Team BYOK models have no alias row; delete_team_model_alias finds nothing. mock_prisma.db.litellm_modeltable = AsyncMock() mock_prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[]) - admin_user = UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) _PS = "litellm.proxy.proxy_server" _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" @@ -2322,9 +2235,7 @@ class TestDeleteTeamBYOKModelGhost: mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() mock_prisma.db.query_raw = AsyncMock(return_value=[]) - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=db_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_teamtable = AsyncMock() @@ -2334,9 +2245,7 @@ class TestDeleteTeamBYOKModelGhost: # No alias row matches -> delete_team_model_alias returns nothing, but it still ran. mock_prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[]) - admin_user = UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) _PS = "litellm.proxy.proxy_server" _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" @@ -2399,25 +2308,17 @@ class TestDeleteTeamBYOKModelGhost: mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() mock_prisma.db.query_raw = AsyncMock(return_value=[]) - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=deleted_row - ) - mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock( - return_value=deleted_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=deleted_row) + mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=deleted_row) # After the deleted replica's row is gone, the sibling still backs the public name. - mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock( - return_value=[sibling_row] - ) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[sibling_row]) mock_prisma.db.litellm_teamtable = AsyncMock() mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=team_row) mock_prisma.db.litellm_modeltable = AsyncMock() mock_prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[]) - admin_user = UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) _PS = "litellm.proxy.proxy_server" _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" @@ -2475,9 +2376,7 @@ class TestDeleteTeamBYOKModelGhost: members_with_roles=[Member(user_id="admin", role="admin")], models=[public_name], ) - alias_row = MagicMock( - id="alias-row-1", model_aliases={public_name: internal_name} - ) + alias_row = MagicMock(id="alias-row-1", model_aliases={public_name: internal_name}) alias_row.team = MagicMock() alias_row.team.team_id = team_id @@ -2485,26 +2384,20 @@ class TestDeleteTeamBYOKModelGhost: mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() mock_prisma.db.query_raw = AsyncMock(return_value=[]) - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=db_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_teamtable = AsyncMock() mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=team_row) mock_prisma.db.litellm_modeltable = AsyncMock() - mock_prisma.db.litellm_modeltable.find_many = AsyncMock( - return_value=[alias_row] - ) + mock_prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[alias_row]) mock_prisma.db.litellm_modeltable.update = AsyncMock() mock_router = MagicMock() mock_router.model_name_to_deployment_indices = {public_name: [0]} - admin_user = UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) _PS = "litellm.proxy.proxy_server" _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" @@ -2567,9 +2460,7 @@ class TestDeleteTeamBYOKModelGhost: mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() mock_prisma.db.query_raw = AsyncMock(return_value=[]) - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=db_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_teamtable = AsyncMock() @@ -2582,9 +2473,7 @@ class TestDeleteTeamBYOKModelGhost: mock_router = MagicMock() mock_router.model_name_to_deployment_indices = {internal_name: [0]} - admin_user = UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) _PS = "litellm.proxy.proxy_server" _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" @@ -2637,9 +2526,7 @@ class TestDeleteModelTeamAuth: mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() mock_prisma.db.query_raw = AsyncMock(return_value=[]) - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=db_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) # The team is gone -> every team lookup returns None. @@ -2661,9 +2548,7 @@ class TestDeleteModelTeamAuth: model_id = "orphaned-byok-1" mock_prisma = self._orphaned_model_mocks(team_id, model_id) - admin_user = UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) _PS = "litellm.proxy.proxy_server" _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" @@ -2699,9 +2584,7 @@ class TestDeleteModelTeamAuth: model_id = "orphaned-byok-2" mock_prisma = self._orphaned_model_mocks(team_id, model_id) - non_admin = UserAPIKeyAuth( - user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER - ) + non_admin = UserAPIKeyAuth(user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER) _PS = "litellm.proxy.proxy_server" _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" @@ -2756,9 +2639,7 @@ class TestDeleteModelTeamAuth: mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() mock_prisma.db.query_raw = AsyncMock(return_value=[]) - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=db_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_teamtable = AsyncMock() @@ -2768,9 +2649,7 @@ class TestDeleteModelTeamAuth: # A team member who is not the team admin: rejected before the delete runs, # so the only team lookup is the single one inside the auth check. - non_admin = UserAPIKeyAuth( - user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER - ) + non_admin = UserAPIKeyAuth(user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER) _PS = "litellm.proxy.proxy_server" _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" @@ -2964,15 +2843,11 @@ class TestDeleteTeamModels: prisma = _TxPrismaClient(rows) router = _RecordingRouter(prisma.events) - await delete_team_models( - team_ids=["team_a", "team_b"], prisma_client=prisma, llm_router=router - ) + await delete_team_models(team_ids=["team_a", "team_b"], prisma_client=prisma, llm_router=router) commit_idx = prisma.events.index(("commit",)) router_indices = [i for i, e in enumerate(prisma.events) if e[0] == "router"] - delete_indices = [ - i for i, e in enumerate(prisma.events) if e[0] == "delete_many" - ] + delete_indices = [i for i, e in enumerate(prisma.events) if e[0] == "delete_many"] assert router_indices, "router was never synced" assert all(i > commit_idx for i in router_indices) assert all(i < commit_idx for i in delete_indices) @@ -2988,9 +2863,7 @@ class TestDeleteTeamModels: prisma = _TxPrismaClient([mine, intruder]) router = _RecordingRouter(prisma.events) - deleted = await delete_team_models( - team_ids=["team_a"], prisma_client=prisma, llm_router=router - ) + deleted = await delete_team_models(team_ids=["team_a"], prisma_client=prisma, llm_router=router) assert deleted == ["a1"] assert router.deleted == ["a1"] @@ -3000,9 +2873,7 @@ class TestDeleteTeamModels: prisma = _TxPrismaClient([]) router = _RecordingRouter(prisma.events) - deleted = await delete_team_models( - team_ids=["team_a"], prisma_client=prisma, llm_router=router - ) + deleted = await delete_team_models(team_ids=["team_a"], prisma_client=prisma, llm_router=router) assert deleted == [] assert router.deleted == [] @@ -3013,9 +2884,7 @@ class TestDeleteTeamModels: rows = [_model_row("a1", "team_a")] prisma = _TxPrismaClient(rows) - deleted = await delete_team_models( - team_ids=["team_a"], prisma_client=prisma, llm_router=None - ) + deleted = await delete_team_models(team_ids=["team_a"], prisma_client=prisma, llm_router=None) assert deleted == ["a1"] assert any(e[0] == "delete_many" for e in prisma.events) @@ -3103,9 +2972,7 @@ class TestUpdateDBModelClearPricing: result = update_db_model( db_model=_build_db_model_with_pricing(), - updated_patch=updateDeployment( - litellm_params=updateLiteLLMParams(input_cost_per_token=None) - ), + updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(input_cost_per_token=None)), ) params = json.loads(result["litellm_params"]) @@ -3124,9 +2991,7 @@ class TestUpdateDBModelClearPricing: result = update_db_model( db_model=_build_db_model_with_pricing(), - updated_patch=updateDeployment( - litellm_params=updateLiteLLMParams(output_cost_per_token=None) - ), + updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(output_cost_per_token=None)), ) params = json.loads(result["litellm_params"]) @@ -3142,9 +3007,7 @@ class TestUpdateDBModelClearPricing: result = update_db_model( db_model=_build_db_model_with_pricing(), - updated_patch=updateDeployment( - litellm_params=updateLiteLLMParams(input_cost_per_token=0.000005) - ), + updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(input_cost_per_token=0.000005)), ) params = json.loads(result["litellm_params"]) @@ -3159,9 +3022,7 @@ class TestUpdateDBModelClearPricing: result = update_db_model( db_model=_build_db_model_with_pricing(), - updated_patch=updateDeployment( - litellm_params=updateLiteLLMParams(output_cost_per_token=0.000007) - ), + updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(output_cost_per_token=0.000007)), ) params = json.loads(result["litellm_params"]) @@ -3196,9 +3057,7 @@ class TestUpdateDBModelClearPricing: # or any other non-pricing field from the merged dict. result = update_db_model( db_model=db_model, - updated_patch=updateDeployment( - litellm_params=updateLiteLLMParams(api_base=None) - ), + updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(api_base=None)), ) info = json.loads(result["model_info"]) @@ -3233,9 +3092,7 @@ class TestUpdateDBModelClearPricing: params = json.loads(result["litellm_params"]) info = json.loads(result["model_info"]) assert "input_cost_per_token" not in params - assert ( - "input_cost_per_token" not in info - ), "model_info passthrough must not resurrect the cleared override" + assert "input_cost_per_token" not in info, "model_info passthrough must not resurrect the cleared override" def test_clear_via_model_info_clears_both_blobs(self): """The mirror works in the reverse direction too: nulling a pricing field @@ -3247,9 +3104,7 @@ class TestUpdateDBModelClearPricing: result = update_db_model( db_model=_build_db_model_with_pricing(), - updated_patch=updateDeployment( - model_info=ModelInfo(id="dep-pricing-0", input_cost_per_token=None) - ), + updated_patch=updateDeployment(model_info=ModelInfo(id="dep-pricing-0", input_cost_per_token=None)), ) params = json.loads(result["litellm_params"]) @@ -3281,9 +3136,7 @@ class TestUpdateDBModelClearPricing: result = update_db_model( db_model=db_model, - updated_patch=updateDeployment( - litellm_params=updateLiteLLMParams(cache_read_input_token_cost=None) - ), + updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(cache_read_input_token_cost=None)), ) params = json.loads(result["litellm_params"]) @@ -3315,9 +3168,7 @@ class TestUpdateDBModelClearPricing: result = update_db_model( db_model=db_model, - updated_patch=updateDeployment( - litellm_params=updateLiteLLMParams(cache_creation_input_token_cost=None) - ), + updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(cache_creation_input_token_cost=None)), ) params = json.loads(result["litellm_params"]) @@ -3351,9 +3202,7 @@ class TestUpdateDBModelClearPricing: result = update_db_model( db_model=db_model, - updated_patch=updateDeployment( - litellm_params=updateLiteLLMParams(cache_read_input_token_cost=None) - ), + updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(cache_read_input_token_cost=None)), ) params = json.loads(result["litellm_params"]) @@ -3419,9 +3268,7 @@ class TestPatchModelBlockedAuthGate: existing_row.model_dump_json.return_value = "{}" mock_prisma = MagicMock() - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=existing_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), @@ -3462,12 +3309,8 @@ class TestPatchModelBlockedAuthGate: updated_row.model_dump_json.return_value = "{}" mock_prisma = MagicMock() - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=existing_row - ) - mock_prisma.db.litellm_proxymodeltable.update = AsyncMock( - return_value=updated_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), @@ -3480,9 +3323,7 @@ class TestPatchModelBlockedAuthGate: ), patch( "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", - new=AsyncMock( - return_value=ReconcileOutcome(still_desired=None, live_after=None) - ), + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): result = await patch_model( @@ -3517,9 +3358,7 @@ class TestPatchModelRowDeletedBeforeWrite: existing_row.model_dump_json.return_value = "{}" mock_prisma = MagicMock() - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=existing_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=None) with ( @@ -3533,9 +3372,7 @@ class TestPatchModelRowDeletedBeforeWrite: ), patch( # test-quality-ok: stubs the cache write so the test observes only the DB result handling "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", - new=AsyncMock( - return_value=ReconcileOutcome(still_desired=None, live_after=None) - ), + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): with pytest.raises(ProxyException) as exc_info: @@ -3619,9 +3456,7 @@ class TestWriteSurfacesReloadDrop: ) with pytest.raises(ProxyException, match="m-gone"): - raise_if_reload_degraded_serving( - before=frozenset(), written_models=[("m-gone", None)], action="update" - ) + raise_if_reload_degraded_serving(before=frozenset(), written_models=[("m-gone", None)], action="update") with pytest.raises(ProxyException, match="m-collateral"): raise_if_reload_degraded_serving( @@ -3736,10 +3571,7 @@ class TestConcurrentModelWritesDoNotEvictEachOther: config = ProxyConfig() await asyncio.gather( - *[ - config.add_deployment(prisma_client=MagicMock(), proxy_logging_obj=MagicMock()) - for _ in range(5) - ] + *[config.add_deployment(prisma_client=MagicMock(), proxy_logging_obj=MagicMock()) for _ in range(5)] ) assert observed_max == 1 @@ -3963,9 +3795,7 @@ class TestDeleteEvictionsHoldTheReconcileLock: ) async def call() -> None: - await delete_team_models( - team_ids=["team-1"], prisma_client=prisma, llm_router=router - ) + await delete_team_models(team_ids=["team-1"], prisma_client=prisma, llm_router=router) await self._assert_evicts_under_lock(monkeypatch, call, model_id) router.delete_deployment.assert_called_once_with(id=model_id) @@ -4570,7 +4400,6 @@ class TestStrategyRouterWriteValidation: heuristic_v2 has its own slot, while custom tier definitions and custom prompts count into one shared customization slot. Renaming built-in tiers through tier_labels claims nothing at all.""" - from fastapi import HTTPException from litellm.proxy.management_endpoints.model_management_endpoints import ( AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY, @@ -4792,7 +4621,6 @@ class TestStrategyRouterWriteValidation: @pytest.mark.asyncio async def test_patch_model_refuses_switching_another_router_to_heuristic_v2(self) -> None: """patch_model relays HTTPException as-is, so the license refusal reaches the client as a plain 403.""" - from fastapi import HTTPException from litellm.proxy.management_endpoints.model_management_endpoints import ( patch_model, @@ -4950,13 +4778,17 @@ class TestAutoRouterClassifierDefaultPrompt: from litellm.router_strategy.complexity_router import ClassificationRubric, classification_system_prompt for preset in ClassificationRubric: - response = await get_auto_router_classifier_default_prompt(context_window_size=5, classification_rubric=preset) + response = await get_auto_router_classifier_default_prompt( + context_window_size=5, classification_rubric=preset + ) assert response.system_prompt == classification_system_prompt(5, classification_rubric=preset) agentic = await get_auto_router_classifier_default_prompt( context_window_size=5, classification_rubric=ClassificationRubric.AGENTIC ) - chat = await get_auto_router_classifier_default_prompt(context_window_size=5, classification_rubric=ClassificationRubric.CHAT) + chat = await get_auto_router_classifier_default_prompt( + context_window_size=5, classification_rubric=ClassificationRubric.CHAT + ) unset = await get_auto_router_classifier_default_prompt(context_window_size=5) assert "Calibration on engineering tasks" in agentic.system_prompt assert "Calibration on engineering tasks" not in chat.system_prompt @@ -5251,6 +5083,759 @@ class TestAutoRouterClassifierDefaultPrompt: assert response.system_prompt == classification_system_prompt(5) +class TestAddModelToDbBlocked: + """`_add_model_to_db` must thread `blocked` into the initial insert, so the wizard can + create a discovered-but-unchecked model already paused instead of active-then-patched.""" + + @staticmethod + def _deployment(blocked): + from litellm.types.router import ModelInfo + + return Deployment( + model_name="anthropic/claude-discovered", + litellm_params=LiteLLM_Params(model="anthropic/claude-discovered"), + model_info=ModelInfo(id="dep-blocked-create-0"), + blocked=blocked, + ) + + @pytest.mark.asyncio + async def test_add_model_to_db_writes_blocked_true(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _add_model_to_db, + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=MagicMock()) + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.master_key", "sk-test-master" + ): # test-quality-ok: the proxy wiring under test is what this patches + await _add_model_to_db( + model_params=self._deployment(True), user_api_key_dict=admin, prisma_client=mock_prisma + ) + + _, kwargs = mock_prisma.db.litellm_proxymodeltable.create.call_args + assert kwargs["data"]["blocked"] is True + + @pytest.mark.asyncio + async def test_add_model_to_db_writes_blocked_false(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _add_model_to_db, + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=MagicMock()) + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.master_key", "sk-test-master" + ): # test-quality-ok: the proxy wiring under test is what this patches + await _add_model_to_db( + model_params=self._deployment(False), user_api_key_dict=admin, prisma_client=mock_prisma + ) + + _, kwargs = mock_prisma.db.litellm_proxymodeltable.create.call_args + assert kwargs["data"]["blocked"] is False + + @pytest.mark.asyncio + async def test_add_model_to_db_omits_blocked_when_not_set(self): + """None means "don't set it" -- the Prisma column defaults to False -- not "explicitly + unblocked", so the key must be absent from the write entirely.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _add_model_to_db, + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=MagicMock()) + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.master_key", "sk-test-master" + ): # test-quality-ok: the proxy wiring under test is what this patches + await _add_model_to_db( + model_params=self._deployment(None), user_api_key_dict=admin, prisma_client=mock_prisma + ) + + _, kwargs = mock_prisma.db.litellm_proxymodeltable.create.call_args + assert "blocked" not in kwargs["data"] + + +class TestAddNewModelBlockedAuthGate: + """Same proxy-admin-only rule patch_model applies to `blocked` must hold at create time + too: a team admin authorized for a team-scoped model must not be able to create it already + paused (or explicitly unpaused) out from under the proxy admin.""" + + @pytest.mark.asyncio + async def test_non_admin_cannot_set_blocked_on_create(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + add_new_model, + ) + + non_admin = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER) + mock_prisma = MagicMock() + + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.store_model_in_db", True + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.premium_user", True + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model( + model_params=Deployment( + model_name="my-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"), + model_info={"id": "blocked-gate-create-0"}, + blocked=True, + ), + user_api_key_dict=non_admin, + ) + assert "proxy admin" in str(exc_info.value.message).lower() + mock_prisma.db.litellm_proxymodeltable.create.assert_not_called() + + @pytest.mark.asyncio + async def test_proxy_admin_can_create_a_blocked_model(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + add_new_model, + ) + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + mock_prisma = MagicMock() + created_row = MagicMock() + created_row.model_id = "blocked-gate-create-1" + created_row.model_dump_json.return_value = "{}" + mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=created_row) + + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.store_model_in_db", True + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.premium_user", True + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.master_key", "sk-test-master" + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.llm_router", + MagicMock(**{"get_model_ids.return_value": ["blocked-gate-create-1"]}), + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.proxy_config", + MagicMock(add_deployment=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None))), + ), + ): + result = await add_new_model( + model_params=Deployment( + model_name="my-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"), + model_info={"id": "blocked-gate-create-1"}, + blocked=True, + ), + user_api_key_dict=admin, + ) + assert result is created_row + _, kwargs = mock_prisma.db.litellm_proxymodeltable.create.call_args + assert kwargs["data"]["blocked"] is True + + +class TestNonAdminCannotPersistWifFieldsOnModel: + """A server-owned Anthropic WIF field (destination, source, or secret reference) chooses + which server-side secret is read and where it is sent. A team admin who is otherwise + authorized for a team-scoped model must not be able to set one via /model/new, + /model/update, or PATCH /model/{id}/update; a proxy admin still can.""" + + @pytest.mark.asyncio + async def test_patch_model_non_admin_cannot_set_wif_field(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + + non_admin = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER) + existing_row = MagicMock() + existing_row.litellm_params = {"model": "anthropic/claude-sonnet-4"} + existing_row.model_dump.return_value = { + "model_name": "claude", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": "m1"}, + } + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", + mock_prisma, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.llm_router", + MagicMock(**{"get_model_ids.return_value": ["m1"]}), + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.store_model_in_db", + True, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.premium_user", + True, + ), + ): + with pytest.raises( + Exception, match="Only proxy admins can modify a deployment configured for workload identity" + ) as exc_info: + await patch_model( + model_id="m1", + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams( + anthropic_keycloak_token_url="https://attacker.example/token", + ) + ), + user_api_key_dict=non_admin, + ) + err = exc_info.value + assert getattr(err, "param", "") == "anthropic_keycloak_token_url" + mock_prisma.db.litellm_proxymodeltable.update.assert_not_called() + + @pytest.mark.asyncio + async def test_patch_model_non_admin_cannot_set_openai_wif_field(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + + non_admin = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER) + existing_row = MagicMock() + existing_row.litellm_params = {"model": "openai/gpt-4o-mini"} + existing_row.model_dump.return_value = { + "model_name": "gpt", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": "m1"}, + } + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", + mock_prisma, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.llm_router", + MagicMock(**{"get_model_ids.return_value": ["m1"]}), + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.store_model_in_db", + True, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.premium_user", + True, + ), + ): + with pytest.raises( + Exception, match="Only proxy admins can modify a deployment configured for workload identity" + ) as exc_info: + await patch_model( + model_id="m1", + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams( + openai_identity_token_file="/var/run/secrets/tokens/attacker", + ) + ), + user_api_key_dict=non_admin, + ) + assert getattr(exc_info.value, "param", "") == "openai_identity_token_file" + mock_prisma.db.litellm_proxymodeltable.update.assert_not_called() + + @pytest.mark.asyncio + async def test_patch_model_admin_can_set_wif_field(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + existing_row = MagicMock() + existing_row.litellm_params = {"model": "anthropic/claude-sonnet-4"} + existing_row.model_dump.return_value = { + "model_name": "claude", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": "m1"}, + } + existing_row.model_dump_json.return_value = "{}" + updated_row = MagicMock() + updated_row.model_dump_json.return_value = "{}" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) + + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", + mock_prisma, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.llm_router", + MagicMock(**{"get_model_ids.return_value": ["m1"]}), + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.store_model_in_db", + True, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.premium_user", + True, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value: value, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + result = await patch_model( + model_id="m1", + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams( + anthropic_keycloak_token_url="https://keycloak.internal/token", + ) + ), + user_api_key_dict=admin, + ) + assert result is updated_row + mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once() + + @pytest.mark.asyncio + async def test_add_new_model_non_admin_cannot_set_wif_field(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + add_new_model, + ) + + non_admin = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER) + mock_prisma = MagicMock() + + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", + mock_prisma, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.store_model_in_db", + True, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.premium_user", + True, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model( + model_params=Deployment( + model_name="my-model", + litellm_params=LiteLLM_Params( + model="anthropic/claude-sonnet-4", + anthropic_keycloak_client_secret_ref="os.environ/LITELLM_MASTER_KEY", + ), + model_info={"id": "wif-gate-create-0"}, + ), + user_api_key_dict=non_admin, + ) + assert "proxy admin" in str(exc_info.value.message).lower() + assert exc_info.value.param == "anthropic_keycloak_client_secret_ref" + mock_prisma.db.litellm_proxymodeltable.create.assert_not_called() + + @pytest.mark.asyncio + async def test_add_new_model_admin_can_set_wif_field(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + add_new_model, + ) + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + mock_prisma = MagicMock() + created_row = MagicMock() + created_row.model_id = "wif-gate-create-1" + created_row.model_dump_json.return_value = "{}" + mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=created_row) + + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", + mock_prisma, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.store_model_in_db", + True, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.premium_user", + True, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.master_key", + "sk-test-master", + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.llm_router", + MagicMock(**{"get_model_ids.return_value": ["wif-gate-create-1"]}), + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.proxy_config", + MagicMock(add_deployment=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None))), + ), + ): + result = await add_new_model( + model_params=Deployment( + model_name="my-model", + litellm_params=LiteLLM_Params( + model="anthropic/claude-sonnet-4", + anthropic_keycloak_client_secret_ref="os.environ/ANTHROPIC_WIF_CLIENT_SECRET", + ), + model_info={"id": "wif-gate-create-1"}, + ), + user_api_key_dict=admin, + ) + assert result is created_row + + @pytest.mark.asyncio + async def test_update_model_non_admin_cannot_set_wif_field(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_model, + ) + + model_id = "wif-gate-update-0" + existing_row = MagicMock() + existing_row.litellm_params = {"model": "anthropic/claude-sonnet-4"} + existing_row.model_dump.return_value = { + "model_name": "claude", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": model_id}, + } + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + non_admin = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER) + + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", + mock_prisma, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.llm_router", + MagicMock(**{"get_model_ids.return_value": [model_id]}), + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.store_model_in_db", + True, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.premium_user", + True, + ), + ): + with pytest.raises( + Exception, match="Only proxy admins can modify a deployment configured for workload identity" + ) as exc_info: + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams( + anthropic_keycloak_client_secret_ref="os.environ/LITELLM_MASTER_KEY", + ), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=non_admin, + ) + assert getattr(exc_info.value, "param", "") == "anthropic_keycloak_client_secret_ref" + mock_prisma.db.litellm_proxymodeltable.update.assert_not_called() + + @pytest.mark.asyncio + async def test_update_model_admin_can_set_wif_field(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_model, + ) + + model_id = "wif-gate-update-1" + existing_row = MagicMock() + existing_row.litellm_params = {"model": "anthropic/claude-sonnet-4"} + existing_row.model_dump.return_value = { + "model_name": "claude", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": model_id}, + } + existing_row.model_dump_json.return_value = "{}" + updated_row = MagicMock() + updated_row.model_dump_json.return_value = "{}" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.prisma_client", + mock_prisma, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.llm_router", + MagicMock(**{"get_model_ids.return_value": [model_id]}), + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.store_model_in_db", + True, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.premium_user", + True, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value: value, + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams( + anthropic_keycloak_client_secret_ref="os.environ/ANTHROPIC_WIF_CLIENT_SECRET", + ), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=admin, + ) + mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once() + written_litellm_params = mock_prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"][ + "litellm_params" + ] + assert "anthropic_keycloak_client_secret_ref" in written_litellm_params + assert "os.environ/ANTHROPIC_WIF_CLIENT_SECRET" in written_litellm_params + + +class TestOneCredentialFeedsManyModelsNoWifCopy: + """Regression: one named WIF credential feeds multiple model rows, and no WIF field is + ever copied onto a model row -- litellm_params carries only `model` and + `litellm_credential_name`, the same shape the wizard's per-row /model/new call produces.""" + + @pytest.mark.asyncio + async def test_two_discovered_models_share_the_credential_reference_only(self): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _add_model_to_db, + ) + from litellm.types.router import ModelInfo + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=MagicMock()) + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with ( + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.proxy_server.master_key", "sk-test-master" + ), + patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.proxy.common_utils.encrypt_decrypt_utils._get_salt_key", return_value="sk-test-master" + ), + ): + for i, discovered_id in enumerate(["claude-a", "claude-b"]): + model_params = Deployment( + model_name=discovered_id, + litellm_params=LiteLLM_Params( + model=f"anthropic/{discovered_id}", litellm_credential_name="anthropic-wif" + ), + model_info=ModelInfo(id=f"dep-shared-{i}"), + blocked=False, + ) + await _add_model_to_db(model_params=model_params, user_api_key_dict=admin, prisma_client=mock_prisma) + + assert mock_prisma.db.litellm_proxymodeltable.create.await_count == 2 + for call in mock_prisma.db.litellm_proxymodeltable.create.await_args_list: + written_litellm_params = json.loads(call.kwargs["data"]["litellm_params"]) + decrypted_credential_name = decrypt_value_helper( + value=written_litellm_params["litellm_credential_name"], key="litellm_credential_name" + ) + assert decrypted_credential_name == "anthropic-wif" + assert "anthropic_federation_rule_id" not in written_litellm_params + assert "anthropic_identity_token" not in written_litellm_params + assert call.kwargs["data"]["blocked"] is False + + +class TestWifBoundaryReadsTheResultingDeployment: + """The proxy-admin rule has to be evaluated against the deployment the write PRODUCES. + Reading only the submitted payload let a team admin keep an existing federated deployment + and change it anyway, because the fields they sent named nothing federated.""" + + @staticmethod + def _existing_wif_row(): + row = MagicMock() + row.litellm_params = { + "model": "anthropic/claude-sonnet-4", + "anthropic_federation_rule_id": "fdrl_admin", + "anthropic_organization_id": "org-admin", + } + row.model_dump.return_value = { + "model_name": "claude", + "litellm_params": row.litellm_params, + "model_info": {"id": "m1"}, + } + return row + + @pytest.mark.asyncio + async def test_non_admin_cannot_retarget_an_existing_wif_deployment_via_api_base(self): + """api_base is not a federation field, so the payload-only check saw nothing to refuse, + and the merged deployment then sent its assertion and minted token to the new host.""" + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + + non_admin = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER) + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=self._existing_wif_row()) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy wiring under test + patch( # test-quality-ok: proxy wiring under test + "litellm.proxy.proxy_server.llm_router", MagicMock(**{"get_model_ids.return_value": ["m1"]}) + ), + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy wiring under test + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy wiring under test + ): + with pytest.raises( + Exception, match="Only proxy admins can modify a deployment configured for workload identity" + ): + await patch_model( + model_id="m1", + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams(api_base="https://gateway.internal") + ), + user_api_key_dict=non_admin, + ) + + @pytest.mark.asyncio + async def test_non_admin_cannot_detach_a_federated_credential_to_escape_the_gate(self): + """Clearing the credential name must not be the way out. A deployment federated through a + named credential carries no federation field of its own, so a patch that sends + litellm_credential_name: null alongside an api_base of the caller's choosing would leave + nothing federated to find, and the write would be allowed.""" + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + + non_admin = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER) + federated_row = MagicMock() + federated_row.litellm_params = { + "model": "anthropic/claude-sonnet-4", + "litellm_credential_name": "admin-wif", + } + federated_row.model_dump.return_value = { + "model_name": "claude", + "litellm_params": federated_row.litellm_params, + "model_info": {"id": "m1"}, + } + + admin_credential_row = { + "credential_name": "admin-wif", + "credential_values": { + "anthropic_federation_rule_id": "fdrl_admin", + "anthropic_organization_id": "org-admin", + }, + "credential_info": {"custom_llm_provider": "anthropic"}, + } + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=federated_row) + mock_prisma.db.litellm_credentialstable.find_unique = AsyncMock(return_value=admin_credential_row) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy wiring under test + patch( # test-quality-ok: proxy wiring under test + "litellm.proxy.proxy_server.llm_router", MagicMock(**{"get_model_ids.return_value": ["m1"]}) + ), + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy wiring under test + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy wiring under test + ): + with pytest.raises( + Exception, match="Only proxy admins can modify a deployment configured for workload identity" + ): + await patch_model( + model_id="m1", + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams( + litellm_credential_name=None, api_base="https://gateway.internal" + ) + ), + user_api_key_dict=non_admin, + ) + + @pytest.mark.asyncio + async def test_non_admin_cannot_attach_a_federated_credential_by_name(self): + """litellm_credential_name names no federation field itself, but request-time hydration + imports whatever the credential holds, so the resulting deployment federates.""" + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + + non_admin = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER) + plain_row = MagicMock() + plain_row.litellm_params = {"model": "anthropic/claude-sonnet-4"} + plain_row.model_dump.return_value = { + "model_name": "claude", + "litellm_params": plain_row.litellm_params, + "model_info": {"id": "m1"}, + } + # The credential is served from the row rather than this pod's memory, which is both the + # multi-pod case and the one the gate must not miss. + admin_credential_row = { + "credential_name": "admin-wif", + "credential_values": { + "anthropic_federation_rule_id": "fdrl_admin", + "anthropic_organization_id": "org-admin", + }, + "credential_info": {"custom_llm_provider": "anthropic"}, + } + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=plain_row) + mock_prisma.db.litellm_credentialstable.find_unique = AsyncMock(return_value=admin_credential_row) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy wiring under test + patch( # test-quality-ok: proxy wiring under test + "litellm.proxy.proxy_server.llm_router", MagicMock(**{"get_model_ids.return_value": ["m1"]}) + ), + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy wiring under test + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy wiring under test + ): + with pytest.raises( + Exception, match="Only proxy admins can modify a deployment configured for workload identity" + ): + await patch_model( + model_id="m1", + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name="admin-wif") + ), + user_api_key_dict=non_admin, + ) + + class TestEnforceRpmTpmOnModelAdd: def test_passes_when_disabled_even_without_limits(self): assert ( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index acb45038df0..3823a14b7c7 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -82,36 +82,28 @@ class TestBaseOpenAIPassThroughHandler: # Test joining base URL with no path and a path base_url = httpx.URL("https://api.example.com") path = "/v1/chat/completions" - result = _join_url_paths( - base_url, path, litellm.LlmProviders.OPENAI.value - ) + result = _join_url_paths(base_url, path, litellm.LlmProviders.OPENAI.value) print(f"Base URL with no path: '{base_url}' + '{path}' → '{result}'") assert str(result) == "https://api.example.com/v1/chat/completions" # Test joining base URL with path and another path base_url = httpx.URL("https://api.example.com/v1") path = "/chat/completions" - result = _join_url_paths( - base_url, path, litellm.LlmProviders.OPENAI.value - ) + result = _join_url_paths(base_url, path, litellm.LlmProviders.OPENAI.value) print(f"Base URL with path: '{base_url}' + '{path}' → '{result}'") assert str(result) == "https://api.example.com/v1/chat/completions" # Test with path not starting with slash base_url = httpx.URL("https://api.example.com/v1") path = "chat/completions" - result = _join_url_paths( - base_url, path, litellm.LlmProviders.OPENAI.value - ) + result = _join_url_paths(base_url, path, litellm.LlmProviders.OPENAI.value) print(f"Path without leading slash: '{base_url}' + '{path}' → '{result}'") assert str(result) == "https://api.example.com/v1/chat/completions" # Test with base URL having trailing slash base_url = httpx.URL("https://api.example.com/v1/") path = "/chat/completions" - result = _join_url_paths( - base_url, path, litellm.LlmProviders.OPENAI.value - ) + result = _join_url_paths(base_url, path, litellm.LlmProviders.OPENAI.value) print(f"Base URL with trailing slash: '{base_url}' + '{path}' → '{result}'") assert str(result) == "https://api.example.com/v1/chat/completions" @@ -130,17 +122,13 @@ class TestBaseOpenAIPassThroughHandler: headers = {"authorization": "Bearer test_key"} # Test with assistants API request - result = BaseOpenAIPassThroughHandler._append_openai_beta_header( - headers, assistants_request - ) + result = BaseOpenAIPassThroughHandler._append_openai_beta_header(headers, assistants_request) print(f"Assistants API request: Added header: {result}") assert result["OpenAI-Beta"] == "assistants=v2" # Test with non-assistants API request headers = {"authorization": "Bearer test_key"} - result = BaseOpenAIPassThroughHandler._append_openai_beta_header( - headers, non_assistants_request - ) + result = BaseOpenAIPassThroughHandler._append_openai_beta_header(headers, non_assistants_request) print(f"Non-assistants API request: Headers: {result}") assert "OpenAI-Beta" not in result @@ -150,9 +138,7 @@ class TestBaseOpenAIPassThroughHandler: assistant_request.url.path = "/v1/assistants/asst_123456" headers = {"authorization": "Bearer test_key"} - result = BaseOpenAIPassThroughHandler._append_openai_beta_header( - headers, assistant_request - ) + result = BaseOpenAIPassThroughHandler._append_openai_beta_header(headers, assistant_request) print(f"Assistant API request: Added header: {result}") assert result["OpenAI-Beta"] == "assistants=v2" @@ -173,9 +159,7 @@ class TestBaseOpenAIPassThroughHandler: "test-header": "value", }, ): - result = BaseOpenAIPassThroughHandler._assemble_headers( - api_key, mock_request - ) + result = BaseOpenAIPassThroughHandler._assemble_headers(api_key, mock_request) print(f"Assembled headers: {result}") assert result["authorization"] == "Bearer test_api_key" assert result["api-key"] == "test_api_key" @@ -215,9 +199,7 @@ class TestBaseOpenAIPassThroughHandler: # Verify create_pass_through_route was called with correct parameters call_args = mock_create_pass_through.call_args[1] - print( - f"create_pass_through_route called with endpoint: {call_args['endpoint']}" - ) + print(f"create_pass_through_route called with endpoint: {call_args['endpoint']}") print(f"create_pass_through_route called with target: {call_args['target']}") assert call_args["endpoint"] == "/chat/completions" assert call_args["target"] == "https://api.openai.com/v1/chat/completions" @@ -269,9 +251,7 @@ class TestVertexAIPassThroughHandler: # Mock request mock_request = Mock() - mock_request.state = ( - None # Prevent Mock from returning a truthy _cached_headers - ) + mock_request.state = None # Prevent Mock from returning a truthy _cached_headers mock_request.method = "POST" mock_request.headers = { "Authorization": "Bearer test-creds", @@ -289,9 +269,7 @@ class TestVertexAIPassThroughHandler: test_token = vertex_credentials with ( - mock.patch( - "litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth" - ) as mock_load_auth, + mock.patch("litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth") as mock_load_auth, mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, @@ -374,9 +352,7 @@ class TestVertexAIPassThroughHandler: # Mock request mock_request = Mock() - mock_request.state = ( - None # Prevent Mock from returning a truthy _cached_headers - ) + mock_request.state = None # Prevent Mock from returning a truthy _cached_headers mock_request.method = "POST" mock_request.headers = { "Authorization": "Bearer test-creds", @@ -394,9 +370,7 @@ class TestVertexAIPassThroughHandler: test_token = vertex_credentials with ( - mock.patch( - "litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth" - ) as mock_load_auth, + mock.patch("litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth") as mock_load_auth, mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, @@ -457,9 +431,7 @@ class TestVertexAIPassThroughHandler: ], ) @pytest.mark.asyncio - async def test_vertex_passthrough_with_default_credentials( - self, monkeypatch, initial_endpoint - ): + async def test_vertex_passthrough_with_default_credentials(self, monkeypatch, initial_endpoint): """ Test that when no passthrough credentials are set, default credentials are used in the request """ @@ -498,9 +470,7 @@ class TestVertexAIPassThroughHandler: mock_response = Response() with ( - mock.patch( - "litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth" - ) as mock_load_auth, + mock.patch("litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth") as mock_load_auth, mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, @@ -642,17 +612,13 @@ class TestVertexAIPassThroughHandler: mock_request.method = "POST" mock_response = Mock() - with patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth" - ) as mock_auth: + with patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth") as mock_auth: mock_auth.return_value = {"api_key": "test-key-123"} with patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_pass_through: - mock_pass_through.return_value = AsyncMock( - return_value={"status": "success"} - ) + mock_pass_through.return_value = AsyncMock(return_value={"status": "success"}) with pytest.raises(HTTPException) as exc_info: await vertex_proxy_route( @@ -712,7 +678,9 @@ class TestVertexAIPassThroughHandler: mock_logging_obj.model_call_details = {} # Test URL with multimodal embedding model - url_route = "/v1/projects/test-project/locations/us-central1/publishers/google/models/multimodalembedding@001:predict" + url_route = ( + "/v1/projects/test-project/locations/us-central1/publishers/google/models/multimodalembedding@001:predict" + ) start_time = datetime.datetime.now() end_time = datetime.datetime.now() @@ -730,19 +698,13 @@ class TestVertexAIPassThroughHandler: mock_embedding_response = EmbeddingResponse( object="list", data=[ - Embedding( - embedding=[0.1, 0.2, 0.3, 0.4, 0.5], index=0, object="embedding" - ), - Embedding( - embedding=[0.6, 0.7, 0.8, 0.9, 1.0], index=1, object="embedding" - ), + Embedding(embedding=[0.1, 0.2, 0.3, 0.4, 0.5], index=0, object="embedding"), + Embedding(embedding=[0.6, 0.7, 0.8, 0.9, 1.0], index=1, object="embedding"), ], model="multimodalembedding@001", usage=Usage(prompt_tokens=0, total_tokens=0, completion_tokens=0), ) - mock_config_instance.transform_embedding_response.return_value = ( - mock_embedding_response - ) + mock_config_instance.transform_embedding_response.return_value = mock_embedding_response # Call the handler result = VertexPassthroughLoggingHandler.vertex_passthrough_handler( @@ -778,26 +740,12 @@ class TestVertexAIPassThroughHandler: ) # Test case 1: Response with textEmbedding should be detected as multimodal - response_with_text_embedding = { - "predictions": [{"textEmbedding": [0.1, 0.2, 0.3]}] - } - assert ( - VertexPassthroughLoggingHandler._is_multimodal_embedding_response( - response_with_text_embedding - ) - is True - ) + response_with_text_embedding = {"predictions": [{"textEmbedding": [0.1, 0.2, 0.3]}]} + assert VertexPassthroughLoggingHandler._is_multimodal_embedding_response(response_with_text_embedding) is True # Test case 2: Response with imageEmbedding should be detected as multimodal - response_with_image_embedding = { - "predictions": [{"imageEmbedding": [0.4, 0.5, 0.6]}] - } - assert ( - VertexPassthroughLoggingHandler._is_multimodal_embedding_response( - response_with_image_embedding - ) - is True - ) + response_with_image_embedding = {"predictions": [{"imageEmbedding": [0.4, 0.5, 0.6]}]} + assert VertexPassthroughLoggingHandler._is_multimodal_embedding_response(response_with_image_embedding) is True # Test case 3: Response with videoEmbeddings should be detected as multimodal response_with_video_embeddings = { @@ -813,43 +761,19 @@ class TestVertexAIPassThroughHandler: } ] } - assert ( - VertexPassthroughLoggingHandler._is_multimodal_embedding_response( - response_with_video_embeddings - ) - is True - ) + assert VertexPassthroughLoggingHandler._is_multimodal_embedding_response(response_with_video_embeddings) is True # Test case 4: Regular text embedding response should NOT be detected as multimodal - regular_embedding_response = { - "predictions": [{"embeddings": {"values": [0.1, 0.2, 0.3]}}] - } - assert ( - VertexPassthroughLoggingHandler._is_multimodal_embedding_response( - regular_embedding_response - ) - is False - ) + regular_embedding_response = {"predictions": [{"embeddings": {"values": [0.1, 0.2, 0.3]}}]} + assert VertexPassthroughLoggingHandler._is_multimodal_embedding_response(regular_embedding_response) is False # Test case 5: Non-embedding response should NOT be detected as multimodal - non_embedding_response = { - "candidates": [{"content": {"parts": [{"text": "Hello world"}]}}] - } - assert ( - VertexPassthroughLoggingHandler._is_multimodal_embedding_response( - non_embedding_response - ) - is False - ) + non_embedding_response = {"candidates": [{"content": {"parts": [{"text": "Hello world"}]}}]} + assert VertexPassthroughLoggingHandler._is_multimodal_embedding_response(non_embedding_response) is False # Test case 6: Empty response should NOT be detected as multimodal empty_response = {} - assert ( - VertexPassthroughLoggingHandler._is_multimodal_embedding_response( - empty_response - ) - is False - ) + assert VertexPassthroughLoggingHandler._is_multimodal_embedding_response(empty_response) is False def test_vertex_passthrough_handler_predict_cost_tracking(self): """ @@ -889,7 +813,9 @@ class TestVertexAIPassThroughHandler: mock_logging_obj.model_call_details = {} # Test URL with /predict endpoint - url_route = "/v1/projects/test-project/locations/us-central1/publishers/google/models/textembedding-gecko@001:predict" + url_route = ( + "/v1/projects/test-project/locations/us-central1/publishers/google/models/textembedding-gecko@001:predict" + ) start_time = datetime.datetime.now() end_time = datetime.datetime.now() @@ -959,7 +885,9 @@ class TestVertexAIPassThroughHandler: mock_logging_obj.litellm_call_id = "test-call-id-embed" mock_logging_obj.model_call_details = {} - url_route = "/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-embedding-001:embedContent" + url_route = ( + "/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-embedding-001:embedContent" + ) start_time = datetime.datetime.now() end_time = datetime.datetime.now() @@ -978,9 +906,7 @@ class TestVertexAIPassThroughHandler: ) assert result is not None - assert ( - result["result"] is not None - ), "result must not be None — logging callbacks need a non-null response" + assert result["result"] is not None, "result must not be None — logging callbacks need a non-null response" assert "kwargs" in result assert result["kwargs"].get("response_cost") == 0.0002 assert result["kwargs"].get("model") == "gemini-embedding-001" @@ -1037,9 +963,7 @@ class TestVertexAIPassThroughHandler: ) assert result is not None - assert ( - result["result"] is not None - ), "result must not be None for batchEmbedContents" + assert result["result"] is not None, "result must not be None for batchEmbedContents" assert result["kwargs"].get("response_cost") == 0.0003 assert result["kwargs"].get("model") == "gemini-embedding-001" assert result["kwargs"].get("custom_llm_provider") == "vertex_ai" @@ -1096,9 +1020,9 @@ class TestVertexAIPassThroughHandler: assert result is not None assert result["result"] is not None - assert ( - result["kwargs"].get("custom_llm_provider") == "gemini" - ), "Google AI Studio embedContent URLs must set custom_llm_provider=gemini, not vertex_ai" + assert result["kwargs"].get("custom_llm_provider") == "gemini", ( + "Google AI Studio embedContent URLs must set custom_llm_provider=gemini, not vertex_ai" + ) assert result["kwargs"].get("model") == "gemini-embedding-2-preview" mock_completion_cost.assert_called_once() @@ -1245,13 +1169,13 @@ class TestVertexAIDiscoveryPassThroughHandler: pass_through_router, ) - endpoint = f"v1/projects/{vertex_project}/locations/{vertex_location}/dataStores/default/servingConfigs/default:search" + endpoint = ( + f"v1/projects/{vertex_project}/locations/{vertex_location}/dataStores/default/servingConfigs/default:search" + ) # Mock request mock_request = Mock() - mock_request.state = ( - None # Prevent Mock from returning a truthy _cached_headers - ) + mock_request.state = None # Prevent Mock from returning a truthy _cached_headers mock_request.method = "POST" mock_request.headers = { "Authorization": "Bearer test-key", @@ -1269,9 +1193,7 @@ class TestVertexAIDiscoveryPassThroughHandler: test_token = "test-auth-token" with ( - mock.patch( - "litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth" - ) as mock_load_auth, + mock.patch("litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth") as mock_load_auth, mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, @@ -1319,10 +1241,7 @@ class TestVertexAIDiscoveryPassThroughHandler: assert test_project in call_args[1]["target"] assert test_location in call_args[1]["target"] assert "Authorization" in call_args[1]["custom_headers"] - assert ( - call_args[1]["custom_headers"]["Authorization"] - == f"Bearer {test_token}" - ) + assert call_args[1]["custom_headers"]["Authorization"] == f"Bearer {test_token}" @pytest.mark.asyncio async def test_vertex_discovery_proxy_route_api_key_auth(self): @@ -1337,17 +1256,13 @@ class TestVertexAIDiscoveryPassThroughHandler: mock_request.method = "POST" mock_response = Mock() - with patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth" - ) as mock_auth: + with patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth") as mock_auth: mock_auth.return_value = {"api_key": "test-key-123"} with patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_pass_through: - mock_pass_through.return_value = AsyncMock( - return_value={"status": "success"} - ) + mock_pass_through.return_value = AsyncMock(return_value={"status": "success"}) with pytest.raises(HTTPException) as exc_info: await vertex_discovery_proxy_route( @@ -1443,9 +1358,7 @@ async def test_mistral_passthrough_accepts_multipart_without_json_parsing(): assert response == {"ok": True} assert captured_kwargs["is_streaming_request"] is False - assert captured_kwargs["custom_headers"] == { - "Authorization": "Bearer mistral-test-key" - } + assert captured_kwargs["custom_headers"] == {"Authorization": "Bearer mistral-test-key"} class TestBedrockLLMProxyRoute: @@ -1457,9 +1370,7 @@ class TestBedrockLLMProxyRoute: mock_user_api_key_dict = Mock() mock_request_body = {"messages": [{"role": "user", "content": "test"}]} mock_processor = Mock() - mock_processor.base_passthrough_process_llm_request = AsyncMock( - return_value="success" - ) + mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success") with ( patch( @@ -1471,9 +1382,10 @@ class TestBedrockLLMProxyRoute: return_value=mock_processor, ), ): - # Test application-inference-profile endpoint - endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/r742sbn2zckd/converse" + endpoint = ( + "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/r742sbn2zckd/converse" + ) result = await bedrock_llm_proxy_route( endpoint=endpoint, @@ -1483,9 +1395,7 @@ class TestBedrockLLMProxyRoute: ) mock_processor.base_passthrough_process_llm_request.assert_called_once() - call_kwargs = ( - mock_processor.base_passthrough_process_llm_request.call_args.kwargs - ) + call_kwargs = mock_processor.base_passthrough_process_llm_request.call_args.kwargs # For application-inference-profile, model should be "arn:aws:bedrock:us-east-1:026090525607:application-inference-profile/r742sbn2zckd" assert ( @@ -1502,9 +1412,7 @@ class TestBedrockLLMProxyRoute: mock_user_api_key_dict = Mock() mock_request_body = {"messages": [{"role": "user", "content": "test"}]} mock_processor = Mock() - mock_processor.base_passthrough_process_llm_request = AsyncMock( - return_value="success" - ) + mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success") with ( patch( @@ -1516,7 +1424,6 @@ class TestBedrockLLMProxyRoute: return_value=mock_processor, ), ): - # Test regular model endpoint endpoint = "model/anthropic.claude-3-sonnet-20240229-v1:0/converse" @@ -1527,9 +1434,7 @@ class TestBedrockLLMProxyRoute: user_api_key_dict=mock_user_api_key_dict, ) mock_processor.base_passthrough_process_llm_request.assert_called_once() - call_kwargs = ( - mock_processor.base_passthrough_process_llm_request.call_args.kwargs - ) + call_kwargs = mock_processor.base_passthrough_process_llm_request.call_args.kwargs # For regular models, model should be just the model ID assert call_kwargs["model"] == "anthropic.claude-3-sonnet-20240229-v1:0" @@ -1552,9 +1457,7 @@ class TestBedrockLLMProxyRoute: # Create a mock httpx.Response for the error mock_error_response = Mock(spec=httpx.Response) mock_error_response.status_code = 400 - mock_error_response.aread = AsyncMock( - return_value=bedrock_error_message.encode("utf-8") - ) + mock_error_response.aread = AsyncMock(return_value=bedrock_error_message.encode("utf-8")) # Create the HTTPStatusError mock_http_error = httpx.HTTPStatusError( @@ -1571,9 +1474,7 @@ class TestBedrockLLMProxyRoute: mock_request.url = MagicMock() mock_request.url.path = "/bedrock/model/test-model/converse" - mock_request_body = { - "messages": [{"role": "user", "content": [{"textaaa": "Hello"}]}] - } + mock_request_body = {"messages": [{"role": "user", "content": [{"textaaa": "Hello"}]}]} mock_llm_router = Mock() @@ -1614,9 +1515,8 @@ class TestBedrockLLMProxyRoute: ) assert exc_info.value.status_code == 400 - assert ( - "ContentBlock object at messages.0.content.0 must set one of the following keys" - in str(exc_info.value.detail) + assert "ContentBlock object at messages.0.content.0 must set one of the following keys" in str( + exc_info.value.detail ) @pytest.mark.asyncio @@ -1694,24 +1594,14 @@ class TestBedrockLLMProxyRoute: deployment_litellm_params = deployment.get("litellm_params", {}) # Verify model-specific credentials are in the deployment - assert ( - deployment_litellm_params.get("aws_access_key_id") == model_access_key - ) - assert ( - deployment_litellm_params.get("aws_secret_access_key") - == model_secret_key - ) + assert deployment_litellm_params.get("aws_access_key_id") == model_access_key + assert deployment_litellm_params.get("aws_secret_access_key") == model_secret_key assert deployment_litellm_params.get("aws_region_name") == model_region - assert ( - deployment_litellm_params.get("aws_session_token") - == model_session_token - ) + assert deployment_litellm_params.get("aws_session_token") == model_session_token # Verify environment variables are NOT in the deployment assert deployment_litellm_params.get("aws_access_key_id") != env_access_key - assert ( - deployment_litellm_params.get("aws_secret_access_key") != env_secret_key - ) + assert deployment_litellm_params.get("aws_secret_access_key") != env_secret_key assert deployment_litellm_params.get("aws_region_name") != env_region # Test 3: Verify credentials are passed through the passthrough route @@ -1722,9 +1612,7 @@ class TestBedrockLLMProxyRoute: captured_kwargs.update(kwargs) mock_response = MagicMock() mock_response.status_code = 200 - mock_response.aread = AsyncMock( - return_value=b'{"content": [{"text": "Hello"}]}' - ) + mock_response.aread = AsyncMock(return_value=b'{"content": [{"text": "Hello"}]}') return mock_response mock_request = MagicMock(spec=Request) @@ -1734,9 +1622,7 @@ class TestBedrockLLMProxyRoute: mock_request.url = MagicMock() mock_request.url.path = "/bedrock/model/claude-opus-4-1/converse" - mock_request_body = { - "messages": [{"role": "user", "content": [{"text": "Hello"}]}] - } + mock_request_body = {"messages": [{"role": "user", "content": [{"text": "Hello"}]}]} mock_user_api_key_dict = Mock() mock_user_api_key_dict.api_key = "test-key" @@ -1757,9 +1643,7 @@ class TestBedrockLLMProxyRoute: # Setup mock response mock_response = MagicMock() mock_response.status_code = 200 - mock_response.aread = AsyncMock( - return_value=b'{"content": [{"text": "Hello"}]}' - ) + mock_response.aread = AsyncMock(return_value=b'{"content": [{"text": "Hello"}]}') mock_process.return_value = mock_response # Call the handler @@ -1966,9 +1850,7 @@ class TestLLMPassthroughFactoryProxyRoute: mock_user_api_key_dict = MagicMock() with ( - patch( - "litellm.utils.ProviderConfigManager.get_provider_model_info" - ) as mock_get_provider, + patch("litellm.utils.ProviderConfigManager.get_provider_model_info") as mock_get_provider, patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials" ) as mock_get_creds, @@ -1978,9 +1860,7 @@ class TestLLMPassthroughFactoryProxyRoute: ): mock_provider_config = MagicMock() mock_provider_config.get_api_base.return_value = "https://example.com/v1" - mock_provider_config.validate_environment.return_value = { - "x-api-key": "dummy" - } + mock_provider_config.validate_environment.return_value = {"x-api-key": "dummy"} mock_get_provider.return_value = mock_provider_config mock_get_creds.return_value = "dummy" @@ -1996,12 +1876,8 @@ class TestLLMPassthroughFactoryProxyRoute: ) assert result == "success" - mock_get_provider.assert_called_once_with( - provider=litellm.LlmProviders(LlmProviders.VLLM), model=None - ) - mock_get_creds.assert_called_once_with( - custom_llm_provider=LlmProviders.VLLM, region_name=None - ) + mock_get_provider.assert_called_once_with(provider=litellm.LlmProviders(LlmProviders.VLLM), model=None) + mock_get_creds.assert_called_once_with(custom_llm_provider=LlmProviders.VLLM, region_name=None) mock_create_route.assert_called_once_with( endpoint="/chat/completions", target="https://example.com/v1/chat/completions", @@ -2399,9 +2275,7 @@ class TestForwardHeaders: # Create a mock request with custom headers mock_request = MagicMock(spec=Request) - mock_request.state = ( - None # Prevent MagicMock from returning a truthy _cached_headers - ) + mock_request.state = None # Prevent MagicMock from returning a truthy _cached_headers mock_request.method = "POST" mock_request.url = MagicMock() mock_request.url.path = "/test/endpoint" @@ -2436,9 +2310,7 @@ class TestForwardHeaders: mock_httpx_response = MagicMock() mock_httpx_response.status_code = 200 mock_httpx_response.headers = {"content-type": "application/json"} - mock_httpx_response.aiter_bytes = AsyncMock( - return_value=[b'{"result": "success"}'] - ) + mock_httpx_response.aiter_bytes = AsyncMock(return_value=[b'{"result": "success"}']) mock_httpx_response.aread = AsyncMock(return_value=b'{"result": "success"}') with ( @@ -2463,9 +2335,7 @@ class TestForwardHeaders: mock_logging_obj.pre_call_hook = AsyncMock(return_value=mock_request_body) mock_logging_obj.post_call_success_hook = AsyncMock() mock_logging_obj.post_call_failure_hook = AsyncMock() - mock_logging_obj.post_call_response_headers_hook = AsyncMock( - return_value={} - ) + mock_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) # Call pass_through_request with forward_headers=True result = await pass_through_request( @@ -2538,9 +2408,7 @@ class TestForwardHeaders: mock_httpx_response = MagicMock() mock_httpx_response.status_code = 200 mock_httpx_response.headers = {"content-type": "application/json"} - mock_httpx_response.aiter_bytes = AsyncMock( - return_value=[b'{"result": "success"}'] - ) + mock_httpx_response.aiter_bytes = AsyncMock(return_value=[b'{"result": "success"}']) mock_httpx_response.aread = AsyncMock(return_value=b'{"result": "success"}') with ( @@ -2565,9 +2433,7 @@ class TestForwardHeaders: mock_logging_obj.pre_call_hook = AsyncMock(return_value=mock_request_body) mock_logging_obj.post_call_success_hook = AsyncMock() mock_logging_obj.post_call_failure_hook = AsyncMock() - mock_logging_obj.post_call_response_headers_hook = AsyncMock( - return_value={} - ) + mock_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) # Call pass_through_request with forward_headers=False (default) result = await pass_through_request( @@ -2625,15 +2491,11 @@ class TestForwardHeaders: mock_httpx_response = MagicMock() mock_httpx_response.status_code = 200 mock_httpx_response.headers = {"content-type": "application/json"} - mock_httpx_response.aiter_bytes = AsyncMock( - return_value=[b'{"result": "success"}'] - ) + mock_httpx_response.aiter_bytes = AsyncMock(return_value=[b'{"result": "success"}']) mock_httpx_response.aread = AsyncMock(return_value=b'{"result": "success"}') with ( - patch( - "litellm.utils.ProviderConfigManager.get_provider_model_info" - ) as mock_get_provider, + patch("litellm.utils.ProviderConfigManager.get_provider_model_info") as mock_get_provider, patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials" ) as mock_get_creds, @@ -2649,9 +2511,7 @@ class TestForwardHeaders: # Setup provider config mock_provider_config = MagicMock() mock_provider_config.get_api_base.return_value = "https://api.openai.com/v1" - mock_provider_config.validate_environment.return_value = { - "authorization": "Bearer sk-test" - } + mock_provider_config.validate_environment.return_value = {"authorization": "Bearer sk-test"} mock_get_provider.return_value = mock_provider_config mock_get_creds.return_value = "sk-test" @@ -2663,9 +2523,7 @@ class TestForwardHeaders: mock_get_client.return_value = mock_client_obj # Setup mock logging object - mock_logging_obj.pre_call_hook = AsyncMock( - return_value={"messages": [{"role": "user", "content": "test"}]} - ) + mock_logging_obj.pre_call_hook = AsyncMock(return_value={"messages": [{"role": "user", "content": "test"}]}) mock_logging_obj.post_call_success_hook = AsyncMock() # This is the key part - when create_pass_through_route is called with _forward_headers=True @@ -2756,24 +2614,16 @@ class TestMilvusProxyRoute: ): # Setup mocks mock_provider_config = MagicMock() - mock_provider_config.get_auth_credentials.return_value = { - "headers": {"Authorization": "Bearer test-token"} - } + mock_provider_config.get_auth_credentials.return_value = {"headers": {"Authorization": "Bearer test-token"}} mock_provider_config.get_complete_url.return_value = api_base mock_get_config.return_value = mock_provider_config mock_index_registry.is_vector_store_index.return_value = True - mock_index_registry.get_vector_store_index_by_name.return_value = ( - mock_index_object - ) + mock_index_registry.get_vector_store_index_by_name.return_value = mock_index_object - mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = ( - mock_vector_store - ) + mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = mock_vector_store - mock_endpoint_func = AsyncMock( - return_value={"results": [{"id": 1, "distance": 0.5}]} - ) + mock_endpoint_func = AsyncMock(return_value={"results": [{"id": 1, "distance": 0.5}]}) mock_create_route.return_value = mock_endpoint_func # Call the route @@ -2786,9 +2636,7 @@ class TestMilvusProxyRoute: # Verify calls mock_get_body.assert_called_once() - mock_index_registry.is_vector_store_index.assert_called_once_with( - vector_store_index_name=collection_name - ) + mock_index_registry.is_vector_store_index.assert_called_once_with(vector_store_index_name=collection_name) mock_is_allowed.assert_called_once() mock_safe_set.assert_called_once() @@ -2800,9 +2648,7 @@ class TestMilvusProxyRoute: mock_create_route.assert_called_once() create_route_args = mock_create_route.call_args[1] assert "vectors/search" in create_route_args["target"] - assert create_route_args["custom_headers"] == { - "Authorization": "Bearer test-token" - } + assert create_route_args["custom_headers"] == {"Authorization": "Bearer test-token"} # Verify endpoint function was called mock_endpoint_func.assert_awaited_once() @@ -2815,7 +2661,6 @@ class TestMilvusProxyRoute: """ from fastapi import HTTPException - mock_request = MagicMock(spec=Request) mock_response = MagicMock(spec=Response) mock_user_api_key_dict = MagicMock() @@ -2849,7 +2694,6 @@ class TestMilvusProxyRoute: """ from fastapi import HTTPException - mock_request = MagicMock(spec=Request) mock_response = MagicMock(spec=Response) mock_user_api_key_dict = MagicMock() @@ -2867,9 +2711,7 @@ class TestMilvusProxyRoute: ) assert exc_info.value.status_code == 500 - assert "Unable to find Milvus vector store config" in str( - exc_info.value.detail - ) + assert "Unable to find Milvus vector store config" in str(exc_info.value.detail) @pytest.mark.asyncio async def test_milvus_proxy_route_no_index_registry(self): @@ -2878,7 +2720,6 @@ class TestMilvusProxyRoute: """ from fastapi import HTTPException - collection_name = "test-collection" mock_request = MagicMock(spec=Request) @@ -2906,9 +2747,7 @@ class TestMilvusProxyRoute: ) assert exc_info.value.status_code == 500 - assert "Unable to find Milvus vector store index registry" in str( - exc_info.value.detail - ) + assert "Unable to find Milvus vector store index registry" in str(exc_info.value.detail) @pytest.mark.asyncio async def test_milvus_proxy_route_not_managed_index(self): @@ -2917,7 +2756,6 @@ class TestMilvusProxyRoute: """ from fastapi import HTTPException - collection_name = "unmanaged-collection" mock_request = MagicMock(spec=Request) @@ -2947,9 +2785,8 @@ class TestMilvusProxyRoute: ) assert exc_info.value.status_code == 400 - assert ( - f"Collection {collection_name} is not a litellm managed vector store index" - in str(exc_info.value.detail) + assert f"Collection {collection_name} is not a litellm managed vector store index" in str( + exc_info.value.detail ) @pytest.mark.asyncio @@ -2981,22 +2818,16 @@ class TestMilvusProxyRoute: patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_allowed_to_call_vector_store_endpoint" ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._safe_set_request_parsed_body" - ), + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._safe_set_request_parsed_body"), patch.object(litellm, "vector_store_index_registry") as mock_index_registry, patch.object(litellm, "vector_store_registry") as mock_vector_registry, ): mock_get_config.return_value = MagicMock() mock_index_registry.is_vector_store_index.return_value = True - mock_index_registry.get_vector_store_index_by_name.return_value = ( - mock_index_object - ) - mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = ( - None - ) + mock_index_registry.get_vector_store_index_by_name.return_value = mock_index_object + mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = None - with pytest.raises(Exception, match='Vector store not found for missing-store') as exc_info: + with pytest.raises(Exception, match="Vector store not found for missing-store") as exc_info: await milvus_proxy_route( endpoint="vectors/search", request=mock_request, @@ -3004,9 +2835,7 @@ class TestMilvusProxyRoute: user_api_key_dict=mock_user_api_key_dict, ) - assert f"Vector store not found for {vector_store_name}" in str( - exc_info.value - ) + assert f"Vector store not found for {vector_store_name}" in str(exc_info.value) @pytest.mark.asyncio async def test_milvus_proxy_route_no_api_base(self): @@ -3039,9 +2868,7 @@ class TestMilvusProxyRoute: patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_allowed_to_call_vector_store_endpoint" ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._safe_set_request_parsed_body" - ), + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._safe_set_request_parsed_body"), patch.object(litellm, "vector_store_index_registry") as mock_index_registry, patch.object(litellm, "vector_store_registry") as mock_vector_registry, ): @@ -3051,14 +2878,10 @@ class TestMilvusProxyRoute: mock_get_config.return_value = mock_provider_config mock_index_registry.is_vector_store_index.return_value = True - mock_index_registry.get_vector_store_index_by_name.return_value = ( - mock_index_object - ) - mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = ( - mock_vector_store - ) + mock_index_registry.get_vector_store_index_by_name.return_value = mock_index_object + mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = mock_vector_store - with pytest.raises(Exception, match='api_base not found in vector store configuration for') as exc_info: + with pytest.raises(Exception, match="api_base not found in vector store configuration for") as exc_info: await milvus_proxy_route( endpoint="vectors/search", request=mock_request, @@ -3066,10 +2889,7 @@ class TestMilvusProxyRoute: user_api_key_dict=mock_user_api_key_dict, ) - assert ( - f"api_base not found in vector store configuration for {vector_store_name}" - in str(exc_info.value) - ) + assert f"api_base not found in vector store configuration for {vector_store_name}" in str(exc_info.value) @pytest.mark.asyncio async def test_milvus_proxy_route_endpoint_without_leading_slash(self): @@ -3103,9 +2923,7 @@ class TestMilvusProxyRoute: patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_allowed_to_call_vector_store_endpoint" ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._safe_set_request_parsed_body" - ), + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._safe_set_request_parsed_body"), patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, @@ -3118,12 +2936,8 @@ class TestMilvusProxyRoute: mock_get_config.return_value = mock_provider_config mock_index_registry.is_vector_store_index.return_value = True - mock_index_registry.get_vector_store_index_by_name.return_value = ( - mock_index_object - ) - mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = ( - mock_vector_store - ) + mock_index_registry.get_vector_store_index_by_name.return_value = mock_index_object + mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = mock_vector_store mock_endpoint_func = AsyncMock(return_value={"status": "success"}) mock_create_route.return_value = mock_endpoint_func @@ -3171,9 +2985,7 @@ class TestOpenAIPassthroughRoute: "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, ): - mock_endpoint_func = AsyncMock( - return_value={"id": "resp_123", "status": "completed"} - ) + mock_endpoint_func = AsyncMock(return_value={"id": "resp_123", "status": "completed"}) mock_create_route.return_value = mock_endpoint_func # Call the route with /v1/responses endpoint @@ -3221,9 +3033,7 @@ class TestOpenAIPassthroughRoute: "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, ): - mock_endpoint_func = AsyncMock( - return_value={"id": "chatcmpl-123", "choices": []} - ) + mock_endpoint_func = AsyncMock(return_value={"id": "chatcmpl-123", "choices": []}) mock_create_route.return_value = mock_endpoint_func result = await openai_proxy_route( @@ -3289,9 +3099,7 @@ class TestOpenAIPassthroughRoute: "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, ): - mock_endpoint_func = AsyncMock( - return_value={"id": "asst_123", "object": "assistant"} - ) + mock_endpoint_func = AsyncMock(return_value={"id": "asst_123", "object": "assistant"}) mock_create_route.return_value = mock_endpoint_func result = await openai_proxy_route( @@ -3396,9 +3204,7 @@ class TestCursorProxyRoute: "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, ): - mock_endpoint_func = AsyncMock( - return_value={"agents": [], "nextCursor": None} - ) + mock_endpoint_func = AsyncMock(return_value={"agents": [], "nextCursor": None}) mock_create_route.return_value = mock_endpoint_func result = await cursor_proxy_route( @@ -3412,12 +3218,8 @@ class TestCursorProxyRoute: call_args = mock_create_route.call_args[1] assert call_args["target"] == "https://api.cursor.com/v0/agents" - expected_auth = base64.b64encode(f"{test_api_key}:".encode("utf-8")).decode( - "ascii" - ) - assert ( - call_args["custom_headers"]["Authorization"] == f"Basic {expected_auth}" - ) + expected_auth = base64.b64encode(f"{test_api_key}:".encode("utf-8")).decode("ascii") + assert call_args["custom_headers"]["Authorization"] == f"Basic {expected_auth}" assert result == {"agents": [], "nextCursor": None} @@ -3441,7 +3243,7 @@ class TestCursorProxyRoute: [], ), ): - with pytest.raises(Exception, match='Cursor API key not found\\. Add Cursor credentials via') as exc_info: + with pytest.raises(Exception, match="Cursor API key not found\\. Add Cursor credentials via") as exc_info: await cursor_proxy_route( endpoint="v0/agents", request=mock_request, @@ -3500,9 +3302,7 @@ class TestCursorProxyRoute: import base64 expected_auth = base64.b64encode(b"crsr_ui_test_key:").decode("ascii") - assert ( - call_args["custom_headers"]["Authorization"] == f"Basic {expected_auth}" - ) + assert call_args["custom_headers"]["Authorization"] == f"Basic {expected_auth}" @pytest.mark.asyncio async def test_cursor_proxy_route_custom_api_base(self): @@ -3515,9 +3315,7 @@ class TestCursorProxyRoute: mock_user_api_key_dict = MagicMock() with ( - patch.dict( - os.environ, {"CURSOR_API_BASE": "https://custom-cursor.example.com"} - ), + patch.dict(os.environ, {"CURSOR_API_BASE": "https://custom-cursor.example.com"}), patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", return_value="test-key", @@ -3597,12 +3395,10 @@ class TestVertexRawPredictStreamingClassification: """ RAW_PREDICT_ENDPOINT = ( - "v1/projects/test-project/locations/us-east5/publishers/anthropic/models/" - "claude-sonnet-4-6:streamRawPredict" + "v1/projects/test-project/locations/us-east5/publishers/anthropic/models/claude-sonnet-4-6:streamRawPredict" ) GENERATE_CONTENT_ENDPOINT = ( - "v1/projects/test-project/locations/us-east5/publishers/google/models/" - "gemini-2.5-flash:streamGenerateContent" + "v1/projects/test-project/locations/us-east5/publishers/google/models/gemini-2.5-flash:streamGenerateContent" ) async def _capture_passthrough_kwargs(self, endpoint: str, body: object) -> dict: @@ -3787,10 +3583,7 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak: """ VKEY = "sk-litellm-victim-key" - ENDPOINT = ( - "v1/projects/my-proj/locations/us-central1/publishers/google/models/" - "gemini-2.5-flash:generateContent" - ) + ENDPOINT = "v1/projects/my-proj/locations/us-central1/publishers/google/models/gemini-2.5-flash:generateContent" async def _run( self, @@ -3918,7 +3711,9 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak: (b"content-type", b"application/json"), ], ) - assert forwarded is None, f"a virtual key echoed as '{scheme} ' in Authorization must be stripped, not forwarded" + assert forwarded is None, ( + f"a virtual key echoed as '{scheme} ' in Authorization must be stripped, not forwarded" + ) assert raised is not None and raised.status_code == 401 @pytest.mark.asyncio @@ -3964,8 +3759,7 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak: @pytest.mark.parametrize( "credential_header", sorted( - SpecialHeaders.litellm_credential_header_names() - - {"authorization", "x-goog-api-key", "x-litellm-api-key"} + SpecialHeaders.litellm_credential_header_names() - {"authorization", "x-goog-api-key", "x-litellm-api-key"} ), ) async def test_every_non_google_credential_header_is_dropped_by_name(self, monkeypatch, credential_header): @@ -4040,7 +3834,9 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak: assert raised is not None and raised.status_code == 401 @pytest.mark.asyncio - async def test_authenticated_authorization_is_stripped_over_a_lower_precedence_pass_through_header(self, monkeypatch): + async def test_authenticated_authorization_is_stripped_over_a_lower_precedence_pass_through_header( + self, monkeypatch + ): with mock.patch.dict( # test-quality-ok: general_settings is the real proxy config surface for pass_through_endpoints; no injection seam exists on this route "litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [{"headers": {"litellm_user_api_key": "x-company-key"}}]}, @@ -4057,7 +3853,9 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak: assert raised is None assert forwarded is not None assert forwarded.get("x-goog-api-key") == "AIza-real-google-api-key" - assert "authorization" not in forwarded, "Authorization authenticated (higher precedence) so its key must be stripped" + assert "authorization" not in forwarded, ( + "Authorization authenticated (higher precedence) so its key must be stripped" + ) assert "x-company-key" not in forwarded assert self.VKEY not in " ".join(f"{name}:{value}" for name, value in forwarded.items()) @@ -4088,7 +3886,9 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak: (b"content-type", b"application/json"), ], ) - assert forwarded is None, "a virtual key in the mapped-route litellm_user_api_key header must be dropped, not forwarded" + assert forwarded is None, ( + "a virtual key in the mapped-route litellm_user_api_key header must be dropped, not forwarded" + ) assert raised is not None and raised.status_code == 401 GOOGLE_OAUTH_TOKEN = "ya29.byo-google-oauth-token" @@ -4785,9 +4585,7 @@ class TestVertexAILiveWebsocketPassthrough: ] ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) - monkeypatch.setattr( - passthrough_module.passthrough_endpoint_router, "default_vertex_config", None - ) + monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None) self._clear_vertex_env(monkeypatch) websocket = self._websocket() ensure_token = AsyncMock(return_value=("token-abc", "proj-db")) @@ -4929,9 +4727,7 @@ class TestVertexAILiveWebsocketPassthrough: ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) - monkeypatch.setattr( - passthrough_module.passthrough_endpoint_router, "default_vertex_config", None - ) + monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None) self._clear_vertex_env(monkeypatch) websocket = self._websocket() ensure_token = AsyncMock(side_effect=Exception("Unable to find your credentials")) @@ -4953,6 +4749,348 @@ class TestVertexAILiveWebsocketPassthrough: assert len(close_kwargs["reason"].encode("utf-8")) <= 123 +class TestAnthropicProxyRoute: + """The /anthropic passthrough route: custom auth headers must not clobber the + client's anthropic-beta, and the WIF tier must mint through the async facade.""" + + def _get_request(self, headers: dict) -> MagicMock: + request = MagicMock(spec=Request) + request.method = "GET" + request.headers = headers + request.query_params = {} + return request + + def _clear_anthropic_env(self, monkeypatch) -> None: + 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_FILE", + "ANTHROPIC_IDENTITY_TOKEN", + ): + monkeypatch.delenv(name, raising=False) + + @pytest.mark.asyncio + async def test_client_anthropic_beta_merged_into_auth_header(self, monkeypatch): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + anthropic_proxy_route, + ) + + self._clear_anthropic_env(monkeypatch) + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-oat01-passthrough-token") + + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + return_value=AsyncMock(return_value={"ok": True}), + ) as mock_create_route: + await anthropic_proxy_route( + endpoint="v1/models", + request=self._get_request({"anthropic-beta": "context-1m-2025-08-07"}), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + ) + + custom_headers = mock_create_route.call_args.kwargs["custom_headers"] + assert custom_headers["authorization"] == "Bearer sk-ant-oat01-passthrough-token" + betas = set(custom_headers["anthropic-beta"].split(",")) + assert {"context-1m-2025-08-07", "oauth-2025-04-20"} <= betas + + @pytest.mark.asyncio + async def test_wif_mint_goes_through_async_facade(self, monkeypatch): + 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 + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + anthropic_proxy_route, + ) + + self._clear_anthropic_env(monkeypatch) + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_route") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-route") + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "route-inline-jwt") + + minted: Final = "sk-ant-oat01-route-minted" + thread_ids: Final = [] + + 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: Final = [] + + 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) + + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + return_value=AsyncMock(return_value={"ok": True}), + ) as mock_create_route: + await anthropic_proxy_route( + endpoint="v1/models", + request=self._get_request({}), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + ) + + custom_headers = mock_create_route.call_args.kwargs["custom_headers"] + assert custom_headers["authorization"] == f"Bearer {minted}" + assert custom_headers["anthropic-beta"] == "oauth-2025-04-20" + assert sync_calls == [] + assert thread_ids and thread_ids[0] != threading.get_ident() + + +class TestAnthropicProxyRouteCallerAuthHeaders: + """Regression for a caller credential riding upstream next to a server-owned one. + + /anthropic forwards the caller's headers, so a caller-supplied ``x-api-key`` used to reach + Anthropic alongside the server-minted ``Authorization: Bearer``. These drive the real relay + (only the httpx client is stubbed) and assert on the bytes actually handed to the upstream. + """ + + _MINTED: Final = "sk-ant-oat01-plan-minted" + + def _clear_anthropic_env(self, monkeypatch) -> None: + 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_FILE", + "ANTHROPIC_IDENTITY_TOKEN", + ): + monkeypatch.delenv(name, raising=False) + + # A sibling test leaving SERVER_ROOT_PATH set re-prefixes the passthrough route, so + # /anthropic/... stops resolving and the request 404s before any header is built. + # Pin it so this class asserts on headers rather than on ambient state. + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + + def _enable_wif(self, monkeypatch) -> None: + from litellm.llms.anthropic import common_utils as anthropic_common_utils + from litellm.llms.anthropic.wif import aget_anthropic_wif_token + from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_plan") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-plan") + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "plan-inline-jwt") + + minted: Final = self._MINTED + + class StubPoster: + def post(self, url, *, content, headers, timeout): + return httpx.Response( + 200, + json={"access_token": minted, "token_type": "Bearer", "expires_in": 3600}, + ) + + engine: Final = JwtBearerTokenExchangeEngine(poster=StubPoster()) + + 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) + + def _request(self, headers: Mapping[str, str]) -> Request: + body: Final = b'{"model":"claude-sonnet-4-5","messages":[]}' + scope: Final = { + "type": "http", + "http_version": "1.1", + "method": "POST", + "scheme": "https", + "path": "/anthropic/v1/messages", + "raw_path": b"/anthropic/v1/messages", + "root_path": "", + "query_string": b"", + "headers": [(name.lower().encode(), value.encode()) for name, value in headers.items()], + "client": ("127.0.0.1", 51234), + "server": ("proxy.local", 4000), + "state": {}, + } + + async def receive() -> dict: + return {"type": "http.request", "body": body, "more_body": False} + + return Request(scope, receive) + + async def _upstream_headers(self, request: Request) -> dict: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + anthropic_proxy_route, + ) + + upstream_response: Final = MagicMock() + upstream_response.status_code = 200 + upstream_response.headers = {"content-type": "application/json"} + upstream_response.aread = AsyncMock(return_value=b'{"ok": true}') + upstream_response.aiter_bytes = AsyncMock(return_value=[b'{"ok": true}']) + + httpx_client: Final = MagicMock() + httpx_client.build_request = MagicMock(return_value=MagicMock()) + httpx_client.send = AsyncMock(return_value=upstream_response) + client_wrapper: Final = MagicMock() + client_wrapper.client = httpx_client + + with ( + patch( # test-quality-ok: stubbing the http client IS the boundary; the test asserts on the bytes handed to it + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client", + return_value=client_wrapper, + ), + patch( # test-quality-ok: the relay calls these hooks, and they need a db this test has no use for + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_logging_obj, + ): + mock_logging_obj.pre_call_hook = AsyncMock(return_value={"model": "claude-sonnet-4-5", "messages": []}) + mock_logging_obj.post_call_success_hook = AsyncMock() + mock_logging_obj.post_call_failure_hook = AsyncMock() + mock_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + await anthropic_proxy_route( + endpoint="v1/messages", + request=request, + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-caller-virtual-key"), + ) + + assert httpx_client.send.called + return {name.lower(): value for name, value in dict(httpx_client.build_request.call_args[1]["headers"]).items()} + + @pytest.mark.asyncio + async def test_wif_credential_drops_caller_supplied_api_key(self, monkeypatch): + self._clear_anthropic_env(monkeypatch) + self._enable_wif(monkeypatch) + + sent: Final = await self._upstream_headers( + self._request( + { + "content-type": "application/json", + "x-api-key": "sk-caller-virtual-key", + "user-agent": "caller/1.0", + } + ) + ) + + assert sent["authorization"] == f"Bearer {self._MINTED}" + assert "x-api-key" not in sent + assert sent["user-agent"] == "caller/1.0" + + @pytest.mark.asyncio + async def test_wif_credential_drops_caller_supplied_authorization(self, monkeypatch): + self._clear_anthropic_env(monkeypatch) + self._enable_wif(monkeypatch) + + sent: Final = await self._upstream_headers( + self._request( + { + "content-type": "application/json", + "authorization": "Bearer sk-caller-virtual-key", + } + ) + ) + + assert sent["authorization"] == f"Bearer {self._MINTED}" + assert all("sk-caller-virtual-key" not in value for value in sent.values()) + + @pytest.mark.asyncio + @pytest.mark.parametrize("header_name", sorted(SpecialHeaders.litellm_credential_header_names())) + async def test_wif_credential_drops_every_proxy_key_header(self, monkeypatch, header_name: str): + """The proxy accepts a LiteLLM key in any SpecialHeaders slot, so the caller's virtual + key must not reach Anthropic from any of them once the server owns the credential.""" + self._clear_anthropic_env(monkeypatch) + self._enable_wif(monkeypatch) + + sent: Final = await self._upstream_headers( + self._request( + { + "content-type": "application/json", + header_name: "sk-caller-virtual-key", + "user-agent": "caller/1.0", + } + ) + ) + + assert sent["authorization"] == f"Bearer {self._MINTED}" + assert header_name == "authorization" or header_name not in sent + assert all("sk-caller-virtual-key" not in value for value in sent.values()) + assert sent["user-agent"] == "caller/1.0" + + @pytest.mark.asyncio + async def test_wif_credential_drops_configured_custom_key_header(self, monkeypatch): + from litellm.proxy import proxy_server + + self._clear_anthropic_env(monkeypatch) + self._enable_wif(monkeypatch) + monkeypatch.setitem(proxy_server.general_settings, "litellm_key_header_name", "X-Tenant-Key") + + sent: Final = await self._upstream_headers( + self._request( + { + "content-type": "application/json", + "x-tenant-key": "sk-caller-virtual-key", + "x-tenant-region": "eu", + } + ) + ) + + assert sent["authorization"] == f"Bearer {self._MINTED}" + assert "x-tenant-key" not in sent + assert all("sk-caller-virtual-key" not in value for value in sent.values()) + assert sent["x-tenant-region"] == "eu" + + @pytest.mark.asyncio + async def test_server_api_key_drops_caller_supplied_authorization(self, monkeypatch): + self._clear_anthropic_env(monkeypatch) + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-server-owned") + + sent: Final = await self._upstream_headers( + self._request( + { + "content-type": "application/json", + "authorization": "Bearer sk-caller-virtual-key", + "x-api-key": "sk-caller-virtual-key", + } + ) + ) + + assert sent["x-api-key"] == "sk-ant-server-owned" + assert "authorization" not in sent + + @pytest.mark.asyncio + async def test_byok_caller_key_still_reaches_upstream(self, monkeypatch): + self._clear_anthropic_env(monkeypatch) + + sent: Final = await self._upstream_headers( + self._request( + { + "content-type": "application/json", + "x-api-key": "sk-ant-caller-owned", + "anthropic-version": "2023-06-01", + } + ) + ) + + assert sent["x-api-key"] == "sk-ant-caller-owned" + assert sent["anthropic-version"] == "2023-06-01" + assert "authorization" not in sent + + class TestPassthroughRouterModelBudgetReservation: """ Router-model passthrough on /vllm and /azure must thread the calling key's diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 2babfe432f3..9fefa93f671 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2406,6 +2406,23 @@ def test_ProxyConfig__add_deployment_resolves_env_refs_on_arbitrary_field(monkey # --------------------------------------------------------------------------- +def test_ProxyConfig_decrypt_credentials_returns_an_encrypted_empty_value_as_empty(monkeypatch): + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-decrypt-credentials-test-salt") + decrypted = ProxyConfig().decrypt_credentials( + { + "credential_name": "openai-wif", + "credential_values": { + "api_base": encrypt_value_helper(""), + "openai_service_account_id": encrypt_value_helper("user-1"), + }, + "credential_info": {"custom_llm_provider": "openai"}, + } + ) + assert decrypted.credential_values == {"api_base": "", "openai_service_account_id": "user-1"} + + def test_ProxyConfig_decrypt_model_list_from_db_returns_decrypted(monkeypatch): monkeypatch.setattr( "litellm.proxy.proxy_server.decrypt_value_helper", diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 9393ec0f8e6..b3cb0286ece 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -12652,6 +12652,57 @@ async def test_load_config_router_authorizes_fallback_targets_against_the_callin assert router.fallback_access_check is router_fallback_access_check +def test_resolve_db_litellm_param_keeps_wif_secret_pointers(monkeypatch): + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setenv("WIF_TEST_KC_SECRET", "kc-secret") + proxy_config = ProxyConfig() + + pointer = proxy_config._resolve_db_litellm_param( + "anthropic_keycloak_client_secret_ref", "os.environ/WIF_TEST_KC_SECRET" + ) + dereferenced = proxy_config._resolve_db_litellm_param("api_key", "os.environ/WIF_TEST_KC_SECRET") + + assert pointer == "os.environ/WIF_TEST_KC_SECRET" + assert dereferenced == "kc-secret" + + +@pytest.mark.asyncio +async def test_load_config_keeps_wif_secret_pointers_on_config_models(tmp_path, monkeypatch): + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setenv("WIF_TEST_SIGNING_KEY", "-----BEGIN PRIVATE KEY-----") + monkeypatch.setenv("WIF_TEST_FDRL", "fdrl_from_env") + config_file = tmp_path / "config.yaml" + config_file.write_text( + yaml.dump( + { + "model_list": [ + { + "model_name": "claude-wif", + "litellm_params": { + "model": "anthropic/claude-haiku-4-5", + "anthropic_federation_rule_id": "os.environ/WIF_TEST_FDRL", + "anthropic_identity_source": "internal_issuer", + "anthropic_issuer_url": "https://litellm.example", + "anthropic_issuer_audience": "https://api.anthropic.com", + "anthropic_issuer_signing_key_ref": "os.environ/WIF_TEST_SIGNING_KEY", + }, + } + ] + } + ) + ) + + _router, model_list, _general_settings = await ProxyConfig().load_config( + router=None, config_file_path=str(config_file) + ) + + litellm_params = model_list[0]["litellm_params"] + assert litellm_params["anthropic_federation_rule_id"] == "fdrl_from_env" + assert litellm_params["anthropic_issuer_signing_key_ref"] == "os.environ/WIF_TEST_SIGNING_KEY" + + def test_docs_redoc_openapi_are_reachable_by_default(): """ LIT-6745: the interactive/machine-readable docs surfaces are on by diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index a4965c49f07..bff1be3ff1d 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -711,9 +711,7 @@ async def test_run_async_fallback_keeps_a_request_override_distinct_from_the_bar with pytest.raises(RuntimeError, match="fallback model also failed"): await run_async_fallback( litellm_router=router, - fallback_model_group=[ - {"model": "already-attempted", "messages": [{"role": "user", "content": "shorter"}]} - ], + fallback_model_group=[{"model": "already-attempted", "messages": [{"role": "user", "content": "shorter"}]}], original_model_group="primary-model", original_exception=RuntimeError("original failed"), max_fallbacks=3, @@ -1075,6 +1073,58 @@ class TestRunAsyncFallbackTriggersCooldown: @pytest.mark.asyncio +async def test_a_stored_fallback_target_cannot_carry_a_federation_field(): + """A dict fallback target is merged into kwargs, and kwargs beat the deployment's own params, + so a stored key/team/global fallback could otherwise set the workspace a federation token is + minted for. The request itself is already forbidden to carry these, and a stored setting is + not a more trusted source than the request.""" + with pytest.raises(ValueError, match="server-owned workload identity federation parameter"): + await run_async_fallback( + litellm_router=FakeRouter(), + fallback_model_group=[{"model": "anthropic-backup", "anthropic_workspace_id": "wrkspc_other"}], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + ) + + +@pytest.mark.asyncio +async def test_a_stored_fallback_target_cannot_carry_an_openai_federation_field(): + """The OpenAI identity trio is server-owned for the same reason: a stored fallback target + naming a token file would pick which workload assertion is exchanged for the bearer.""" + with pytest.raises(ValueError, match="openai_identity_token_file"): + await run_async_fallback( + litellm_router=FakeRouter(), + fallback_model_group=[ + {"model": "openai-backup", "openai_identity_token_file": "/var/run/secrets/tokens/other"} + ], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + ) + + +@pytest.mark.asyncio +async def test_the_refusal_is_not_swallowed_as_a_fallback_error(): + """Checked before the per-target loop on purpose: inside it, the refusal would be caught as + that target's failure and the run would quietly continue to the next one.""" + with pytest.raises(ValueError, match="anthropic_issuer_signing_key_ref"): + await run_async_fallback( + litellm_router=FakeRouter(), + fallback_model_group=[ + {"model": "anthropic-backup", "anthropic_issuer_signing_key_ref": "os.environ/ADMIN_KEY"}, + "a-perfectly-fine-model", + ], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + include_fallback_errors=True, + ) + + async def test_run_async_fallback_stamps_fallback_info_into_metadata(): """Spend logs are built from the request metadata of the nested call, so the fallback signal has to be stamped there before recursing.""" diff --git a/tests/test_litellm/test_anthropic_skills_transformation.py b/tests/test_litellm/test_anthropic_skills_transformation.py index 1b917f08ca9..1d6be70dd7a 100644 --- a/tests/test_litellm/test_anthropic_skills_transformation.py +++ b/tests/test_litellm/test_anthropic_skills_transformation.py @@ -27,9 +27,7 @@ FAKE_API_KEY = "sk-ant-test-key-1234" FAKE_API_BASE = "https://api.anthropic.com" -def _make_mock_response( - json_data: dict, status_code: int = 200, method: str = "POST" -) -> httpx.Response: +def _make_mock_response(json_data: dict, status_code: int = 200, method: str = "POST") -> httpx.Response: return httpx.Response( status_code=status_code, json=json_data, @@ -111,9 +109,7 @@ class TestAnthropicSkillsConfigHeaderValidation: "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", return_value=FAKE_API_KEY, ): - headers = self.config.validate_environment( - headers={}, litellm_params=self._make_litellm_params() - ) + headers = self.config.validate_environment(headers={}, litellm_params=self._make_litellm_params()) assert headers["x-api-key"] == FAKE_API_KEY def test_sets_anthropic_version_header(self): @@ -121,9 +117,7 @@ class TestAnthropicSkillsConfigHeaderValidation: "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", return_value=FAKE_API_KEY, ): - headers = self.config.validate_environment( - headers={}, litellm_params=self._make_litellm_params() - ) + headers = self.config.validate_environment(headers={}, litellm_params=self._make_litellm_params()) assert headers["anthropic-version"] == "2023-06-01" def test_sets_skills_beta_header(self): @@ -131,12 +125,12 @@ class TestAnthropicSkillsConfigHeaderValidation: "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", return_value=FAKE_API_KEY, ): - headers = self.config.validate_environment( - headers={}, litellm_params=self._make_litellm_params() - ) + headers = self.config.validate_environment(headers={}, litellm_params=self._make_litellm_params()) assert headers["anthropic-beta"] == ANTHROPIC_SKILLS_API_BETA_VERSION - def test_merges_existing_beta_header_string(self): + def test_merges_existing_beta_header_into_string(self): + """The merged value must stay a comma-separated string: a list value makes + httpx.Headers raise TypeError when the request is built.""" with patch( "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", return_value=FAKE_API_KEY, @@ -145,21 +139,26 @@ class TestAnthropicSkillsConfigHeaderValidation: headers={"anthropic-beta": "other-beta-2024-01-01"}, litellm_params=self._make_litellm_params(), ) - assert isinstance(headers["anthropic-beta"], list) - assert "other-beta-2024-01-01" in headers["anthropic-beta"] - assert ANTHROPIC_SKILLS_API_BETA_VERSION in headers["anthropic-beta"] + assert isinstance(headers["anthropic-beta"], str) + betas = set(headers["anthropic-beta"].split(",")) + assert {"other-beta-2024-01-01", ANTHROPIC_SKILLS_API_BETA_VERSION} <= betas + httpx.Headers(headers) - def test_merges_existing_beta_header_list(self): + def test_oauth_key_beta_merges_without_crashing_httpx(self): + """Regression: an sk-ant-oat/WIF auth header carries its own anthropic-beta; + the old list-building merge produced a Python list that crashed httpx.""" with patch( "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", - return_value=FAKE_API_KEY, + return_value="sk-ant-oat01-fake-skills-token", ): headers = self.config.validate_environment( - headers={"anthropic-beta": ["other-beta-2024-01-01"]}, - litellm_params=self._make_litellm_params(), + headers={}, litellm_params=self._make_litellm_params(api_key=None) ) - assert ANTHROPIC_SKILLS_API_BETA_VERSION in headers["anthropic-beta"] - assert "other-beta-2024-01-01" in headers["anthropic-beta"] + assert headers["authorization"] == "Bearer sk-ant-oat01-fake-skills-token" + assert isinstance(headers["anthropic-beta"], str) + betas = set(headers["anthropic-beta"].split(",")) + assert {"oauth-2025-04-20", ANTHROPIC_SKILLS_API_BETA_VERSION} <= betas + httpx.Headers(headers) def test_does_not_duplicate_beta_header(self): with patch( @@ -170,11 +169,7 @@ class TestAnthropicSkillsConfigHeaderValidation: headers={"anthropic-beta": ANTHROPIC_SKILLS_API_BETA_VERSION}, litellm_params=self._make_litellm_params(), ) - beta = headers["anthropic-beta"] - if isinstance(beta, list): - assert beta.count(ANTHROPIC_SKILLS_API_BETA_VERSION) == 1 - else: - assert beta == ANTHROPIC_SKILLS_API_BETA_VERSION + assert headers["anthropic-beta"] == ANTHROPIC_SKILLS_API_BETA_VERSION def test_raises_without_api_key(self): with patch( @@ -182,9 +177,7 @@ class TestAnthropicSkillsConfigHeaderValidation: return_value=None, ): with pytest.raises(ValueError, match="ANTHROPIC_API_KEY"): - self.config.validate_environment( - headers={}, litellm_params=self._make_litellm_params(api_key=None) - ) + self.config.validate_environment(headers={}, litellm_params=self._make_litellm_params(api_key=None)) class TestAnthropicSkillsConfigCreateRequestTransformation: @@ -275,9 +268,7 @@ class TestAnthropicSkillsConfigResponseTransformation: def test_create_skill_response_parses_skill(self): payload = _make_skill_payload() raw = _make_mock_response(payload) - skill = self.config.transform_create_skill_response( - raw_response=raw, logging_obj=self.logging_obj - ) + skill = self.config.transform_create_skill_response(raw_response=raw, logging_obj=self.logging_obj) assert isinstance(skill, Skill) assert skill.id == "skill_abc123" assert skill.source == "custom" @@ -286,9 +277,7 @@ class TestAnthropicSkillsConfigResponseTransformation: def test_get_skill_response_parses_skill(self): payload = _make_skill_payload(id="skill_xyz", display_title="Another") raw = _make_mock_response(payload, method="GET") - skill = self.config.transform_get_skill_response( - raw_response=raw, logging_obj=self.logging_obj - ) + skill = self.config.transform_get_skill_response(raw_response=raw, logging_obj=self.logging_obj) assert isinstance(skill, Skill) assert skill.id == "skill_xyz" assert skill.display_title == "Another" @@ -300,9 +289,7 @@ class TestAnthropicSkillsConfigResponseTransformation: "next_page": None, } raw = _make_mock_response(payload, method="GET") - result = self.config.transform_list_skills_response( - raw_response=raw, logging_obj=self.logging_obj - ) + result = self.config.transform_list_skills_response(raw_response=raw, logging_obj=self.logging_obj) assert isinstance(result, ListSkillsResponse) assert len(result.data) == 2 assert result.data[0].id == "skill_abc123" @@ -316,18 +303,14 @@ class TestAnthropicSkillsConfigResponseTransformation: "next_page": "page_token_xyz", } raw = _make_mock_response(payload, method="GET") - result = self.config.transform_list_skills_response( - raw_response=raw, logging_obj=self.logging_obj - ) + result = self.config.transform_list_skills_response(raw_response=raw, logging_obj=self.logging_obj) assert result.has_more is True assert result.next_page == "page_token_xyz" def test_delete_skill_response_parses_correctly(self): payload = {"id": "skill_abc123", "type": "skill_deleted"} raw = _make_mock_response(payload, method="DELETE") - result = self.config.transform_delete_skill_response( - raw_response=raw, logging_obj=self.logging_obj - ) + result = self.config.transform_delete_skill_response(raw_response=raw, logging_obj=self.logging_obj) assert isinstance(result, DeleteSkillResponse) assert result.id == "skill_abc123" assert result.type == "skill_deleted" @@ -341,8 +324,6 @@ class TestAnthropicSkillsConfigResponseTransformation: "type": "skill", } raw = _make_mock_response(payload) - skill = self.config.transform_create_skill_response( - raw_response=raw, logging_obj=self.logging_obj - ) + skill = self.config.transform_create_skill_response(raw_response=raw, logging_obj=self.logging_obj) assert skill.display_title is None assert skill.latest_version is None diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index ffd6c5f97ce..f779abaf870 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -4859,6 +4859,40 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params(): assert credentials.get(key) == value, key +def test_get_deployment_credentials_with_provider_preserves_anthropic_wif_params(): + """ + Test that get_deployment_credentials_with_provider preserves a litellm_params-configured + Anthropic workload identity federation setup (both the legacy token_file fields and the + Phase 1 internal_issuer/keycloak identity-source fields) so files/batches/passthrough + deployments using WIF do not silently fall back to a missing credential. + """ + wif_params = { + "anthropic_federation_rule_id": "fdrl_deployment", + "anthropic_organization_id": "org-deployment", + "anthropic_identity_source": "keycloak", + "anthropic_keycloak_token_url": "https://keycloak.internal.example/realms/r/protocol/openid-connect/token", + "anthropic_keycloak_client_id": "litellm", + "anthropic_keycloak_client_secret_ref": "oidc/env/KEYCLOAK_CLIENT_SECRET", + } + router = litellm.Router( + model_list=[ + { + "model_name": "anthropic-wif-model", + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5", + **wif_params, + }, + } + ], + ) + + credentials = router.get_deployment_credentials_with_provider(model_id="anthropic-wif-model") + + assert credentials is not None + for key, value in wif_params.items(): + assert credentials.get(key) == value, key + + def _team_wildcard_model(api_key: str, model_id: str = "team-wildcard-id") -> dict: return { "model_name": f"model_name_team-1_{model_id}", @@ -12593,6 +12627,31 @@ class TestTierParamsTheTargetAccepts: assert accepted == {"reasoning_effort": "max"} +def test_router_keeps_wif_secret_pointers_unresolved(monkeypatch): + monkeypatch.setenv("WIF_TEST_KC_SECRET", "kc-secret") + monkeypatch.setenv("WIF_TEST_FDRL", "fdrl_from_env") + router = Router( + model_list=[ + { + "model_name": "claude-wif", + "litellm_params": { + "model": "anthropic/claude-haiku-4-5", + "anthropic_federation_rule_id": "os.environ/WIF_TEST_FDRL", + "anthropic_identity_source": "keycloak", + "anthropic_keycloak_token_url": "https://keycloak.example/token", + "anthropic_keycloak_client_id": "litellm", + "anthropic_keycloak_client_secret_ref": "os.environ/WIF_TEST_KC_SECRET", + }, + } + ] + ) + + litellm_params = router.get_model_list()[0]["litellm_params"] + + assert litellm_params["anthropic_federation_rule_id"] == "fdrl_from_env" + assert litellm_params["anthropic_keycloak_client_secret_ref"] == "os.environ/WIF_TEST_KC_SECRET" + + class TestRequestReasoningEffortOverride: def test_drop_effort_from_nested_carrier_preserves_other_nested_values(self): params: dict[str, object] = {"output_config": {"effort": "high", "format": "json"}} diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index accd3b32a0d..5f9fb6c57b5 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -2,11 +2,22 @@ import pytest from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, + CredentialLiteLLMParams, Deployment, LiteLLM_Params, ModelInfo, + holds_secret_pointer, + reject_server_owned_wif_params, + server_owned_wif_fields_named, + server_owned_wif_fields_present, +) +from litellm.types.utils import ( + CustomPricingLiteLLMParams, + MirroredPricingParams, + anthropic_wif_litellm_params, + openai_wif_litellm_params, + server_owned_wif_litellm_params, ) -from litellm.types.utils import CustomPricingLiteLLMParams, MirroredPricingParams def test_model_info_declares_mirrored_pricing_fields(): @@ -87,5 +98,105 @@ def test_pricing_strings_are_coerced_to_float(): def test_invalid_pricing_is_rejected(): - with pytest.raises(ValueError, match='validation error for ModelInfo'): + with pytest.raises(ValueError, match="validation error for ModelInfo"): ModelInfo(id="x", input_cost_per_token="free") + + +def test_credential_litellm_params_declares_every_anthropic_wif_field(): + """Without these, get_deployment_credentials_with_provider round-trips litellm_params + through a strict Pydantic dump and silently drops every WIF field before files/batches/ + passthrough callers see it -- the same #30235-shaped gap azure_ad_token closed above.""" + for field in anthropic_wif_litellm_params: + assert field in CredentialLiteLLMParams.model_fields, field + + +def test_anthropic_wif_fields_round_trip_through_model_dump(): + values = {field: f"value-for-{field}" for field in anthropic_wif_litellm_params} + values["anthropic_issuer_ttl_seconds"] = 300 + values["anthropic_disable_workload_identity_federation"] = True + + dumped = CredentialLiteLLMParams(**values).model_dump(exclude_none=True) + + for field, value in values.items(): + assert dumped[field] == value, field + + +def test_server_owned_wif_fields_present_reports_only_set_fields(): + assert server_owned_wif_fields_present({}) == () + assert server_owned_wif_fields_present({"model": "gpt-4o"}) == () + assert server_owned_wif_fields_present( + {"anthropic_keycloak_token_url": "https://idp.example/token", "model": "gpt-4o"} + ) == ("anthropic_keycloak_token_url",) + + +def test_server_owned_wif_fields_present_is_derived_from_the_shared_list(): + """A non-admin persistence gate built on this must automatically cover a field added + later to server_owned_wif_litellm_params, not just the fields known when the gate was + written -- so this must read the shared list rather than a hand-copied one.""" + values = {field: "set" for field in server_owned_wif_litellm_params} + assert set(server_owned_wif_fields_present(values)) == set(server_owned_wif_litellm_params) + + +def test_server_owned_wif_fields_named_reports_keys_whatever_their_value(): + """The credential write gates must see a key a caller sets to ``None``: the federation + resolver reacts to the key's presence, not its value, so ``{"anthropic_issuer_url": None}`` + wedges every deployment referencing the credential once persisted.""" + assert server_owned_wif_fields_named({}) == () + assert server_owned_wif_fields_named({"model": "gpt-4o"}) == () + assert server_owned_wif_fields_named({"anthropic_issuer_url": None}) == ("anthropic_issuer_url",) + assert server_owned_wif_fields_present({"anthropic_issuer_url": None}) == () + assert server_owned_wif_fields_named(("anthropic_keycloak_token_url", "api_key")) == ( + "anthropic_keycloak_token_url", + ) + + +def test_server_owned_wif_fields_named_is_derived_from_the_shared_list(): + assert set(server_owned_wif_fields_named(frozenset(server_owned_wif_litellm_params))) == set( + server_owned_wif_litellm_params + ) + + +@pytest.mark.parametrize("param_name", ["anthropic_issuer_signing_key_ref", "anthropic_keycloak_client_secret_ref"]) +def test_wif_ref_fields_hold_secret_pointers(param_name: str): + assert holds_secret_pointer(param_name) + + +@pytest.mark.parametrize("param_name", ["api_key", "anthropic_federation_rule_id", "anthropic_identity_token"]) +def test_dereferenced_fields_do_not_hold_secret_pointers(param_name: str): + assert not holds_secret_pointer(param_name) + + +def test_credential_litellm_params_declares_every_openai_wif_field(): + for field in openai_wif_litellm_params: + assert field in CredentialLiteLLMParams.model_fields, field + + +def test_openai_wif_fields_round_trip_through_model_dump(): + values = {field: f"value-for-{field}" for field in openai_wif_litellm_params} + + dumped = CredentialLiteLLMParams(**values).model_dump(exclude_none=True) + + for field, value in values.items(): + assert dumped[field] == value, field + + +def test_server_owned_registry_is_anthropic_plus_openai(): + assert server_owned_wif_litellm_params == anthropic_wif_litellm_params + openai_wif_litellm_params + assert set(openai_wif_litellm_params) == { + "openai_identity_provider_id", + "openai_service_account_id", + "openai_identity_token_file", + } + + +def test_server_owned_wif_fields_present_reports_openai_fields(): + assert server_owned_wif_fields_present( + {"openai_identity_token_file": "/var/run/secrets/tokens/openai", "model": "gpt-4o"} + ) == ("openai_identity_token_file",) + assert server_owned_wif_fields_named({"openai_service_account_id": None}) == ("openai_service_account_id",) + + +@pytest.mark.parametrize("param_name", openai_wif_litellm_params) +def test_reject_server_owned_wif_params_names_each_openai_field(param_name: str): + with pytest.raises(ValueError, match=param_name): + reject_server_owned_wif_params({param_name: "client-supplied"}) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 14485a0a115..676fb4c45ef 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -3463,6 +3463,28 @@ export interface paths { patch: operations["update_credential_credentials__credential_name__patch"]; trace?: never; }; + "/credentials/{credential_name}/jwks": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Credential Internal Issuer Jwks + * @description 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. + */ + get: operations["get_credential_internal_issuer_jwks_credentials__credential_name__jwks_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/cursor/chat/completions": { parameters: { query?: never; @@ -26520,7 +26542,22 @@ export interface components { search_tool: components["schemas"]["SearchTool"]; }; /** CredentialItem */ - CredentialItem: { + "CredentialItem-Input": { + /** Credential Info */ + credential_info: { + [key: string]: unknown; + }; + /** Credential Name */ + credential_name: string; + /** Credential Values */ + credential_values: { + [key: string]: unknown; + }; + /** Credential Values To Delete */ + credential_values_to_delete?: string[] | null; + }; + /** CredentialItem */ + "CredentialItem-Output": { /** Credential Info */ credential_info: { [key: string]: unknown; @@ -26892,6 +26929,8 @@ export interface components { }; /** Deployment */ Deployment: { + /** Blocked */ + blocked?: boolean | null; litellm_params: components["schemas"]["LiteLLM_Params"]; model_info: components["schemas"]["litellm__types__router__ModelInfo"]; /** Model Name */ @@ -29265,6 +29304,42 @@ export interface components { allow_client_keepalive_override: boolean | null; /** Annotation Cost Per Page */ annotation_cost_per_page?: number | null; + /** Anthropic Disable Workload Identity Federation */ + anthropic_disable_workload_identity_federation?: boolean | null; + /** Anthropic Federation Rule Id */ + anthropic_federation_rule_id?: string | null; + /** Anthropic Identity Source */ + anthropic_identity_source?: string | null; + /** Anthropic Identity Token */ + anthropic_identity_token?: string | null; + /** Anthropic Identity Token File */ + anthropic_identity_token_file?: string | null; + /** Anthropic Issuer Audience */ + anthropic_issuer_audience?: string | null; + /** Anthropic Issuer Signing Key Ref */ + anthropic_issuer_signing_key_ref?: string | null; + /** Anthropic Issuer Subject */ + anthropic_issuer_subject?: string | null; + /** Anthropic Issuer Ttl Seconds */ + anthropic_issuer_ttl_seconds?: number | null; + /** Anthropic Issuer Url */ + anthropic_issuer_url?: string | null; + /** Anthropic Keycloak Auth Method */ + anthropic_keycloak_auth_method?: string | null; + /** Anthropic Keycloak Client Id */ + anthropic_keycloak_client_id?: string | null; + /** Anthropic Keycloak Client Secret Ref */ + anthropic_keycloak_client_secret_ref?: string | null; + /** Anthropic Keycloak Scope */ + anthropic_keycloak_scope?: string | null; + /** Anthropic Keycloak Token Url */ + anthropic_keycloak_token_url?: string | null; + /** Anthropic Organization Id */ + anthropic_organization_id?: string | null; + /** Anthropic Service Account Id */ + anthropic_service_account_id?: string | null; + /** Anthropic Workspace Id */ + anthropic_workspace_id?: string | null; /** Api Base */ api_base?: string | null; /** Api Key */ @@ -29470,6 +29545,12 @@ export interface components { ocr_cost_per_credit?: number | null; /** Ocr Cost Per Page */ ocr_cost_per_page?: number | null; + /** Openai Identity Provider Id */ + openai_identity_provider_id?: string | null; + /** Openai Identity Token File */ + openai_identity_token_file?: string | null; + /** Openai Service Account Id */ + openai_service_account_id?: string | null; /** Organization */ organization?: string | null; /** Otpm */ @@ -39395,6 +39476,42 @@ export interface components { allow_client_keepalive_override: boolean | null; /** Annotation Cost Per Page */ annotation_cost_per_page?: number | null; + /** Anthropic Disable Workload Identity Federation */ + anthropic_disable_workload_identity_federation?: boolean | null; + /** Anthropic Federation Rule Id */ + anthropic_federation_rule_id?: string | null; + /** Anthropic Identity Source */ + anthropic_identity_source?: string | null; + /** Anthropic Identity Token */ + anthropic_identity_token?: string | null; + /** Anthropic Identity Token File */ + anthropic_identity_token_file?: string | null; + /** Anthropic Issuer Audience */ + anthropic_issuer_audience?: string | null; + /** Anthropic Issuer Signing Key Ref */ + anthropic_issuer_signing_key_ref?: string | null; + /** Anthropic Issuer Subject */ + anthropic_issuer_subject?: string | null; + /** Anthropic Issuer Ttl Seconds */ + anthropic_issuer_ttl_seconds?: number | null; + /** Anthropic Issuer Url */ + anthropic_issuer_url?: string | null; + /** Anthropic Keycloak Auth Method */ + anthropic_keycloak_auth_method?: string | null; + /** Anthropic Keycloak Client Id */ + anthropic_keycloak_client_id?: string | null; + /** Anthropic Keycloak Client Secret Ref */ + anthropic_keycloak_client_secret_ref?: string | null; + /** Anthropic Keycloak Scope */ + anthropic_keycloak_scope?: string | null; + /** Anthropic Keycloak Token Url */ + anthropic_keycloak_token_url?: string | null; + /** Anthropic Organization Id */ + anthropic_organization_id?: string | null; + /** Anthropic Service Account Id */ + anthropic_service_account_id?: string | null; + /** Anthropic Workspace Id */ + anthropic_workspace_id?: string | null; /** Api Base */ api_base?: string | null; /** Api Key */ @@ -39600,6 +39717,12 @@ export interface components { ocr_cost_per_credit?: number | null; /** Ocr Cost Per Page */ ocr_cost_per_page?: number | null; + /** Openai Identity Provider Id */ + openai_identity_provider_id?: string | null; + /** Openai Identity Token File */ + openai_identity_token_file?: string | null; + /** Openai Service Account Id */ + openai_service_account_id?: string | null; /** Organization */ organization?: string | null; /** Otpm */ @@ -44818,7 +44941,7 @@ export interface operations { [name: string]: unknown; }; content: { - "application/json": components["schemas"]["CredentialItem"]; + "application/json": components["schemas"]["CredentialItem-Output"]; }; }; /** @description Validation Error */ @@ -44850,7 +44973,7 @@ export interface operations { [name: string]: unknown; }; content: { - "application/json": components["schemas"]["CredentialItem"]; + "application/json": components["schemas"]["CredentialItem-Output"]; }; }; /** @description Validation Error */ @@ -44960,7 +45083,7 @@ export interface operations { }; requestBody: { content: { - "application/json": components["schemas"]["CredentialItem"]; + "application/json": components["schemas"]["CredentialItem-Input"]; }; }; responses: { @@ -44984,6 +45107,38 @@ export interface operations { }; }; }; + get_credential_internal_issuer_jwks_credentials__credential_name__jwks_get: { + parameters: { + query?: never; + header?: never; + path: { + /** @description The credential name, percent-decoded; may contain slashes */ + credential_name: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; cursor_chat_completions_cursor_chat_completions_post: { parameters: { query?: never;