From f06fdab04bd5001f76604636d4edf0d9a5766d11 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 17:11:16 +0000 Subject: [PATCH 001/107] fix(bedrock): stop leaking thinking/reasoning_effort into additionalModelRequestFields for DeepSeek --- .../bedrock/chat/converse_transformation.py | 29 +++++++---- .../chat/test_converse_transformation.py | 52 +++++++++++++++++++ 2 files changed, 70 insertions(+), 11 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index c38b3593465..08f44d5b1b1 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -493,6 +493,23 @@ class AmazonConverseConfig(BaseConfig): ) thinking["budget_tokens"] = BEDROCK_MIN_THINKING_BUDGET_TOKENS + def _model_accepts_anthropic_thinking_param(self, model: str, base_model: str) -> bool: + """Whether the model accepts the Anthropic-shaped ``thinking`` / ``reasoning_effort`` request field. + + The Converse mapping serializes ``thinking`` into ``additionalModelRequestFields`` in Anthropic's + shape. Only Anthropic Claude reasoning models accept that field; DeepSeek models reason natively + and reject it (a 400 when it leaks through), even though they advertise ``supports_reasoning``. + """ + if "deepseek" in model or "deepseek" in base_model: + return False + return ( + "claude-3-7" in model + or "claude-sonnet-4" in model + or "claude-opus-4" in model + or supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider) + or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider) + ) + def get_supported_openai_params(self, model: str) -> List[str]: from litellm.utils import supports_function_calling @@ -553,17 +570,7 @@ class AmazonConverseConfig(BaseConfig): # Nova 2 models support reasoning_effort (transformed to reasoningConfig) # These models use a different reasoning structure than Anthropic's thinking parameter supported_params.append("reasoning_effort") - elif ( - "claude-3-7" in model - or "claude-sonnet-4" in model - or "claude-opus-4" in model - or "deepseek.r1" in model - or supports_reasoning( - model=model, - custom_llm_provider=self.custom_llm_provider, - ) - or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider) - ): + elif self._model_accepts_anthropic_thinking_param(model=model, base_model=base_model): supported_params.append("thinking") supported_params.append("reasoning_effort") diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index fc12ead36a1..f7e2716bd19 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -553,6 +553,58 @@ def test_get_supported_openai_params(): assert "reasoning_effort" in supported_params +@pytest.mark.parametrize( + "model", + [ + "bedrock/us.deepseek.r1-v1:0", + "bedrock/converse/us.deepseek.r1-v1:0", + "bedrock/deepseek.v3-v1:0", + "bedrock/deepseek.v3.2", + ], +) +def test_bedrock_deepseek_does_not_advertise_thinking(model): + """DeepSeek reasons natively on Bedrock and rejects the Anthropic-shaped + `thinking`/`reasoning_effort` field, so it must not be advertised as supported + (otherwise it leaks into additionalModelRequestFields and Bedrock 400s).""" + config = AmazonConverseConfig() + supported_params = config.get_supported_openai_params(model=model) + assert "thinking" not in supported_params + assert "reasoning_effort" not in supported_params + + +def test_bedrock_deepseek_r1_thinking_raises_without_drop_params(): + """Passing `thinking` to Bedrock DeepSeek R1 must fail client-side with a clear + UnsupportedParamsError instead of leaking through and hitting a Bedrock 400.""" + with pytest.raises(litellm.UnsupportedParamsError): + litellm.utils.get_optional_params( + model="us.deepseek.r1-v1:0", + custom_llm_provider="bedrock", + thinking={"type": "enabled", "budget_tokens": 1024}, + ) + + +def test_bedrock_deepseek_r1_thinking_dropped_does_not_leak_into_request(): + """With drop_params, `thinking` is dropped rather than forwarded into + additionalModelRequestFields for Bedrock DeepSeek R1.""" + optional_params = litellm.utils.get_optional_params( + model="us.deepseek.r1-v1:0", + custom_llm_provider="bedrock", + thinking={"type": "enabled", "budget_tokens": 1024}, + drop_params=True, + ) + assert "thinking" not in optional_params + + config = AmazonConverseConfig() + request = config._transform_request( + model="bedrock/converse/us.deepseek.r1-v1:0", + messages=[{"role": "user", "content": "Say hi in one word."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + assert "thinking" not in request.get("additionalModelRequestFields", {}) + + def test_get_supported_openai_params_bedrock_converse(): """ Test that all documented bedrock converse models have the same set of supported openai params when using From f4ad6b0f0bd13bad1fd8ee68e6b5463730b50e0e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 17:20:03 +0000 Subject: [PATCH 002/107] fix(bedrock): drop thinking/reasoning_effort in converse mapping for DeepSeek --- .../bedrock/chat/converse_transformation.py | 15 +++++++++- .../chat/test_converse_transformation.py | 28 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 08f44d5b1b1..d7ee15cedd7 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -857,6 +857,9 @@ class AmazonConverseConfig(BaseConfig): drop_params: bool, ) -> dict: is_thinking_enabled = self.is_thinking_enabled(non_default_params) + accepts_thinking_param = self._model_accepts_anthropic_thinking_param( + model=model, base_model=BedrockModelInfo.get_base_model(model) + ) for param, value in non_default_params.items(): if param == "response_format" and isinstance(value, dict): @@ -906,7 +909,12 @@ class AmazonConverseConfig(BaseConfig): optional_params["_parallel_tool_use_config"] = { "tool_choice": {"disable_parallel_tool_use": disable_parallel} } - if param == "thinking": + if param == "thinking" and not accepts_thinking_param: + verbose_logger.debug( + "Dropping unsupported `thinking` param for Bedrock model=%s; it reasons natively and rejects it.", + model, + ) + elif param == "thinking": if ( isinstance(value, dict) and value.get("type") == "adaptive" @@ -929,6 +937,11 @@ class AmazonConverseConfig(BaseConfig): litellm.verbose_logger.warning(DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, model) else: optional_params["thinking"] = value + elif param == "reasoning_effort" and isinstance(value, str) and not accepts_thinking_param: + verbose_logger.debug( + "Dropping unsupported `reasoning_effort` param for Bedrock model=%s; it reasons natively and rejects it.", + model, + ) elif param == "reasoning_effort" and isinstance(value, str): self._handle_reasoning_effort_parameter( model=model, reasoning_effort=value, optional_params=optional_params diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index f7e2716bd19..93455f00143 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -605,6 +605,34 @@ def test_bedrock_deepseek_r1_thinking_dropped_does_not_leak_into_request(): assert "thinking" not in request.get("additionalModelRequestFields", {}) +@pytest.mark.parametrize("param", ["thinking", "reasoning_effort"]) +def test_bedrock_deepseek_r1_reasoning_params_not_forwarded_by_map(param): + """Even when map_openai_params is called directly (bypassing the supported-params + gate), DeepSeek R1 must not forward the Anthropic-shaped thinking/reasoning_effort + into additionalModelRequestFields, since Bedrock rejects it with a 400.""" + config = AmazonConverseConfig() + model = "bedrock/converse/us.deepseek.r1-v1:0" + value = {"type": "enabled", "budget_tokens": 1024} if param == "thinking" else "high" + + optional_params = config.map_openai_params( + non_default_params={param: value, "max_tokens": 100}, + optional_params={}, + model=model, + drop_params=False, + ) + assert "thinking" not in optional_params + assert "reasoning_effort" not in optional_params + + request = config._transform_request( + model=model, + messages=[{"role": "user", "content": "Say hi in one word."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + assert request.get("additionalModelRequestFields") is None + + def test_get_supported_openai_params_bedrock_converse(): """ Test that all documented bedrock converse models have the same set of supported openai params when using From ab1f739188cf4096ae7882a6b7848e702480ed02 Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 23 Jul 2026 03:29:45 +0000 Subject: [PATCH 003/107] fix(bedrock): only drop reasoning params for DeepSeek, not all non-allowlisted models The DeepSeek leak fix gated the drop of thinking/reasoning_effort on a positive allowlist (_model_accepts_anthropic_thinking_param). Its negation dropped the params for any model the allowlist and ARN introspection missed, so Claude behind application-inference-profile ARNs, gpt-oss-safeguard (absent from the cost map so supports_reasoning is False), and Nova 2 custom-import ARNs all silently lost reasoning in prod. Gate the drop on a narrow denylist instead: only DeepSeek reasons natively and 400s on the request field, so drop only there and leave every other model (including opaque ARNs and gpt-oss/Nova 2 which route reasoning_effort through their own shapes) untouched. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/chat/converse_transformation.py | 18 ++++++-- .../chat/test_converse_transformation.py | 46 +++++++++++++++++++ 2 files changed, 61 insertions(+), 3 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index d7ee15cedd7..617ce423d51 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -510,6 +510,18 @@ class AmazonConverseConfig(BaseConfig): or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider) ) + def _model_reasons_natively_and_rejects_request_param(self, model: str, base_model: str) -> bool: + """Whether the model reasons natively and rejects any reasoning request field on Converse. + + The Converse mapping serializes ``thinking`` into ``additionalModelRequestFields`` in Anthropic's + shape and ``reasoning_effort`` into a provider-specific shape. DeepSeek reasons on its own and + returns a 400 when either field is sent, even though it advertises ``supports_reasoning``, so both + must be dropped for it. Every other model either accepts one of those shapes (Claude ``thinking``, + gpt-oss / Nova 2 ``reasoning_effort``) or is an opaque ARN we can't introspect, so we leave those + untouched rather than silently degrading reasoning. + """ + return "deepseek" in model or "deepseek" in base_model + def get_supported_openai_params(self, model: str) -> List[str]: from litellm.utils import supports_function_calling @@ -857,7 +869,7 @@ class AmazonConverseConfig(BaseConfig): drop_params: bool, ) -> dict: is_thinking_enabled = self.is_thinking_enabled(non_default_params) - accepts_thinking_param = self._model_accepts_anthropic_thinking_param( + drop_reasoning_request_param = self._model_reasons_natively_and_rejects_request_param( model=model, base_model=BedrockModelInfo.get_base_model(model) ) @@ -909,7 +921,7 @@ class AmazonConverseConfig(BaseConfig): optional_params["_parallel_tool_use_config"] = { "tool_choice": {"disable_parallel_tool_use": disable_parallel} } - if param == "thinking" and not accepts_thinking_param: + if param == "thinking" and drop_reasoning_request_param: verbose_logger.debug( "Dropping unsupported `thinking` param for Bedrock model=%s; it reasons natively and rejects it.", model, @@ -937,7 +949,7 @@ class AmazonConverseConfig(BaseConfig): litellm.verbose_logger.warning(DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, model) else: optional_params["thinking"] = value - elif param == "reasoning_effort" and isinstance(value, str) and not accepts_thinking_param: + elif param == "reasoning_effort" and isinstance(value, str) and drop_reasoning_request_param: verbose_logger.debug( "Dropping unsupported `reasoning_effort` param for Bedrock model=%s; it reasons natively and rejects it.", model, diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 93455f00143..87a172f2409 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -633,6 +633,52 @@ def test_bedrock_deepseek_r1_reasoning_params_not_forwarded_by_map(param): assert request.get("additionalModelRequestFields") is None +@pytest.mark.parametrize( + "model, param, value, kept_key", + [ + ( + "bedrock/us.anthropic.claude-opus-4-20250514-v1:0", + "thinking", + {"type": "enabled", "budget_tokens": 1024}, + "thinking", + ), + ( + "bedrock/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + "thinking", + {"type": "enabled", "budget_tokens": 1024}, + "thinking", + ), + ( + "bedrock/openai.gpt-oss-safeguard-20b-1:0", + "reasoning_effort", + "high", + "reasoning_effort", + ), + ( + "bedrock/us.amazon.nova-2-lite-v1:0", + "reasoning_effort", + "high", + "reasoningConfig", + ), + ], +) +def test_bedrock_non_deepseek_reasoning_params_preserved(model, param, value, kept_key): + """The DeepSeek leak fix must only drop reasoning request params for DeepSeek. + + Claude behind an application-inference-profile ARN, gpt-oss-safeguard (absent from the + cost map so `supports_reasoning` is False), and Nova 2 all reason via a request param and + must keep it. Regression guard against gating the drop on a positive allowlist, which + silently degraded reasoning for anything the allowlist/ARN introspection missed.""" + config = AmazonConverseConfig() + optional_params = config.map_openai_params( + non_default_params={param: value, "max_tokens": 100}, + optional_params={}, + model=model, + drop_params=False, + ) + assert kept_key in optional_params + + def test_get_supported_openai_params_bedrock_converse(): """ Test that all documented bedrock converse models have the same set of supported openai params when using From 512e730f2fbbb4a832a42ca65b8b29c41c77508d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 21:30:48 -0700 Subject: [PATCH 004/107] fix(azure_ai): add passthrough config so router-model relays reach the deployment's own endpoint Every /azure_ai// relay failed with HTTP 500 because azure_ai had no passthrough config. The new AzureAIPassthroughConfig strips the router-model prefix from the relayed path, forwards to the deployment's api_base with its own credential (api-key on Foundry and Azure OpenAI hosts, Bearer elsewhere, Entra as the fallback), and delegates chat/completions cost logging to the Azure passthrough config. The router's provider inference now receives the deployment's api_base so an OpenAI-family model on a Foundry resource stays azure_ai instead of flipping to azure through the AZURE_AI_API_BASE env var. --- basedpyright-code-budget.json | 6 +- litellm/llms/azure_ai/chat/transformation.py | 12 +- litellm/llms/azure_ai/common_utils.py | 7 + .../azure_ai/passthrough/transformation.py | 106 +++++++++++ .../base_llm/passthrough/transformation.py | 3 +- litellm/router.py | 2 + litellm/utils.py | 6 + ruff-strict-budget.json | 2 +- ...est_azure_ai_passthrough_transformation.py | 175 ++++++++++++++++++ .../passthrough/test_passthrough_main.py | 115 ++++++++++++ type-discipline-budget.json | 2 +- 11 files changed, 423 insertions(+), 13 deletions(-) create mode 100644 litellm/llms/azure_ai/passthrough/transformation.py create mode 100644 tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 669107bb5b1..f31d620557d 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15284 + "limit": 15283 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,10 +105,10 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38309 + "limit": 38308 }, "reportUnknownParameterType": { - "limit": 19621 + "limit": 19620 }, "reportUnknownVariableType": { "limit": 29844 diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index f2d405e9a17..4c0a2e1f546 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -2,7 +2,6 @@ import copy import enum import re from typing import TYPE_CHECKING, Final, cast -from urllib.parse import urlparse import httpx from httpx import Response @@ -15,7 +14,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, ) from litellm.llms.azure.common_utils import BaseAzureLLM -from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base +from litellm.llms.azure_ai.common_utils import ( + api_key_header_for_base, + is_foundry_model_inference_base, +) from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error from litellm.llms.openai.openai import OpenAIConfig @@ -99,11 +101,7 @@ class AzureAIStudioConfig(OpenAIConfig): """ Returns True if the request should use `api-key` header for authentication. """ - parsed_url: Final = urlparse(api_base) - host: Final = parsed_url.hostname - if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")): - return True - return False + return api_key_header_for_base(api_base) == "api-key" def get_complete_url( self, diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index aa34bab5b2e..0665b3f64c5 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -19,6 +19,13 @@ def is_foundry_model_inference_base(api_base: str) -> bool: return "/openai/deployments" not in parsed.path +def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader: + host: Final = urlparse(api_base).hostname if api_base else None + if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")): + return "api-key" + return "Authorization" + + def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None: """ Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment. diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py new file mode 100644 index 00000000000..afe0b814524 --- /dev/null +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + api_key_header_for_base, + get_azure_ai_auth_headers, +) +from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +from litellm.types.llms.openai import AllMessageValues + +if TYPE_CHECKING: + from httpx import URL, Response + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CostResponseTypes + + +def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str: + path: Final = endpoint.lstrip("/") + for model_name in model_names: + if not model_name: + continue + if path == model_name: + return "" + if path.startswith(f"{model_name}/"): + return path[len(model_name) + 1 :] + return path + + +class PassthroughMetadata(BaseModel): + model_config = ConfigDict(extra="ignore") + + model_group: str = "" + + +def model_group_from(litellm_params: Mapping[str, object]) -> str: + try: + return PassthroughMetadata.model_validate(litellm_params.get("litellm_metadata")).model_group + except ValidationError: + return "" + + +class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): + def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool: + return request_data.get("stream") is True + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + endpoint: str, + request_query_params: Mapping[str, object] | None, + litellm_params: Mapping[str, object], + ) -> tuple[URL, str]: + base_target_url: Final = self.get_api_base(api_base) + if base_target_url is None: + raise ValueError("Azure AI api base not found: set `api_base` on the deployment or AZURE_AI_API_BASE") + + native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params))) + return ( + self.format_url(native_endpoint, base_target_url, request_query_params), + base_target_url, + ) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: base class contract returns dict for httpx + auth_headers: Final = get_azure_ai_auth_headers( + api_key=api_key, + litellm_params=litellm_params, + api_key_header=api_key_header_for_base(api_base), + ) + return {**headers, **auth_headers} # mutable-ok: base class contract returns dict for httpx + + def logging_non_streaming_response( + self, + model: str, + custom_llm_provider: str, + httpx_response: Response, + request_data: Mapping[str, object], + logging_obj: Logging, + endpoint: str, + ) -> CostResponseTypes | None: + from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig + + return AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict + model=model, + custom_llm_provider=custom_llm_provider, + httpx_response=httpx_response, + request_data=dict(request_data), # mutable-ok: AzurePassthroughConfig wants a dict + logging_obj=logging_obj, + endpoint=endpoint, + ) diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index adbf2e126fb..34533e4c603 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -1,4 +1,5 @@ from abc import abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, Union from ..base_utils import BaseLLMModelInfo @@ -23,7 +24,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): self, endpoint: str, base_target_url: str, - request_query_params: dict | None, + request_query_params: Mapping[str, object] | None, ) -> "URL": """ Helper function to add query params to the url diff --git a/litellm/router.py b/litellm/router.py index 6c7611c6236..b84f790287f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5051,6 +5051,7 @@ class Router: _, inferred_custom_llm_provider, _, _ = get_llm_provider( model=data["model"], custom_llm_provider=custom_llm_provider, + api_base=data.get("api_base"), ) custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider except Exception: @@ -5570,6 +5571,7 @@ class Router: _, inferred_custom_llm_provider, _, _ = get_llm_provider( model=data["model"], custom_llm_provider=custom_llm_provider, + api_base=data.get("api_base"), ) custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider except Exception: diff --git a/litellm/utils.py b/litellm/utils.py index 9d20d32d147..0f055308e3d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8831,6 +8831,12 @@ class ProviderConfigManager: ) return AzurePassthroughConfig() + elif LlmProviders.AZURE_AI == provider: + from litellm.llms.azure_ai.passthrough.transformation import ( + AzureAIPassthroughConfig, + ) + + return AzureAIPassthroughConfig() elif LlmProviders.GIGACHAT == provider: from litellm.llms.gigachat.passthrough.transformation import ( GigaChatPassthroughConfig, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 70408ea022b..a011230c2d3 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 119 + "limit": 118 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py new file mode 100644 index 00000000000..e1abbe1a7e1 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -0,0 +1,175 @@ +import json +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.llms.azure_ai.passthrough.transformation import AzureAIPassthroughConfig +from litellm.types.utils import LlmProviders, ModelResponse +from litellm.utils import ProviderConfigManager + +FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" + + +@pytest.fixture(autouse=True) +def clear_azure_ai_env(monkeypatch): + for env_var in ("AZURE_AI_API_BASE", "AZURE_AI_API_KEY", "AZURE_AD_TOKEN", "AZURE_API_KEY"): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.setattr(litellm, "api_key", None) + + +def test_provider_config_manager_resolves_azure_ai_passthrough_config(): + config = ProviderConfigManager.get_provider_passthrough_config(model="Cohere-parse-v5", provider=LlmProviders.AZURE_AI) + + assert isinstance(config, AzureAIPassthroughConfig) + + +def test_router_model_prefix_is_stripped_and_native_path_kept_verbatim(): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key=None, + model="Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + assert base == FOUNDRY_BASE + + +def test_model_group_prefix_is_stripped_when_router_metadata_names_it(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key=None, + model="Cohere-parse-v5", + endpoint="/parse-alias/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={"litellm_metadata": {"model_group": "parse-alias"}}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + + +def test_model_inside_the_path_stays_and_query_params_are_forwarded(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/", + api_key=None, + model="gpt-5.4-mini", + endpoint="openai/deployments/gpt-5.4-mini/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/openai/deployments/gpt-5.4-mini/chat/completions?api-version=2024-10-21" + + +def test_missing_api_base_raises_instead_of_building_a_relative_url(): + with pytest.raises(ValueError, match="AZURE_AI_API_BASE"): + AzureAIPassthroughConfig().get_complete_url( + api_base=None, + api_key=None, + model="Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={}, + ) + + +def _auth_headers(api_key: str | None, api_base: str, litellm_params: dict | None = None) -> dict: + return AzureAIPassthroughConfig().validate_environment( + headers={"content-type": "application/json"}, + model="Cohere-parse-v5", + messages=[], + optional_params={}, + litellm_params=litellm_params or {}, + api_key=api_key, + api_base=api_base, + ) + + +def test_foundry_host_gets_the_api_key_header(): + headers = _auth_headers(api_key="deployment-key", api_base=FOUNDRY_BASE) + + assert headers == {"content-type": "application/json", "api-key": "deployment-key"} + + +def test_serverless_host_gets_a_bearer_token(): + headers = _auth_headers(api_key="deployment-key", api_base="https://cohere-parse.eastus.models.ai.azure.com") + + assert headers["Authorization"] == "Bearer deployment-key" + assert "api-key" not in headers + + +def test_entra_token_is_used_when_the_deployment_has_no_api_key(): + headers = _auth_headers(api_key=None, api_base=FOUNDRY_BASE, litellm_params={"azure_ad_token": "entra-token"}) + + assert headers["Authorization"] == "Bearer entra-token" + + +def test_no_credentials_at_all_raises(): + with pytest.raises(ValueError, match="Missing Azure AI credentials"): + _auth_headers(api_key=None, api_base=FOUNDRY_BASE) + + +@pytest.mark.parametrize( + "request_data, expected", + [({"stream": True}, True), ({"stream": False}, False), ({}, False)], +) +def test_is_streaming_request_reads_the_stream_flag(request_data, expected): + assert AzureAIPassthroughConfig().is_streaming_request(endpoint="models/chat/completions", request_data=request_data) is expected + + +def _chat_completion_response() -> httpx.Response: + body = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-5.4-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}, + } + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request("POST", f"{FOUNDRY_BASE}/models/chat/completions"), + ) + + +def test_chat_completions_relay_yields_a_model_response_for_cost_tracking(): + result = AzureAIPassthroughConfig().logging_non_streaming_response( + model="gpt-5.4-mini", + custom_llm_provider="azure_ai", + httpx_response=_chat_completion_response(), + request_data={"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "hi"}]}, + logging_obj=MagicMock(), + endpoint="models/chat/completions", + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "hi" + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 8 + + +def test_non_chat_relay_yields_no_cost_response(): + parse_response = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"id":"parse-1","pages":[]}', + request=httpx.Request("POST", f"{FOUNDRY_BASE}/providers/cohere/v2/parse"), + ) + + result = AzureAIPassthroughConfig().logging_non_streaming_response( + model="Cohere-parse-v5", + custom_llm_provider="azure_ai", + httpx_response=parse_response, + request_data={"model": "Cohere-parse-v5"}, + logging_obj=MagicMock(), + endpoint="providers/cohere/v2/parse", + ) + + assert result is None diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index 1950c37a12e..546cff18b5d 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -873,3 +873,118 @@ def test_llm_passthrough_route_propagates_allm_passthrough_route_to_logging_obj( assert captured_litellm_params.get("allm_passthrough_route") is True assert LitellmLogging._is_sync_litellm_request(captured_litellm_params) is False + + +FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" + + +def _foundry_parse_response() -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"id":"parse-1","pages":[]}', + request=httpx.Request("POST", f"{FOUNDRY_BASE}/providers/cohere/v2/parse"), + ) + + +def test_azure_ai_relay_reaches_the_deployment_with_its_own_credential(): + """ + Regression for LIT-7022: azure_ai had no passthrough config, so every + /azure_ai// relay raised "Provider azure_ai not found" + before a request was built. + """ + client = HTTPHandler() + + with patch.object(client.client, "send", return_value=_foundry_parse_response()) as mock_send: + response = llm_passthrough_route( + model="azure_ai/Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + method="POST", + custom_llm_provider="azure_ai", + api_base=FOUNDRY_BASE, + api_key="deployment-key", + json={"model": "Cohere-parse-v5", "document": {"type": "image_url", "image_url": "https://x/y.png"}}, + client=client, + litellm_logging_obj=MagicMock(), + ) + + sent = mock_send.call_args.kwargs["request"] + assert str(sent.url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + assert sent.headers["api-key"] == "deployment-key" + assert json.loads(sent.content)["model"] == "Cohere-parse-v5" + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_router_relays_azure_ai_model_through_the_deployment_api_base(): + router = litellm.Router( + model_list=[ + { + "model_name": "foundry-parse", + "litellm_params": { + "model": "azure_ai/Cohere-parse-v5", + "api_base": FOUNDRY_BASE, + "api_key": "deployment-key", + }, + } + ] + ) + async_client = AsyncHTTPHandler() + + with patch.object(async_client.client, "send", AsyncMock(return_value=_foundry_parse_response())) as mock_send: + response = await router.allm_passthrough_route( + model="foundry-parse", + method="POST", + endpoint="foundry-parse/providers/cohere/v2/parse", + json={"model": "foundry-parse", "document": {"type": "image_url", "image_url": "https://x/y.png"}}, + client=async_client, + ) + + sent = mock_send.call_args.kwargs["request"] + assert str(sent.url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + assert sent.headers["api-key"] == "deployment-key" + assert json.loads(sent.content)["model"] == "Cohere-parse-v5" + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_router_relays_an_openai_model_on_a_foundry_base_as_azure_ai(monkeypatch): + monkeypatch.setenv("AZURE_AI_API_BASE", "https://unrelated.openai.azure.com") + router = litellm.Router( + model_list=[ + { + "model_name": "foundry-gpt", + "litellm_params": { + "model": "azure_ai/gpt-5.4-mini", + "api_base": FOUNDRY_BASE, + "api_key": "deployment-key", + }, + } + ] + ) + async_client = AsyncHTTPHandler() + upstream = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=( + b'{"id":"chatcmpl-1","object":"chat.completion","model":"gpt-5.4-mini",' + b'"choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"hi"}}],' + b'"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}' + ), + request=httpx.Request("POST", f"{FOUNDRY_BASE}/models/chat/completions"), + ) + + with patch.object(async_client.client, "send", AsyncMock(return_value=upstream)) as mock_send: + await router.allm_passthrough_route( + model="foundry-gpt", + method="POST", + endpoint="foundry-gpt/models/chat/completions", + request_query_params={"api-version": "2024-05-01-preview"}, + json={"model": "foundry-gpt", "messages": [{"role": "user", "content": "hi"}]}, + client=async_client, + ) + + sent = mock_send.call_args.kwargs["request"] + assert str(sent.url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + assert sent.headers["api-key"] == "deployment-key" + assert json.loads(sent.content)["model"] == "gpt-5.4-mini" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 3d01c08e8eb..7d12721175f 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22326 + "limit": 22325 }, "LIT002": { "limit": 26748 From a276690ce200bebe3547c110ac9a830b32052173 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 23:32:26 -0700 Subject: [PATCH 005/107] fix(router): keep a model's own provider prefix for generic SDK calls Generic passthrough calls inferred the provider from the bare model name, so an azure_ai/gpt-* deployment on an Azure OpenAI host flipped to azure and get_llm_provider re-prefixed the deployment name into azure_ai/gpt-5.4-mini, a 404 DeploymentNotFound. provider_for_generic_call takes the declared custom_llm_provider first, then the model's own prefix, and only infers for unprefixed models --- litellm/router.py | 24 ++--------- litellm/router_utils/common_utils.py | 27 ++++++++++++ .../test_router_utils_common_utils.py | 18 ++++++++ tests/test_litellm/test_router.py | 43 +++++++++++++++++++ 4 files changed, 91 insertions(+), 21 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b84f790287f..9b7a7ee7ce8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -135,6 +135,7 @@ from litellm.router_utils.common_utils import ( _is_proxy_admin_request, filter_team_based_models, filter_web_search_deployments, + provider_for_generic_call, resolve_model_group_alias, truncate_fallback_error_detail, warn_on_provider_credential_mismatch, @@ -5045,17 +5046,7 @@ class Router: kwargs=kwargs, model=model, model_name=model_name ) - # Get custom_llm_provider from deployment params - try: - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - api_base=data.get("api_base"), - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - except Exception: - custom_llm_provider = None + custom_llm_provider: Final = provider_for_generic_call(data) response_kwargs: Final = { **data, @@ -5566,16 +5557,7 @@ class Router: # Perform pre-call checks for routing strategy self.routing_strategy_pre_call_checks(deployment=deployment) - try: - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - api_base=data.get("api_base"), - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - except Exception: - custom_llm_provider = None + custom_llm_provider: Final = provider_for_generic_call(data) response: Final = original_function( **{ diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 280a7defcf8..ec3ffa1247e 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject +import litellm from litellm._logging import verbose_logger, verbose_router_logger from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError @@ -244,6 +245,32 @@ PROVIDER_SCOPED_CREDENTIAL_PARAMS: Final[Mapping[str, frozenset[str]]] = Mapping ) +def provider_for_generic_call(litellm_params: Mapping[str, object]) -> str | None: + """ + The provider the router hands a deployment's generic SDK call, or None when it cannot be resolved. + + A model that carries its own provider prefix keeps that prefix even where get_llm_provider + would resolve it to a sibling provider (azure_ai/ on an Azure OpenAI host + resolves to azure): the SDK call still receives the prefixed model, and an explicit provider + that contradicts the prefix makes get_llm_provider re-prefix it into a deployment name that + does not exist upstream. + """ + declared: Final = litellm_params.get("custom_llm_provider") + if isinstance(declared, str) and declared: + return declared + model: Final = litellm_params.get("model") + if not isinstance(model, str) or not model: + return None + prefix: Final = model.split("/", 1)[0] + if "/" in model and prefix in litellm.provider_list: + return prefix + try: + _, inferred, _, _ = get_llm_provider(model=model) + except BadRequestError: + return None + return inferred + + def warn_on_provider_credential_mismatch(model_name: str, litellm_params: Mapping[str, object]) -> str | None: """ Warn when a deployment carries one provider's credentials but resolves to another. diff --git a/tests/test_litellm/router_utils/test_router_utils_common_utils.py b/tests/test_litellm/router_utils/test_router_utils_common_utils.py index 30f658d7ea2..ac18b4889dd 100644 --- a/tests/test_litellm/router_utils/test_router_utils_common_utils.py +++ b/tests/test_litellm/router_utils/test_router_utils_common_utils.py @@ -12,6 +12,7 @@ from litellm.router_utils.common_utils import ( add_model_file_id_mappings, filter_team_based_models, filter_web_search_deployments, + provider_for_generic_call, resolve_model_group_alias, truncate_fallback_error_detail, PROVIDER_SCOPED_CREDENTIAL_PARAMS, @@ -756,3 +757,20 @@ class TestWarnOnProviderCredentialMismatch: ) is None ) + + +@pytest.mark.parametrize( + ("litellm_params", "expected"), + [ + ({"model": "azure_ai/gpt-5.4-mini", "custom_llm_provider": "azure"}, "azure"), + ({"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.openai.azure.com"}, "azure_ai"), + ({"model": "cohere/command-r"}, "cohere"), + ({"model": "gpt-5.4-mini"}, "openai"), + ({"model": "no-provider-knows-this-model"}, None), + ({"api_base": "https://my-resource.openai.azure.com"}, None), + ], + ids=["declared_wins", "prefix_beats_host_flip", "prefix_beats_cohere_chat_flip", "unprefixed_inferred", "unknown", "no_model"], +) +def test_provider_for_generic_call(litellm_params, expected, monkeypatch): + monkeypatch.setenv("AZURE_AI_API_BASE", "https://unrelated.openai.azure.com") + assert provider_for_generic_call(litellm_params) == expected diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 31eb46f1458..b220b23c338 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -12938,3 +12938,46 @@ async def test_router_retry_policy_controls_upstream_attempt_count( await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) assert upstream.call_count == expected_upstream_calls + + +@pytest.mark.asyncio +async def test_generic_call_keeps_the_deployment_name_of_an_azure_ai_model_on_an_azure_openai_host(monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + { + "model_name": "aoai-gpt", + "litellm_params": { + "model": "azure_ai/gpt-5.4-mini", + "api_base": "https://my-resource.openai.azure.com", + "api_key": "deployment-key", + }, + } + ] + ) + + with respx.mock(assert_all_called=True) as respx_mock: + upstream = respx_mock.post(host="my-resource.openai.azure.com", path__regex=r"^/openai/.*responses$").mock( + return_value=httpx.Response( + 200, + json={ + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + }, + ) + ) + await router.aresponses(model="aoai-gpt", input="hi") + + assert json.loads(upstream.calls.last.request.content)["model"] == "gpt-5.4-mini" From cd25eb918995ef115aa67e8fd39337e4786bdcd2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 23:32:27 -0700 Subject: [PATCH 006/107] fix(azure_ai): cost streaming relays and return upstream errors from router relays Streaming chat relays on Azure and azure_ai deployments rebuild the response from the SSE chunks through the OpenAI passthrough assembler, so the spend log carries usage. The router relays keep the JSON body when the Content-Type carries a charset, return the upstream status and body instead of a 500 when the deployment rejects the call, and fall back to the caller's api-version when the deployment sets none. Lint budgets ratcheted to the measured totals --- basedpyright-code-budget.json | 8 +- .../llms/azure/passthrough/transformation.py | 24 ++++- .../azure_ai/passthrough/transformation.py | 18 ++++ .../proxy/common_utils/http_parsing_utils.py | 4 +- .../llm_passthrough_endpoints.py | 53 ++++++----- .../openai_passthrough_logging_handler.py | 3 +- ruff-strict-budget.json | 4 +- .../test_azure_passthrough_transformation.py | 69 +++++++++++++++ ...est_azure_ai_passthrough_transformation.py | 21 +++++ .../test_llm_pass_through_endpoints.py | 87 +++++++++++++++++++ type-discipline-budget.json | 4 +- 11 files changed, 260 insertions(+), 35 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index b446c79f49b..ef1bd675c5e 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15279 + "limit": 15278 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38281 + "limit": 38276 }, "reportUnknownParameterType": { - "limit": 19582 + "limit": 19581 }, "reportUnknownVariableType": { - "limit": 29829 + "limit": 29825 }, "reportUnnecessaryCast": { "limit": 110 diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 898852e645f..7654958fc0b 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from typing import TYPE_CHECKING, Final, Optional import httpx @@ -43,7 +44,7 @@ class AzurePassthroughConfig(BasePassthroughConfig): api_base=base_target_url, litellm_params=litellm_params, route=endpoint, - default_api_version=litellm_params.get("api_version"), + default_api_version=request_query_params.get("api-version") if request_query_params else None, ) return ( httpx.URL(complete_url), @@ -116,3 +117,24 @@ class AzurePassthroughConfig(BasePassthroughConfig): ) return litellm_model_response + + def handle_logging_collected_chunks( + self, + all_chunks: Sequence[str], + litellm_logging_obj: Logging, + model: str, + custom_llm_provider: str, + endpoint: str, + ) -> Optional["CostResponseTypes"]: + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, + ) + + if "chat/completions" not in endpoint: + return None + + return OpenAIPassthroughLoggingHandler()._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only OpenAI SSE-to-ModelResponse assembler; reimplementing it would fork the parser + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + ) diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index afe0b814524..5948c865bad 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -104,3 +104,21 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): logging_obj=logging_obj, endpoint=endpoint, ) + + def handle_logging_collected_chunks( + self, + all_chunks: Sequence[str], + litellm_logging_obj: Logging, + model: str, + custom_llm_provider: str, + endpoint: str, + ) -> CostResponseTypes | None: + from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig + + return AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + custom_llm_provider=custom_llm_provider, + endpoint=endpoint, + ) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 54a0f18fd63..635cbf3ba9b 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -39,7 +39,7 @@ def _is_form_content_type(content_type: str) -> bool: return _normalize_media_type(content_type) in _FORM_CONTENT_TYPES -def _is_json_content_type(content_type: str) -> bool: +def is_json_content_type(content_type: str) -> bool: """True iff the body should be parsed as JSON.""" return _normalize_media_type(content_type) == "application/json" @@ -406,7 +406,7 @@ async def get_request_body(request: Request) -> dict[str, Any]: """ if request.method == "POST": content_type: Final = request.headers.get("content-type", "") - if _is_json_content_type(content_type): + if is_json_content_type(content_type): return await _read_request_body(request) elif _is_form_content_type(content_type): return await get_form_data(request) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 6b1d6405a6a..cf7b7fdf07a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -49,6 +49,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_set_request_parsed_body, get_form_data, get_request_body, + is_json_content_type, ) from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, @@ -408,7 +409,7 @@ async def vllm_proxy_route( content=None, data=None, files=None, - json=(request_body if request.headers.get("content-type") == "application/json" else None), + json=(request_body if is_json_content_type(request.headers.get("content-type", "")) else None), params=None, headers=None, cookies=None, @@ -1492,6 +1493,14 @@ async def _relay_upstream_bytes(upstream: AsyncGenerator[bytes, bytes]) -> Async await upstream.aclose() +async def _relay_upstream_response(upstream: httpx.Response) -> Response: + return Response( + content=await upstream.aread(), + status_code=upstream.status_code, + headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None), + ) + + async def _relay_azure_router_model( llm_router: litellm.Router, model: str, @@ -1501,30 +1510,28 @@ async def _relay_azure_router_model( is_streaming_request: bool, user_api_key_dict: UserAPIKeyAuth, ) -> Response: - result: Final = await llm_router.allm_passthrough_route( - model=model, - method=request.method, - endpoint=endpoint, - request_query_params=request.query_params, - request_headers=_safe_get_request_headers(request), - stream=is_streaming_request, - content=None, - data=None, - files=None, - json=(request_body if request.headers.get("content-type") == "application/json" else None), - params=None, - headers=None, - cookies=None, - litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), - ) + try: + result: Final = await llm_router.allm_passthrough_route( + model=model, + method=request.method, + endpoint=endpoint, + request_query_params=request.query_params, + request_headers=_safe_get_request_headers(request), + stream=is_streaming_request, + content=None, + data=None, + files=None, + json=(request_body if is_json_content_type(request.headers.get("content-type", "")) else None), + params=None, + headers=None, + cookies=None, + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), + ) + except httpx.HTTPStatusError as upstream_error: + return await _relay_upstream_response(upstream_error.response) if not is_streaming_request: - upstream: Final = cast(httpx.Response, result) - return Response( - content=await upstream.aread(), - status_code=upstream.status_code, - headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None), - ) + return await _relay_upstream_response(cast(httpx.Response, result)) if inspect.isasyncgen(result): sse_headers: Final = {"content-type": "text/event-stream"} diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 5f6489a69ca..a39ea342f0a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -4,6 +4,7 @@ OpenAI Passthrough Logging Handler Handles cost tracking and logging for OpenAI passthrough endpoints, specifically /chat/completions. """ +from collections.abc import Sequence from datetime import datetime from typing import Final from urllib.parse import urlparse @@ -512,7 +513,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): def _build_complete_streaming_response( self, - all_chunks: list[str], + all_chunks: Sequence[str], litellm_logging_obj: LiteLLMLoggingObj, model: str, ) -> ModelResponse | TextCompletionResponse | None: diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 694ce6e86e0..3354610bfcf 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2916 + "limit": 2914 }, "C401": { "limit": 8 @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 116 + "limit": 115 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 29b74c2ee4a..68fbc99b5d5 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -92,3 +92,72 @@ def test_azure_passthrough_logging_non_streaming_response_unknown_endpoint_retur ) assert result is None + + +def _sse_line(payload: dict) -> str: + return "data: " + json.dumps(payload) + + +def _azure_chat_completion_chunks() -> list[str]: + head = {"id": "chatcmpl-abc123", "object": "chat.completion.chunk", "created": 1700000000, "model": "gpt-4.1-mini"} + return [ + _sse_line({**head, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello!"}, "finish_reason": None}]}), + _sse_line({**head, "choices": [{"index": 0, "delta": {"content": " How can I assist?"}, "finish_reason": None}]}), + _sse_line({**head, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}), + _sse_line({**head, "choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}}), + "data: [DONE]", + ] + + +def test_azure_passthrough_streaming_chat_chunks_build_the_complete_response(): + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_chat_completion_chunks(), + litellm_logging_obj=MagicMock(), + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "Hello! How can I assist?" + assert response.usage.prompt_tokens == 10 + assert response.usage.completion_tokens == 8 + + +def test_azure_passthrough_streaming_chunks_for_unknown_endpoint_return_none(): + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_chat_completion_chunks(), + litellm_logging_obj=MagicMock(), + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/embeddings", + ) + + assert response is None + + +def _complete_url(request_query_params: dict, litellm_params: dict) -> httpx.URL: + url, _ = AzurePassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com", + api_key="key", + model="gpt-4.1-mini", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + request_query_params=request_query_params, + litellm_params=litellm_params, + ) + return url + + +def test_azure_passthrough_url_falls_back_to_the_callers_api_version(): + url = _complete_url(request_query_params={"api-version": "2025-04-01-preview"}, litellm_params={}) + + assert url.path == "/openai/deployments/gpt-4.1-mini/chat/completions" + assert url.params["api-version"] == "2025-04-01-preview" + + +def test_azure_passthrough_url_prefers_the_deployments_api_version(): + url = _complete_url( + request_query_params={"api-version": "2025-04-01-preview"}, litellm_params={"api_version": "2024-10-21"} + ) + + assert url.params["api-version"] == "2024-10-21" diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index e1abbe1a7e1..67851d2d58e 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -173,3 +173,24 @@ def test_non_chat_relay_yields_no_cost_response(): ) assert result is None + + +def test_streaming_chat_completion_chunks_are_costed_like_azure(): + head = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-5.4-mini"} + chunks = [ + "data: " + json.dumps({**head, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}]}), + "data: " + json.dumps({**head, "choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}}), + "data: [DONE]", + ] + + response = AzureAIPassthroughConfig().handle_logging_collected_chunks( + all_chunks=chunks, + litellm_logging_obj=MagicMock(), + model="gpt-5.4-mini", + custom_llm_provider="azure_ai", + endpoint="chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "hi" + assert response.usage.total_tokens == 4 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..a2043808b6b 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 @@ -5206,3 +5206,90 @@ class TestAzureRouterModelStreamingKeepalive: assert result.headers["x-upstream"] == "kept" assert chunks == [b"data: hello\n\n"] + + +class TestRouterModelRelayUpstreamContract: + def _request(self, content_type: str) -> MagicMock: + request = MagicMock(spec=Request) + request.method = "POST" + request.headers = {"content-type": content_type} + request.query_params = {} + return request + + def _install_router(self, monkeypatch, router, body: dict) -> None: + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + import litellm.proxy.proxy_server as proxy_server + + async def fake_get_request_body(_request): + return body + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) + monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + + def _recording_router(self, captured: list[dict]): + class RecordingRouter: + async def allm_passthrough_route(self, **kwargs): + captured.append(kwargs) + return httpx.Response(200, json={"ok": True}) + + return RecordingRouter() + + @pytest.mark.asyncio + async def test_azure_relay_keeps_the_json_body_when_the_content_type_carries_a_charset(self, monkeypatch): + body = {"model": "gpt-5", "messages": [{"role": "user", "content": "hi"}]} + captured: list[dict] = [] + self._install_router(monkeypatch, self._recording_router(captured), body) + + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json; charset=utf-8"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert captured[0]["json"] == body + + @pytest.mark.asyncio + async def test_vllm_relay_keeps_the_json_body_when_the_content_type_carries_a_charset(self, monkeypatch): + body = {"model": "router-model", "messages": [{"role": "user", "content": "hi"}]} + captured: list[dict] = [] + self._install_router(monkeypatch, self._recording_router(captured), body) + + await vllm_proxy_route( + endpoint="/chat/completions", + request=self._request("application/json; charset=utf-8"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert captured[0]["json"] == body + + @pytest.mark.asyncio + async def test_azure_relay_returns_the_upstream_status_and_body_when_the_deployment_rejects_the_call( + self, monkeypatch + ): + upstream_body = {"error": {"code": "DeploymentNotFound", "message": "The API deployment does not exist."}} + + class RejectingRouter: + async def allm_passthrough_route(self, **kwargs): + upstream_request = httpx.Request( + "POST", "https://my-azure.openai.azure.com/openai/deployments/gpt-5/chat/completions" + ) + upstream = httpx.Response( + 404, json=upstream_body, headers={"x-ms-request-id": "req-1"}, request=upstream_request + ) + raise httpx.HTTPStatusError("404", request=upstream_request, response=upstream) + + self._install_router(monkeypatch, RejectingRouter(), {"model": "gpt-5", "stream": False}) + + result = await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert result.status_code == 404 + assert json.loads(result.body) == upstream_body + assert result.headers["x-ms-request-id"] == "req-1" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index bfa055b43e5..30bd580c4e2 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22180 + "limit": 22178 }, "LIT002": { "limit": 26745 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16464 + "limit": 16458 }, "LIT011": { "limit": 5506 From 9ea9aa2e7b1ef97f80e2141f0bc299498e351143 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:00:23 -0700 Subject: [PATCH 007/107] fix(azure_ai): count prompt tokens for streaming relays that carry no usage chunk --- basedpyright-code-budget.json | 8 ++--- .../streaming_chunk_builder_utils.py | 4 +-- .../llms/azure/passthrough/transformation.py | 20 ++++++++++- litellm/main.py | 4 +-- .../openai_passthrough_logging_handler.py | 7 ++-- ruff-strict-budget.json | 4 +-- .../test_azure_passthrough_transformation.py | 33 +++++++++++++++++-- type-discipline-budget.json | 4 +-- 8 files changed, 67 insertions(+), 17 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index ef1bd675c5e..4dc7a0d8671 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15278 + "limit": 15277 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38276 + "limit": 38271 }, "reportUnknownParameterType": { - "limit": 19581 + "limit": 19580 }, "reportUnknownVariableType": { - "limit": 29825 + "limit": 29821 }, "reportUnnecessaryCast": { "limit": 110 diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0e01577b20e..1beb06f52e7 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -209,7 +209,7 @@ def apply_grounding_request_counts( class ChunkProcessor: - def __init__(self, chunks: list, messages: list | None = None): + def __init__(self, chunks: list, messages: Sequence | None = None): self.chunks = self._sort_chunks(chunks) self.messages = messages self.first_chunk = chunks[0] @@ -992,7 +992,7 @@ class ChunkProcessor: chunks: Sequence["_UsageBearingChunk | ModelResponse"], model: str, completion_output: str, - messages: list | None = None, + messages: Sequence | None = None, reasoning_tokens: int | None = None, ) -> Usage: """ diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 7654958fc0b..0b8b216faae 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,8 +1,9 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional import httpx from httpx import Response +from pydantic import BaseModel, ValidationError from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure.common_utils import BaseAzureLLM @@ -17,6 +18,22 @@ if TYPE_CHECKING: from litellm.types.utils import CostResponseTypes +class RelayedChatRequest(BaseModel): + messages: Sequence[Mapping[str, object]] | None = None + + +class RelayedCallDetails(BaseModel): + request_data: RelayedChatRequest | None = None + + +def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, object]] | None: + try: + details: Final = RelayedCallDetails.model_validate(litellm_logging_obj.model_call_details) + except ValidationError: + return None + return details.request_data.messages if details.request_data else None + + class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return "stream" in request_data @@ -137,4 +154,5 @@ class AzurePassthroughConfig(BasePassthroughConfig): all_chunks=all_chunks, litellm_logging_obj=litellm_logging_obj, model=model, + messages=_relayed_messages(litellm_logging_obj), ) diff --git a/litellm/main.py b/litellm/main.py index 2929790f2bd..a7fa5560d30 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8553,7 +8553,7 @@ def config_completion(**kwargs): ) -def stream_chunk_builder_text_completion(chunks: list, messages: list | None = None) -> TextCompletionResponse: +def stream_chunk_builder_text_completion(chunks: list, messages: Sequence | None = None) -> TextCompletionResponse: id: Final = chunks[0]["id"] object: Final = chunks[0]["object"] created: Final = chunks[0]["created"] @@ -8670,7 +8670,7 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o def stream_chunk_builder( chunks: list, - messages: list | None = None, + messages: Sequence | None = None, start_time=None, end_time=None, logging_obj: Optional["Logging"] = None, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index a39ea342f0a..9805c7d7fae 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -4,7 +4,7 @@ OpenAI Passthrough Logging Handler Handles cost tracking and logging for OpenAI passthrough endpoints, specifically /chat/completions. """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime from typing import Final from urllib.parse import urlparse @@ -516,6 +516,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): all_chunks: Sequence[str], litellm_logging_obj: LiteLLMLoggingObj, model: str, + messages: Sequence[Mapping[str, object]] | None = None, ) -> ModelResponse | TextCompletionResponse | None: """ Builds complete response from raw chunks for OpenAI streaming responses. @@ -559,7 +560,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): return None # Build complete response from chunks - complete_streaming_response: Final = litellm.stream_chunk_builder(chunks=all_openai_chunks) + complete_streaming_response: Final = litellm.stream_chunk_builder( + chunks=all_openai_chunks, messages=messages + ) return complete_streaming_response diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 3354610bfcf..cfe737867f1 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2914 + "limit": 2912 }, "C401": { "limit": 8 @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 115 + "limit": 114 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 68fbc99b5d5..d0657679983 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -3,6 +3,8 @@ from unittest.mock import MagicMock import httpx +import litellm + from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig from litellm.types.utils import ModelResponse @@ -101,8 +103,15 @@ def _sse_line(payload: dict) -> str: def _azure_chat_completion_chunks() -> list[str]: head = {"id": "chatcmpl-abc123", "object": "chat.completion.chunk", "created": 1700000000, "model": "gpt-4.1-mini"} return [ - _sse_line({**head, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello!"}, "finish_reason": None}]}), - _sse_line({**head, "choices": [{"index": 0, "delta": {"content": " How can I assist?"}, "finish_reason": None}]}), + _sse_line( + { + **head, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello!"}, "finish_reason": None}], + } + ), + _sse_line( + {**head, "choices": [{"index": 0, "delta": {"content": " How can I assist?"}, "finish_reason": None}]} + ), _sse_line({**head, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}), _sse_line({**head, "choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}}), "data: [DONE]", @@ -124,6 +133,26 @@ def test_azure_passthrough_streaming_chat_chunks_build_the_complete_response(): assert response.usage.completion_tokens == 8 +def test_azure_passthrough_streaming_chunks_without_usage_count_prompt_tokens_from_the_relayed_request(): + messages = [{"role": "user", "content": "Say hi in three words"}] + logging_obj = MagicMock() + logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True}} + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=[chunk for chunk in _azure_chat_completion_chunks() if '"usage"' not in chunk], + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "Hello! How can I assist?" + assert response.usage.prompt_tokens > 0 + assert response.usage.prompt_tokens == litellm.token_counter(model="gpt-4.1-mini", messages=messages) + assert response.usage.completion_tokens > 0 + + def test_azure_passthrough_streaming_chunks_for_unknown_endpoint_return_none(): response = AzurePassthroughConfig().handle_logging_collected_chunks( all_chunks=_azure_chat_completion_chunks(), diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 30bd580c4e2..74b54148418 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22178 + "limit": 22172 }, "LIT002": { "limit": 26745 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16458 + "limit": 16452 }, "LIT011": { "limit": 5506 From bbbdccb82d88f3489a1a311109f86d87bb1782e8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:30:13 -0700 Subject: [PATCH 008/107] fix(azure_ai): count relayed image prompt tokens without fetching the image --- basedpyright-code-budget.json | 8 +++--- .../streaming_chunk_builder_utils.py | 5 +++- litellm/main.py | 3 ++ .../openai_passthrough_logging_handler.py | 2 +- ruff-strict-budget.json | 4 +-- .../test_azure_passthrough_transformation.py | 28 +++++++++++++++++++ type-discipline-budget.json | 4 +-- 7 files changed, 44 insertions(+), 10 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 4dc7a0d8671..f63fb2a497a 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15277 + "limit": 15276 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38271 + "limit": 38266 }, "reportUnknownParameterType": { - "limit": 19580 + "limit": 19579 }, "reportUnknownVariableType": { - "limit": 29821 + "limit": 29817 }, "reportUnnecessaryCast": { "limit": 110 diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 1beb06f52e7..1dca64f1a80 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -994,6 +994,7 @@ class ChunkProcessor: completion_output: str, messages: Sequence | None = None, reasoning_tokens: int | None = None, + use_default_image_token_count: bool = False, ) -> Usage: """ Calculate usage for the given chunks. @@ -1018,7 +1019,9 @@ class ChunkProcessor: cost: Final[float | None] = calculated_usage_per_chunk["cost"] try: - returned_usage.prompt_tokens = prompt_tokens or token_counter(model=model, messages=messages) + returned_usage.prompt_tokens = prompt_tokens or token_counter( + model=model, messages=messages, use_default_image_token_count=use_default_image_token_count + ) except Exception: # don't allow this failing to block a complete streaming response from being returned print_verbose("token_counter failed, assuming prompt tokens is 0") returned_usage.prompt_tokens = 0 diff --git a/litellm/main.py b/litellm/main.py index a7fa5560d30..d125f1c1d7f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8674,6 +8674,7 @@ def stream_chunk_builder( start_time=None, end_time=None, logging_obj: Optional["Logging"] = None, + use_default_image_token_count: bool = False, ) -> ModelResponse | TextCompletionResponse | None: try: if chunks is None: @@ -8747,6 +8748,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=0, + use_default_image_token_count=use_default_image_token_count, ) setattr(response, "usage", usage) @@ -8924,6 +8926,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=reasoning_tokens, + use_default_image_token_count=use_default_image_token_count, ) setattr(response, "usage", usage) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 9805c7d7fae..e21105f760c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -561,7 +561,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): # Build complete response from chunks complete_streaming_response: Final = litellm.stream_chunk_builder( - chunks=all_openai_chunks, messages=messages + chunks=all_openai_chunks, messages=messages, use_default_image_token_count=True ) return complete_streaming_response diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index cfe737867f1..907317e81b1 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2912 + "limit": 2910 }, "C401": { "limit": 8 @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 114 + "limit": 113 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index d0657679983..b07b7915277 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -153,6 +153,34 @@ def test_azure_passthrough_streaming_chunks_without_usage_count_prompt_tokens_fr assert response.usage.completion_tokens > 0 +def test_azure_passthrough_streaming_chunks_count_remote_image_prompt_tokens_without_fetching_the_image(): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": "http://127.0.0.1:9/doc.png"}}, + ], + } + ] + logging_obj = MagicMock() + logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True}} + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=[chunk for chunk in _azure_chat_completion_chunks() if '"usage"' not in chunk], + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.usage.prompt_tokens > 0 + assert response.usage.prompt_tokens == litellm.token_counter( + model="gpt-4.1-mini", messages=messages, use_default_image_token_count=True + ) + + def test_azure_passthrough_streaming_chunks_for_unknown_endpoint_return_none(): response = AzurePassthroughConfig().handle_logging_collected_chunks( all_chunks=_azure_chat_completion_chunks(), diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 74b54148418..78cd84f5b35 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22172 + "limit": 22166 }, "LIT002": { "limit": 26745 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16452 + "limit": 16446 }, "LIT011": { "limit": 5506 From 3f695846b3cbb688d0d581bc64c4da5258546240 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:09:25 -0700 Subject: [PATCH 009/107] fix(router): rewrite the passthrough model group as a whole path segment --- .../base_llm/passthrough/transformation.py | 20 ++++++++++++++-- litellm/router.py | 3 ++- tests/test_litellm/test_router.py | 23 +++++++++++++++++++ 3 files changed, 43 insertions(+), 3 deletions(-) diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 34533e4c603..d97ecddc75d 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -8,11 +8,27 @@ if TYPE_CHECKING: from httpx import URL, Headers, Response from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.types.utils import CostResponseTypes + from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject from ..chat.transformation import BaseLLMException +def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str: + path: Final = endpoint.lstrip("/") + for model_name in model_names: + if not model_name: + continue + if path == model_name: + return "" + if path.startswith(f"{model_name}/"): + return path[len(model_name) + 1 :] + return path + + +def replace_path_segment(endpoint: str, segment: str, replacement: str) -> str: + return "/".join(replacement if part == segment else part for part in endpoint.split("/")) + + class BasePassthroughConfig(BaseLLMModelInfo): @abstractmethod def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: @@ -104,7 +120,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): request_data: dict, logging_obj: "LiteLLMLoggingObj", endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> Optional["CostResponseTypes | StandardPassThroughResponseObject"]: pass def handle_logging_collected_chunks( diff --git a/litellm/router.py b/litellm/router.py index 9b7a7ee7ce8..273c1f9d99c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -86,6 +86,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_credentials_in_payload, mask_sensitive_structure, ) +from litellm.llms.base_llm.passthrough.transformation import replace_path_segment from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, vector_store_request_metadata, @@ -5010,7 +5011,7 @@ class Router: # If get_llm_provider fails, fall back to using model_name as-is replacement_model_name = model_name - kwargs["endpoint"] = kwargs["endpoint"].replace(model, replacement_model_name) + kwargs["endpoint"] = replace_path_segment(kwargs["endpoint"], model, replacement_model_name) return kwargs async def _ageneric_api_call_with_fallbacks_helper(self, model: str, original_generic_function: Callable, **kwargs): diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b220b23c338..0f4e199694d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -4469,6 +4469,29 @@ def test_get_deployment_model_info_base_model_merge_priority(): print("✓ Base model merge priority test passed!") +def test_add_deployment_model_to_endpoint_rewrites_whole_path_segments_only(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt", + "litellm_params": { + "model": "azure_ai/gpt-5.4-mini", + "api_base": "https://my-resource.services.ai.azure.com", + "api_key": "key", + }, + } + ], + ) + + result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( + kwargs={"endpoint": "gpt/openai/deployments/gpt-4o/chat/completions", "custom_llm_provider": "azure_ai"}, + model="gpt", + model_name="azure_ai/gpt-5.4-mini", + ) + + assert result["endpoint"] == "gpt-5.4-mini/openai/deployments/gpt-4o/chat/completions" + + def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): """ Test that _add_deployment_model_to_endpoint_for_llm_passthrough_route correctly strips bedrock provider prefix From df6fb9e5d95e97b7b207cd278d107e82b134fe45 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:09:26 -0700 Subject: [PATCH 010/107] fix(azure_ai): relay from the Foundry root and log non-chat relays for spend tracking --- .../llms/azure/passthrough/transformation.py | 14 +-- .../azure_ai/passthrough/transformation.py | 69 +++++++++++---- ...est_azure_ai_passthrough_transformation.py | 87 +++++++++++++++++-- 3 files changed, 140 insertions(+), 30 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 0b8b216faae..983033a8b14 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -7,7 +7,11 @@ from pydantic import BaseModel, ValidationError from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure.common_utils import BaseAzureLLM -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +from litellm.llms.base_llm.passthrough.transformation import ( + BasePassthroughConfig, + replace_path_segment, + strip_leading_model_segment, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams @@ -36,7 +40,7 @@ def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, obj class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: - return "stream" in request_data + return bool(request_data.get("stream")) def get_complete_url( self, @@ -54,13 +58,13 @@ class AzurePassthroughConfig(BasePassthroughConfig): litellm_metadata: Final = litellm_params.get("litellm_metadata") or {} model_group: Final = litellm_metadata.get("model_group") - if model_group and model_group in endpoint: - endpoint = endpoint.replace(model_group, model) + routed_endpoint: Final = replace_path_segment(endpoint, model_group, model) if model_group else endpoint + native_endpoint: Final = strip_leading_model_segment(routed_endpoint, (model,)) complete_url: Final = BaseAzureLLM._get_base_azure_url( api_base=base_target_url, litellm_params=litellm_params, - route=endpoint, + route=native_endpoint, default_api_version=request_query_params.get("api-version") if request_query_params else None, ) return ( diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index 5948c865bad..4776bbe0755 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -1,17 +1,20 @@ from __future__ import annotations from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final -from pydantic import BaseModel, ConfigDict, ValidationError +import httpx +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.llms.azure_ai.common_utils import ( AzureFoundryModelInfo, api_key_header_for_base, get_azure_ai_auth_headers, ) -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, strip_leading_model_segment from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import StandardPassThroughResponseObject if TYPE_CHECKING: from httpx import URL, Response @@ -20,16 +23,7 @@ if TYPE_CHECKING: from litellm.types.utils import CostResponseTypes -def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str: - path: Final = endpoint.lstrip("/") - for model_name in model_names: - if not model_name: - continue - if path == model_name: - return "" - if path.startswith(f"{model_name}/"): - return path[len(model_name) + 1 :] - return path +EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({}) class PassthroughMetadata(BaseModel): @@ -45,9 +39,44 @@ def model_group_from(litellm_params: Mapping[str, object]) -> str: return "" +def api_version_from(litellm_params: Mapping[str, object]) -> str | None: + try: + return TypeAdapter(str | None).validate_python(litellm_params.get("api_version")) + except ValidationError: + return None + + +def foundry_root(api_base: str) -> str: + url: Final = httpx.URL(api_base) + segments: Final = tuple(segment for segment in url.path.split("/") if segment) + root_segments: Final = segments[: segments.index("models")] if "models" in segments else segments + return str(url.copy_with(path="/" + "/".join(root_segments), query=None)).rstrip("/") + + +def relay_query_params( + request_query_params: Mapping[str, object] | None, + deployment_api_version: str | None, + api_base: str, +) -> Mapping[str, object] | None: + if request_query_params and "api-version" in request_query_params: + return request_query_params + api_version: Final = deployment_api_version or httpx.URL(api_base).params.get("api-version") + if api_version is None: + return request_query_params + return MappingProxyType({**(request_query_params or EMPTY_QUERY), "api-version": api_version}) + + +def relayed_body(httpx_response: Response) -> str | dict: + try: + body: Final[object] = httpx_response.json() + except ValueError: + return httpx_response.text + return body if isinstance(body, dict) else httpx_response.text + + class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool: - return request_data.get("stream") is True + return bool(request_data.get("stream")) def get_complete_url( self, @@ -62,11 +91,12 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): if base_target_url is None: raise ValueError("Azure AI api base not found: set `api_base` on the deployment or AZURE_AI_API_BASE") + root: Final = foundry_root(base_target_url) native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params))) - return ( - self.format_url(native_endpoint, base_target_url, request_query_params), - base_target_url, + query_params: Final = relay_query_params( + request_query_params, api_version_from(litellm_params), base_target_url ) + return (self.format_url(native_endpoint, root, query_params), root) def validate_environment( self, @@ -93,10 +123,10 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): request_data: Mapping[str, object], logging_obj: Logging, endpoint: str, - ) -> CostResponseTypes | None: + ) -> CostResponseTypes | StandardPassThroughResponseObject | None: from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig - return AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict + chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict model=model, custom_llm_provider=custom_llm_provider, httpx_response=httpx_response, @@ -104,6 +134,9 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): logging_obj=logging_obj, endpoint=endpoint, ) + if chat_result is not None: + return chat_result + return StandardPassThroughResponseObject(response=relayed_body(httpx_response)) def handle_logging_collected_chunks( self, diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 67851d2d58e..e343d05c5df 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -66,6 +66,72 @@ def test_model_inside_the_path_stays_and_query_params_are_forwarded(): assert str(url) == f"{FOUNDRY_BASE}/openai/deployments/gpt-5.4-mini/chat/completions?api-version=2024-10-21" +def test_api_base_that_already_ends_in_models_is_cut_back_to_the_foundry_root(): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/models", + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params={"api-version": "2024-05-01-preview"}, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + assert base == FOUNDRY_BASE + + +def test_parse_relay_under_a_models_api_base_targets_the_foundry_root(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/models", + api_key="key", + model="Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + + +def test_deployment_api_version_fills_in_when_the_caller_sends_none(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params=None, + litellm_params={"api_version": "2024-05-01-preview"}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + + +def test_callers_api_version_beats_the_deployments(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params={"api-version": "2025-04-01-preview"}, + litellm_params={"api_version": "2024-05-01-preview"}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2025-04-01-preview" + + +def test_api_version_on_the_configured_api_base_is_the_last_fallback(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview", + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params=None, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + + def test_missing_api_base_raises_instead_of_building_a_relative_url(): with pytest.raises(ValueError, match="AZURE_AI_API_BASE"): AzureAIPassthroughConfig().get_complete_url( @@ -116,7 +182,7 @@ def test_no_credentials_at_all_raises(): @pytest.mark.parametrize( "request_data, expected", - [({"stream": True}, True), ({"stream": False}, False), ({}, False)], + [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], ) def test_is_streaming_request_reads_the_stream_flag(request_data, expected): assert AzureAIPassthroughConfig().is_streaming_request(endpoint="models/chat/completions", request_data=request_data) is expected @@ -155,15 +221,14 @@ def test_chat_completions_relay_yields_a_model_response_for_cost_tracking(): assert result.usage.completion_tokens == 8 -def test_non_chat_relay_yields_no_cost_response(): +def _non_chat_logging_result(content: bytes, content_type: str): parse_response = httpx.Response( status_code=200, - headers={"content-type": "application/json"}, - content=b'{"id":"parse-1","pages":[]}', + headers={"content-type": content_type}, + content=content, request=httpx.Request("POST", f"{FOUNDRY_BASE}/providers/cohere/v2/parse"), ) - - result = AzureAIPassthroughConfig().logging_non_streaming_response( + return AzureAIPassthroughConfig().logging_non_streaming_response( model="Cohere-parse-v5", custom_llm_provider="azure_ai", httpx_response=parse_response, @@ -172,7 +237,15 @@ def test_non_chat_relay_yields_no_cost_response(): endpoint="providers/cohere/v2/parse", ) - assert result is None + +def test_non_chat_relay_logs_the_parsed_body_so_spend_tracking_sees_the_call(): + result = _non_chat_logging_result(b'{"id":"parse-1","pages":[],"meta":{"billed_units":{"pages":1}}}', "application/json") + + assert result == {"response": {"id": "parse-1", "pages": [], "meta": {"billed_units": {"pages": 1}}}} + + +def test_non_chat_relay_with_a_non_json_body_logs_the_raw_text(): + assert _non_chat_logging_result(b"page one", "text/plain") == {"response": "page one"} def test_streaming_chat_completion_chunks_are_costed_like_azure(): From 2f981d14e4262db3e573891ba94f0e66209b216d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:09:26 -0700 Subject: [PATCH 011/107] refactor(passthrough): inject the streaming prompt-token counter instead of a default-image flag --- .../streaming_chunk_builder_utils.py | 8 ++-- litellm/main.py | 8 ++-- .../openai_passthrough_logging_handler.py | 44 +++++++++++++++++- ruff-strict-budget.json | 4 +- .../test_azure_passthrough_transformation.py | 44 ++++++++++++++++-- ...test_openai_passthrough_logging_handler.py | 46 +++++++++++++++++++ type-discipline-budget.json | 6 +-- 7 files changed, 142 insertions(+), 18 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 1dca64f1a80..f30a8e5f8bd 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1,6 +1,6 @@ import base64 import time -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from itertools import groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast @@ -994,7 +994,7 @@ class ChunkProcessor: completion_output: str, messages: Sequence | None = None, reasoning_tokens: int | None = None, - use_default_image_token_count: bool = False, + count_prompt_tokens: Callable[[], int] | None = None, ) -> Usage: """ Calculate usage for the given chunks. @@ -1019,8 +1019,8 @@ class ChunkProcessor: cost: Final[float | None] = calculated_usage_per_chunk["cost"] try: - returned_usage.prompt_tokens = prompt_tokens or token_counter( - model=model, messages=messages, use_default_image_token_count=use_default_image_token_count + returned_usage.prompt_tokens = prompt_tokens or ( + count_prompt_tokens() if count_prompt_tokens else token_counter(model=model, messages=messages) ) except Exception: # don't allow this failing to block a complete streaming response from being returned print_verbose("token_counter failed, assuming prompt tokens is 0") diff --git a/litellm/main.py b/litellm/main.py index d125f1c1d7f..a3830192e7a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -19,7 +19,7 @@ import random import sys import time import traceback -from collections.abc import AsyncIterator, Coroutine, Iterable, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Mapping, Sequence from concurrent import futures from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from copy import deepcopy @@ -8674,7 +8674,7 @@ def stream_chunk_builder( start_time=None, end_time=None, logging_obj: Optional["Logging"] = None, - use_default_image_token_count: bool = False, + count_prompt_tokens: Callable[[], int] | None = None, ) -> ModelResponse | TextCompletionResponse | None: try: if chunks is None: @@ -8748,7 +8748,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=0, - use_default_image_token_count=use_default_image_token_count, + count_prompt_tokens=count_prompt_tokens, ) setattr(response, "usage", usage) @@ -8926,7 +8926,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=reasoning_tokens, - use_default_image_token_count=use_default_image_token_count, + count_prompt_tokens=count_prompt_tokens, ) setattr(response, "usage", usage) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index e21105f760c..6a142d08be5 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -13,6 +13,7 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, @@ -97,6 +98,45 @@ def _is_openai_compatible_url(url_route: str | None) -> bool: return False +def _is_remote_high_detail_image(part: object) -> bool: + if not isinstance(part, Mapping) or part.get("type") != "image_url": + return False + image_url: Final = part.get("image_url") + if not isinstance(image_url, Mapping): + return False + url: Final = image_url.get("url") + return isinstance(url, str) and url.startswith(("http://", "https://")) and image_url.get("detail") == "high" + + +def _content_parts(message: Mapping[str, object]) -> Sequence[object]: + content: Final = message.get("content") + return content if isinstance(content, list) else () + + +def _without_remote_high_detail_images(message: Mapping[str, object]) -> Mapping[str, object]: + if not isinstance(message.get("content"), list): + return message + kept_parts: Final = [ # mutable-ok: token_counter reads message content only when it is a list + part for part in _content_parts(message) if not _is_remote_high_detail_image(part) + ] + return {**message, "content": kept_parts} # mutable-ok: token_counter rejects any message that is not a dict + + +def count_relayed_prompt_tokens(model: str, messages: Sequence[Mapping[str, object]] | None) -> int: + if messages is None: + return 0 + remote_high_detail_images: Final = sum( + 1 for message in messages for part in _content_parts(message) if _is_remote_high_detail_image(part) + ) + local_messages: Final = [ # mutable-ok: token_counter takes a list of messages + _without_remote_high_detail_images(message) for message in messages + ] + return ( + litellm.token_counter(model=model, messages=local_messages) + + DEFAULT_IMAGE_TOKEN_COUNT * remote_high_detail_images + ) + + class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): """ OpenAI-specific passthrough logging handler that provides cost tracking for /chat/completions endpoints. @@ -561,7 +601,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): # Build complete response from chunks complete_streaming_response: Final = litellm.stream_chunk_builder( - chunks=all_openai_chunks, messages=messages, use_default_image_token_count=True + chunks=all_openai_chunks, + messages=messages, + count_prompt_tokens=lambda: count_relayed_prompt_tokens(model, messages), ) return complete_streaming_response diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 907317e81b1..574bf01cc34 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2910 + "limit": 2904 }, "C401": { "limit": 8 @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 113 + "limit": 110 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index b07b7915277..5d9e4312b5e 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -2,8 +2,10 @@ import json from unittest.mock import MagicMock import httpx +import pytest import litellm +from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig @@ -159,7 +161,7 @@ def test_azure_passthrough_streaming_chunks_count_remote_image_prompt_tokens_wit "role": "user", "content": [ {"type": "text", "text": "Describe this"}, - {"type": "image_url", "image_url": {"url": "http://127.0.0.1:9/doc.png"}}, + {"type": "image_url", "image_url": {"url": "http://127.0.0.1:9/doc.png", "detail": "high"}}, ], } ] @@ -174,10 +176,10 @@ def test_azure_passthrough_streaming_chunks_count_remote_image_prompt_tokens_wit endpoint="openai/deployments/gpt-4.1-mini/chat/completions", ) + text_only_messages = [{"role": "user", "content": [{"type": "text", "text": "Describe this"}]}] assert isinstance(response, ModelResponse) - assert response.usage.prompt_tokens > 0 - assert response.usage.prompt_tokens == litellm.token_counter( - model="gpt-4.1-mini", messages=messages, use_default_image_token_count=True + assert response.usage.prompt_tokens == ( + litellm.token_counter(model="gpt-4.1-mini", messages=text_only_messages) + DEFAULT_IMAGE_TOKEN_COUNT ) @@ -218,3 +220,37 @@ def test_azure_passthrough_url_prefers_the_deployments_api_version(): ) assert url.params["api-version"] == "2024-10-21" + + +def test_azure_passthrough_url_strips_the_leading_router_model_segment(): + url, _ = AzurePassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com", + api_key="key", + model="gpt-4.1-mini", + endpoint="gpt-4.1-mini/openai/deployments/gpt-4.1-mini/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={}, + ) + + assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + + +def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment(): + url, _ = AzurePassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com", + api_key="key", + model="gpt-4.1-mini", + endpoint="gpt/openai/deployments/gpt-4o/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={"litellm_metadata": {"model_group": "gpt"}}, + ) + + assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21" + + +@pytest.mark.parametrize( + "request_data, expected", + [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], +) +def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_data, expected): + assert AzurePassthroughConfig().is_streaming_request(endpoint="openai/deployments/x/chat/completions", request_data=request_data) is expected diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 6f9142c85df..23c26a0b8a3 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -8,9 +8,11 @@ import pytest import litellm +from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, + count_relayed_prompt_tokens, ) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, @@ -2037,3 +2039,47 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: if __name__ == "__main__": pytest.main([__file__]) + + +ONE_PIXEL_PNG_DATA_URL = ( + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" +) +UNREACHABLE_IMAGE_URL = "http://127.0.0.1:9/doc.png" +TEXT_ONLY_MESSAGES = [{"role": "user", "content": [{"type": "text", "text": "Describe this"}]}] + + +def _image_messages(url: str, detail: str) -> list[dict]: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": url, "detail": detail}}, + ], + } + ] + + +def test_count_relayed_prompt_tokens_counts_a_data_url_image_exactly(): + messages = _image_messages(ONE_PIXEL_PNG_DATA_URL, "high") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == litellm.token_counter( + model="gpt-4.1-mini", messages=messages + ) + + +def test_count_relayed_prompt_tokens_keeps_a_low_detail_remote_image_at_the_base_count(): + messages = _image_messages(UNREACHABLE_IMAGE_URL, "low") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == litellm.token_counter( + model="gpt-4.1-mini", messages=messages + ) + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) < DEFAULT_IMAGE_TOKEN_COUNT + + +def test_count_relayed_prompt_tokens_estimates_only_the_remote_high_detail_image(): + messages = _image_messages(UNREACHABLE_IMAGE_URL, "high") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == ( + litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + DEFAULT_IMAGE_TOKEN_COUNT + ) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 78cd84f5b35..3a4abfbc20a 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22166 + "limit": 22156 }, "LIT002": { "limit": 26745 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16446 + "limit": 16434 }, "LIT011": { - "limit": 5506 + "limit": 5504 }, "LIT012": { "limit": 4486 From 4706acef958a726f1407ad5413ea2db896d6c825 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 03:48:06 -0700 Subject: [PATCH 012/107] fix(passthrough): charge remote high-detail images at the high-detail upper bound --- litellm/litellm_core_utils/token_counter.py | 7 ++++++ .../openai_passthrough_logging_handler.py | 4 ++-- ruff-strict-budget.json | 4 ++-- .../litellm_core_utils/test_token_counter.py | 22 ++++++++++++++++++- .../test_azure_passthrough_transformation.py | 4 ++-- ...test_openai_passthrough_logging_handler.py | 8 +++---- type-discipline-budget.json | 6 ++--- 7 files changed, 41 insertions(+), 14 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 3732ffd734c..4c66d878f14 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -172,6 +172,13 @@ def calculate_tiles_needed( return total_tiles +def high_detail_image_token_upper_bound(base_tokens: int = 85) -> int: + largest_tile_count: Final = calculate_tiles_needed( + MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES, MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES + ) + return base_tokens + (base_tokens * 2) * largest_tile_count + + def _unpack_ints(fmt: str, buffer: bytes) -> tuple[int, ...]: return struct.unpack(fmt, buffer) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 6a142d08be5..206002d13f7 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -13,11 +13,11 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger -from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, ) +from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.llms.openai.openai import OpenAIConfig from litellm.llms.openai.openai import OpenAIConfig as OpenAIConfigType from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig @@ -133,7 +133,7 @@ def count_relayed_prompt_tokens(model: str, messages: Sequence[Mapping[str, obje ] return ( litellm.token_counter(model=model, messages=local_messages) - + DEFAULT_IMAGE_TOKEN_COUNT * remote_high_detail_images + + high_detail_image_token_upper_bound() * remote_high_detail_images ) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 574bf01cc34..134eec578bc 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2904 + "limit": 2900 }, "C401": { "limit": 8 @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 110 + "limit": 108 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 4694fa8fbed..1898fd57220 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1,5 +1,6 @@ #### What this tests #### # This tests litellm.token_counter.token_counter() function +import base64 import importlib import time import traceback @@ -14,7 +15,11 @@ import litellm from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants -from litellm.litellm_core_utils.token_counter import _get_tiktoken_count_function +from litellm.litellm_core_utils.token_counter import ( + _get_tiktoken_count_function, + calculate_img_tokens, + high_detail_image_token_upper_bound, +) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from tests.large_text import text from tests.test_litellm.litellm_core_utils.messages_with_counts import ( @@ -1412,3 +1417,18 @@ def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): assert _count_user_content([prompt, named]) == _count_user_content( [prompt, {"type": "text", "text": "report.pdf"}] ) + + +def _png_data_url(width: int, height: int) -> str: + ihdr = b"\x89PNG\r\n\x1a\n" + (13).to_bytes(4, "big") + b"IHDR" + width.to_bytes(4, "big") + height.to_bytes(4, "big") + return "data:image/png;base64," + base64.b64encode(ihdr + b"\x08\x06\x00\x00\x00").decode() + + +@pytest.mark.parametrize(("width", "height"), [(1, 1), (768, 768), (2000, 768), (768, 2000), (4096, 4096), (8000, 3072)]) +def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int, height: int) -> None: + assert calculate_img_tokens(_png_data_url(width, height), mode="high") <= high_detail_image_token_upper_bound() + + +def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_image() -> None: + assert calculate_img_tokens(_png_data_url(2000, 768), mode="high") == high_detail_image_token_upper_bound() + assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound() diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 5d9e4312b5e..8ffd632fe85 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -5,9 +5,9 @@ import httpx import pytest import litellm -from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT +from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig from litellm.types.utils import ModelResponse @@ -179,7 +179,7 @@ def test_azure_passthrough_streaming_chunks_count_remote_image_prompt_tokens_wit text_only_messages = [{"role": "user", "content": [{"type": "text", "text": "Describe this"}]}] assert isinstance(response, ModelResponse) assert response.usage.prompt_tokens == ( - litellm.token_counter(model="gpt-4.1-mini", messages=text_only_messages) + DEFAULT_IMAGE_TOKEN_COUNT + litellm.token_counter(model="gpt-4.1-mini", messages=text_only_messages) + high_detail_image_token_upper_bound() ) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 23c26a0b8a3..0ddafd20e68 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -8,8 +8,8 @@ import pytest import litellm -from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, count_relayed_prompt_tokens, @@ -2074,12 +2074,12 @@ def test_count_relayed_prompt_tokens_keeps_a_low_detail_remote_image_at_the_base assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == litellm.token_counter( model="gpt-4.1-mini", messages=messages ) - assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) < DEFAULT_IMAGE_TOKEN_COUNT + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) < high_detail_image_token_upper_bound() -def test_count_relayed_prompt_tokens_estimates_only_the_remote_high_detail_image(): +def test_count_relayed_prompt_tokens_charges_only_the_remote_high_detail_image_at_the_upper_bound(): messages = _image_messages(UNREACHABLE_IMAGE_URL, "high") assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == ( - litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + DEFAULT_IMAGE_TOKEN_COUNT + litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + high_detail_image_token_upper_bound() ) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 3a4abfbc20a..b4c17ffd47b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22156 + "limit": 22146 }, "LIT002": { "limit": 26745 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16434 + "limit": 16422 }, "LIT011": { - "limit": 5504 + "limit": 5502 }, "LIT012": { "limit": 4486 From 808787659b05ac986a01892938e7387fe88d1a46 Mon Sep 17 00:00:00 2001 From: Mihidum Hettiyahandi <55163074+mihidumh@users.noreply.github.com> Date: Fri, 21 Aug 2026 09:42:08 +1000 Subject: [PATCH 013/107] fix(azure_ai): stop accepting MAI image params the endpoint cannot honour MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two params were advertised for the MAI image models and dropped downstream, so the caller got a 200 that did not match the request, or an opaque provider 400. n: get_supported_openai_params returns ["n", "size"], so n passes validation and is forwarded. The MAI endpoint (/mai/v1/images/generations) has no count field at all — its documented body is model/prompt/width/height, plus image for edits — and ignores both `n` and the native `sampleCount`. Measured against MAI-Image-2.5 and MAI-Image-2.5-Flash: n=2 and n=4 each return HTTP 200 with exactly one image, billed as one, with nothing in the response saying the request was reduced. A caller balancing cost against image count cannot see it. n=1 still passes through; n>1 now raises unless drop_params is set, which is the existing opt-in for silently dropping a param. size: _map_size_param's table offered five sizes, of which one is usable. MAI requires width and height >= 768px and width*height <= 1048576, so 512x512 and 256x256 are under the per-side minimum and 1792x1024 / 1024x1792 are over the pixel budget — all four 400 at the provider with "Model does not support request parameter value supplied: 'width' must be at least 768 pixels." Only 1024x1024 works. The bounds are now checked where the size is mapped, so the error names the constraint instead of arriving from Azure. width/height are deliberately left unchecked: they pass through unmapped, so a future MAI model with different bounds stays reachable without a code change. Verified on a live Azure AI Foundry deployment of MAI-Image-2.5 and MAI-Image-2.5-Flash (2026-08-17). One existing test asserted the 1792x1024 mapping; its size is changed to a size the provider accepts, keeping what it was testing. Co-Authored-By: Claude Opus 5 --- .../image_generation/mai_transformation.py | 49 +++++++- .../test_mai_image_generation.py | 106 +++++++++++++----- 2 files changed, 122 insertions(+), 33 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 64f81956ad7..3421cc0a4a1 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -21,6 +21,19 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_WIDTH = 1024 DEFAULT_HEIGHT = 1024 + # The MAI endpoint produces exactly one image per request. Its documented + # body is model/prompt/width/height (plus `image` for edits) — there is no + # count field, and `n` (or the native `sampleCount`) is accepted and + # ignored, so a request for more silently comes back with one. + MAX_IMAGES_PER_REQUEST: Final = 1 + + # Provider-side bounds on the generated image. Both are enforced by the + # MAI endpoint, which 400s with "'width' must be at least 768 pixels." + # Only `size` is checked against them: `width`/`height` pass through + # unmapped, which keeps a future model with different bounds reachable. + MIN_DIMENSION_PX: Final = 768 + MAX_TOTAL_PX: Final = 1024 * 1024 + @staticmethod def get_mai_image_generation_url( api_base: str | None, @@ -146,6 +159,15 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if k in supported_params: if k == "size" and v: self._map_size_param(v, optional_params) + elif k == "n" and v is not None and v > self.MAX_IMAGES_PER_REQUEST: + if not drop_params: + raise ValueError( + f"n={v} is not supported for model {model}. The Azure AI MAI image " + f"endpoint returns exactly {self.MAX_IMAGES_PER_REQUEST} image per " + "request and ignores any count, so a larger value would silently " + "return fewer images than requested. Send one request per image, or " + "set drop_params=True to drop n." + ) else: optional_params[k] = v elif k in ("width", "height"): @@ -176,13 +198,9 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if size in size_mapping: width, height = size_mapping[size] - optional_params["width"] = width - optional_params["height"] = height elif "x" in size: try: width, height = map(int, size.lower().split("x")) - optional_params["width"] = width - optional_params["height"] = height except ValueError: raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').") else: @@ -191,6 +209,29 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." ) + self._validate_dimensions(size=size, width=width, height=height) + optional_params["width"] = width + optional_params["height"] = height + + def _validate_dimensions(self, size: str, width: int, height: int) -> None: + """Reject a `size` the MAI endpoint would 400 on. + + Several OpenAI-standard sizes are outside MAI's bounds: 512x512 and + 256x256 fall under the per-side minimum, and 1792x1024 / 1024x1792 + exceed the total pixel budget. Checking here turns an opaque provider + 400 into an error that names the constraint. + """ + if width < self.MIN_DIMENSION_PX or height < self.MIN_DIMENSION_PX: + raise ValueError( + f"Unsupported size value: '{size}'. Azure AI MAI image models require width and " + f"height of at least {self.MIN_DIMENSION_PX} pixels." + ) + if width * height > self.MAX_TOTAL_PX: + raise ValueError( + f"Unsupported size value: '{size}'. Azure AI MAI image models accept at most " + f"{self.MAX_TOTAL_PX} total pixels ({width}x{height} is {width * height})." + ) + def transform_image_generation_response( self, model: str, diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py index 2a44e77ce09..6cbc79fff2b 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -1,10 +1,8 @@ -import os from unittest.mock import MagicMock import httpx import pytest - import litellm from litellm.llms.azure.azure import AzureChatCompletion from litellm.llms.azure.image_generation import get_azure_image_generation_config @@ -30,9 +28,7 @@ from litellm.utils import get_optional_params_image_gen class TestAzureMAIImageGeneration: def test_is_mai_model(self): assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5") - assert AzureFoundryMAIImageGenerationConfig.is_mai_model( - "azure_ai/MAI-Image-2.5" - ) + assert AzureFoundryMAIImageGenerationConfig.is_mai_model("azure_ai/MAI-Image-2.5") assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5-Flash") assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2e") assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro") @@ -62,16 +58,10 @@ class TestAzureMAIImageGeneration: api_base="https://my-resource.services.ai.azure.com", api_version="preview", ) - assert ( - url - == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" - ) + assert url == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" def test_get_mai_image_generation_url_preserves_full_path(self): - api = ( - "https://my-resource.services.ai.azure.com/mai/v1/images/generations" - "?api-version=preview" - ) + api = "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( api_base=api, api_version="preview", @@ -83,10 +73,7 @@ class TestAzureMAIImageGeneration: api_base="https://my-resource.services.ai.azure.com/mai/v1", api_version="preview", ) - assert ( - url - == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" - ) + assert url == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" def test_get_azure_ai_image_generation_config_returns_mai(self): config = get_azure_ai_image_generation_config("MAI-Image-2.5") @@ -124,13 +111,13 @@ class TestAzureMAIImageGeneration: config = AzureFoundryMAIImageGenerationConfig() optional_params = get_optional_params_image_gen( model="MAI-Image-2.5", - size="1792x1024", + size="1024x1024", n=1, custom_llm_provider="azure_ai", provider_config=config, drop_params=True, ) - assert optional_params["width"] == 1792 + assert optional_params["width"] == 1024 assert optional_params["height"] == 1024 assert "size" not in optional_params @@ -147,10 +134,7 @@ class TestAzureMAIImageGeneration: assert "api-version=preview" in url def test_mai_json_body_keeps_model(self): - api = ( - "https://my-resource.services.ai.azure.com/mai/v1/images/generations" - "?api-version=preview" - ) + api = "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" data = { "model": "MAI-Image-2.5", "prompt": "A photograph of a red fox", @@ -214,6 +198,74 @@ class TestAzureMAIImageGeneration: drop_params=True, ) + @pytest.mark.parametrize("size", ["512x512", "256x256", "700x1400"]) + def test_map_openai_params_size_below_minimum_dimension_raises(self, size): + """MAI requires >= 768px per side; the OpenAI size table offered smaller ones.""" + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(ValueError, match="at least 768 pixels"): + config.map_openai_params( + non_default_params={"size": size}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + + @pytest.mark.parametrize("size", ["1792x1024", "1024x1792"]) + def test_map_openai_params_size_over_total_pixel_budget_raises(self, size): + """MAI caps total pixels at 1024*1024, so both landscape/portrait sizes 400 upstream.""" + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(ValueError, match="at most 1048576 total pixels"): + config.map_openai_params( + non_default_params={"size": size}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + + def test_map_openai_params_explicit_width_height_not_range_checked(self): + """width/height pass through unmapped, so a future model's bounds stay reachable.""" + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"width": 1792, "height": 1024}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == 1792 + assert optional_params["height"] == 1024 + + @pytest.mark.parametrize("n", [2, 4]) + def test_map_openai_params_multi_image_n_raises(self, n): + """The MAI endpoint returns one image and ignores any count, so n>1 must not pass silently.""" + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(ValueError, match="returns exactly 1 image per request"): + config.map_openai_params( + non_default_params={"n": n}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + + def test_map_openai_params_multi_image_n_dropped_with_drop_params(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"n": 4}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert "n" not in optional_params + + def test_map_openai_params_single_image_n_still_passes_through(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"n": 1}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + assert optional_params["n"] == 1 + def test_map_openai_params_unsupported_param_raises(self): config = AzureFoundryMAIImageGenerationConfig() with pytest.raises(ValueError, match="Parameter quality is not supported"): @@ -363,16 +415,12 @@ class TestAzureMAIImageGeneration: litellm.model_cost = litellm.get_model_cost_map(url="") model = "azure_ai/MAI-Image-2.5" model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") - image_response = ImageResponse( - data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")] - ) + image_response = ImageResponse(data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]) cost = azure_ai_image_cost_calculator( model=model, image_response=image_response, ) - assert ( - cost == len(image_response.data or []) * model_info["output_cost_per_image"] - ) + assert cost == len(image_response.data or []) * model_info["output_cost_per_image"] assert cost > 0 From 1d8650413d4237f9dfce6aec838035618a858d8b Mon Sep 17 00:00:00 2001 From: jesus Date: Mon, 7 Sep 2026 11:51:35 +0000 Subject: [PATCH 014/107] fix(responses): keep background polling alive after the client disconnects Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../response_polling/background_streaming.py | 17 +++++- .../test_response_polling_handler.py | 4 +- .../proxy/response_polling/__init__.py | 0 .../test_background_streaming.py | 52 +++++++++++++++++++ 4 files changed, 70 insertions(+), 3 deletions(-) create mode 100644 tests/test_litellm/proxy/response_polling/__init__.py create mode 100644 tests/test_litellm/proxy/response_polling/test_background_streaming.py diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 0fd242f2bc1..fac45d4391c 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Final, TypeAlias from fastapi import Request, Response from fastapi.responses import StreamingResponse +from starlette.types import Message from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -74,6 +75,20 @@ class _StreamEventParser: parse: Callable[[str], _StreamEvent] = staticmethod(json.loads) +async def _never_receive() -> Message: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + +def detach_request_from_client(request: Request) -> Request: + """Same scope (headers, parsed body, auth) but a receive() that never yields http.disconnect. + + The polling client closes its connection right after getting the polling id, so the + upstream call must not be cancelled by the client-disconnect guards. + """ + return Request(request.scope, _never_receive) + + async def background_streaming_task( polling_id: str, data: dict[str, object], @@ -123,7 +138,7 @@ async def background_streaming_task( # Pre-call checks (rate limits, guardrails, budget) were already run # before polling ID creation, so skip them here to avoid double-counting. response: Final[StreamingResponse] = await processor.base_process_llm_request( - request=request, + request=detach_request_from_client(request), fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, route_type="aresponses", diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/proxy_unit_tests/test_response_polling_handler.py index 772d3622745..7a4e0ca89c0 100644 --- a/tests/proxy_unit_tests/test_response_polling_handler.py +++ b/tests/proxy_unit_tests/test_response_polling_handler.py @@ -20,7 +20,7 @@ from typing import Any, Dict, Optional from unittest.mock import AsyncMock, Mock, patch import pytest - +from fastapi import Request from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler @@ -1414,7 +1414,7 @@ def _make_background_streaming_kwargs( polling_id=polling_id, data={"model": "gpt-4o", "stream": False, "background": True}, polling_handler=polling_handler, - request=Mock(), + request=Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}), fastapi_response=Mock(), user_api_key_dict=Mock(), general_settings={}, diff --git a/tests/test_litellm/proxy/response_polling/__init__.py b/tests/test_litellm/proxy/response_polling/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/response_polling/test_background_streaming.py b/tests/test_litellm/proxy/response_polling/test_background_streaming.py new file mode 100644 index 00000000000..98d914b1848 --- /dev/null +++ b/tests/test_litellm/proxy/response_polling/test_background_streaming.py @@ -0,0 +1,52 @@ +import asyncio +from collections.abc import AsyncGenerator +from typing import Final + +import pytest +from fastapi import Request +from fastapi.responses import JSONResponse, StreamingResponse +from starlette.types import Message + +from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED +from litellm.proxy.common_request_processing import create_response +from litellm.proxy.response_polling.background_streaming import detach_request_from_client + + +def _request_whose_client_already_left() -> Request: + async def receive() -> Message: + return {"type": "http.disconnect"} + + scope: Final = { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [(b"x-litellm-call-id", b"call-123")], + "query_string": b"", + } + return Request(scope, receive) + + +async def _slow_first_chunk() -> AsyncGenerator[str, None]: + await asyncio.sleep(0.05) + yield 'data: {"type": "response.created"}\n\n' + + +@pytest.mark.asyncio +async def test_detached_request_survives_client_disconnect_before_first_chunk(): + original: Final = _request_whose_client_already_left() + + cancelled: Final = await create_response(_slow_first_chunk(), "text/event-stream", {}, request=original) + assert isinstance(cancelled, JSONResponse) + assert cancelled.status_code == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED + + detached: Final = detach_request_from_client(original) + kept_alive: Final = await create_response(_slow_first_chunk(), "text/event-stream", {}, request=detached) + assert isinstance(kept_alive, StreamingResponse) + assert kept_alive.status_code == 200 + + +def test_detached_request_keeps_scope(): + original: Final = _request_whose_client_already_left() + detached: Final = detach_request_from_client(original) + assert detached.headers["x-litellm-call-id"] == "call-123" + assert detached.scope is original.scope From 9c8caae1e4a793d17493f4de6caebf2cfa1ad651 Mon Sep 17 00:00:00 2001 From: jesus Date: Mon, 7 Sep 2026 11:53:50 +0000 Subject: [PATCH 015/107] ci: add tests/test_litellm/proxy/response_polling to the proxy-endpoints shard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 + Makefile | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 33245ec5b5f..761960145eb 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -155,6 +155,7 @@ jobs: tests/test_litellm/proxy/vector_store_files_endpoints tests/test_litellm/proxy/video_endpoints tests/test_litellm/proxy/response_api_endpoints + tests/test_litellm/proxy/response_polling tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/ocr_endpoints tests/test_litellm/proxy/vector_store_endpoints diff --git a/Makefile b/Makefile index e17fdba3c85..b14aaaf896f 100644 --- a/Makefile +++ b/Makefile @@ -298,7 +298,7 @@ test-unit-proxy-core: install-test-deps $(UV_RUN) pytest tests/test_litellm/proxy/auth tests/test_litellm/proxy/client tests/test_litellm/proxy/db tests/test_litellm/proxy/hooks tests/test_litellm/proxy/policy_engine --tb=short -vv -n 4 --durations=20 test-unit-proxy-misc: install-test-deps - $(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/response_polling tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20 test-unit-integrations: install-test-deps $(UV_RUN) pytest tests/test_litellm/integrations --tb=short -vv -n 4 --durations=20 From e589a0ed82023d6fa550e6ddfa005c5ebbf4951e Mon Sep 17 00:00:00 2001 From: jesus Date: Mon, 7 Sep 2026 12:07:31 +0000 Subject: [PATCH 016/107] test(responses): cover the polling disconnect regression through background_streaming_task Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 - Makefile | 2 +- .../test_response_polling_handler.py | 58 +++++++++++++++++++ .../proxy/response_polling/__init__.py | 0 .../test_background_streaming.py | 52 ----------------- 5 files changed, 59 insertions(+), 54 deletions(-) delete mode 100644 tests/test_litellm/proxy/response_polling/__init__.py delete mode 100644 tests/test_litellm/proxy/response_polling/test_background_streaming.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 761960145eb..33245ec5b5f 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -155,7 +155,6 @@ jobs: tests/test_litellm/proxy/vector_store_files_endpoints tests/test_litellm/proxy/video_endpoints tests/test_litellm/proxy/response_api_endpoints - tests/test_litellm/proxy/response_polling tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/ocr_endpoints tests/test_litellm/proxy/vector_store_endpoints diff --git a/Makefile b/Makefile index b14aaaf896f..e17fdba3c85 100644 --- a/Makefile +++ b/Makefile @@ -298,7 +298,7 @@ test-unit-proxy-core: install-test-deps $(UV_RUN) pytest tests/test_litellm/proxy/auth tests/test_litellm/proxy/client tests/test_litellm/proxy/db tests/test_litellm/proxy/hooks tests/test_litellm/proxy/policy_engine --tb=short -vv -n 4 --durations=20 test-unit-proxy-misc: install-test-deps - $(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/response_polling tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20 test-unit-integrations: install-test-deps $(UV_RUN) pytest tests/test_litellm/integrations --tb=short -vv -n 4 --durations=20 diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/proxy_unit_tests/test_response_polling_handler.py index 7a4e0ca89c0..f5335508ce4 100644 --- a/tests/proxy_unit_tests/test_response_polling_handler.py +++ b/tests/proxy_unit_tests/test_response_polling_handler.py @@ -14,6 +14,7 @@ These tests ensure the polling handler correctly manages response state following the OpenAI Response API format. """ +import asyncio import json from datetime import datetime, timezone from typing import Any, Dict, Optional @@ -1663,6 +1664,63 @@ class TestBackgroundStreamingTerminalEvents: final_call = handler.update_state.call_args_list[-1] assert final_call.kwargs["status"] == "completed" + @pytest.mark.asyncio + async def test_polling_client_disconnect_does_not_cancel_upstream_call(self): + """The polling client hangs up right after getting its polling id. The detached task + must still stream the upstream response through the client-disconnect guards.""" + from litellm.proxy.common_request_processing import create_response + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + async def client_already_left(): + return {"type": "http.disconnect"} + + async def slow_upstream_stream(): + await asyncio.sleep(0.05) + for event in ( + {"type": "response.in_progress"}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "status": "completed", + "usage": {"input_tokens": 13, "output_tokens": 10}, + "model": "gpt-4o", + "output": [{"id": "item_1", "type": "message"}], + }, + }, + ): + yield f"data: {json.dumps(event)}\n\n" + + async def upstream_call_behind_disconnect_guard(**kwargs): + return await create_response( + slow_upstream_stream(), "text/event-stream", {}, request=kwargs["request"] + ) + + handler = AsyncMock(spec=ResponsePollingHandler) + kwargs = _make_background_streaming_kwargs("poll_7", handler) + kwargs["request"] = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [(b"x-litellm-call-id", b"call-123")], + "query_string": b"", + }, + client_already_left, + ) + + with patch( + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = upstream_call_behind_disconnect_guard + await background_streaming_task(**kwargs) + + final_call = handler.update_state.call_args_list[-1] + assert final_call.kwargs["status"] == "completed" + assert final_call.kwargs["usage"] == {"input_tokens": 13, "output_tokens": 10} + class TestEdgeCases: """Test edge cases and error scenarios""" diff --git a/tests/test_litellm/proxy/response_polling/__init__.py b/tests/test_litellm/proxy/response_polling/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/proxy/response_polling/test_background_streaming.py b/tests/test_litellm/proxy/response_polling/test_background_streaming.py deleted file mode 100644 index 98d914b1848..00000000000 --- a/tests/test_litellm/proxy/response_polling/test_background_streaming.py +++ /dev/null @@ -1,52 +0,0 @@ -import asyncio -from collections.abc import AsyncGenerator -from typing import Final - -import pytest -from fastapi import Request -from fastapi.responses import JSONResponse, StreamingResponse -from starlette.types import Message - -from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED -from litellm.proxy.common_request_processing import create_response -from litellm.proxy.response_polling.background_streaming import detach_request_from_client - - -def _request_whose_client_already_left() -> Request: - async def receive() -> Message: - return {"type": "http.disconnect"} - - scope: Final = { - "type": "http", - "method": "POST", - "path": "/v1/responses", - "headers": [(b"x-litellm-call-id", b"call-123")], - "query_string": b"", - } - return Request(scope, receive) - - -async def _slow_first_chunk() -> AsyncGenerator[str, None]: - await asyncio.sleep(0.05) - yield 'data: {"type": "response.created"}\n\n' - - -@pytest.mark.asyncio -async def test_detached_request_survives_client_disconnect_before_first_chunk(): - original: Final = _request_whose_client_already_left() - - cancelled: Final = await create_response(_slow_first_chunk(), "text/event-stream", {}, request=original) - assert isinstance(cancelled, JSONResponse) - assert cancelled.status_code == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED - - detached: Final = detach_request_from_client(original) - kept_alive: Final = await create_response(_slow_first_chunk(), "text/event-stream", {}, request=detached) - assert isinstance(kept_alive, StreamingResponse) - assert kept_alive.status_code == 200 - - -def test_detached_request_keeps_scope(): - original: Final = _request_whose_client_already_left() - detached: Final = detach_request_from_client(original) - assert detached.headers["x-litellm-call-id"] == "call-123" - assert detached.scope is original.scope From 0d27ec952aa6b1df6697d952b90e9ea4569c06e5 Mon Sep 17 00:00:00 2001 From: jesus Date: Mon, 7 Sep 2026 12:14:36 +0000 Subject: [PATCH 017/107] test(responses): annotate the processor patch for the test quality gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/proxy_unit_tests/test_response_polling_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/proxy_unit_tests/test_response_polling_handler.py index f5335508ce4..467c1332325 100644 --- a/tests/proxy_unit_tests/test_response_polling_handler.py +++ b/tests/proxy_unit_tests/test_response_polling_handler.py @@ -1711,7 +1711,7 @@ class TestBackgroundStreamingTerminalEvents: client_already_left, ) - with patch( + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: MockProcessor.return_value.base_process_llm_request = upstream_call_behind_disconnect_guard From 5c23d296d31c3eba4a362fd0fa0e8cca96f51426 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 16:51:56 -0700 Subject: [PATCH 018/107] fix(azure_ai): cost OCR relays per page so non-chat relays debit budgets --- .../azure_ai/passthrough/transformation.py | 45 +++++- .../base_llm/passthrough/transformation.py | 3 +- ...est_azure_ai_passthrough_transformation.py | 134 +++++++++++++++++- 3 files changed, 172 insertions(+), 10 deletions(-) diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index 4776bbe0755..4cb93b4f8f9 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -1,25 +1,28 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from litellm._logging import verbose_logger from litellm.llms.azure_ai.common_utils import ( AzureFoundryModelInfo, api_key_header_for_base, get_azure_ai_auth_headers, ) +from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, strip_leading_model_segment from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import StandardPassThroughResponseObject +from litellm.types.utils import CallTypes, StandardPassThroughResponseObject if TYPE_CHECKING: from httpx import URL, Response from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse from litellm.types.utils import CostResponseTypes @@ -75,6 +78,10 @@ def relayed_body(httpx_response: Response) -> str | dict: class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): + def __init__(self, ocr_config_for: Callable[[str], BaseOCRConfig | None] = get_azure_ai_ocr_config) -> None: + super().__init__() + self.ocr_config_for: Final = ocr_config_for + def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool: return bool(request_data.get("stream")) @@ -123,7 +130,7 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): request_data: Mapping[str, object], logging_obj: Logging, endpoint: str, - ) -> CostResponseTypes | StandardPassThroughResponseObject | None: + ) -> CostResponseTypes | OCRResponse | StandardPassThroughResponseObject | None: from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict @@ -136,8 +143,40 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): ) if chat_result is not None: return chat_result + ocr_result: Final = self.logged_ocr_response(model, httpx_response, logging_obj, endpoint) + if ocr_result is not None: + return ocr_result return StandardPassThroughResponseObject(response=relayed_body(httpx_response)) + def logged_ocr_response( + self, model: str, httpx_response: Response, logging_obj: Logging, endpoint: str + ) -> OCRResponse | None: + ocr_config: Final = self.ocr_config_for(model) + if ocr_config is None or httpx_response.status_code != 200: + return None + relayed_url: Final = httpx_response.request.url + relayed_origin: Final = str(relayed_url.copy_with(path="/", query=None, fragment=None)).rstrip("/") + ocr_url: Final = httpx.URL( + ocr_config.get_complete_url( + api_base=relayed_origin, + model=model, + optional_params={}, # mutable-ok: BaseOCRConfig wants a dict + ) + ) + known_prefixes: Final = (model, model_group_from(logging_obj.litellm_params)) + native_endpoint: Final = strip_leading_model_segment(endpoint, known_prefixes) + if f"/{native_endpoint.strip('/')}" != ocr_url.path: + return None + try: + ocr_response: Final = ocr_config.transform_ocr_response( + model=model, raw_response=httpx_response, logging_obj=logging_obj + ) + except (ValueError, AttributeError) as error: + verbose_logger.warning("azure_ai passthrough: OCR body from %s is not costable: %s", ocr_url, error) + return None + logging_obj.call_type = CallTypes.aocr.value # rebind-ok: routes cost calculation to the per-page OCR path + return ocr_response + def handle_logging_collected_chunks( self, all_chunks: Sequence[str], diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index d97ecddc75d..464c121ecc1 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -11,6 +11,7 @@ if TYPE_CHECKING: from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject from ..chat.transformation import BaseLLMException + from ..ocr.transformation import OCRResponse def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str: @@ -120,7 +121,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): request_data: dict, logging_obj: "LiteLLMLoggingObj", endpoint: str, - ) -> Optional["CostResponseTypes | StandardPassThroughResponseObject"]: + ) -> Optional["CostResponseTypes | OCRResponse | StandardPassThroughResponseObject"]: pass def handle_logging_collected_chunks( diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index e343d05c5df..644316c2e37 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -1,11 +1,14 @@ import json +from datetime import datetime from unittest.mock import MagicMock import httpx import pytest import litellm +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure_ai.passthrough.transformation import AzureAIPassthroughConfig +from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.types.utils import LlmProviders, ModelResponse from litellm.utils import ProviderConfigManager @@ -238,16 +241,135 @@ def _non_chat_logging_result(content: bytes, content_type: str): ) -def test_non_chat_relay_logs_the_parsed_body_so_spend_tracking_sees_the_call(): - result = _non_chat_logging_result(b'{"id":"parse-1","pages":[],"meta":{"billed_units":{"pages":1}}}', "application/json") - - assert result == {"response": {"id": "parse-1", "pages": [], "meta": {"billed_units": {"pages": 1}}}} - - def test_non_chat_relay_with_a_non_json_body_logs_the_raw_text(): assert _non_chat_logging_result(b"page one", "text/plain") == {"response": "page one"} +def _relay_logging_obj(model: str, api_base: str) -> Logging: + logging_obj = Logging( + model=model, + messages=[], + stream=False, + call_type="allm_passthrough_route", + start_time=datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.update_environment_variables( + model=model, + litellm_params={"api_base": api_base, "custom_llm_provider": "azure_ai"}, + optional_params={}, + custom_llm_provider="azure_ai", + ) + return logging_obj + + +def _relay_logging_result( + config: AzureAIPassthroughConfig, model: str, native_path: str, body, api_base: str = FOUNDRY_BASE, status_code: int = 200 +): + relayed_url = f"{FOUNDRY_BASE}/{native_path}?api-version=2024-05-01-preview" + logging_obj = _relay_logging_obj(model, api_base) + response = httpx.Response( + status_code=status_code, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request("POST", relayed_url), + ) + result = config.logging_non_streaming_response( + model=model, + custom_llm_provider="azure_ai", + httpx_response=response, + request_data={"model": model}, + logging_obj=logging_obj, + endpoint=f"{model}/{native_path}", + ) + return result, logging_obj + + +MISTRAL_OCR_BODY = { + "pages": [{"index": 0, "markdown": "page one"}, {"index": 1, "markdown": "page two"}], + "model": "mistral-document-ai-2512", + "usage_info": {"pages_processed": 2, "doc_size_bytes": 4321}, +} + + +def test_mistral_document_ai_relay_is_costed_per_page(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", MISTRAL_OCR_BODY + ) + per_page = litellm.get_model_info("azure_ai/mistral-document-ai-2512")["ocr_cost_per_page"] + + assert isinstance(result, OCRResponse) + assert result.usage_info.pages_processed == 2 + assert per_page > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(2 * per_page) + + +def test_ocr_route_under_a_models_api_base_is_still_recognised(): + result, _ = _relay_logging_result( + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + MISTRAL_OCR_BODY, + api_base=f"{FOUNDRY_BASE}/models", + ) + + assert isinstance(result, OCRResponse) + + +def test_relay_to_a_non_ocr_route_keeps_the_passthrough_object_and_call_type(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "mistral-document-ai-2512", "models/info", {"name": "mistral-document-ai-2512"} + ) + + assert result == {"response": {"name": "mistral-document-ai-2512"}} + assert logging_obj.call_type == "allm_passthrough_route" + + +COHERE_PARSE_BODY = {"id": "parse-1", "pages": [], "meta": {"billed_units": {"pages": 3}}} + + +def test_cohere_parse_relay_is_costed_per_billed_page(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "Cohere-parse-v5", "providers/cohere/v2/parse", COHERE_PARSE_BODY + ) + per_page = litellm.get_model_info("azure_ai/Cohere-parse-v5")["ocr_cost_per_page"] + + assert isinstance(result, OCRResponse) + assert result.usage_info.pages_processed == 3 + assert logging_obj.call_type == "aocr" + assert per_page > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(3 * per_page) + + +def test_deployment_without_an_ocr_config_is_never_costed_as_ocr(): + config = AzureAIPassthroughConfig(ocr_config_for=lambda model: None) + result, logging_obj = _relay_logging_result( + config, "mistral-document-ai-2512", "providers/mistral/azure/ocr", MISTRAL_OCR_BODY + ) + + assert result == {"response": MISTRAL_OCR_BODY} + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_accepted_ocr_job_without_a_result_body_is_not_costed(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", {"status": "running"}, status_code=202 + ) + + assert result == {"response": {"status": "running"}} + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_unparseable_ocr_body_falls_back_to_the_passthrough_object(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", ["not", "an", "ocr", "body"] + ) + + assert result == {"response": '["not", "an", "ocr", "body"]'} + assert logging_obj.call_type == "allm_passthrough_route" + + def test_streaming_chat_completion_chunks_are_costed_like_azure(): head = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-5.4-mini"} chunks = [ From 6f99917b332a1038ac83f5f5a5f8614fa573b1e6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:57:54 -0700 Subject: [PATCH 019/107] fix(router): rewrite multi-segment model groups as whole passthrough path segments --- .../base_llm/passthrough/transformation.py | 4 +- tests/test_litellm/test_router.py | 46 ++++++++++++------- 2 files changed, 32 insertions(+), 18 deletions(-) diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 464c121ecc1..2c21df01b5b 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -1,3 +1,4 @@ +import re from abc import abstractmethod from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, Union @@ -27,7 +28,8 @@ def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> def replace_path_segment(endpoint: str, segment: str, replacement: str) -> str: - return "/".join(replacement if part == segment else part for part in endpoint.split("/")) + bounded_segment: Final = re.compile(rf"(? Date: Mon, 7 Sep 2026 18:00:48 -0700 Subject: [PATCH 020/107] fix(azure): let the caller's api-version win over the deployment's on passthrough relays --- litellm/llms/azure/passthrough/transformation.py | 4 ++-- .../test_azure_passthrough_transformation.py | 10 ++++++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 983033a8b14..062c60d7231 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -61,11 +61,11 @@ class AzurePassthroughConfig(BasePassthroughConfig): routed_endpoint: Final = replace_path_segment(endpoint, model_group, model) if model_group else endpoint native_endpoint: Final = strip_leading_model_segment(routed_endpoint, (model,)) + caller_api_version: Final = request_query_params.get("api-version") if request_query_params else None complete_url: Final = BaseAzureLLM._get_base_azure_url( api_base=base_target_url, - litellm_params=litellm_params, + litellm_params={**litellm_params, "api_version": caller_api_version or litellm_params.get("api_version")}, route=native_endpoint, - default_api_version=request_query_params.get("api-version") if request_query_params else None, ) return ( httpx.URL(complete_url), diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 8ffd632fe85..262a281aeb3 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -207,18 +207,24 @@ def _complete_url(request_query_params: dict, litellm_params: dict) -> httpx.URL return url -def test_azure_passthrough_url_falls_back_to_the_callers_api_version(): +def test_azure_passthrough_url_forwards_the_callers_api_version(): url = _complete_url(request_query_params={"api-version": "2025-04-01-preview"}, litellm_params={}) assert url.path == "/openai/deployments/gpt-4.1-mini/chat/completions" assert url.params["api-version"] == "2025-04-01-preview" -def test_azure_passthrough_url_prefers_the_deployments_api_version(): +def test_azure_passthrough_url_prefers_the_callers_api_version_over_the_deployments(): url = _complete_url( request_query_params={"api-version": "2025-04-01-preview"}, litellm_params={"api_version": "2024-10-21"} ) + assert url.params["api-version"] == "2025-04-01-preview" + + +def test_azure_passthrough_url_fills_in_the_deployments_api_version_when_the_caller_sends_none(): + url = _complete_url(request_query_params={}, litellm_params={"api_version": "2024-10-21"}) + assert url.params["api-version"] == "2024-10-21" From 170fece7db9b4997fab48cd716811975e43f11f3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:08:25 -0700 Subject: [PATCH 021/107] fix(token_counter): release the GIL for HuggingFace counts and cap exact counting per string Both proxy token counting endpoints already count in a worker thread, but the HuggingFace tokenizer's encode holds the GIL for the whole call, so a 600k-token count on a Claude model still froze the event loop for up to 0.8 s and every other request with it. Count through encode_batch_fast, which releases the GIL, and tokenize at most TOKEN_COUNTER_MAX_EXACT_CHARS characters of any one string (default 4,000,000), scaling the exact count of that prefix by the string's length above it so the largest payloads stay bounded. --- litellm/constants.py | 6 ++ litellm/litellm_core_utils/token_counter.py | 25 ++++++- .../litellm_core_utils/test_token_counter.py | 71 ++++++++++++++++++- tests/test_litellm/proxy/test_proxy_server.py | 34 ++++++++- 4 files changed, 132 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index d53686e5e5b..a9112c90a54 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -391,6 +391,12 @@ TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS: Final = get_env_int_in_range( minimum=1, maximum=TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS, ) +TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range( + "TOKEN_COUNTER_MAX_EXACT_CHARS", + default=4_000_000, + minimum=1, + maximum=1_000_000_000, +) MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512)) MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512)) OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 3732ffd734c..d526f0f2996 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -8,6 +8,7 @@ from typing import Final, Literal, cast import httpx import tiktoken +from tokenizers import Tokenizer import litellm from litellm import verbose_logger @@ -21,6 +22,7 @@ from litellm.constants import ( MAX_TILE_HEIGHT, MAX_TILE_WIDTH, TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS, + TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.default_encoding import encoding as default_encoding from litellm.litellm_core_utils.url_utils import safe_get @@ -538,9 +540,28 @@ def _count_extra( return num_tokens +def _get_extrapolating_count_function( + count_exactly: TokenCounterFunction, + max_exact_chars: int = TOKEN_COUNTER_MAX_EXACT_CHARS, +) -> TokenCounterFunction: + def count_tokens(text: str) -> int: + if len(text) <= max_exact_chars: + return count_exactly(text) + return round(count_exactly(text[:max_exact_chars]) * len(text) / max_exact_chars) + + return count_tokens + + def _get_count_function( model: str | None, custom_tokenizer: dict | SelectTokenizerResponse | None = None, +) -> TokenCounterFunction: + return _get_extrapolating_count_function(_get_exact_count_function(model, custom_tokenizer)) + + +def _get_exact_count_function( + model: str | None, + custom_tokenizer: dict | SelectTokenizerResponse | None = None, ) -> TokenCounterFunction: """ Get the function to count tokens based on the model and custom tokenizer.""" @@ -549,10 +570,10 @@ def _get_count_function( if model is not None or custom_tokenizer is not None: tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model) if tokenizer_json["type"] == "huggingface_tokenizer": + tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"] def count_tokens(text: str) -> int: - enc: Final = tokenizer_json["tokenizer"].encode(text) - return len(enc.ids) + return len(tokenizer.encode_batch_fast([text])[0]) return count_tokens elif tokenizer_json["type"] == "openai_tokenizer": diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 4694fa8fbed..fddeacc37f9 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1,8 +1,10 @@ #### What this tests #### # This tests litellm.token_counter.token_counter() function +import asyncio import importlib import time import traceback +from typing import Final from unittest.mock import MagicMock import pytest @@ -14,7 +16,12 @@ import litellm from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants -from litellm.litellm_core_utils.token_counter import _get_tiktoken_count_function +from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.token_counter import ( + _get_exact_count_function, + _get_extrapolating_count_function, + _get_tiktoken_count_function, +) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from tests.large_text import text from tests.test_litellm.litellm_core_utils.messages_with_counts import ( @@ -120,6 +127,68 @@ def test_valid_chunk_size_config_is_honoured(monkeypatch): importlib.reload(litellm.constants) +async def _loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]: + async def wake_lag() -> float: + started: Final = time.perf_counter() + await asyncio.sleep(0.001) + return time.perf_counter() - started - 0.001 + + return tuple([await wake_lag() for _ in iter(until.is_set, True)]) + + +async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free(): + counted: Final = asyncio.Event() + + async def count_off_the_loop() -> tuple[int, float]: + started: Final = time.perf_counter() + try: + tokens: Final = await asyncify(token_counter_new)(model="claude-fable-5", text=text * 100) + return tokens, time.perf_counter() - started + finally: + counted.set() + + (tokens, took), lags = await asyncio.gather(count_off_the_loop(), _loop_wake_lags(counted)) + + assert tokens > 0 + assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" + + +@pytest.mark.parametrize( + ("max_exact_chars", "expected"), + [(1_000, 10_000), (5_000, 6_000), (10_000, 6_000)], +) +def test_count_above_the_cap_scales_the_exact_count_of_the_prefix(max_exact_chars, expected): + def count_exactly(chunk: str) -> int: + return chunk.count("a") + len(chunk) + + count_tokens: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars) + + assert count_tokens("a" * 1_000 + "b" * 4_000) == expected + + +def test_token_counter_applies_the_default_cap(): + max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS + prefix: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] + over_the_cap: Final = prefix + "a" * 200_000 + scaled: Final = round(token_counter_new(model="gpt-5.6", text=prefix) * len(over_the_cap) / max_exact_chars) + + assert token_counter_new(model="gpt-5.6", text=over_the_cap) == scaled + assert _get_exact_count_function("gpt-5.6")(over_the_cap) != scaled + + +@pytest.mark.parametrize( + ("configured", "expected"), + [("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)], +) +def test_max_exact_chars_config_is_honoured(monkeypatch, configured, expected): + monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured) + try: + assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected + finally: + monkeypatch.delenv("TOKEN_COUNTER_MAX_EXACT_CHARS") + importlib.reload(litellm.constants) + + def test_token_counter_with_prefix(): messages = [ {"role": "user", "content": "Who won the world cup in 2022?"}, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index b7bb58378d4..20b41cf1530 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -6,6 +6,7 @@ import os import re import socket import subprocess +import time import types from datetime import datetime, timedelta, timezone from pathlib import Path @@ -28,7 +29,7 @@ from litellm.caching.caching import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app, initialize from litellm.utils import _invalidate_model_cost_lowercase_map @@ -12857,3 +12858,34 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough() import litellm.proxy.proxy_server as ps assert ps.general_settings["enable_openai_websocket_passthrough"] is False + + +async def _loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]: + async def wake_lag() -> float: + started: Final = time.perf_counter() + await asyncio.sleep(0.001) + return time.perf_counter() - started - 0.001 + + return tuple([await wake_lag() for _ in iter(until.is_set, True)]) + + +async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch): + from tests.large_text import text + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + counted: Final = asyncio.Event() + + async def count_off_the_loop() -> tuple[int, float]: + started: Final = time.perf_counter() + try: + response: Final = await proxy_server_module.token_counter( + TokenCountRequest(model="claude-fable-5", prompt=text * 100) + ) + return response.total_tokens, time.perf_counter() - started + finally: + counted.set() + + (tokens, took), lags = await asyncio.gather(count_off_the_loop(), _loop_wake_lags(counted)) + + assert tokens > 0 + assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" From 17b70035924b7bb2e97065a7f3512c540e105d66 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:11:17 -0700 Subject: [PATCH 022/107] fix(azure_ai): cost embeddings, responses, images, and rerank relays instead of logging zero --- .../llms/azure/passthrough/transformation.py | 45 ++++++- .../azure_ai/passthrough/transformation.py | 32 ++++- .../base_llm/passthrough/transformation.py | 20 ++- .../test_azure_passthrough_transformation.py | 120 +++++++++++++++--- ...est_azure_ai_passthrough_transformation.py | 68 +++++++++- 5 files changed, 261 insertions(+), 24 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 062c60d7231..d7078d94a0b 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,4 +1,5 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Optional import httpx @@ -9,12 +10,14 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.passthrough.transformation import ( BasePassthroughConfig, + relayed_json_object, replace_path_segment, strip_leading_model_segment, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse if TYPE_CHECKING: from httpx import URL @@ -38,6 +41,40 @@ def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, obj return details.request_data.messages if details.request_data else None +@dataclass(frozen=True, slots=True) +class OpenAIRelayShape: + path_suffix: str + call_type: CallTypes + parse: Callable[[Mapping[str, object]], EmbeddingResponse | ImageResponse | ResponsesAPIResponse] + + +OPENAI_RELAY_SHAPES: Final = ( + OpenAIRelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate), + OpenAIRelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate), + OpenAIRelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate), +) + + +def logged_openai_response( + httpx_response: Response, logging_obj: Logging, endpoint: str +) -> EmbeddingResponse | ImageResponse | ResponsesAPIResponse | None: + relayed_path: Final = f"/{endpoint.strip('/')}" + shape: Final = next( + (candidate for candidate in OPENAI_RELAY_SHAPES if relayed_path.endswith(candidate.path_suffix)), None + ) + body: Final = relayed_json_object(httpx_response) if shape else None + if shape is None or body is None: + return None + try: + parsed: Final = shape.parse(body) + except ValidationError: + return None + logging_obj.call_type = ( + shape.call_type.value + ) # rebind-ok: routes cost calculation to the relayed shape's pricing path + return parsed + + class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return bool(request_data.get("stream")) @@ -114,13 +151,13 @@ class AzurePassthroughConfig(BasePassthroughConfig): request_data: dict, logging_obj: Logging, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> Optional["CostResponseTypes | ResponsesAPIResponse"]: from litellm import encoding from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.utils import ModelResponse if "chat/completions" not in endpoint: - return None + return logged_openai_response(httpx_response, logging_obj, endpoint) openai_chat_config: Final = OpenAIGPTConfig() diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index 4cb93b4f8f9..c49e10566ed 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -14,8 +14,13 @@ from litellm.llms.azure_ai.common_utils import ( get_azure_ai_auth_headers, ) from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, strip_leading_model_segment +from litellm.llms.base_llm.passthrough.transformation import ( + BasePassthroughConfig, + relayed_json_object, + strip_leading_model_segment, +) from litellm.types.llms.openai import AllMessageValues +from litellm.types.rerank import RerankResponse from litellm.types.utils import CallTypes, StandardPassThroughResponseObject if TYPE_CHECKING: @@ -23,6 +28,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse + from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import CostResponseTypes @@ -77,6 +83,18 @@ def relayed_body(httpx_response: Response) -> str | dict: return body if isinstance(body, dict) else httpx_response.text +def logged_rerank_response(httpx_response: Response, logging_obj: Logging, endpoint: str) -> RerankResponse | None: + body: Final = relayed_json_object(httpx_response) if f"/{endpoint.strip('/')}".endswith("/rerank") else None + if body is None: + return None + try: + rerank_response: Final = RerankResponse.model_validate(body) + except ValidationError: + return None + logging_obj.call_type = CallTypes.arerank.value # rebind-ok: routes cost calculation to the per-query rerank path + return rerank_response + + class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): def __init__(self, ocr_config_for: Callable[[str], BaseOCRConfig | None] = get_azure_ai_ocr_config) -> None: super().__init__() @@ -130,7 +148,14 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): request_data: Mapping[str, object], logging_obj: Logging, endpoint: str, - ) -> CostResponseTypes | OCRResponse | StandardPassThroughResponseObject | None: + ) -> ( + CostResponseTypes + | OCRResponse + | RerankResponse + | ResponsesAPIResponse + | StandardPassThroughResponseObject + | None + ): from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict @@ -146,6 +171,9 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): ocr_result: Final = self.logged_ocr_response(model, httpx_response, logging_obj, endpoint) if ocr_result is not None: return ocr_result + rerank_result: Final = logged_rerank_response(httpx_response, logging_obj, endpoint) + if rerank_result is not None: + return rerank_result return StandardPassThroughResponseObject(response=relayed_body(httpx_response)) def logged_ocr_response( diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 2c21df01b5b..975b74137db 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -3,18 +3,25 @@ from abc import abstractmethod from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, Union +from pydantic import TypeAdapter, ValidationError + from ..base_utils import BaseLLMModelInfo if TYPE_CHECKING: from httpx import URL, Headers, Response from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.rerank import RerankResponse from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject from ..chat.transformation import BaseLLMException from ..ocr.transformation import OCRResponse +RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) + + def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str: path: Final = endpoint.lstrip("/") for model_name in model_names: @@ -32,6 +39,15 @@ def replace_path_segment(endpoint: str, segment: str, replacement: str) -> str: return bounded_segment.sub(lambda _: replacement, endpoint) +def relayed_json_object(httpx_response: "Response") -> Mapping[str, object] | None: + if httpx_response.status_code != 200: + return None + try: + return RELAYED_JSON_OBJECT.validate_python(httpx_response.json()) + except (ValueError, ValidationError): + return None + + class BasePassthroughConfig(BaseLLMModelInfo): @abstractmethod def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: @@ -123,7 +139,9 @@ class BasePassthroughConfig(BaseLLMModelInfo): request_data: dict, logging_obj: "LiteLLMLoggingObj", endpoint: str, - ) -> Optional["CostResponseTypes | OCRResponse | StandardPassThroughResponseObject"]: + ) -> Optional[ + "CostResponseTypes | OCRResponse | RerankResponse | ResponsesAPIResponse | StandardPassThroughResponseObject" + ]: pass def handle_logging_collected_chunks( diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 262a281aeb3..2723a13c3fc 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -1,15 +1,16 @@ import json +from datetime import datetime from unittest.mock import MagicMock import httpx import pytest import litellm - - +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig -from litellm.types.utils import ModelResponse +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import EmbeddingResponse, ModelResponse def _azure_chat_completion_body(): @@ -77,25 +78,112 @@ def test_azure_passthrough_logging_non_streaming_response_chat_completions(): assert result.usage.total_tokens == 18 -def test_azure_passthrough_logging_non_streaming_response_unknown_endpoint_returns_none(): - """ - Endpoints other than chat/completions (responses, messages, images) fall - through to None — matches base-class behavior and Bedrock's "unknown - endpoint" handling. Not a regression; just scoping. - """ - config = AzurePassthroughConfig() - logging_obj = MagicMock() - - result = config.logging_non_streaming_response( - model="gpt-4.1-mini", +def _relay_logging_obj(model: str) -> Logging: + logging_obj = Logging( + model=model, + messages=[], + stream=False, + call_type="allm_passthrough_route", + start_time=datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.update_environment_variables( + model=model, + litellm_params={"api_base": "https://my-resource.openai.azure.com", "custom_llm_provider": "azure"}, + optional_params={}, custom_llm_provider="azure", - httpx_response=_make_httpx_response(_azure_chat_completion_body()), + ) + return logging_obj + + +def _relay_logging_result(model: str, endpoint: str, body, status_code: int = 200): + logging_obj = _relay_logging_obj(model) + response = httpx.Response( + status_code=status_code, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request("POST", f"https://my-resource.openai.azure.com/{endpoint}?api-version=2025-04-01-preview"), + ) + result = AzurePassthroughConfig().logging_non_streaming_response( + model=model, + custom_llm_provider="azure", + httpx_response=response, request_data={}, logging_obj=logging_obj, - endpoint="openai/responses", + endpoint=endpoint, + ) + return result, logging_obj + + +EMBEDDINGS_BODY = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 1000, "total_tokens": 1000}, +} + +RESPONSES_BODY = { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4.1-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, +} + + +def test_azure_passthrough_embeddings_relay_is_costed_per_input_token(): + result, logging_obj = _relay_logging_result( + "text-embedding-3-small", "openai/deployments/text-embedding-3-small/embeddings", EMBEDDINGS_BODY + ) + per_token = litellm.get_model_info("azure/text-embedding-3-small")["input_cost_per_token"] + + assert isinstance(result, EmbeddingResponse) + assert logging_obj.call_type == "aembedding" + assert per_token > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(1000 * per_token) + + +def test_azure_passthrough_responses_relay_is_costed_per_token(): + result, logging_obj = _relay_logging_result("gpt-4.1-mini", "openai/responses", RESPONSES_BODY) + info = litellm.get_model_info("azure/gpt-4.1-mini") + + assert isinstance(result, ResponsesAPIResponse) + assert logging_obj.call_type == "aresponses" + assert logging_obj._response_cost_calculator(result=result) == pytest.approx( + 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] + ) + + +def test_azure_passthrough_failed_embeddings_relay_is_not_costed(): + result, logging_obj = _relay_logging_result( + "text-embedding-3-small", + "openai/deployments/text-embedding-3-small/embeddings", + {"error": {"code": "429", "message": "rate limited"}}, + status_code=429, ) assert result is None + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_azure_passthrough_logging_non_streaming_response_unknown_endpoint_returns_none(): + result, logging_obj = _relay_logging_result( + "gpt-4o-mini-tts", "openai/deployments/gpt-4o-mini-tts/audio/speech", {"audio": "..."} + ) + + assert result is None + assert logging_obj.call_type == "allm_passthrough_route" def _sse_line(payload: dict) -> str: diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 644316c2e37..788aced50f8 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -9,7 +9,8 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure_ai.passthrough.transformation import AzureAIPassthroughConfig from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.types.utils import LlmProviders, ModelResponse +from litellm.types.rerank import RerankResponse +from litellm.types.utils import EmbeddingResponse, ImageResponse, LlmProviders, ModelResponse from litellm.utils import ProviderConfigManager FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" @@ -370,6 +371,71 @@ def test_unparseable_ocr_body_falls_back_to_the_passthrough_object(): assert logging_obj.call_type == "allm_passthrough_route" +EMBEDDINGS_BODY = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "embed-v-4-0", + "usage": {"prompt_tokens": 1200, "total_tokens": 1200}, +} + +RERANK_BODY = { + "id": "rerank-1", + "results": [{"index": 1, "relevance_score": 0.9}, {"index": 0, "relevance_score": 0.2}], + "meta": {"api_version": {"version": "2"}, "billed_units": {"search_units": 2}}, +} + +IMAGE_BODY = {"created": 1, "data": [{"b64_json": "AAAA"}]} + + +def test_foundry_embeddings_relay_is_costed_per_input_token(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "embed-v-4-0", "models/embeddings", EMBEDDINGS_BODY + ) + per_token = litellm.get_model_info("azure_ai/embed-v-4-0")["input_cost_per_token"] + + assert isinstance(result, EmbeddingResponse) + assert logging_obj.call_type == "aembedding" + assert per_token > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(1200 * per_token) + + +def test_cohere_rerank_relay_is_costed_per_search_unit(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "cohere-rerank-v4.0-fast", "providers/cohere/v2/rerank", RERANK_BODY + ) + per_query = litellm.get_model_info("azure_ai/cohere-rerank-v4.0-fast")["input_cost_per_query"] + + assert isinstance(result, RerankResponse) + assert logging_obj.call_type == "arerank" + assert per_query > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(2 * per_query) + + +def test_image_generation_relay_is_costed_per_image(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "FLUX.2-pro", "openai/deployments/FLUX.2-pro/images/generations", IMAGE_BODY + ) + per_image = litellm.get_model_info("azure_ai/FLUX.2-pro")["output_cost_per_image"] + + assert isinstance(result, ImageResponse) + assert logging_obj.call_type == "aimage_generation" + assert per_image > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(per_image) + + +def test_rejected_rerank_relay_keeps_the_passthrough_object_and_call_type(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), + "cohere-rerank-v4.0-fast", + "providers/cohere/v2/rerank", + {"message": "invalid request"}, + status_code=400, + ) + + assert result == {"response": {"message": "invalid request"}} + assert logging_obj.call_type == "allm_passthrough_route" + + def test_streaming_chat_completion_chunks_are_costed_like_azure(): head = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-5.4-mini"} chunks = [ From 3829ebdce54a0330ca2c9d60a5bbee94c2b09382 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:26:23 -0700 Subject: [PATCH 023/107] fix(token_counter): estimate over-cap strings from evenly spaced samples instead of the prefix A string above TOKEN_COUNTER_MAX_EXACT_CHARS was counted from its first cap characters and scaled, so a string whose start tokenizes unlike its end got a skewed count, and that count reaches fallback billing when the provider sends no usage. The estimate now tokenizes 16 evenly spaced samples that together total the cap and scales their sum by the string's length, keeping the same bound on work while tracking the whole string --- litellm/litellm_core_utils/token_counter.py | 16 ++++++++- .../litellm_core_utils/test_token_counter.py | 36 +++++++++++-------- 2 files changed, 37 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index d526f0f2996..e3a0d4c785a 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -315,6 +315,8 @@ def calculate_img_tokens( TokenCounterFunction = Callable[[str], int] + +EXTRAPOLATION_SAMPLES: Final = 16 """ Type for a function that counts tokens in a string. """ @@ -547,11 +549,23 @@ def _get_extrapolating_count_function( def count_tokens(text: str) -> int: if len(text) <= max_exact_chars: return count_exactly(text) - return round(count_exactly(text[:max_exact_chars]) * len(text) / max_exact_chars) + samples: Final = _evenly_spaced_samples(text, max_exact_chars) + sampled_chars: Final = sum(len(sample) for sample in samples) + return round(sum(count_exactly(sample) for sample in samples) * len(text) / sampled_chars) return count_tokens +def _evenly_spaced_samples(text: str, total_chars: int) -> tuple[str, ...]: + sample_count: Final = min(EXTRAPOLATION_SAMPLES, total_chars) + sample_chars: Final = total_chars // sample_count + last_start: Final = len(text) - sample_chars + return tuple( + text[start : start + sample_chars] + for start in (last_start * index // max(sample_count - 1, 1) for index in range(sample_count)) + ) + + def _get_count_function( model: str | None, custom_tokenizer: dict | SelectTokenizerResponse | None = None, diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index fddeacc37f9..5da8413ac9b 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -153,27 +153,35 @@ async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free() assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" -@pytest.mark.parametrize( - ("max_exact_chars", "expected"), - [(1_000, 10_000), (5_000, 6_000), (10_000, 6_000)], -) -def test_count_above_the_cap_scales_the_exact_count_of_the_prefix(max_exact_chars, expected): - def count_exactly(chunk: str) -> int: - return chunk.count("a") + len(chunk) +@pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500]) +def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars): + count_exactly: Final = MagicMock(side_effect=lambda chunk: chunk.count("a") + len(chunk)) + front_heavy: Final = "a" * 1_000 + "b" * 4_000 + exact: Final = 1_000 + len(front_heavy) - count_tokens: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars) + estimate: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars)(front_heavy) - assert count_tokens("a" * 1_000 + "b" * 4_000) == expected + assert abs(estimate - exact) <= exact // 100 + assert sum(len(call.args[0]) for call in count_exactly.call_args_list) <= max_exact_chars + + +def test_count_at_or_below_the_cap_is_exact(): + count_exactly: Final = MagicMock(side_effect=len) + + assert _get_extrapolating_count_function(count_exactly, max_exact_chars=5_000)("a" * 5_000) == 5_000 + assert count_exactly.call_args_list == [(("a" * 5_000,),)] def test_token_counter_applies_the_default_cap(): max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS - prefix: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] - over_the_cap: Final = prefix + "a" * 200_000 - scaled: Final = round(token_counter_new(model="gpt-5.6", text=prefix) * len(over_the_cap) / max_exact_chars) + prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] + over_the_cap: Final = prose + "a" * 200_000 + exact: Final = _get_exact_count_function("gpt-5.6")(over_the_cap) - assert token_counter_new(model="gpt-5.6", text=over_the_cap) == scaled - assert _get_exact_count_function("gpt-5.6")(over_the_cap) != scaled + estimate: Final = token_counter_new(model="gpt-5.6", text=over_the_cap) + + assert estimate != exact + assert abs(estimate - exact) <= exact // 100 @pytest.mark.parametrize( From 0d89873daf9e82a692825820adca2a47c0aab8dc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:36:09 -0700 Subject: [PATCH 024/107] test(token_counter): type the parametrized cap test arguments --- tests/test_litellm/litellm_core_utils/test_token_counter.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 5da8413ac9b..db87bc52600 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -154,7 +154,7 @@ async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free() @pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500]) -def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars): +def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars: int): count_exactly: Final = MagicMock(side_effect=lambda chunk: chunk.count("a") + len(chunk)) front_heavy: Final = "a" * 1_000 + "b" * 4_000 exact: Final = 1_000 + len(front_heavy) @@ -188,7 +188,7 @@ def test_token_counter_applies_the_default_cap(): ("configured", "expected"), [("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)], ) -def test_max_exact_chars_config_is_honoured(monkeypatch, configured, expected): +def test_max_exact_chars_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured) try: assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected From 12d9413860e96f1cf2b0d138466b34fcae53f925 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:37:58 -0700 Subject: [PATCH 025/107] refactor(passthrough): share the relay shape table and price FLUX 2 provider relays --- .../llms/azure/passthrough/transformation.py | 44 +++----------- .../azure_ai/passthrough/transformation.py | 36 ++++-------- .../base_llm/passthrough/transformation.py | 58 ++++++++++++++----- ...est_azure_ai_passthrough_transformation.py | 49 +++++++++++++--- 4 files changed, 106 insertions(+), 81 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index d7078d94a0b..bdb854bb87a 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,5 +1,4 @@ -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional import httpx @@ -10,7 +9,8 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.passthrough.transformation import ( BasePassthroughConfig, - relayed_json_object, + RelayShape, + logged_relay_shape, replace_path_segment, strip_leading_model_segment, ) @@ -22,6 +22,7 @@ from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse if TYPE_CHECKING: from httpx import URL + from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse from litellm.types.utils import CostResponseTypes @@ -41,40 +42,13 @@ def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, obj return details.request_data.messages if details.request_data else None -@dataclass(frozen=True, slots=True) -class OpenAIRelayShape: - path_suffix: str - call_type: CallTypes - parse: Callable[[Mapping[str, object]], EmbeddingResponse | ImageResponse | ResponsesAPIResponse] - - OPENAI_RELAY_SHAPES: Final = ( - OpenAIRelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate), - OpenAIRelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate), - OpenAIRelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate), + RelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate), + RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate), + RelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate), ) -def logged_openai_response( - httpx_response: Response, logging_obj: Logging, endpoint: str -) -> EmbeddingResponse | ImageResponse | ResponsesAPIResponse | None: - relayed_path: Final = f"/{endpoint.strip('/')}" - shape: Final = next( - (candidate for candidate in OPENAI_RELAY_SHAPES if relayed_path.endswith(candidate.path_suffix)), None - ) - body: Final = relayed_json_object(httpx_response) if shape else None - if shape is None or body is None: - return None - try: - parsed: Final = shape.parse(body) - except ValidationError: - return None - logging_obj.call_type = ( - shape.call_type.value - ) # rebind-ok: routes cost calculation to the relayed shape's pricing path - return parsed - - class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return bool(request_data.get("stream")) @@ -151,13 +125,13 @@ class AzurePassthroughConfig(BasePassthroughConfig): request_data: dict, logging_obj: Logging, endpoint: str, - ) -> Optional["CostResponseTypes | ResponsesAPIResponse"]: + ) -> Optional["LoggedRelayResponse"]: from litellm import encoding from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.utils import ModelResponse if "chat/completions" not in endpoint: - return logged_openai_response(httpx_response, logging_obj, endpoint) + return logged_relay_shape(OPENAI_RELAY_SHAPES, httpx_response, logging_obj, endpoint) openai_chat_config: Final = OpenAIGPTConfig() diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index c49e10566ed..e854874eaee 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -16,19 +16,20 @@ from litellm.llms.azure_ai.common_utils import ( from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config from litellm.llms.base_llm.passthrough.transformation import ( BasePassthroughConfig, - relayed_json_object, + RelayShape, + logged_relay_shape, strip_leading_model_segment, ) from litellm.types.llms.openai import AllMessageValues from litellm.types.rerank import RerankResponse -from litellm.types.utils import CallTypes, StandardPassThroughResponseObject +from litellm.types.utils import CallTypes, ImageResponse, StandardPassThroughResponseObject if TYPE_CHECKING: from httpx import URL, Response from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse - from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse from litellm.types.utils import CostResponseTypes @@ -83,16 +84,10 @@ def relayed_body(httpx_response: Response) -> str | dict: return body if isinstance(body, dict) else httpx_response.text -def logged_rerank_response(httpx_response: Response, logging_obj: Logging, endpoint: str) -> RerankResponse | None: - body: Final = relayed_json_object(httpx_response) if f"/{endpoint.strip('/')}".endswith("/rerank") else None - if body is None: - return None - try: - rerank_response: Final = RerankResponse.model_validate(body) - except ValidationError: - return None - logging_obj.call_type = CallTypes.arerank.value # rebind-ok: routes cost calculation to the per-query rerank path - return rerank_response +FOUNDRY_RELAY_SHAPES: Final = ( + RelayShape("/rerank", CallTypes.arerank, RerankResponse.model_validate), + RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, ImageResponse.model_validate), +) class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): @@ -148,14 +143,7 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): request_data: Mapping[str, object], logging_obj: Logging, endpoint: str, - ) -> ( - CostResponseTypes - | OCRResponse - | RerankResponse - | ResponsesAPIResponse - | StandardPassThroughResponseObject - | None - ): + ) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None: from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict @@ -171,9 +159,9 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): ocr_result: Final = self.logged_ocr_response(model, httpx_response, logging_obj, endpoint) if ocr_result is not None: return ocr_result - rerank_result: Final = logged_rerank_response(httpx_response, logging_obj, endpoint) - if rerank_result is not None: - return rerank_result + foundry_result: Final = logged_relay_shape(FOUNDRY_RELAY_SHAPES, httpx_response, logging_obj, endpoint) + if foundry_result is not None: + return foundry_result return StandardPassThroughResponseObject(response=relayed_body(httpx_response)) def logged_ocr_response( diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 975b74137db..93ec5a09e62 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -1,10 +1,15 @@ +from __future__ import annotations + import re from abc import abstractmethod -from collections.abc import Mapping -from typing import TYPE_CHECKING, Final, Optional, Union +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, TypeAlias from pydantic import TypeAdapter, ValidationError +from litellm.types.utils import CallTypes + from ..base_utils import BaseLLMModelInfo if TYPE_CHECKING: @@ -18,6 +23,8 @@ if TYPE_CHECKING: from ..chat.transformation import BaseLLMException from ..ocr.transformation import OCRResponse + LoggedRelayResponse: TypeAlias = CostResponseTypes | RerankResponse | ResponsesAPIResponse + RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) @@ -39,7 +46,7 @@ def replace_path_segment(endpoint: str, segment: str, replacement: str) -> str: return bounded_segment.sub(lambda _: replacement, endpoint) -def relayed_json_object(httpx_response: "Response") -> Mapping[str, object] | None: +def relayed_json_object(httpx_response: Response) -> Mapping[str, object] | None: if httpx_response.status_code != 200: return None try: @@ -48,6 +55,31 @@ def relayed_json_object(httpx_response: "Response") -> Mapping[str, object] | No return None +@dataclass(frozen=True, slots=True) +class RelayShape: + path_suffix: str + call_type: CallTypes + parse: Callable[[Mapping[str, object]], LoggedRelayResponse] + + +def logged_relay_shape( + shapes: Sequence[RelayShape], httpx_response: Response, logging_obj: LiteLLMLoggingObj, endpoint: str +) -> LoggedRelayResponse | None: + relayed_path: Final = f"/{endpoint.strip('/')}" + shape: Final = next((candidate for candidate in shapes if relayed_path.endswith(candidate.path_suffix)), None) + body: Final = relayed_json_object(httpx_response) if shape else None + if shape is None or body is None: + return None + try: + parsed: Final = shape.parse(body) + except ValidationError: + return None + logging_obj.call_type = ( + shape.call_type.value + ) # rebind-ok: routes cost calculation to the relayed shape's pricing path + return parsed + + class BasePassthroughConfig(BaseLLMModelInfo): @abstractmethod def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: @@ -60,7 +92,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): endpoint: str, base_target_url: str, request_query_params: Mapping[str, object] | None, - ) -> "URL": + ) -> URL: """ Helper function to add query params to the url Args: @@ -94,7 +126,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): endpoint: str, request_query_params: dict | None, litellm_params: dict, - ) -> tuple["URL", str]: + ) -> tuple[URL, str]: """ Get the complete url for the request Returns: @@ -124,9 +156,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): """ return headers, None - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, "Headers"] - ) -> "BaseLLMException": + def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException: from litellm.llms.base_llm.chat.transformation import BaseLLMException return BaseLLMException(status_code=status_code, message=error_message, headers=headers) @@ -135,23 +165,21 @@ class BasePassthroughConfig(BaseLLMModelInfo): self, model: str, custom_llm_provider: str, - httpx_response: "Response", + httpx_response: Response, request_data: dict, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, endpoint: str, - ) -> Optional[ - "CostResponseTypes | OCRResponse | RerankResponse | ResponsesAPIResponse | StandardPassThroughResponseObject" - ]: + ) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None: pass def handle_logging_collected_chunks( self, all_chunks: list[str], - litellm_logging_obj: "LiteLLMLoggingObj", + litellm_logging_obj: LiteLLMLoggingObj, model: str, custom_llm_provider: str, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> CostResponseTypes | None: return None def _convert_raw_bytes_to_str_lines(self, raw_bytes: list[bytes]) -> list[str]: diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 788aced50f8..8b6c76952d6 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -25,7 +25,9 @@ def clear_azure_ai_env(monkeypatch): def test_provider_config_manager_resolves_azure_ai_passthrough_config(): - config = ProviderConfigManager.get_provider_passthrough_config(model="Cohere-parse-v5", provider=LlmProviders.AZURE_AI) + config = ProviderConfigManager.get_provider_passthrough_config( + model="Cohere-parse-v5", provider=LlmProviders.AZURE_AI + ) assert isinstance(config, AzureAIPassthroughConfig) @@ -189,7 +191,10 @@ def test_no_credentials_at_all_raises(): [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], ) def test_is_streaming_request_reads_the_stream_flag(request_data, expected): - assert AzureAIPassthroughConfig().is_streaming_request(endpoint="models/chat/completions", request_data=request_data) is expected + assert ( + AzureAIPassthroughConfig().is_streaming_request(endpoint="models/chat/completions", request_data=request_data) + is expected + ) def _chat_completion_response() -> httpx.Response: @@ -266,7 +271,12 @@ def _relay_logging_obj(model: str, api_base: str) -> Logging: def _relay_logging_result( - config: AzureAIPassthroughConfig, model: str, native_path: str, body, api_base: str = FOUNDRY_BASE, status_code: int = 200 + config: AzureAIPassthroughConfig, + model: str, + native_path: str, + body, + api_base: str = FOUNDRY_BASE, + status_code: int = 200, ): relayed_url = f"{FOUNDRY_BASE}/{native_path}?api-version=2024-05-01-preview" logging_obj = _relay_logging_obj(model, api_base) @@ -355,7 +365,11 @@ def test_deployment_without_an_ocr_config_is_never_costed_as_ocr(): def test_accepted_ocr_job_without_a_result_body_is_not_costed(): result, logging_obj = _relay_logging_result( - AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", {"status": "running"}, status_code=202 + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + {"status": "running"}, + status_code=202, ) assert result == {"response": {"status": "running"}} @@ -364,7 +378,10 @@ def test_accepted_ocr_job_without_a_result_body_is_not_costed(): def test_unparseable_ocr_body_falls_back_to_the_passthrough_object(): result, logging_obj = _relay_logging_result( - AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", ["not", "an", "ocr", "body"] + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + ["not", "an", "ocr", "body"], ) assert result == {"response": '["not", "an", "ocr", "body"]'} @@ -423,6 +440,17 @@ def test_image_generation_relay_is_costed_per_image(): assert logging_obj._response_cost_calculator(result=result) == pytest.approx(per_image) +def test_flux_2_relay_through_the_provider_route_is_costed_per_image(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "FLUX.2-pro", "providers/blackforestlabs/v1/flux-2-pro", IMAGE_BODY + ) + per_image = litellm.get_model_info("azure_ai/FLUX.2-pro")["output_cost_per_image"] + + assert isinstance(result, ImageResponse) + assert logging_obj.call_type == "aimage_generation" + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(per_image) + + def test_rejected_rerank_relay_keeps_the_passthrough_object_and_call_type(): result, logging_obj = _relay_logging_result( AzureAIPassthroughConfig(), @@ -439,8 +467,15 @@ def test_rejected_rerank_relay_keeps_the_passthrough_object_and_call_type(): def test_streaming_chat_completion_chunks_are_costed_like_azure(): head = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-5.4-mini"} chunks = [ - "data: " + json.dumps({**head, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}]}), - "data: " + json.dumps({**head, "choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}}), + "data: " + + json.dumps( + { + **head, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + } + ), + "data: " + + json.dumps({**head, "choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}}), "data: [DONE]", ] From 02b18b4fc638dfba8fd77e70db796b3ea38c8757 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 19:10:06 -0700 Subject: [PATCH 026/107] fix(passthrough): match image URL schemes case-insensitively when counting relayed prompt tokens --- .../openai_passthrough_logging_handler.py | 4 +++- .../test_openai_passthrough_logging_handler.py | 9 +++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 206002d13f7..93fe3c5b31b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -105,7 +105,9 @@ def _is_remote_high_detail_image(part: object) -> bool: if not isinstance(image_url, Mapping): return False url: Final = image_url.get("url") - return isinstance(url, str) and url.startswith(("http://", "https://")) and image_url.get("detail") == "high" + return ( + isinstance(url, str) and url.lower().startswith(("http://", "https://")) and image_url.get("detail") == "high" + ) def _content_parts(message: Mapping[str, object]) -> Sequence[object]: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 0ddafd20e68..a3d3ae32169 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -2083,3 +2083,12 @@ def test_count_relayed_prompt_tokens_charges_only_the_remote_high_detail_image_a assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == ( litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + high_detail_image_token_upper_bound() ) + + +@pytest.mark.parametrize("scheme", ["HTTPS://", "Http://"]) +def test_count_relayed_prompt_tokens_charges_an_uppercase_scheme_remote_high_detail_image_at_the_upper_bound(scheme): + messages = _image_messages(scheme + UNREACHABLE_IMAGE_URL.split("://", 1)[1], "high") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == ( + litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + high_detail_image_token_upper_bound() + ) From dcd38ab9f0235c18fc429eeb5802519fde26c264 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 19:15:36 -0700 Subject: [PATCH 027/107] fix(proxy): count failure and rate-limit input tokens off the event loop The failure hook's usage estimate and the project ITPM reservation both called litellm.token_counter inline on the event loop, so a large request that failed or hit the limiter stalled the gateway the same way the count_tokens endpoints did. Both now run through asyncify. The loop-lag probe the existing tests used moves into a shared helper that warms the tokenizer first, and two new tests fail when either count runs inline --- .../hooks/parallel_request_limiter_v3.py | 3 +- litellm/proxy/utils.py | 3 +- .../litellm_core_utils/event_loop_lag.py | 39 +++++++++++++++++++ .../litellm_core_utils/test_token_counter.py | 30 +++++--------- .../proxy/hooks/test_tpm_concurrent.py | 37 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 34 ++++++---------- tests/test_litellm/proxy/test_proxy_utils.py | 38 ++++++++++++++++++ 7 files changed, 139 insertions(+), 45 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/event_loop_lag.py diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 31437af7770..74266801e50 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -28,6 +28,7 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -3303,7 +3304,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit=min_configured_otpm_limit, call_type=call_type, ) - raw_estimated_input_tokens: Final = self._estimate_precise_input_tokens( + raw_estimated_input_tokens: Final = await asyncify(self._estimate_precise_input_tokens)( data=data, model=requested_model, call_type=call_type ) estimated_input_tokens: Final = max(raw_estimated_input_tokens, 1) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ddf31cb1d8a..1014c2b9b6a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,6 +91,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, independent_snapshot, @@ -2569,7 +2570,7 @@ class ProxyLogging: original_exception=original_exception, ) - request_data.update(_failure_fields_to_lift(request_data)) + request_data.update(await asyncify(_failure_fields_to_lift)(request_data)) # Remove before callbacks iterate — not serialisable request_data.pop("litellm_logging_obj", None) diff --git a/tests/test_litellm/litellm_core_utils/event_loop_lag.py b/tests/test_litellm/litellm_core_utils/event_loop_lag.py new file mode 100644 index 00000000000..c2bbe2204c7 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/event_loop_lag.py @@ -0,0 +1,39 @@ +import asyncio +import time +from collections.abc import Awaitable, Callable +from typing import Final, TypeVar + +import litellm + +T = TypeVar("T") + + +def warm_tokenizer(model: str) -> None: + litellm.token_counter(model=model, text="load the tokenizer before anything is timed") + + +async def loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]: + async def wake_lag() -> float: + started: Final = time.perf_counter() + await asyncio.sleep(0.001) + return time.perf_counter() - started - 0.001 + + return tuple([await wake_lag() for _ in iter(until.is_set, True)]) + + +async def timed_with_loop_lags(run: Callable[[], Awaitable[T]]) -> tuple[T, float, tuple[float, ...]]: + finished: Final = asyncio.Event() + + async def timed() -> tuple[T, float]: + started: Final = time.perf_counter() + try: + return await run(), time.perf_counter() - started + finally: + finished.set() + + (result, took), lags = await asyncio.gather(timed(), loop_wake_lags(finished)) + return result, took, lags + + +def assert_loop_stayed_free(took: float, lags: tuple[float, ...]) -> None: + assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index db87bc52600..e3dd5feb18d 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -24,6 +24,11 @@ from litellm.litellm_core_utils.token_counter import ( ) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from tests.large_text import text +from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, +) from tests.test_litellm.litellm_core_utils.messages_with_counts import ( MESSAGES_TEXT, MESSAGES_WITH_IMAGES, @@ -127,30 +132,15 @@ def test_valid_chunk_size_config_is_honoured(monkeypatch): importlib.reload(litellm.constants) -async def _loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]: - async def wake_lag() -> float: - started: Final = time.perf_counter() - await asyncio.sleep(0.001) - return time.perf_counter() - started - 0.001 - - return tuple([await wake_lag() for _ in iter(until.is_set, True)]) - - async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free(): - counted: Final = asyncio.Event() + warm_tokenizer("claude-fable-5") - async def count_off_the_loop() -> tuple[int, float]: - started: Final = time.perf_counter() - try: - tokens: Final = await asyncify(token_counter_new)(model="claude-fable-5", text=text * 100) - return tokens, time.perf_counter() - started - finally: - counted.set() - - (tokens, took), lags = await asyncio.gather(count_off_the_loop(), _loop_wake_lags(counted)) + tokens, took, lags = await timed_with_loop_lags( + lambda: asyncify(token_counter_new)(model="claude-fable-5", text=text * 100) + ) assert tokens > 0 - assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" + assert_loop_stayed_free(took, lags) @pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500]) diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 2839acab6b0..398abf60b54 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -3670,5 +3670,42 @@ async def test_post_call_success_hook_contains_header_merge_failures( ) +@pytest.mark.asyncio +async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loop(rate_limiter): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + handler, _cache = rate_limiter + stash = get_or_create_request_stash() + warm_tokenizer("claude-fable-5") + data: Dict[str, Any] = { + "model": "claude-fable-5", + "messages": [{"role": "user", "content": text * 100}], + } + itpm_descriptor = { + "key": PROJECT_ITPM_DESCRIPTOR_KEY, + "value": "proj-loop:claude-fable-5", + "rate_limit": {"tokens_per_unit": 10_000_000, "window_size": 60}, + } + + _, took, lags = await timed_with_loop_lags( + lambda: handler._reserve_project_io_tokens_or_raise( + descriptors=[itpm_descriptor], + data=data, + requested_model="claude-fable-5", + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-itpm-loop"), project_id="proj-loop"), + tpm_reservation_scopes=[], + tpm_reservation_amount=0, + ) + ) + + assert stash.rate_limit_response is not None + assert_loop_stayed_free(took, lags) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 20b41cf1530..7fa3d5dd246 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -12860,32 +12860,20 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough() assert ps.general_settings["enable_openai_websocket_passthrough"] is False -async def _loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]: - async def wake_lag() -> float: - started: Final = time.perf_counter() - await asyncio.sleep(0.001) - return time.perf_counter() - started - 0.001 - - return tuple([await wake_lag() for _ in iter(until.is_set, True)]) - - async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch): from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) - counted: Final = asyncio.Event() + warm_tokenizer("claude-fable-5") - async def count_off_the_loop() -> tuple[int, float]: - started: Final = time.perf_counter() - try: - response: Final = await proxy_server_module.token_counter( - TokenCountRequest(model="claude-fable-5", prompt=text * 100) - ) - return response.total_tokens, time.perf_counter() - started - finally: - counted.set() + response, took, lags = await timed_with_loop_lags( + lambda: proxy_server_module.token_counter(TokenCountRequest(model="claude-fable-5", prompt=text * 100)) + ) - (tokens, took), lags = await asyncio.gather(count_off_the_loop(), _loop_wake_lags(counted)) - - assert tokens > 0 - assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" + assert response.total_tokens > 0 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 9462f2c8eb0..b78ec7dcff6 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1823,6 +1823,44 @@ def test_a_dispatched_failure_lifts_the_four_fields_the_spend_log_needs(): assert lifted["standard_logging_object"] == {"id": "log-1"} +@pytest.mark.asyncio +async def test_a_dispatched_failure_is_counted_off_the_event_loop(): + from unittest.mock import AsyncMock, patch + + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("claude-fable-5") + request_data = { + "litellm_logging_obj": _LoggingObj( + { + "first_api_call_start_time": 1700000000.0, + "call_type": "acompletion", + "model": "claude-fable-5", + "messages": [{"role": "user", "content": text * 100}], + } + ), + "metadata": {}, + } + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + _, took, lags = await timed_with_loop_lags( + lambda: proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=Exception("boom"), + user_api_key_dict=UserAPIKeyAuth(), + ) + ) + + assert request_data["combined_usage_object"].prompt_tokens > 0 + assert_loop_stayed_free(took, lags) + + @pytest.mark.asyncio async def test_proxy_only_error_expected_4xx_skips_traceback_for_both_handlers(monkeypatch): """Regression for LIT-6043: an expected 4xx must not format a traceback for From 83d500bac9c375221120b4ae09e0e3cbd55f96b1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:29:26 -0700 Subject: [PATCH 028/107] fix(router): count deployment itpm reservations off the event loop --- litellm/litellm_core_utils/token_counter.py | 4 +-- .../io_token_rate_limit_check.py | 3 +- .../test_router/test_io_token_rate_limits.py | 28 +++++++++++++++++++ 3 files changed, 32 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index e3a0d4c785a..e3af64fd3bf 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -315,12 +315,12 @@ def calculate_img_tokens( TokenCounterFunction = Callable[[str], int] - -EXTRAPOLATION_SAMPLES: Final = 16 """ Type for a function that counts tokens in a string. """ +EXTRAPOLATION_SAMPLES: Final = 16 + def _get_tiktoken_count_function( encode_length: Callable[[str], int], diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 01d42627001..dca5a2be845 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -21,6 +21,7 @@ import litellm from litellm import token_counter from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.litellm_core_utils.asyncify import asyncify from litellm.types.router import RouterCacheEnum, RouterErrors from litellm.utils import get_utc_datetime @@ -466,7 +467,7 @@ async def async_io_token_pre_call_check( request_kwargs: Final = get_io_token_rate_limit_request_kwargs() _model: Final = (deployment.get("litellm_params") or {}).get("model") or "" - estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model) + estimated_input: Final = await asyncify(_estimate_input_tokens)(request_kwargs, model=_model) max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment) dt: Final = get_utc_datetime() diff --git a/tests/test_litellm/test_router/test_io_token_rate_limits.py b/tests/test_litellm/test_router/test_io_token_rate_limits.py index a5a68271111..3cef1c7bb63 100644 --- a/tests/test_litellm/test_router/test_io_token_rate_limits.py +++ b/tests/test_litellm/test_router/test_io_token_rate_limits.py @@ -1039,3 +1039,31 @@ class TestContextSlotRetention: assert deployment is not None router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) assert get_io_token_rate_limit_request_kwargs() is kwargs + + +@pytest.mark.asyncio +async def test_the_deployment_itpm_reservation_counts_the_request_off_the_event_loop(): + from litellm.utils import get_utc_datetime + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + warm_tokenizer("anthropic/claude-fable-5") + deployment = { + "litellm_params": {"model": "anthropic/claude-fable-5", "itpm": 10_000_000}, + "model_info": {"id": "io-loop-id"}, + "model_name": "claude", + } + set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": text * 100}], "metadata": {}}) + + _, took, lags = await timed_with_loop_lags(lambda: check.async_pre_call_check(deployment)) + + minute = get_utc_datetime().strftime("%H-%M") + reserved = await dual_cache.async_get_cache(key=f"global_router:io-loop-id:anthropic/claude-fable-5:itpm:{minute}") + assert reserved > 100_000 + assert_loop_stayed_free(took, lags) From 1de369a1a446b7ecc1ed60a0c04d26bed1601c3a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:54:37 -0700 Subject: [PATCH 029/107] fix(passthrough): bound azure relays to the key's model group and reject foreign deployment segments --- litellm/proxy/auth/auth_utils.py | 10 ++ .../llm_passthrough_endpoints.py | 41 ++++++ .../test_azure_passthrough_transformation.py | 4 +- .../proxy/auth/test_auth_utils.py | 32 +++++ .../test_llm_pass_through_endpoints.py | 134 +++++++++++++++++- tests/test_litellm/test_router.py | 4 +- 6 files changed, 217 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 1e4836654a1..dd7df2b3fd9 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1985,9 +1985,19 @@ def get_model_from_request( bedrock_model: Final = _model_from_bedrock_route(route) return model if bedrock_model is None else bedrock_model + if route.lower().startswith(("/azure/", "/azure_ai/")): + azure_model: Final = _router_model_from_azure_route(route, llm_router) + return model if azure_model is None else azure_model + return model +def _router_model_from_azure_route(route: str, llm_router: Router | None) -> str | None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import azure_router_model_in_endpoint + + return azure_router_model_in_endpoint(re.sub(r"^/azure(?:_ai)?/", "", route, flags=re.IGNORECASE), llm_router) + + def _model_from_bedrock_route(route: str) -> str | None: from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _extract_model_from_bedrock_endpoint, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 839c9d47c6f..47b3c56e66a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -78,6 +78,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, ) from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials +from litellm.types.router import LiteLLMParamsTypedDict from litellm.types.utils import LlmProviders from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager @@ -120,6 +121,37 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li return False +def azure_router_model_in_endpoint(endpoint: str, llm_router: litellm.Router | None) -> str | None: + parts: Final = endpoint.split("/") + if len(parts) < 2: + return None + return next((part for part in parts if is_known_model(part, llm_router)), None) + + +AZURE_DEPLOYMENT_SEGMENT: Final = re.compile(r"(? str: + model: Final = litellm_params.get("model", "") + try: + return get_llm_provider(model=model, custom_llm_provider=litellm_params.get("custom_llm_provider"))[0] + except litellm.BadRequestError: + return model + + +def foreign_azure_deployment(endpoint: str, model_group: str, llm_router: litellm.Router) -> str | None: + match: Final = AZURE_DEPLOYMENT_SEGMENT.search(endpoint) + if match is None: + return None + deployment: Final = match.group(1) + if deployment == model_group: + return None + served: Final = frozenset( + _deployment_model_name(row["litellm_params"]) for row in llm_router.get_model_list(model_name=model_group) or () + ) + return None if deployment in served else deployment + + def is_passthrough_request_streaming(request_body: object) -> bool: """ Returns True if the request is streaming. @@ -1523,6 +1555,15 @@ async def _relay_azure_router_model( is_streaming_request: bool, user_api_key_dict: UserAPIKeyAuth, ) -> Response: + foreign_deployment: Final = foreign_azure_deployment(endpoint, model, llm_router) + if foreign_deployment is not None: + raise HTTPException( + status_code=400, + detail={ + "error": f"deployment '{foreign_deployment}' in the path is not served by model group '{model}'; " + "put the model group name in the deployments segment" + }, + ) try: result: Final = await llm_router.allm_passthrough_route( model=model, diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 2723a13c3fc..54e27b34a59 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -334,12 +334,12 @@ def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment( api_base="https://my-resource.openai.azure.com", api_key="key", model="gpt-4.1-mini", - endpoint="gpt/openai/deployments/gpt-4o/chat/completions", + endpoint="gpt/openai/deployments/gpt-4.1-mini/chat/completions", request_query_params={"api-version": "2024-10-21"}, litellm_params={"litellm_metadata": {"model_group": "gpt"}}, ) - assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21" + assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" @pytest.mark.parametrize( diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index a996de4d40c..ccb707a8e52 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -528,6 +528,38 @@ def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model(): ) +def _azure_relay_router(): + from litellm.router import Router + + return Router( + model_list=[ + { + "model_name": "gpt", + "litellm_params": {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://a.services.ai.azure.com", "api_key": "k"}, + }, + { + "model_name": "other-group", + "litellm_params": {"model": "azure/gpt-5.4", "api_base": "https://b.openai.azure.com", "api_key": "k"}, + }, + ] + ) + + +@pytest.mark.parametrize( + "route, request_data, expected", + [ + ("/azure_ai/other-group/openai/deployments/other-group/chat/completions", {"model": "gpt"}, "other-group"), + ("/azure_ai/other-group/models/chat/completions", {}, "other-group"), + ("/azure/openai/deployments/gpt/chat/completions", {"model": "other-group"}, "gpt"), + ("/azure/openai/deployments/gpt/chat/completions", {}, "gpt"), + ("/azure/openai/deployments/my-azure-deployment/chat/completions", {"model": "gpt"}, "gpt"), + ("/azure_ai/gpt", {"model": "other-group"}, "other-group"), + ], +) +def test_get_model_from_request_azure_relay_routes_use_the_model_group_in_the_path(route, request_data, expected): + assert get_model_from_request(request_data=request_data, route=route, llm_router=_azure_relay_router()) == expected + + def test_get_model_from_request_includes_file_endpoint_header_model(): assert ( get_model_from_request( 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 6ff14157468..fe4400df704 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 @@ -4999,7 +4999,11 @@ class TestPassthroughRouterModelBudgetReservation: monkeypatch.setattr(proxy_server, "llm_router", RecordingRouter()) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) return captured def _assert_metadata_carries_attribution(self, captured: list[dict], user_api_key_dict: UserAPIKeyAuth) -> None: @@ -5098,7 +5102,11 @@ class TestAzureRouterModelStreamingDispatch: monkeypatch.setattr(proxy_server, "llm_router", StreamingRouter()) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) request = MagicMock(spec=Request) request.method = "POST" @@ -5158,7 +5166,11 @@ class TestAzureRouterModelStreamingKeepalive: monkeypatch.setattr(proxy_server, "llm_router", StreamingRouter()) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) request = MagicMock(spec=Request) request.method = "POST" @@ -5225,7 +5237,11 @@ class TestRouterModelRelayUpstreamContract: monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) def _recording_router(self, captured: list[dict]): class RecordingRouter: @@ -5327,3 +5343,113 @@ async def test_bedrock_count_tokens_error_forwards_provider_headers(): assert exc_info.value.status_code == 500 assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-count-tokens-500" + + +class _AzureGroupRouter: + def __init__(self, captured: list[dict]) -> None: + self.captured = captured + + def get_model_names(self, team_id=None): + return ["gpt", "other-group"] + + def get_model_list(self, model_name=None, team_id=None): + rows = [ + {"model_name": "gpt", "litellm_params": {"model": "azure_ai/gpt-5.4-mini", "api_key": "k"}}, + {"model_name": "other-group", "litellm_params": {"model": "azure/gpt-5.4", "api_key": "k"}}, + ] + return [row for row in rows if model_name is None or row["model_name"] == model_name] + + async def allm_passthrough_route(self, **kwargs): + self.captured.append(kwargs) + return httpx.Response(200, json={"ok": True}) + + +class TestAzureRelayDeploymentSegment: + """A key allowed one model group must not reach another deployment by naming it in the + ``openai/deployments/`` segment while the group segment picks the credential.""" + + @pytest.mark.parametrize( + "endpoint, expected", + [ + ("gpt/openai/deployments/gpt/chat/completions", None), + ("openai/deployments/gpt/chat/completions", None), + ("gpt/openai/deployments/gpt-5.4-mini/chat/completions", None), + ("gpt/models/chat/completions", None), + ("gpt/openai/deployments/gpt-5.4/chat/completions", "gpt-5.4"), + ("gpt/openai/deployments/other-group/chat/completions", "other-group"), + ("openai/deployments/victim/gpt/chat/completions", "victim"), + ], + ) + def test_foreign_azure_deployment_names_a_segment_outside_the_group(self, endpoint, expected): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import foreign_azure_deployment + + assert foreign_azure_deployment(endpoint, "gpt", _AzureGroupRouter([])) == expected + + @pytest.mark.parametrize( + "endpoint, expected", + [ + ("other-group/openai/deployments/other-group/chat/completions", "other-group"), + ("openai/deployments/gpt/chat/completions", "gpt"), + ("openai/deployments/my-azure-deployment/chat/completions", None), + ("gpt", None), + ], + ) + def test_azure_router_model_in_endpoint_matches_the_relay_decision(self, endpoint, expected): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import azure_router_model_in_endpoint + + assert azure_router_model_in_endpoint(endpoint, _AzureGroupRouter([])) == expected + + def _install(self, monkeypatch, body: dict) -> list[dict]: + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + import litellm.proxy.proxy_server as proxy_server + + captured: list[dict] = [] + + async def fake_get_request_body(_request): + return body + + monkeypatch.setattr(proxy_server, "llm_router", _AzureGroupRouter(captured)) + monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) + return captured + + def _request(self) -> Request: + request = MagicMock(spec=Request) + request.method = "POST" + request.headers = {"content-type": "application/json"} + request.query_params = {} + return request + + @pytest.mark.asyncio + async def test_azure_relay_rejects_a_deployment_the_group_does_not_serve(self, monkeypatch): + from fastapi import HTTPException + + captured = self._install(monkeypatch, {"model": "gpt", "messages": []}) + + with pytest.raises(HTTPException) as exc_info: + await azure_proxy_route( + endpoint="gpt/openai/deployments/gpt-5.4/chat/completions", + request=self._request(), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token", models=["gpt"]), + ) + + assert exc_info.value.status_code == 400 + assert "gpt-5.4" in exc_info.value.detail["error"] + assert captured == [] + + @pytest.mark.asyncio + async def test_azure_relay_dispatches_the_group_and_its_own_deployment_name(self, monkeypatch): + captured = self._install(monkeypatch, {"model": "gpt", "messages": []}) + + for endpoint in ( + "gpt/openai/deployments/gpt/chat/completions", + "gpt/openai/deployments/gpt-5.4-mini/chat/completions", + ): + await azure_proxy_route( + endpoint=endpoint, + request=self._request(), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token", models=["gpt"]), + ) + + assert [call["model"] for call in captured] == ["gpt", "gpt"] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 204c6f109ff..a766f0c9db1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5053,8 +5053,8 @@ def test_get_deployment_model_info_base_model_merge_priority(): ( "gpt", {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.services.ai.azure.com", "api_key": "key"}, - "gpt/openai/deployments/gpt-4o/chat/completions", - "gpt-5.4-mini/openai/deployments/gpt-4o/chat/completions", + "gpt/openai/deployments/gpt-5.4-mini/chat/completions", + "gpt-5.4-mini/openai/deployments/gpt-5.4-mini/chat/completions", ), ( "aws/anthropic/bedrock-claude", From 52e9cfe254355b2ffdb75fe314c2b62b467f5018 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:57:00 -0700 Subject: [PATCH 030/107] fix(azure): price streamed Responses relays and keep a full-URL api_base from doubling the native path --- litellm/litellm_core_utils/litellm_logging.py | 5 +- .../llms/azure/passthrough/transformation.py | 23 +++++++- .../azure_ai/passthrough/transformation.py | 5 +- .../base_llm/passthrough/transformation.py | 2 +- .../test_azure_passthrough_transformation.py | 47 +++++++++++++++ ...est_azure_ai_passthrough_transformation.py | 57 +++++++++++++++++++ 6 files changed, 129 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ca2cca5360f..ec988b200d0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -119,7 +119,6 @@ from litellm.types.utils import ( CachingDetails, CallTypes, CostBreakdown, - CostResponseTypes, CustomPricingLiteLLMParams, DynamicPromptManagementParamLiteral, EmbeddingResponse, @@ -201,7 +200,7 @@ if TYPE_CHECKING: from mcp.types import EmbeddedResource, ImageContent, TextContent from litellm.integrations.otel.logger import OpenTelemetryV2 - from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, LoggedRelayResponse try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -2363,7 +2362,7 @@ class Logging(LiteLLMLoggingBaseClass): self, raw_bytes: list[bytes], provider_config: "BasePassthroughConfig", - ) -> Optional["CostResponseTypes"]: + ) -> Optional["LoggedRelayResponse"]: all_chunks: Final = provider_config._convert_raw_bytes_to_str_lines(raw_bytes) complete_streaming_response: Final = provider_config.handle_logging_collected_chunks( all_chunks=all_chunks, diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index bdb854bb87a..ab7fe412628 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -23,7 +23,6 @@ if TYPE_CHECKING: from httpx import URL from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse - from litellm.types.utils import CostResponseTypes class RelayedChatRequest(BaseModel): @@ -42,13 +41,29 @@ def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, obj return details.request_data.messages if details.request_data else None +RESPONSES_RELAY_SHAPE: Final = RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate) + OPENAI_RELAY_SHAPES: Final = ( RelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate), - RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate), + RESPONSES_RELAY_SHAPE, RelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate), ) +def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesAPIResponse | None: + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + + terminal_response: Final = OpenAIResponsesAPIConfig.parse_terminal_response_from_stream_chunks( + all_chunks=list(all_chunks) + ) + if terminal_response is None: + return None + logging_obj.call_type = ( + RESPONSES_RELAY_SHAPE.call_type.value + ) # rebind-ok: routes cost calculation to the relayed shape's pricing path + return terminal_response + + class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return bool(request_data.get("stream")) @@ -157,11 +172,13 @@ class AzurePassthroughConfig(BasePassthroughConfig): model: str, custom_llm_provider: str, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> Optional["LoggedRelayResponse"]: from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, ) + if f"/{endpoint.strip('/')}".endswith(RESPONSES_RELAY_SHAPE.path_suffix): + return logged_responses_stream(all_chunks, litellm_logging_obj) if "chat/completions" not in endpoint: return None diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index e854874eaee..1689755b042 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -30,7 +30,6 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse - from litellm.types.utils import CostResponseTypes EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({}) @@ -111,8 +110,8 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): if base_target_url is None: raise ValueError("Azure AI api base not found: set `api_base` on the deployment or AZURE_AI_API_BASE") - root: Final = foundry_root(base_target_url) native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params))) + root: Final = foundry_root(base_target_url).removesuffix(f"/{native_endpoint.strip('/')}") query_params: Final = relay_query_params( request_query_params, api_version_from(litellm_params), base_target_url ) @@ -200,7 +199,7 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): model: str, custom_llm_provider: str, endpoint: str, - ) -> CostResponseTypes | None: + ) -> LoggedRelayResponse | None: from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig return AzurePassthroughConfig().handle_logging_collected_chunks( diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 93ec5a09e62..dc857dfc808 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -179,7 +179,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): model: str, custom_llm_provider: str, endpoint: str, - ) -> CostResponseTypes | None: + ) -> LoggedRelayResponse | None: return None def _convert_raw_bytes_to_str_lines(self, raw_bytes: list[bytes]) -> list[str]: diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 54e27b34a59..42dd65ea661 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -283,6 +283,53 @@ def test_azure_passthrough_streaming_chunks_for_unknown_endpoint_return_none(): assert response is None +def _azure_responses_stream_chunks(terminal_event: str | None = "response.completed") -> list[str]: + in_progress = {**RESPONSES_BODY, "status": "in_progress", "output": [], "usage": None} + events = [ + ("response.created", {"type": "response.created", "sequence_number": 0, "response": in_progress}), + ( + "response.output_text.delta", + {"type": "response.output_text.delta", "sequence_number": 1, "item_id": "msg_1", "delta": "hi"}, + ), + ] + ([(terminal_event, {"type": terminal_event, "sequence_number": 2, "response": RESPONSES_BODY})] if terminal_event else []) + return [line for name, payload in events for line in (f"event: {name}", _sse_line(payload))] + + +def test_azure_passthrough_streaming_responses_chunks_are_costed_per_token(): + logging_obj = _relay_logging_obj("gpt-4.1-mini") + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_responses_stream_chunks(), + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/responses", + ) + info = litellm.get_model_info("azure/gpt-4.1-mini") + + assert isinstance(response, ResponsesAPIResponse) + assert response.usage.input_tokens == 1000 + assert logging_obj.call_type == "aresponses" + assert logging_obj._response_cost_calculator(result=response) == pytest.approx( + 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] + ) + + +def test_azure_passthrough_streaming_responses_without_a_terminal_event_are_not_costed(): + logging_obj = _relay_logging_obj("gpt-4.1-mini") + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_responses_stream_chunks(terminal_event=None), + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/responses", + ) + + assert response is None + assert logging_obj.call_type == "allm_passthrough_route" + + def _complete_url(request_query_params: dict, litellm_params: dict) -> httpx.URL: url, _ = AzurePassthroughConfig().get_complete_url( api_base="https://my-resource.openai.azure.com", diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 8b6c76952d6..0b23da984e8 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -86,6 +86,22 @@ def test_api_base_that_already_ends_in_models_is_cut_back_to_the_foundry_root(): assert base == FOUNDRY_BASE +def test_full_url_api_base_that_already_ends_with_the_native_path_is_not_doubled(): + model_router_url = "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router/chat/completions" + + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{model_router_url}?api-version=2025-01-01-preview", + api_key="key", + model="model_router/model-router", + endpoint="model-router/chat/completions", + request_query_params=None, + litellm_params={"litellm_metadata": {"model_group": "model-router"}}, + ) + + assert str(url) == f"{model_router_url}?api-version=2025-01-01-preview" + assert base == "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router" + + def test_parse_relay_under_a_models_api_base_targets_the_foundry_root(): url, _ = AzureAIPassthroughConfig().get_complete_url( api_base=f"{FOUNDRY_BASE}/models", @@ -490,3 +506,44 @@ def test_streaming_chat_completion_chunks_are_costed_like_azure(): assert isinstance(response, ModelResponse) assert response.choices[0].message.content == "hi" assert response.usage.total_tokens == 4 + + +def test_streaming_responses_chunks_through_a_router_relay_are_costed_like_azure(): + completed = { + "type": "response.completed", + "sequence_number": 2, + "response": { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, + }, + } + logging_obj = _relay_logging_obj("gpt-5.4-mini", FOUNDRY_BASE) + + response = AzureAIPassthroughConfig().handle_logging_collected_chunks( + all_chunks=["event: response.completed", "data: " + json.dumps(completed)], + litellm_logging_obj=logging_obj, + model="gpt-5.4-mini", + custom_llm_provider="azure_ai", + endpoint="gpt/openai/responses", + ) + info = litellm.get_model_info("azure_ai/gpt-5.4-mini") + + assert response is not None + assert response.usage.output_tokens == 100 + assert logging_obj.call_type == "aresponses" + assert logging_obj._response_cost_calculator(result=response) == pytest.approx( + 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] + ) From d202885f8b4ce1b0d56b7757ebb0c5849aeb3cb8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:21:45 -0700 Subject: [PATCH 031/107] fix(proxy): run prompt caching counts and custom tokenizer loads off the event loop --- litellm/proxy/proxy_server.py | 4 +- .../prompt_caching_deployment_check.py | 5 +- litellm/utils.py | 10 +-- .../test_custom_tokenizer_bug.py | 10 ++- .../litellm_core_utils/event_loop_lag.py | 1 + tests/test_litellm/proxy/test_proxy_server.py | 37 +++++++++++ .../test_prompt_caching_deployment_check.py | 62 +++++++++++++++++++ 7 files changed, 111 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0915b8dd1b9..77354e9885f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12713,7 +12713,9 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = litellm.utils._select_tokenizer(model=model_to_use, custom_tokenizer=custom_tokenizer) + _tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)( + model=model_to_use, custom_tokenizer=custom_tokenizer + ) tokenizer_used: Final = str(_tokenizer_used["type"]) system_message: Final = _system_message(system) diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 70362e60495..9ea528b9700 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -14,6 +14,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.asyncify import asyncify from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt @@ -61,7 +62,7 @@ class PromptCachingDeploymentCheck(CustomLogger): if request_kwargs is not None and request_kwargs.get("_target_order") is not None: return healthy_deployments - if messages is not None and is_prompt_caching_valid_prompt( + if messages is not None and await asyncify(is_prompt_caching_valid_prompt)( messages=messages, model=model, min_token_count=_get_min_token_count_for_deployments(healthy_deployments), @@ -139,7 +140,7 @@ class PromptCachingDeploymentCheck(CustomLogger): return ## PROMPT CACHING - cache model id, if prompt caching valid prompt + provider - if is_prompt_caching_valid_prompt( + if await asyncify(is_prompt_caching_valid_prompt)( model=model, messages=cast(list[AllMessageValues], messages), ): diff --git a/litellm/utils.py b/litellm/utils.py index d0e11bc9551..ab905381b22 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2284,15 +2284,7 @@ def create_pretrained_tokenizer(identifier: str, revision="main", auth_token: st dict: A dictionary with the tokenizer and its type. """ - try: - tokenizer = Tokenizer.from_pretrained( - identifier, - revision=revision, - auth_token=auth_token, - ) - except Exception as e: - verbose_logger.error("Error creating pretrained tokenizer: %s. Defaulting to version without 'auth_token'.", e) - tokenizer = Tokenizer.from_pretrained(identifier, revision=revision) + tokenizer: Final = Tokenizer.from_pretrained(identifier, revision=revision, token=auth_token) return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} diff --git a/tests/proxy_unit_tests/test_custom_tokenizer_bug.py b/tests/proxy_unit_tests/test_custom_tokenizer_bug.py index 89899d3e762..c4b1f4f3afd 100644 --- a/tests/proxy_unit_tests/test_custom_tokenizer_bug.py +++ b/tests/proxy_unit_tests/test_custom_tokenizer_bug.py @@ -23,9 +23,9 @@ from litellm.proxy.proxy_server import token_counter def _fake_hf_tokenizer(num_tokens: int) -> MagicMock: encoding = MagicMock() - encoding.ids = list(range(num_tokens)) + encoding.__len__.return_value = num_tokens tokenizer = MagicMock() - tokenizer.encode.return_value = encoding + tokenizer.encode_batch_fast.return_value = [encoding] return tokenizer @@ -68,13 +68,11 @@ async def test_custom_tokenizer_from_model_info_is_used(monkeypatch): ) ) - mock_tokenizer_cls.from_pretrained.assert_called_once_with( - "my-org/custom-tokenizer", revision="v2", auth_token=None - ) + mock_tokenizer_cls.from_pretrained.assert_called_once_with("my-org/custom-tokenizer", revision="v2", token=None) assert response.tokenizer_type == "huggingface_tokenizer" assert response.request_model == "my-embedding-model" assert response.model_used == "self-hosted-embedder" - assert response.total_tokens > 0 + assert response.total_tokens >= 7 @pytest.mark.asyncio diff --git a/tests/test_litellm/litellm_core_utils/event_loop_lag.py b/tests/test_litellm/litellm_core_utils/event_loop_lag.py index c2bbe2204c7..1cac0365547 100644 --- a/tests/test_litellm/litellm_core_utils/event_loop_lag.py +++ b/tests/test_litellm/litellm_core_utils/event_loop_lag.py @@ -25,6 +25,7 @@ async def timed_with_loop_lags(run: Callable[[], Awaitable[T]]) -> tuple[T, floa finished: Final = asyncio.Event() async def timed() -> tuple[T, float]: + await asyncio.sleep(0) started: Final = time.perf_counter() try: return await run(), time.perf_counter() - started diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7fa3d5dd246..6fa99587dbf 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -12877,3 +12877,40 @@ async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_coun assert response.total_tokens > 0 assert_loop_stayed_free(took, lags) + + +async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch): + from tokenizers import Tokenizer + + from litellm import Router + from tests.test_litellm.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags + + claude_tokenizer: Final = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"] + + class SlowHubTokenizer: + @staticmethod + def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: + time.sleep(0.3) + return claude_tokenizer + + monkeypatch.setattr(litellm.utils, "Tokenizer", SlowHubTokenizer) + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + Router( + model_list=[ + { + "model_name": "self-hosted", + "litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"}, + "model_info": {"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}}, + } + ] + ), + ) + + response, took, lags = await timed_with_loop_lags( + lambda: proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me off the loop")) + ) + + assert response.tokenizer_type == "huggingface_tokenizer" + assert response.total_tokens > 0 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 030bdfe03e9..333e7b2ff31 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -477,3 +477,65 @@ async def test_wildcard_route_resolves_underlying_model_minimum(local_model_cost assert deployments[0]["litellm_params"]["model"] == "anthropic/claude-opus-4-6" assert _get_min_token_count_for_deployments(deployments) == 4096 + + +@pytest.mark.asyncio +async def test_async_filter_deployments_counts_the_prompt_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("anthropic/claude-fable-5") + check = PromptCachingDeploymentCheck(cache=DualCache()) + deployments = _deployments("anthropic/claude-fable-5") + messages = cast(List[AllMessageValues], [{"role": "user", "content": text * 100}]) + + result, took, lags = await timed_with_loop_lags( + lambda: check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=messages + ) + ) + + assert result == deployments + assert_loop_stayed_free(took, lags) + + +@pytest.mark.asyncio +async def test_async_log_success_event_counts_the_prompt_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("anthropic/claude-fable-5") + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + messages = cast( + List[AllMessageValues], + [{"role": "user", "content": [{"type": "text", "text": text * 100, "cache_control": {"type": "ephemeral"}}]}], + ) + standard_logging_object = { + "call_type": "acompletion", + "model": "anthropic/claude-fable-5", + "messages": messages, + "model_id": "dep-1", + } + + _, took, lags = await timed_with_loop_lags( + lambda: check.async_log_success_event( + kwargs={"standard_logging_object": standard_logging_object}, + response_obj=None, + start_time=None, + end_time=None, + ) + ) + + assert await PromptCachingCache(cache=cache).async_get_model_id(messages=messages, tools=None) == { + "model_id": "dep-1" + } + assert_loop_stayed_free(took, lags) From 13db79da46f41b9775e781ac0df9c1aa14c29f9a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:38:41 -0700 Subject: [PATCH 032/107] fix(azure_ai): price streamed Responses relays from their terminal event A streamed Responses API relay handed the success handler a bare ResponsesAPIResponse, which the streaming assembly step drops, so the relay never reached the spend callbacks. Hand it the terminal response.completed event instead, which the assembly step already converts, and cover the whole flush path with a regression test that fails on the previous tip. --- .../llms/azure/passthrough/transformation.py | 11 +-- .../base_llm/passthrough/transformation.py | 4 +- .../llms/openai/responses/transformation.py | 9 +- litellm/types/llms/openai.py | 3 + .../test_azure_passthrough_transformation.py | 35 ++++++-- ...est_azure_ai_passthrough_transformation.py | 90 +++++++++++++------ 6 files changed, 107 insertions(+), 45 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index ab7fe412628..85647ae02c9 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -15,7 +15,7 @@ from litellm.llms.base_llm.passthrough.transformation import ( strip_leading_model_segment, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse +from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse, ResponsesTerminalEvent from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse @@ -50,18 +50,19 @@ OPENAI_RELAY_SHAPES: Final = ( ) -def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesAPIResponse | None: +def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesTerminalEvent | None: + """A streaming logging object assembles the logged response from the terminal event, not from its body.""" from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig - terminal_response: Final = OpenAIResponsesAPIConfig.parse_terminal_response_from_stream_chunks( + terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks( all_chunks=list(all_chunks) ) - if terminal_response is None: + if terminal_event is None: return None logging_obj.call_type = ( RESPONSES_RELAY_SHAPE.call_type.value ) # rebind-ok: routes cost calculation to the relayed shape's pricing path - return terminal_response + return terminal_event class AzurePassthroughConfig(BasePassthroughConfig): diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index dc857dfc808..20180c5cfa2 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -16,14 +16,14 @@ if TYPE_CHECKING: from httpx import URL, Headers, Response from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesTerminalEvent from litellm.types.rerank import RerankResponse from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject from ..chat.transformation import BaseLLMException from ..ocr.transformation import OCRResponse - LoggedRelayResponse: TypeAlias = CostResponseTypes | RerankResponse | ResponsesAPIResponse + LoggedRelayResponse: TypeAlias = CostResponseTypes | RerankResponse | ResponsesAPIResponse | ResponsesTerminalEvent RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 926de3e8854..cde399065fd 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -620,15 +620,20 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return event_pydantic_model.model_construct(**parsed_chunk) @staticmethod - def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None: + def parse_terminal_event_from_stream_chunks(all_chunks: list[str]) -> ResponsesTerminalEvent | None: for chunk_str in reversed(all_chunks): for event_model in (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent): try: - return event_model.model_validate_json(chunk_str.removeprefix("data: ")).response + return event_model.model_validate_json(chunk_str.removeprefix("data: ")) except ValueError: continue return None + @staticmethod + def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None: + terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks) + return None if terminal_event is None else terminal_event.response + @staticmethod def get_event_model_class(event_type: str) -> type[BaseLiteLLMOpenAIResponseObject]: """ diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b6da9490e01..b7c4371f32f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1564,6 +1564,9 @@ class ResponseIncompleteEvent(BaseLiteLLMOpenAIResponseObject): response: ResponsesAPIResponse +ResponsesTerminalEvent: TypeAlias = ResponseCompletedEvent | ResponseIncompleteEvent | ResponseFailedEvent + + class ResponsePartAddedEvent(BaseLiteLLMOpenAIResponseObject): type: Literal[ResponsesAPIStreamEvents.RESPONSE_PART_ADDED] item_id: str diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 42dd65ea661..6e85b0cdca6 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -9,7 +9,7 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import EmbeddingResponse, ModelResponse @@ -103,7 +103,9 @@ def _relay_logging_result(model: str, endpoint: str, body, status_code: int = 20 status_code=status_code, headers={"content-type": "application/json"}, content=json.dumps(body).encode("utf-8"), - request=httpx.Request("POST", f"https://my-resource.openai.azure.com/{endpoint}?api-version=2025-04-01-preview"), + request=httpx.Request( + "POST", f"https://my-resource.openai.azure.com/{endpoint}?api-version=2025-04-01-preview" + ), ) result = AzurePassthroughConfig().logging_non_streaming_response( model=model, @@ -291,7 +293,11 @@ def _azure_responses_stream_chunks(terminal_event: str | None = "response.comple "response.output_text.delta", {"type": "response.output_text.delta", "sequence_number": 1, "item_id": "msg_1", "delta": "hi"}, ), - ] + ([(terminal_event, {"type": terminal_event, "sequence_number": 2, "response": RESPONSES_BODY})] if terminal_event else []) + ] + ( + [(terminal_event, {"type": terminal_event, "sequence_number": 2, "response": RESPONSES_BODY})] + if terminal_event + else [] + ) return [line for name, payload in events for line in (f"event: {name}", _sse_line(payload))] @@ -307,10 +313,10 @@ def test_azure_passthrough_streaming_responses_chunks_are_costed_per_token(): ) info = litellm.get_model_info("azure/gpt-4.1-mini") - assert isinstance(response, ResponsesAPIResponse) - assert response.usage.input_tokens == 1000 + assert isinstance(response, ResponseCompletedEvent) + assert response.response.usage.input_tokens == 1000 assert logging_obj.call_type == "aresponses" - assert logging_obj._response_cost_calculator(result=response) == pytest.approx( + assert logging_obj._response_cost_calculator(result=response.response) == pytest.approx( 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] ) @@ -373,7 +379,10 @@ def test_azure_passthrough_url_strips_the_leading_router_model_segment(): litellm_params={}, ) - assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + assert ( + str(url) + == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + ) def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment(): @@ -386,7 +395,10 @@ def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment( litellm_params={"litellm_metadata": {"model_group": "gpt"}}, ) - assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + assert ( + str(url) + == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + ) @pytest.mark.parametrize( @@ -394,4 +406,9 @@ def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment( [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], ) def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_data, expected): - assert AzurePassthroughConfig().is_streaming_request(endpoint="openai/deployments/x/chat/completions", request_data=request_data) is expected + assert ( + AzurePassthroughConfig().is_streaming_request( + endpoint="openai/deployments/x/chat/completions", request_data=request_data + ) + is expected + ) diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 0b23da984e8..064be954518 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -6,6 +6,7 @@ import httpx import pytest import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure_ai.passthrough.transformation import AzureAIPassthroughConfig from litellm.llms.base_llm.ocr.transformation import OCRResponse @@ -14,6 +15,36 @@ from litellm.types.utils import EmbeddingResponse, ImageResponse, LlmProviders, from litellm.utils import ProviderConfigManager FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" +RESPONSES_COMPLETED_EVENT = { + "type": "response.completed", + "sequence_number": 2, + "response": { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, + }, +} + + +class _SpendProbe(CustomLogger): + logged_call_type: str | None = None + logged_cost: float | None = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_call_type = kwargs["call_type"] + self.logged_cost = kwargs["response_cost"] @pytest.fixture(autouse=True) @@ -87,7 +118,9 @@ def test_api_base_that_already_ends_in_models_is_cut_back_to_the_foundry_root(): def test_full_url_api_base_that_already_ends_with_the_native_path_is_not_doubled(): - model_router_url = "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router/chat/completions" + model_router_url = ( + "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router/chat/completions" + ) url, base = AzureAIPassthroughConfig().get_complete_url( api_base=f"{model_router_url}?api-version=2025-01-01-preview", @@ -267,21 +300,29 @@ def test_non_chat_relay_with_a_non_json_body_logs_the_raw_text(): assert _non_chat_logging_result(b"page one", "text/plain") == {"response": "page one"} -def _relay_logging_obj(model: str, api_base: str) -> Logging: +def _relay_logging_obj( + model: str, + api_base: str, + stream: bool = False, + callbacks: list[CustomLogger] | None = None, + endpoint: str = "", +) -> Logging: logging_obj = Logging( model=model, messages=[], - stream=False, + stream=stream, call_type="allm_passthrough_route", start_time=datetime.now(), litellm_call_id="call-1", function_id="fn-1", + dynamic_async_success_callbacks=callbacks, ) logging_obj.update_environment_variables( model=model, litellm_params={"api_base": api_base, "custom_llm_provider": "azure_ai"}, optional_params={}, custom_llm_provider="azure_ai", + endpoint=endpoint, ) return logging_obj @@ -509,31 +550,10 @@ def test_streaming_chat_completion_chunks_are_costed_like_azure(): def test_streaming_responses_chunks_through_a_router_relay_are_costed_like_azure(): - completed = { - "type": "response.completed", - "sequence_number": 2, - "response": { - "id": "resp_1", - "object": "response", - "created_at": 1, - "status": "completed", - "model": "gpt-5.4-mini", - "output": [ - { - "type": "message", - "id": "msg_1", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "hi", "annotations": []}], - } - ], - "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, - }, - } logging_obj = _relay_logging_obj("gpt-5.4-mini", FOUNDRY_BASE) response = AzureAIPassthroughConfig().handle_logging_collected_chunks( - all_chunks=["event: response.completed", "data: " + json.dumps(completed)], + all_chunks=["event: response.completed", "data: " + json.dumps(RESPONSES_COMPLETED_EVENT)], litellm_logging_obj=logging_obj, model="gpt-5.4-mini", custom_llm_provider="azure_ai", @@ -542,8 +562,24 @@ def test_streaming_responses_chunks_through_a_router_relay_are_costed_like_azure info = litellm.get_model_info("azure_ai/gpt-5.4-mini") assert response is not None - assert response.usage.output_tokens == 100 + assert response.response.usage.output_tokens == 100 assert logging_obj.call_type == "aresponses" - assert logging_obj._response_cost_calculator(result=response) == pytest.approx( + assert logging_obj._response_cost_calculator(result=response.response) == pytest.approx( 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] ) + + +async def test_streaming_responses_relay_flush_reaches_the_success_callbacks_with_a_price(): + probe = _SpendProbe() + logging_obj = _relay_logging_obj( + "gpt-5.4-mini", FOUNDRY_BASE, stream=True, callbacks=[probe], endpoint="gpt/openai/responses" + ) + stream = "event: response.completed\ndata: " + json.dumps(RESPONSES_COMPLETED_EVENT) + "\n\n" + + await logging_obj.async_flush_passthrough_collected_chunks( + raw_bytes=[stream.encode()], provider_config=AzureAIPassthroughConfig() + ) + info = litellm.get_model_info("azure_ai/gpt-5.4-mini") + + assert probe.logged_call_type == "allm_passthrough_route" + assert probe.logged_cost == pytest.approx(1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"]) From 17dd6afbaafed87e6e7380278ca13a543046160a Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 13:51:58 +0000 Subject: [PATCH 033/107] fix: keep reasoning_effort for mode: responses bridge deployments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/main.py | 12 +++++++++++- tests/test_litellm/test_main.py | 34 +++++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/litellm/main.py b/litellm/main.py index 75b7f7f10a5..062ddb9a293 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5398,6 +5398,16 @@ def completion( if dynamic_api_key is not None: api_key = dynamic_api_key # check if user passed in any of the OpenAI optional params + allowed_openai_params: Final[list[str] | None] = cast( + list[str] | None, + ( + [*(kwargs.get("allowed_openai_params") or []), "reasoning_effort"] + if responses_api_model_info.get("mode") == "responses" + and not skip_responses_api_bridge + and "reasoning_effort" not in (kwargs.get("allowed_openai_params") or []) + else kwargs.get("allowed_openai_params") + ), + ) optional_param_args: Final = { "functions": functions, "function_call": function_call, @@ -5442,7 +5452,7 @@ def completion( "service_tier": service_tier, "store": store, "prompt_cache_key": prompt_cache_key, - "allowed_openai_params": kwargs.get("allowed_openai_params"), + "allowed_openai_params": allowed_openai_params, "base_model": base_model, } optional_params = get_optional_params(**optional_param_args, **non_default_params) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 038df3656fe..9ce2bb8400d 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -1367,6 +1367,40 @@ def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict( } +@pytest.mark.parametrize("reasoning_effort", ["high", {"effort": "high"}]) +@patch("litellm.completion_extras.responses_api_bridge.completion") +def test_responses_bridge_preserves_reasoning_effort_with_drop_params( + mock_responses_completion, reasoning_effort, restore_model_registry +): + mock_responses_completion.return_value = MagicMock() + model = "perplexity/test-responses-bridge" + litellm.register_model( + { + model: { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_reasoning": True, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + }, + persist_across_reloads=False, + ) + + with patch.object(litellm, "supports_reasoning", return_value=False): + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + drop_params=True, + api_key="fake-key", + api_base="https://api.perplexity.ai", + ) + + optional_params = mock_responses_completion.call_args.kwargs["optional_params"] + assert optional_params["reasoning_effort"] == reasoning_effort + + @pytest.mark.parametrize( "model, model_info, expected_model_param, expected_base_model_param", [ From 5520c3bb1a768263bbc4458b14da7d87e3696a41 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 13:53:02 +0000 Subject: [PATCH 034/107] refactor: simplify allowed_openai_params bridge override Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/main.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 062ddb9a293..01fe6a7c674 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5398,15 +5398,13 @@ def completion( if dynamic_api_key is not None: api_key = dynamic_api_key # check if user passed in any of the OpenAI optional params - allowed_openai_params: Final[list[str] | None] = cast( - list[str] | None, - ( - [*(kwargs.get("allowed_openai_params") or []), "reasoning_effort"] - if responses_api_model_info.get("mode") == "responses" - and not skip_responses_api_bridge - and "reasoning_effort" not in (kwargs.get("allowed_openai_params") or []) - else kwargs.get("allowed_openai_params") - ), + bridges_to_responses_api: Final = ( + responses_api_model_info.get("mode") == "responses" and not skip_responses_api_bridge + ) + allowed_openai_params: Final[list[str] | None] = ( + [*(kwargs.get("allowed_openai_params") or []), "reasoning_effort"] + if bridges_to_responses_api + else kwargs.get("allowed_openai_params") ) optional_param_args: Final = { "functions": functions, From 6f3b3c7957a67d45542538bcda6dc689752b7316 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 14:27:12 +0000 Subject: [PATCH 035/107] test: cover responses bridge reasoning at HTTP boundary Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_main.py | 70 +++++++++++++++++++++++++-------- 1 file changed, 54 insertions(+), 16 deletions(-) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 9ce2bb8400d..8a9d5e07258 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -1368,18 +1368,57 @@ def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict( @pytest.mark.parametrize("reasoning_effort", ["high", {"effort": "high"}]) -@patch("litellm.completion_extras.responses_api_bridge.completion") def test_responses_bridge_preserves_reasoning_effort_with_drop_params( - mock_responses_completion, reasoning_effort, restore_model_registry + reasoning_effort, + restore_model_registry, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, ): - mock_responses_completion.return_value = MagicMock() - model = "perplexity/test-responses-bridge" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + response_body: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "test-responses-bridge", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": None, + "temperature": None, + "tool_choice": "auto", + "tools": [], + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "truncation": None, + "user": None, + } + response_route: Final = respx_mock.post("https://api.perplexity.ai/v1/responses").respond(json=response_body) + model: Final = "perplexity/test-responses-bridge" litellm.register_model( { model: { "litellm_provider": "perplexity", "mode": "responses", - "supports_reasoning": True, + "supports_reasoning": False, "input_cost_per_token": 0.0, "output_cost_per_token": 0.0, } @@ -1387,18 +1426,17 @@ def test_responses_bridge_preserves_reasoning_effort_with_drop_params( persist_across_reloads=False, ) - with patch.object(litellm, "supports_reasoning", return_value=False): - litellm.completion( - model=model, - messages=[{"role": "user", "content": "hello"}], - reasoning_effort=reasoning_effort, - drop_params=True, - api_key="fake-key", - api_base="https://api.perplexity.ai", - ) + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + drop_params=True, + api_key="fake-key", + api_base="https://api.perplexity.ai", + ) - optional_params = mock_responses_completion.call_args.kwargs["optional_params"] - assert optional_params["reasoning_effort"] == reasoning_effort + request_body: Final = json.loads(response_route.calls[0].request.content) + assert request_body["reasoning"] == {"effort": "high"} @pytest.mark.parametrize( From 06497ab42a1ad055e8cbb7695d778cd7a4158ea5 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 15:13:53 +0000 Subject: [PATCH 036/107] fix(ui): guard playground cost metric against null and NaN Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat_ui/ResponseMetrics.test.tsx | 19 ++++++++++++ .../components/chat_ui/ResponseMetrics.tsx | 2 +- .../llm_calls/chat_completion.test.tsx | 6 ++++ .../components/llm_calls/chat_completion.tsx | 5 +++- .../llm_calls/responses_api.test.tsx | 29 +++++++++++++++++++ .../components/llm_calls/responses_api.tsx | 5 +++- 6 files changed, 63 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx index f31e1839739..4ba6d8d50b4 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx @@ -51,4 +51,23 @@ describe("ResponseMetrics prompt cache chips", () => { expect(screen.queryByText(/Response Cache/)).not.toBeInTheDocument(); }); + + it("does not render the Cost chip when a persisted cost is null", () => { + render(); + + expect(screen.queryByText(/Cost:/)).not.toBeInTheDocument(); + expect(screen.getByText("In: 1")).toBeInTheDocument(); + }); + + it("does not render the Cost chip for NaN", () => { + render(); + + expect(screen.queryByText(/Cost:/)).not.toBeInTheDocument(); + }); + + it("renders the Cost chip for a finite cost", () => { + render(); + + expect(screen.getByText("Cost: $0.000063")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx index ec62d0618d7..bb7debc7aee 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx @@ -159,7 +159,7 @@ const ResponseMetrics: React.FC = ({ timeToFirstToken, tot /> )} - {usage?.cost !== undefined && ( + {typeof usage?.cost === "number" && Number.isFinite(usage.cost) && ( { expect(usageData).not.toHaveProperty("cacheReadTokens"); expect(usageData).not.toHaveProperty("cacheCreationTokens"); }); + + it("omits cost when the provider reports a non-numeric value", async () => { + const usageData = await captureUsage({ cost: "not-a-number" }); + + expect(usageData).toEqual(expect.not.objectContaining({ cost: expect.anything() })); + }); }); describe("chat_completion response cache", () => { diff --git a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx index cd0852bc06e..b813a35f687 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx @@ -245,7 +245,10 @@ export async function makeOpenAIChatCompletionRequest( // Extract cost from usage object if available if (chunkWithUsage.usage.cost !== undefined && chunkWithUsage.usage.cost !== null) { - usageData.cost = parseFloat(chunkWithUsage.usage.cost); + const parsedCost = parseFloat(chunkWithUsage.usage.cost); + if (Number.isFinite(parsedCost)) { + usageData.cost = parsedCost; + } } onUsageData(usageData); diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx index 0b94c093acd..7b13f0ac2e3 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx @@ -231,6 +231,35 @@ describe("responses_api", () => { expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ cost: expect.anything() }), ""); }); + it("should omit cost when the proxy reports a non-numeric cost", async () => { + async function* streamWithNonNumericCost() { + yield { + type: "response.completed", + response: { + id: "resp_non_numeric_cost", + usage: { output_tokens: 12, input_tokens: 12, total_tokens: 24, cost: "not-a-number" }, + }, + }; + } + mockResponsesCreate.mockResolvedValueOnce(streamWithNonNumericCost()); + + const onUsageData = vi.fn(); + + await makeOpenAIResponsesRequest( + messages, + mockUpdateTextUI, + "gpt-4", + "test-token", + undefined, + undefined, + undefined, + undefined, + onUsageData, + ); + + expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ cost: expect.anything() }), ""); + }); + it("should replay MCP output items as events for a non-streaming response", async () => { mockResponsesCreate.mockReturnValueOnce( nonStreamingResponse({ diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx index 7ab76488504..63085a81b94 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx @@ -312,7 +312,10 @@ export async function makeOpenAIResponsesRequest( } if (usage.cost !== undefined && usage.cost !== null) { - usageData.cost = Number(usage.cost); + const parsedCost = Number(usage.cost); + if (Number.isFinite(parsedCost)) { + usageData.cost = parsedCost; + } } onUsageData(usageData, mcpToolUsed); From 89c3a8216be6c42bc9d327e312c05d62c99118b9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 11:54:43 -0700 Subject: [PATCH 037/107] fix(bedrock): sign requests off the event loop on every async path SigV4 signing resolves AWS credentials, and botocore refreshes expiring credentials inside that signing with a blocking HTTP call. Every async Bedrock path that still signed on the event loop (/v1/messages, Converse, count tokens, the agent-runtime and Comprehend Medical pass-throughs, async-invoke status polling, realtime, AgentCore, SQS, S3) now signs on a worker thread, so one Bedrock request no longer stalls the whole worker. Fixes #40165 --- .../providers/bedrock_agentcore/handler.py | 7 +- litellm/integrations/s3_v2.py | 4 +- litellm/integrations/sqs.py | 2 +- litellm/llms/bedrock/base_aws_llm.py | 19 ++++- litellm/llms/bedrock/chat/converse_handler.py | 7 +- litellm/llms/bedrock/count_tokens/handler.py | 9 ++- litellm/llms/bedrock/embed/embedding.py | 48 ++++++------ litellm/llms/bedrock/realtime/handler.py | 5 +- litellm/llms/custom_httpx/llm_http_handler.py | 25 ++++-- .../llm_passthrough_endpoints.py | 78 +++++++++++-------- .../chat/test_bedrock_converse_handler.py | 25 +++++- .../test_bedrock_count_tokens_handler.py | 50 ++++++++++++ .../llms/bedrock/event_loop_probe.py | 52 +++++++++++++ .../llms/bedrock/test_base_aws_llm.py | 24 ++++++ .../custom_httpx/test_llm_http_handler.py | 60 ++++++++++++++ .../test_llm_pass_through_endpoints.py | 27 ++++++- 16 files changed, 363 insertions(+), 79 deletions(-) create mode 100644 tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py create mode 100644 tests/test_litellm/llms/bedrock/event_loop_probe.py diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index a4e6fa50901..e01e215b5f0 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -5,6 +5,7 @@ Sends JSON-RPC envelopes directly to AgentCore endpoints, bypassing the completion bridge that would otherwise strip the envelope. """ +import asyncio import json from collections.abc import AsyncIterator, Mapping from typing import Any, Final @@ -45,7 +46,8 @@ class BedrockAgentCoreA2AHandler: Returns: A2A JSON-RPC response dict from the AgentCore agent """ - url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + url, headers, body = await asyncio.to_thread( + BedrockAgentCoreA2ATransformation.get_url_and_signed_request, request_id=request_id, params=params, litellm_params=litellm_params, @@ -91,7 +93,8 @@ class BedrockAgentCoreA2AHandler: Yields: A2A streaming response events from the AgentCore agent """ - url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + url, headers, body = await asyncio.to_thread( + BedrockAgentCoreA2ATransformation.get_url_and_signed_request, request_id=request_id, params=params, litellm_params=litellm_params, diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 712ce41d09e..e3297cfd27c 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -366,7 +366,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Sign the request aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers) aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) - S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) + await asyncio.to_thread(S3SigV4Auth(credentials, "s3", aws_region_name).add_auth, aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) @@ -597,7 +597,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Sign the request aws_request: Final = AWSRequest(method="GET", url=url, headers=headers) - S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request) + await asyncio.to_thread(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 2b4c8c9928d..31cc20f5995 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -295,7 +295,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): data=prepped.body, headers=prepped.headers, ) - SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth(aws_request) + await asyncio.to_thread(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request) signed_headers: Final = dict(aws_request.headers.items()) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index a197702921f..bbd4988b204 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1,3 +1,4 @@ +import asyncio import base64 import hashlib import json @@ -7,7 +8,7 @@ import urllib.parse from collections.abc import Callable, Mapping from datetime import datetime from threading import Lock -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args, overload +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload import httpx from pydantic import BaseModel, ValidationError @@ -1668,3 +1669,19 @@ class BaseAWSLLM: request_headers_dict["Authorization"] = incoming_authorization return request_headers_dict, request.body + + +_SignParams = ParamSpec("_SignParams") +_SignedRequest = TypeVar("_SignedRequest") + + +async def sign_request_off_loop_if_aws( + provider_config: object, + sign_request: Callable[_SignParams, _SignedRequest], + /, + *args: _SignParams.args, + **kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature +) -> _SignedRequest: + if isinstance(provider_config, BaseAWSLLM): + return await asyncio.to_thread(sign_request, *args, **kwargs) + return sign_request(*args, **kwargs) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 984ba371898..3e8e1226a55 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -1,3 +1,4 @@ +import asyncio import json from collections.abc import Mapping from types import MappingProxyType @@ -136,7 +137,8 @@ class BedrockConverseLLM(BaseAWSLLM): ) data: Final = json.dumps(request_data) - prepped: Final = self.get_request_headers( + prepped: Final = await asyncio.to_thread( + self.get_request_headers, credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", extra_headers=headers, @@ -206,7 +208,8 @@ class BedrockConverseLLM(BaseAWSLLM): ) data: Final = json.dumps(request_data) - prepped: Final = self.get_request_headers( + prepped: Final = await asyncio.to_thread( + self.get_request_headers, credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", extra_headers=headers, diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 1fb53f6ff0a..cae2315e744 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -4,6 +4,7 @@ AWS Bedrock CountTokens API handler. Simplified handler leveraging existing LiteLLM Bedrock infrastructure. """ +import asyncio from typing import Any, Final import httpx @@ -12,7 +13,7 @@ import litellm from litellm._logging import verbose_logger from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client class BedrockCountTokensHandler(BedrockCountTokensConfig): @@ -27,6 +28,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): request_data: dict[str, Any], litellm_params: dict[str, Any], resolved_model: str, + client: AsyncHTTPHandler | None = None, ) -> dict[str, Any]: """ Handle a CountTokens request using existing LiteLLM patterns. @@ -75,7 +77,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): # Extract api_key for bearer token auth if provided api_key: Final = litellm_params.get("api_key", None) headers: Final = {"Content-Type": "application/json"} - signed_headers, signed_body = self._sign_request( + signed_headers, signed_body = await asyncio.to_thread( + self._sign_request, service_name="bedrock", headers=headers, optional_params=litellm_params, @@ -85,7 +88,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): api_key=api_key, ) - async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) + async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) response: Final = await async_client.post( endpoint_url, diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index d3725434498..e5f94e51e6d 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -2,10 +2,11 @@ Handles embedding calls to Bedrock's `/invoke` endpoint """ +import asyncio import copy import json import urllib.parse -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Final, get_args, overload import httpx @@ -26,7 +27,7 @@ from litellm.types.llms.bedrock import ( ) from litellm.types.utils import EmbeddingResponse, LlmProviders -from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token +from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token from ..common_utils import BedrockError from .amazon_nova_transformation import AmazonNovaEmbeddingConfig from .amazon_titan_g1_transformation import AmazonTitanG1Config @@ -41,6 +42,20 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +def _sign_get_request( + credentials: Credentials, url: str, headers: Mapping[str, str], aws_region_name: str +) -> AWSPreparedRequest: + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + + request: Final = AWSRequest(method="GET", url=url, data=None, headers=dict(headers)) + SigV4Auth(credentials, "bedrock", aws_region_name).add_auth(request) + return request.prepare() + + class BedrockEmbedding(BaseAWSLLM): @overload def _load_credentials( @@ -599,9 +614,6 @@ class BedrockEmbedding(BaseAWSLLM): dict: Status response from AWS Bedrock """ - # Get AWS credentials using the same method as other Bedrock methods - credentials, _ = self._load_credentials(kwargs) - # Get the runtime endpoint endpoint_url, _ = self.get_runtime_endpoint( api_base=None, @@ -618,27 +630,13 @@ class BedrockEmbedding(BaseAWSLLM): # Prepare headers for GET request headers: Final = {"Content-Type": "application/json"} - # Use AWSRequest directly for GET requests (get_request_headers hardcodes POST) - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - except ImportError: - raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + def sign_status_request() -> AWSPreparedRequest: + credentials, _ = self._load_credentials(kwargs) + return _sign_get_request( + credentials=credentials, url=status_url, headers=headers, aws_region_name=aws_region_name + ) - # Create AWSRequest with GET method and encoded URL - request: Final = AWSRequest( - method="GET", - url=status_url, - data=None, # GET request, no body - headers=headers, - ) - - # Sign the request - SigV4Auth will create canonical string from request URL - sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name) - sigv4.add_auth(request) - - # Prepare the request - prepped: Final = request.prepare() + prepped: Final = await asyncio.to_thread(sign_status_request) # LOGGING if logging_obj is not None: diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 42fe8941443..792e3169a7a 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -149,7 +149,8 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Connecting to %s with model %s", endpoint_uri, model) - credentials: Final = self.get_credentials( + credentials: Final = await asyncio.to_thread( + self.get_credentials, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, aws_session_token=aws_session_token, @@ -169,7 +170,7 @@ class BedrockRealtime(BaseAWSLLM): "or configure credentials in the environment" ), ) - frozen_credentials: Final = credentials.get_frozen_credentials() + frozen_credentials: Final = await asyncio.to_thread(credentials.get_frozen_credentials) # Initialize Bedrock client with aws_sdk_bedrock_runtime config: Final = Config( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2f561809940..48dc009d116 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -77,6 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws from litellm.llms.custom_httpx.container_handler import raise_for_error_status from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -1961,7 +1962,9 @@ class BaseLLMHTTPHandler: api_key=api_key, ) - signed_headers, signed_json_body = provider_config.sign_request( + signed_headers, signed_json_body = await sign_request_off_loop_if_aws( + provider_config, + provider_config.sign_request, headers=headers, optional_params=optional_params, request_data=data, @@ -2062,7 +2065,9 @@ class BaseLLMHTTPHandler: max_attempts, ) provider_config.transform_anthropic_messages_request_on_http_error(e=e, request_data=request_body) - headers, signed_json_body = provider_config.sign_request( + headers, signed_json_body = await sign_request_off_loop_if_aws( + provider_config, + provider_config.sign_request, headers=headers, optional_params=optional_params_dict, request_data=request_body, @@ -2222,7 +2227,9 @@ class BaseLLMHTTPHandler: stream=stream, ) - headers, signed_json_body = anthropic_messages_provider_config.sign_request( + headers, signed_json_body = await sign_request_off_loop_if_aws( + anthropic_messages_provider_config, + anthropic_messages_provider_config.sign_request, headers=headers, optional_params=dict(litellm_params), # dynamic aws_* params are passed under litellm_params request_data=request_body, @@ -2898,7 +2905,9 @@ class BaseLLMHTTPHandler: fake_stream=fake_stream, ) - headers, signed_body = responses_api_provider_config.sign_request( + headers, signed_body = await sign_request_off_loop_if_aws( + responses_api_provider_config, + responses_api_provider_config.sign_request, headers=headers, optional_params=dict(litellm_params), request_data=data, @@ -4606,7 +4615,9 @@ class BaseLLMHTTPHandler: ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) - headers, signed_body = responses_api_provider_config.sign_request( + headers, signed_body = await sign_request_off_loop_if_aws( + responses_api_provider_config, + responses_api_provider_config.sign_request, headers=headers, optional_params=dict(litellm_params), request_data=data, @@ -9833,7 +9844,9 @@ class BaseLLMHTTPHandler: ) all_optional_params: Final[dict[str, object]] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) - headers, signed_json_body = vector_store_provider_config.sign_request( + headers, signed_json_body = await sign_request_off_loop_if_aws( + vector_store_provider_config, + vector_store_provider_config.sign_request, headers=headers, optional_params=all_optional_params, request_data=request_body, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 2ea46b740a8..6eb7b5d8541 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -8,6 +8,7 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc. from __future__ import annotations +import asyncio import hmac import inspect import json @@ -15,6 +16,7 @@ import os import re from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from dataclasses import dataclass +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast @@ -84,6 +86,9 @@ from litellm.utils import ProviderConfigManager from .passthrough_endpoint_router import PassthroughEndpointRouter if TYPE_CHECKING: + from botocore.awsrequest import AWSPreparedRequest + from botocore.credentials import Credentials + from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig from litellm.router import Router @@ -1099,13 +1104,6 @@ async def bedrock_proxy_route( """ create_request_copy(request) - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - from botocore.credentials import Credentials - except ImportError: - raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") - aws_region_name: Final = get_secret_str(secret_name="AWS_REGION_NAME") if not _is_bedrock_agent_runtime_route(endpoint=endpoint): return await bedrock_llm_proxy_route( @@ -1139,17 +1137,20 @@ async def bedrock_proxy_route( from litellm.llms.bedrock.chat import BedrockConverseLLM bedrock_llm: Final = BedrockConverseLLM() - credentials: Final[Credentials] = bedrock_llm.get_credentials() - sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name) - headers: Final = {"Content-Type": "application/json"} # Assuming the body contains JSON data, parse it try: data: Final = await _json_request_body(request) except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) - _request: Final = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers) - sigv4.add_auth(_request) - prepped: Final = _request.prepare() + prepped: Final = await asyncio.to_thread( + _sign_aws_json_post, + get_credentials=bedrock_llm.get_credentials, + service_name="bedrock", + aws_region_name=aws_region_name, + url=str(updated_url), + body=json.dumps(data), + headers=MappingProxyType({"Content-Type": "application/json"}), + ) ## check for streaming is_streaming_request = False @@ -1177,6 +1178,25 @@ async def bedrock_proxy_route( return received_value +def _sign_aws_json_post( + get_credentials: Callable[[], Credentials], + service_name: str, + aws_region_name: str | None, + url: str, + body: str, + headers: Mapping[str, str], +) -> AWSPreparedRequest: + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.") + + aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=dict(headers)) + SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request) + return aws_request.prepare() + + COMPREHEND_MEDICAL_TARGET_PREFIX: Final = "ComprehendMedical_20181030" @@ -1207,13 +1227,6 @@ async def comprehend_medical_proxy_route( [Docs](https://docs.litellm.ai/docs/pass_through/comprehend_medical) """ - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - from botocore.credentials import Credentials - except ImportError: - raise ImportError("Missing boto3 to call comprehendmedical. Run 'pip install boto3'.") - from .llm_provider_handlers.comprehend_medical_passthrough_logging_handler import ( COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS, ) @@ -1246,18 +1259,21 @@ async def comprehend_medical_proxy_route( from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM - credentials: Final[Credentials] = BaseAWSLLM().get_credentials(aws_region_name=aws_region_name) - sigv4: Final = SigV4Auth(credentials, "comprehendmedical", aws_region_name) - headers: Final = MappingProxyType( - { - "Content-Type": "application/x-amz-json-1.1", - "X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}", - } - ) target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" - _request: Final = AWSRequest(method="POST", url=target_url, data=json.dumps(data), headers=headers) - sigv4.add_auth(_request) - prepped: Final = _request.prepare() + prepped: Final = await asyncio.to_thread( + _sign_aws_json_post, + get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name), + service_name="comprehendmedical", + aws_region_name=aws_region_name, + url=target_url, + body=json.dumps(data), + headers=MappingProxyType( + { + "Content-Type": "application/x-amz-json-1.1", + "X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}", + } + ), + ) endpoint_func: Final = create_pass_through_route( endpoint=operation, diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index f34b8eb1fb9..8831fd9061c 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -6,6 +6,7 @@ extension, and AWS credential resolution is stubbed so nothing reaches STS. from __future__ import annotations +import asyncio from unittest.mock import MagicMock, patch import httpx @@ -16,6 +17,7 @@ from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.rust_bridge import chat_completions as bridge from litellm.types.utils import ModelResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe RUST_RESPONSE = { "created": 1_700_000_000, @@ -308,7 +310,9 @@ CONVERSE_RESPONSE = { } -async def _drive_async_completion(*, skip_pre_call_logging: bool, logging_obj): +async def _drive_async_completion( + *, skip_pre_call_logging: bool, logging_obj, credentials: Credentials = RESOLVED_CREDENTIALS +): """Run the real `async_completion` with a stubbed transport.""" import httpx as _httpx @@ -335,7 +339,7 @@ async def _drive_async_completion(*, skip_pre_call_logging: bool, logging_obj): stream=None, optional_params={"maxTokens": 16}, litellm_params={"aws_region_name": "us-west-2"}, - credentials=RESOLVED_CREDENTIALS, + credentials=credentials, headers={}, client=client, skip_pre_call_logging=skip_pre_call_logging, @@ -357,6 +361,23 @@ async def test_async_completion_logs_pre_call_by_default(): assert logging_obj.pre_call.call_count == 1 +@pytest.mark.asyncio +async def test_async_completion_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: botocore refreshes expiring credentials inside SigV4 signing with a + blocking HTTP call, so `async_completion` must sign on a worker thread to keep the loop serving.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + + response = await _drive_async_completion( + skip_pre_call_logging=False, logging_obj=MagicMock(), credentials=probe.credentials() + ) + await release + + assert response.choices[0].message.content == "hi" + assert probe.served_during_refresh is True + + def _sync_client_returning_converse_response(): client = MagicMock() client.post.side_effect = lambda **_kwargs: httpx.Response( diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py new file mode 100644 index 00000000000..31e65bb6e27 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -0,0 +1,50 @@ +import asyncio +from unittest.mock import AsyncMock + +import httpx +import pytest + +from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe + + +class _ProbedCountTokensHandler(BedrockCountTokensHandler): + def __init__(self, probe: EventLoopProbe) -> None: + super().__init__() + self._probe = probe + + def get_credentials(self, **kwargs): + return self._probe.credentials() + + +@pytest.mark.asyncio +async def test_handle_count_tokens_request_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the count_tokens handler signed on the loop, so botocore's blocking + credential refresh inside SigV4 stalled every other request on the worker.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + client = AsyncMock(spec=AsyncHTTPHandler) + client.post = AsyncMock( + return_value=httpx.Response( + 200, + json={"inputTokens": 7}, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com/"), + ) + ) + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + + result = await _ProbedCountTokensHandler(probe).handle_count_tokens_request( + request_data={ + "model": "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "messages": [{"role": "user", "content": "hi"}], + }, + litellm_params={"aws_region_name": "us-west-2"}, + resolved_model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + client=client, + ) + await release + + assert result == {"input_tokens": 7} + assert client.post.call_args.kwargs["headers"]["Authorization"].startswith("AWS4-HMAC-SHA256") + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/bedrock/event_loop_probe.py b/tests/test_litellm/llms/bedrock/event_loop_probe.py new file mode 100644 index 00000000000..68acce308b8 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/event_loop_probe.py @@ -0,0 +1,52 @@ +"""Refreshable credentials whose refresh only completes while the event loop keeps serving.""" + +from __future__ import annotations + +import asyncio +import threading +from datetime import datetime, timedelta, timezone +from typing import Final + +from botocore.credentials import RefreshableCredentials + +REFRESH_RELEASE_TIMEOUT_SECONDS: Final = 2.0 + + +class EventLoopProbe: + """Blocks inside botocore's credential refresh until a coroutine on the loop releases it. + + Signing on the event loop thread can never be released, so `served_during_refresh` reads False there + and True only when the refresh ran on another thread while the loop stayed responsive. + """ + + def __init__(self) -> None: + self.refresh_started: Final = threading.Event() + self.loop_served: Final = threading.Event() + self.served_during_refresh: bool | None = None + + def refresh(self) -> dict[str, str | None]: + self.refresh_started.set() + served: Final = self.loop_served.wait(timeout=REFRESH_RELEASE_TIMEOUT_SECONDS) + if self.served_during_refresh is None: + self.served_during_refresh = served + return { + "access_key": "AKIAREFRESHED", + "secret_key": "refreshed-secret", + "token": None, + "expiry_time": (datetime.now(timezone.utc) + timedelta(hours=1)).isoformat(), + } + + def credentials(self) -> RefreshableCredentials: + return RefreshableCredentials( + access_key="AKIASTALE", + secret_key="stale-secret", + token=None, + expiry_time=datetime.now(timezone.utc) + timedelta(seconds=60), + refresh_using=self.refresh, + method="event-loop-probe", + ) + + async def release_refresh_from_the_loop(self) -> None: + while not self.refresh_started.is_set(): + await asyncio.sleep(0.005) + self.loop_served.set() diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index f854d806bdc..b3fc8dbe959 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1,3 +1,4 @@ +import asyncio import json import os import threading @@ -22,7 +23,9 @@ from litellm.llms.bedrock.base_aws_llm import ( AwsAuthError, BaseAWSLLM, Boto3CredentialsInfo, + sign_request_off_loop_if_aws, ) +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe # Global variable for the base_aws_llm.py file path @@ -3215,3 +3218,24 @@ class TestGetRequestHeadersResign: extra_headers={"Authorization": "Bearer foo"}, ) assert prepped.headers["Authorization"] == "Bearer foo" + + +@pytest.mark.asyncio +async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credentials_refresh(): + """Regression for issue #40165: an AWS provider's signing (and the botocore credential refresh + inside it) must run off the event loop, so other requests keep being served meanwhile.""" + probe = EventLoopProbe() + + def sign(headers: dict[str, str]) -> dict[str, str]: + request = AWSRequest( + method="POST", url="https://bedrock-runtime.us-west-2.amazonaws.com/", data="{}", headers=headers + ) + SigV4Auth(probe.credentials(), "bedrock", "us-west-2").add_auth(request) + return dict(request.headers) + + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + signed = await sign_request_off_loop_if_aws(BaseAWSLLM(), sign, headers={"Content-Type": "application/json"}) + await release + + assert "Authorization" in signed + assert probe.served_during_refresh is True 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 e16855da8cb..6799850c4ae 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 @@ -29,10 +29,14 @@ from litellm.llms.custom_httpx.llm_http_handler import ( _rust_responses_websocket_enabled, ) from litellm.llms.azure.videos.transformation import AzureVideoConfig +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, +) from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" @@ -749,6 +753,62 @@ async def test_anthropic_messages_streaming_response_aclose_closes_agentic_upstr assert tracker.closed is True +class _ProbedBedrockMessagesConfig(AmazonAnthropicClaudeMessagesConfig): + def __init__(self, probe: EventLoopProbe) -> None: + super().__init__() + self._probe = probe + + def get_credentials(self, **kwargs): + return self._probe.credentials() + + +@pytest.mark.asyncio +async def test_async_anthropic_messages_handler_signs_bedrock_off_the_event_loop(monkeypatch): + """Regression for issue #40165: /v1/messages on Bedrock signed on the loop, so botocore's blocking + credential refresh inside SigV4 stalled every other request on the worker.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + handler = BaseLLMHTTPHandler() + upstream_response = httpx.Response( + 200, + json={ + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + "model": "claude-haiku-4-5-20251001", + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com/"), + ) + mock_client = AsyncMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=upstream_response) + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + mock_logging_obj.dynamic_success_callbacks = None + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + + await handler.async_anthropic_messages_handler( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_provider_config=_ProbedBedrockMessagesConfig(probe), + anthropic_messages_optional_request_params={"max_tokens": 16}, + custom_llm_provider="bedrock", + litellm_params=GenericLiteLLMParams(aws_region_name="us-west-2"), + logging_obj=mock_logging_obj, + client=mock_client, + stream=False, + kwargs={}, + ) + await release + + sent_headers = mock_client.post.call_args.kwargs["headers"] + assert sent_headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert probe.served_during_refresh is True + + @pytest.mark.asyncio async def test_async_anthropic_messages_handler_passes_litellm_metadata(): """Ensure litellm_metadata from kwargs is forwarded via update_from_kwargs. 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 d9969dd1dc9..a1cfd973b75 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 @@ -1,3 +1,4 @@ +import asyncio import base64 import contextlib import json @@ -19,6 +20,7 @@ from starlette.datastructures import FormData import litellm from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( BaseOpenAIPassThroughHandler, @@ -1852,11 +1854,11 @@ class TestBedrockAgentRuntimePassthroughToggle: return request @contextlib.contextmanager - def _patched_dispatch(self, general_settings: Mapping[str, object]): + def _patched_dispatch(self, general_settings: Mapping[str, object], credentials: object | None = None): from botocore.credentials import Credentials bedrock_llm: Final = Mock() - bedrock_llm.get_credentials = Mock(return_value=Credentials("ak", "sk")) + bedrock_llm.get_credentials = Mock(return_value=credentials or Credentials("ak", "sk")) forwarder: Final = AsyncMock(return_value="forwarded") with ( @@ -1891,6 +1893,27 @@ class TestBedrockAgentRuntimePassthroughToggle: forwarder.assert_awaited_once() assert "bedrock-agent-runtime.us-east-1.amazonaws.com" in create_route.call_args.kwargs["target"] + @pytest.mark.asyncio + async def test_agent_runtime_dispatch_signs_off_the_event_loop(self, monkeypatch): + """Regression for issue #40165: the agent-runtime pass-through signed on the loop, so botocore's + blocking credential refresh inside SigV4 stalled every other request on the worker.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe: Final = EventLoopProbe() + release: Final = asyncio.create_task(probe.release_refresh_from_the_loop()) + + with self._patched_dispatch(MappingProxyType({}), credentials=probe.credentials()) as (create_route, forwarder): + result: Final = await bedrock_proxy_route( + endpoint=self.AGENT_RUNTIME_ENDPOINT, + request=self._mock_request(), + fastapi_response=Mock(), + user_api_key_dict=UserAPIKeyAuth(), + ) + await release + + assert result == "forwarded" + assert create_route.call_args.kwargs["custom_headers"]["Authorization"].startswith("AWS4-HMAC-SHA256") + assert probe.served_during_refresh is True + @pytest.mark.asyncio @pytest.mark.parametrize("value", (True, "true", "True")) async def test_agent_runtime_dispatch_rejected_when_disabled(self, value: bool | str): From db59e99932f9438ac364c731cb1584092a76c94d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:21:34 -0700 Subject: [PATCH 038/107] fix(azure): let the caller's api-version override a full-URL api_base on relays --- .../llms/azure/passthrough/transformation.py | 9 +++++- .../test_azure_passthrough_transformation.py | 32 +++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 85647ae02c9..dc71c53e126 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -65,6 +65,12 @@ def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> return terminal_event +def without_api_version(api_base: str) -> str: + url: Final = httpx.URL(api_base) + kept_params: Final = tuple((key, value) for key, value in url.params.multi_items() if key != "api-version") + return str(url.copy_with(params=httpx.QueryParams(kept_params))) + + class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return bool(request_data.get("stream")) @@ -89,8 +95,9 @@ class AzurePassthroughConfig(BasePassthroughConfig): native_endpoint: Final = strip_leading_model_segment(routed_endpoint, (model,)) caller_api_version: Final = request_query_params.get("api-version") if request_query_params else None + relay_base: Final = without_api_version(base_target_url) if caller_api_version else base_target_url complete_url: Final = BaseAzureLLM._get_base_azure_url( - api_base=base_target_url, + api_base=relay_base, litellm_params={**litellm_params, "api_version": caller_api_version or litellm_params.get("api_version")}, route=native_endpoint, ) diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 6e85b0cdca6..a3f5c415417 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -369,6 +369,38 @@ def test_azure_passthrough_url_fills_in_the_deployments_api_version_when_the_cal assert url.params["api-version"] == "2024-10-21" +FULL_URL_API_BASE = ( + "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" +) + + +def _full_url_complete_url(request_query_params: dict) -> httpx.URL: + url, _ = AzurePassthroughConfig().get_complete_url( + api_base=FULL_URL_API_BASE, + api_key="key", + model="gpt-4.1-mini", + endpoint="chat/completions", + request_query_params=request_query_params, + litellm_params={}, + ) + return url + + +def test_azure_passthrough_url_prefers_the_callers_api_version_over_a_full_url_api_bases(): + url = _full_url_complete_url(request_query_params={"api-version": "2025-04-01-preview"}) + + assert str(url) == ( + "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions" + "?api-version=2025-04-01-preview" + ) + + +def test_azure_passthrough_url_keeps_a_full_url_api_bases_api_version_when_the_caller_sends_none(): + url = _full_url_complete_url(request_query_params={}) + + assert url.params["api-version"] == "2024-10-21" + + def test_azure_passthrough_url_strips_the_leading_router_model_segment(): url, _ = AzurePassthroughConfig().get_complete_url( api_base="https://my-resource.openai.azure.com", From 5e056a264ecff102a72ee36fed333e5e2ecd9c37 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:45:05 -0700 Subject: [PATCH 039/107] refactor(azure): move the passthrough deployment-segment helpers under llms --- .../llms/azure/passthrough/transformation.py | 23 +++++++++++- litellm/proxy/auth/auth_utils.py | 8 +++-- .../llm_passthrough_endpoints.py | 26 ++++---------- .../test_azure_passthrough_transformation.py | 35 ++++++++++++++++++- .../test_llm_pass_through_endpoints.py | 33 +++-------------- 5 files changed, 71 insertions(+), 54 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index dc71c53e126..51a409c61da 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,4 +1,5 @@ -from collections.abc import Mapping, Sequence +import re +from collections.abc import Callable, Collection, Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional import httpx @@ -65,6 +66,26 @@ def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> return terminal_event +AZURE_DEPLOYMENT_SEGMENT: Final = re.compile(r"(? str | None: + parts: Final = endpoint.split("/") + if len(parts) < 2: + return None + return next((part for part in parts if part in router_models), None) + + +def foreign_azure_deployment( + endpoint: str, model_group: str, served_models: Callable[[], Collection[str]] +) -> str | None: + match: Final = AZURE_DEPLOYMENT_SEGMENT.search(endpoint) + if match is None: + return None + deployment: Final = match.group(1) + return None if deployment == model_group or deployment in served_models() else deployment + + def without_api_version(api_base: str) -> str: url: Final = httpx.URL(api_base) kept_params: Final = tuple((key, value) for key, value in url.params.multi_items() if key != "api-version") diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index bb6ebddf77a..aa77835e858 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -27,6 +27,7 @@ from litellm.litellm_core_utils.url_utils import ( provider_url_destination_candidates, validate_url, ) +from litellm.llms.azure.passthrough.transformation import azure_router_model_in_endpoint from litellm.proxy._types import * from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_metadata from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -2011,9 +2012,10 @@ def get_model_from_request( def _router_model_from_azure_route(route: str, llm_router: Router | None) -> str | None: - from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import azure_router_model_in_endpoint - - return azure_router_model_in_endpoint(re.sub(r"^/azure(?:_ai)?/", "", route, flags=re.IGNORECASE), llm_router) + if llm_router is None: + return None + endpoint: Final = re.sub(r"^/azure(?:_ai)?/", "", route, flags=re.IGNORECASE) + return azure_router_model_in_endpoint(endpoint, frozenset(llm_router.get_model_names())) def _model_from_bedrock_route(route: str) -> str | None: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 47b3c56e66a..66987042962 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -33,6 +33,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.azure.passthrough.transformation import foreign_azure_deployment 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 @@ -121,16 +122,6 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li return False -def azure_router_model_in_endpoint(endpoint: str, llm_router: litellm.Router | None) -> str | None: - parts: Final = endpoint.split("/") - if len(parts) < 2: - return None - return next((part for part in parts if is_known_model(part, llm_router)), None) - - -AZURE_DEPLOYMENT_SEGMENT: Final = re.compile(r"(? str: model: Final = litellm_params.get("model", "") try: @@ -139,17 +130,10 @@ def _deployment_model_name(litellm_params: LiteLLMParamsTypedDict) -> str: return model -def foreign_azure_deployment(endpoint: str, model_group: str, llm_router: litellm.Router) -> str | None: - match: Final = AZURE_DEPLOYMENT_SEGMENT.search(endpoint) - if match is None: - return None - deployment: Final = match.group(1) - if deployment == model_group: - return None - served: Final = frozenset( +def _models_served_by_group(llm_router: litellm.Router, model_group: str) -> frozenset[str]: + return frozenset( _deployment_model_name(row["litellm_params"]) for row in llm_router.get_model_list(model_name=model_group) or () ) - return None if deployment in served else deployment def is_passthrough_request_streaming(request_body: object) -> bool: @@ -1555,7 +1539,9 @@ async def _relay_azure_router_model( is_streaming_request: bool, user_api_key_dict: UserAPIKeyAuth, ) -> Response: - foreign_deployment: Final = foreign_azure_deployment(endpoint, model, llm_router) + foreign_deployment: Final = foreign_azure_deployment( + endpoint, model, lambda: _models_served_by_group(llm_router, model) + ) if foreign_deployment is not None: raise HTTPException( status_code=400, diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index a3f5c415417..69fdd59a748 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -8,7 +8,11 @@ import pytest import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound -from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig +from litellm.llms.azure.passthrough.transformation import ( + AzurePassthroughConfig, + azure_router_model_in_endpoint, + foreign_azure_deployment, +) from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import EmbeddingResponse, ModelResponse @@ -444,3 +448,32 @@ def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_da ) is expected ) + + +@pytest.mark.parametrize( + "endpoint, expected", + [ + ("gpt/openai/deployments/gpt/chat/completions", None), + ("openai/deployments/gpt/chat/completions", None), + ("gpt/openai/deployments/gpt-5.4-mini/chat/completions", None), + ("gpt/models/chat/completions", None), + ("gpt/openai/deployments/gpt-5.4/chat/completions", "gpt-5.4"), + ("gpt/openai/deployments/other-group/chat/completions", "other-group"), + ("openai/deployments/victim/gpt/chat/completions", "victim"), + ], +) +def test_foreign_azure_deployment_names_a_segment_outside_the_group(endpoint, expected): + assert foreign_azure_deployment(endpoint, "gpt", lambda: frozenset({"gpt-5.4-mini"})) == expected + + +@pytest.mark.parametrize( + "endpoint, expected", + [ + ("other-group/openai/deployments/other-group/chat/completions", "other-group"), + ("openai/deployments/gpt/chat/completions", "gpt"), + ("openai/deployments/my-azure-deployment/chat/completions", None), + ("gpt", None), + ], +) +def test_azure_router_model_in_endpoint_picks_the_first_router_model_segment(endpoint, expected): + assert azure_router_model_in_endpoint(endpoint, frozenset({"gpt", "other-group"})) == expected 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 fe4400df704..d1f9e5c4c1d 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 @@ -5368,36 +5368,11 @@ class TestAzureRelayDeploymentSegment: """A key allowed one model group must not reach another deployment by naming it in the ``openai/deployments/`` segment while the group segment picks the credential.""" - @pytest.mark.parametrize( - "endpoint, expected", - [ - ("gpt/openai/deployments/gpt/chat/completions", None), - ("openai/deployments/gpt/chat/completions", None), - ("gpt/openai/deployments/gpt-5.4-mini/chat/completions", None), - ("gpt/models/chat/completions", None), - ("gpt/openai/deployments/gpt-5.4/chat/completions", "gpt-5.4"), - ("gpt/openai/deployments/other-group/chat/completions", "other-group"), - ("openai/deployments/victim/gpt/chat/completions", "victim"), - ], - ) - def test_foreign_azure_deployment_names_a_segment_outside_the_group(self, endpoint, expected): - from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import foreign_azure_deployment + def test_models_served_by_group_resolves_each_deployment_to_its_model_name(self): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _models_served_by_group - assert foreign_azure_deployment(endpoint, "gpt", _AzureGroupRouter([])) == expected - - @pytest.mark.parametrize( - "endpoint, expected", - [ - ("other-group/openai/deployments/other-group/chat/completions", "other-group"), - ("openai/deployments/gpt/chat/completions", "gpt"), - ("openai/deployments/my-azure-deployment/chat/completions", None), - ("gpt", None), - ], - ) - def test_azure_router_model_in_endpoint_matches_the_relay_decision(self, endpoint, expected): - from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import azure_router_model_in_endpoint - - assert azure_router_model_in_endpoint(endpoint, _AzureGroupRouter([])) == expected + assert _models_served_by_group(_AzureGroupRouter([]), "gpt") == frozenset({"gpt-5.4-mini"}) + assert _models_served_by_group(_AzureGroupRouter([]), "missing-group") == frozenset() def _install(self, monkeypatch, body: dict) -> list[dict]: import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep From 71a4a6e912624c31fde159fad5f6508e6b2a217a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:20:50 -0700 Subject: [PATCH 040/107] refactor: keep the AWS JSON signing helper under llms and bound the probe's refresh wait --- litellm/llms/bedrock/base_aws_llm.py | 19 ++++++++++++ .../llm_passthrough_endpoints.py | 29 +++---------------- .../llms/bedrock/event_loop_probe.py | 5 ++++ 3 files changed, 28 insertions(+), 25 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index bbd4988b204..1c7268ffa22 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1671,6 +1671,25 @@ class BaseAWSLLM: return request_headers_dict, request.body +def sign_aws_json_post( + get_credentials: Callable[[], Credentials], + service_name: str, + aws_region_name: str | None, + url: str, + body: str, + headers: Mapping[str, str], +) -> AWSPreparedRequest: + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.") + + aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers) + SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request) + return aws_request.prepare() + + _SignParams = ParamSpec("_SignParams") _SignedRequest = TypeVar("_SignedRequest") diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 6085bc2bfce..037f663044f 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -86,9 +86,6 @@ from litellm.utils import ProviderConfigManager from .passthrough_endpoint_router import PassthroughEndpointRouter if TYPE_CHECKING: - from botocore.awsrequest import AWSPreparedRequest - from botocore.credentials import Credentials - from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig from litellm.router import Router @@ -1134,6 +1131,7 @@ async def bedrock_proxy_route( ) # Add or update query parameters + from litellm.llms.bedrock.base_aws_llm import sign_aws_json_post from litellm.llms.bedrock.chat import BedrockConverseLLM bedrock_llm: Final = BedrockConverseLLM() @@ -1143,7 +1141,7 @@ async def bedrock_proxy_route( except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) prepped: Final = await asyncio.to_thread( - _sign_aws_json_post, + sign_aws_json_post, get_credentials=bedrock_llm.get_credentials, service_name="bedrock", aws_region_name=aws_region_name, @@ -1178,25 +1176,6 @@ async def bedrock_proxy_route( return received_value -def _sign_aws_json_post( - get_credentials: Callable[[], Credentials], - service_name: str, - aws_region_name: str | None, - url: str, - body: str, - headers: Mapping[str, str], -) -> AWSPreparedRequest: - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - except ImportError: - raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.") - - aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers) - SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request) - return aws_request.prepare() - - COMPREHEND_MEDICAL_TARGET_PREFIX: Final = "ComprehendMedical_20181030" @@ -1257,11 +1236,11 @@ async def comprehend_medical_proxy_route( if "stream" in data: raise HTTPException(status_code=400, detail="'stream' is not a Comprehend Medical request member") - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, sign_aws_json_post target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" prepped: Final = await asyncio.to_thread( - _sign_aws_json_post, + sign_aws_json_post, get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name), service_name="comprehendmedical", aws_region_name=aws_region_name, diff --git a/tests/test_litellm/llms/bedrock/event_loop_probe.py b/tests/test_litellm/llms/bedrock/event_loop_probe.py index 68acce308b8..c347247ec32 100644 --- a/tests/test_litellm/llms/bedrock/event_loop_probe.py +++ b/tests/test_litellm/llms/bedrock/event_loop_probe.py @@ -4,12 +4,14 @@ from __future__ import annotations import asyncio import threading +import time from datetime import datetime, timedelta, timezone from typing import Final from botocore.credentials import RefreshableCredentials REFRESH_RELEASE_TIMEOUT_SECONDS: Final = 2.0 +REFRESH_START_TIMEOUT_SECONDS: Final = 10.0 class EventLoopProbe: @@ -47,6 +49,9 @@ class EventLoopProbe: ) async def release_refresh_from_the_loop(self) -> None: + deadline: Final = time.monotonic() + REFRESH_START_TIMEOUT_SECONDS while not self.refresh_started.is_set(): + if time.monotonic() > deadline: + raise TimeoutError("signing finished without ever starting a credential refresh") await asyncio.sleep(0.005) self.loop_served.set() From 0f224dd4ed1fb4bed9ff67d8d8a095e0d032fcae Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:59:34 -0700 Subject: [PATCH 041/107] fix: sign Bedrock guardrail, embedding, and Mantle requests off the event loop Bedrock Guardrails resolved credentials and signed inline on the loop in its three async paths, Titan embeddings signed each batch item inline in the async loop, and Bedrock Mantle requests slipped past the off-loop gate because BedrockMantleAuthMixin composes a BaseAWSLLM instead of inheriting from it. Introduce the SignsRequestsWithAWS marker that both BaseAWSLLM and the Mantle mixin carry so sign_request_off_loop_if_aws covers Mantle, and move the guardrail and embedding signing into asyncio.to_thread. Every new test fails at the previous tip. --- litellm/llms/bedrock/base_aws_llm.py | 8 +++- litellm/llms/bedrock/embed/embedding.py | 3 +- litellm/llms/bedrock_mantle/common_utils.py | 4 +- .../guardrail_hooks/bedrock_guardrails.py | 14 +++++-- .../test_bedrock_async_invoke_embedding.py | 35 ++++++++++++++++ .../bedrock/embed/test_bedrock_embedding.py | 42 +++++++++++++++++++ .../test_bedrock_mantle_transformation.py | 28 +++++++++++++ .../test_bedrock_guardrails.py | 36 ++++++++++++++++ .../test_bedrock_invoke_guardrail_checks.py | 33 +++++++++++++++ 9 files changed, 194 insertions(+), 9 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 1c7268ffa22..eddd4d6f826 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -81,7 +81,11 @@ class AwsAuthError(Exception): super().__init__(self.message) # Call the base class constructor with the parameters it needs -class BaseAWSLLM: +class SignsRequestsWithAWS: + pass + + +class BaseAWSLLM(SignsRequestsWithAWS): # Process-wide IAM credential cache (shared across instances — Bedrock passthrough is per-request). # Storage is in-process memory only: no Redis backend unless attached elsewhere. Entry TTL: static # access-key + secret + region use ``_get_default_ttl_for_boto3_credentials`` (~59 minutes); ambient @@ -1701,6 +1705,6 @@ async def sign_request_off_loop_if_aws( *args: _SignParams.args, **kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature ) -> _SignedRequest: - if isinstance(provider_config, BaseAWSLLM): + if isinstance(provider_config, SignsRequestsWithAWS): return await asyncio.to_thread(sign_request, *args, **kwargs) return sign_request(*args, **kwargs) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index d0e8d734c5e..2c66fa1ef06 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -357,7 +357,8 @@ class BedrockEmbedding(BaseAWSLLM): if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - prepped = self.get_request_headers( + prepped = await asyncio.to_thread( + self.get_request_headers, credentials=credentials, aws_region_name=aws_region_name, extra_headers=extra_headers, diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index 850738bc320..ac94d9cb922 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -23,7 +23,7 @@ from botocore.exceptions import ( ProfileNotFound, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS from litellm.secret_managers.main import get_secret_str BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1" @@ -55,7 +55,7 @@ def resolve_mantle_region(params: Mapping[str, object]) -> str: ) -class BedrockMantleAuthMixin: +class BedrockMantleAuthMixin(SignsRequestsWithAWS): _aws_signer: BaseAWSLLM @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 7204839f6d3..0ae2c48a3b4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -917,7 +917,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source, ) return BedrockGuardrailResponse() - credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key)) + credentials, aws_region_name = await asyncio.to_thread( + self._load_credentials, bearer_token=bedrock_bearer_token(api_key) + ) allow_chunking: Final = not self._content_uses_contextual_grounding(content) completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator @@ -1178,7 +1180,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): **base_request_data, "content": content, } # mutable-ok: outbound JSON request body - prepared_request: Final = self._prepare_request( + prepared_request: Final = await asyncio.to_thread( + self._prepare_request, credentials=credentials, data=bedrock_request_data, optional_params=self.optional_params, @@ -1875,10 +1878,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return BedrockGuardrailResponse() api_key: Final[str | None] = request_data.get("api_key") if request_data else None - credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key)) + credentials, aws_region_name = await asyncio.to_thread( + self._load_credentials, bearer_token=bedrock_bearer_token(api_key) + ) body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks} - prepared_request: Final = self._prepare_request( + prepared_request: Final = await asyncio.to_thread( + self._prepare_request, credentials=credentials, data=body, optional_params=self.optional_params, diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py index 74a55cc1ef2..d8dfaaeac99 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py @@ -1,11 +1,15 @@ import json +import asyncio from unittest.mock import Mock, patch +import httpx import pytest +import respx import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.base import HiddenParams +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe # Mock async invoke responses async_invoke_response = { @@ -383,3 +387,34 @@ class TestBedrockAsyncInvokeEmbedding: async_endpoint == "https://bedrock-runtime.us-east-1.amazonaws.com/async-invoke" ) + + +@pytest.mark.asyncio +async def test_async_invoke_status_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the GetAsyncInvoke poll is a signed GET, and botocore refreshes + expiring credentials inside that signing with a blocking HTTP call, so it must run on a worker + thread to keep the loop serving other requests.""" + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + embedder = BedrockEmbedding() + probe = EventLoopProbe() + + with ( + patch.object(embedder, "_load_credentials", return_value=(probe.credentials(), "us-east-1")), + respx.mock, + ): + route = respx.get(url__regex=r"https://bedrock-runtime\.us-east-1\.amazonaws\.com/async-invoke/.*").mock( + return_value=httpx.Response(200, json=async_invoke_status_response) + ) + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + status = await embedder._get_async_invoke_status( + invocation_arn=async_invoke_status_response["invocationArn"], aws_region_name="us-east-1" + ) + await release + + assert status["status"] == "InProgress" + assert "Authorization" in route.calls.last.request.headers + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 50f8bbcf584..4ebf08dc709 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -1,11 +1,16 @@ import json +import asyncio import os from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, MagicMock import pytest +import httpx import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.bedrock.embed.embedding import BedrockEmbedding +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe # Mock responses for different embedding models titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10} @@ -1059,3 +1064,40 @@ def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(mo assert response.data[0]["embedding"] == titan_embedding_response["embedding"] assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +@pytest.mark.asyncio +async def test_async_single_func_embeddings_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: Titan, Nova, and TwelveLabs embeddings sign one SigV4 request per + input, and botocore refreshes expiring credentials inside that signing with a blocking HTTP call, + so each signing must run on a worker thread to keep the loop serving other requests.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + client = MagicMock() + client.__class__ = AsyncHTTPHandler + client.post = AsyncMock( + return_value=httpx.Response( + 200, + json=titan_embedding_response, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"), + ) + ) + + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + response = await BedrockEmbedding()._async_single_func_embeddings( + client=client, + timeout=None, + batch_data=[{"inputText": test_input}], + credentials=probe.credentials(), + extra_headers=None, + endpoint_url="https://bedrock-runtime.us-west-2.amazonaws.com/model/amazon.titan-embed-text-v1/invoke", + aws_region_name="us-west-2", + model="amazon.titan-embed-text-v1", + logging_obj=MagicMock(), + provider="amazon", + ) + await release + + assert response.data[0]["embedding"] == titan_embedding_response["embedding"] + assert "Authorization" in client.post.call_args.kwargs["headers"] + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 1be94d4daa2..e83a844c87e 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -6,15 +6,20 @@ API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.ht """ import json +import asyncio from unittest.mock import patch import httpx import pytest +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest import litellm from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig +from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws from litellm.types.utils import LlmProviders +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe @pytest.fixture @@ -710,3 +715,26 @@ def test_gemma_4_models_register_under_bedrock_mantle(local_cost_map, model_id): resolved_model, provider, _, _ = litellm.get_llm_provider(full_model_name) assert provider == "bedrock_mantle" assert resolved_model == model_id + + +@pytest.mark.asyncio +async def test_mantle_signing_runs_off_the_event_loop(): + """Regression for issue #40165: Mantle signs with SigV4 through a composed BaseAWSLLM, so the + off-loop gate must recognise it too, or its credential refresh blocks the loop like Bedrock's did.""" + probe = EventLoopProbe() + + def sign(headers: dict[str, str]) -> dict[str, str]: + request = AWSRequest( + method="POST", url="https://bedrock-mantle.us-east-1.api.aws/v1/responses", data="{}", headers=headers + ) + SigV4Auth(probe.credentials(), "bedrock", "us-east-1").add_auth(request) + return dict(request.headers) + + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + signed = await sign_request_off_loop_if_aws( + BedrockMantleChatConfig(), sign, headers={"Content-Type": "application/json"} + ) + await release + + assert "Authorization" in signed + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 7da55f22bda..c0762edec92 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -3,6 +3,8 @@ Unit tests for Bedrock Guardrails """ import json +import asyncio +from datetime import datetime, timezone import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -28,6 +30,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockTextContent, ) from litellm.types.utils import CallTypes, ModelResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe @pytest.mark.asyncio @@ -5842,3 +5845,36 @@ async def test_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): assert response["action"] == "NONE" assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +@pytest.mark.asyncio +async def test_apply_guardrail_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the ApplyGuardrail request is signed with SigV4, and botocore + refreshes expiring credentials inside that signing with a blocking HTTP call, so it must run + on a worker thread to keep the loop serving other requests.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") + probe = EventLoopProbe() + allowed = httpx.Response( + 200, + json={"action": "NONE", "outputs": [], "assessments": []}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com"), + ) + + with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=allowed)): + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + response = await guardrail._post_apply_guardrail_content( + content=[{"text": {"text": "hello"}}], + base_request_data={"source": "INPUT"}, + credentials=probe.credentials(), + aws_region_name="us-east-1", + api_key=None, + request_data={}, + event_type=GuardrailEventHooks.pre_call, + start_time=datetime.now(timezone.utc), + completed_chunk_usages=[], + ) + await release + + assert response["action"] == "NONE" + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py index f4af77d2e40..bbc8fd539a3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py @@ -5,10 +5,12 @@ All Bedrock HTTP calls are mocked; no real AWS calls are made. """ import json +import asyncio import logging from unittest.mock import AsyncMock, MagicMock, patch import pytest +import httpx from fastapi import HTTPException @@ -21,6 +23,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrailResponse, ) from litellm.types.utils import Choices, Message, ModelResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe CONTENT_FILTER_CHECKS = {"contentFilter": {"categories": [{"category": "VIOLENCE"}]}} @@ -861,3 +864,33 @@ async def test_checks_bearer_token_never_runs_the_sigv4_credential_chain(monkeyp {"check": "contentFilter", "category": "VIOLENCE", "severityScore": 0.8} ] assert post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +@pytest.mark.asyncio +async def test_invoke_guardrail_checks_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the checks request is signed with SigV4, and botocore refreshes + expiring credentials inside that signing with a blocking HTTP call, so it must run on a worker + thread to keep the loop serving other requests.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + g = BedrockGuardrail(checks=CONTENT_FILTER_CHECKS, content_filter_threshold=0.5) + probe = EventLoopProbe() + allowed = httpx.Response( + 200, + json={"results": {"contentFilter": {"results": [{"category": "VIOLENCE", "severityScore": 0.1}]}}}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com"), + ) + + with ( + patch.object(g, "_load_credentials", return_value=(probe.credentials(), "us-east-1")), + patch.object(g.async_handler, "post", new=AsyncMock(return_value=allowed)), + ): + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + response = await g.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={"messages": []}, + ) + await release + + assert response == BedrockGuardrailResponse() + assert probe.served_during_refresh is True From a892e67c40c3d84064f4c9f3d83c74bb1e4295a5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:05:06 -0700 Subject: [PATCH 042/107] fix(azure): match deployment segments case-insensitively and keep relay helpers immutable --- litellm/llms/azure/passthrough/transformation.py | 15 ++++++++++----- litellm/llms/openai/responses/transformation.py | 2 +- .../llm_passthrough_endpoints.py | 16 +++++++++------- .../test_azure_passthrough_transformation.py | 10 ++++++++++ 4 files changed, 30 insertions(+), 13 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 51a409c61da..69767bb29a3 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,5 +1,6 @@ import re from collections.abc import Callable, Collection, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional import httpx @@ -55,9 +56,7 @@ def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> """A streaming logging object assembles the logged response from the terminal event, not from its body.""" from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig - terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks( - all_chunks=list(all_chunks) - ) + terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks=all_chunks) if terminal_event is None: return None logging_obj.call_type = ( @@ -83,7 +82,11 @@ def foreign_azure_deployment( if match is None: return None deployment: Final = match.group(1) - return None if deployment == model_group or deployment in served_models() else deployment + folded: Final = deployment.casefold() + if folded == model_group.casefold(): + return None + served: Final = frozenset(name.casefold() for name in served_models()) + return None if folded in served else deployment def without_api_version(api_base: str) -> str: @@ -119,7 +122,9 @@ class AzurePassthroughConfig(BasePassthroughConfig): relay_base: Final = without_api_version(base_target_url) if caller_api_version else base_target_url complete_url: Final = BaseAzureLLM._get_base_azure_url( api_base=relay_base, - litellm_params={**litellm_params, "api_version": caller_api_version or litellm_params.get("api_version")}, + litellm_params=MappingProxyType( + {**litellm_params, "api_version": caller_api_version or litellm_params.get("api_version")} + ), route=native_endpoint, ) return ( diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index cde399065fd..97e7bcc60b7 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -620,7 +620,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return event_pydantic_model.model_construct(**parsed_chunk) @staticmethod - def parse_terminal_event_from_stream_chunks(all_chunks: list[str]) -> ResponsesTerminalEvent | None: + def parse_terminal_event_from_stream_chunks(all_chunks: Sequence[str]) -> ResponsesTerminalEvent | None: for chunk_str in reversed(all_chunks): for event_model in (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent): try: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 66987042962..086e1654f70 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -122,6 +122,10 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li return False +class RelayRejection(TypedDict): + error: ReadOnly[str] + + def _deployment_model_name(litellm_params: LiteLLMParamsTypedDict) -> str: model: Final = litellm_params.get("model", "") try: @@ -1543,13 +1547,11 @@ async def _relay_azure_router_model( endpoint, model, lambda: _models_served_by_group(llm_router, model) ) if foreign_deployment is not None: - raise HTTPException( - status_code=400, - detail={ - "error": f"deployment '{foreign_deployment}' in the path is not served by model group '{model}'; " - "put the model group name in the deployments segment" - }, - ) + rejection: Final[RelayRejection] = { + "error": f"deployment '{foreign_deployment}' in the path is not served by model group '{model}'; " + "put the model group name in the deployments segment" + } + raise HTTPException(status_code=400, detail=rejection) try: result: Final = await llm_router.allm_passthrough_route( model=model, diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 69fdd59a748..798ef1f5a4f 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -456,16 +456,26 @@ def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_da ("gpt/openai/deployments/gpt/chat/completions", None), ("openai/deployments/gpt/chat/completions", None), ("gpt/openai/deployments/gpt-5.4-mini/chat/completions", None), + ("gpt/openai/deployments/GPT-5.4-MINI/chat/completions", None), + ("gpt/openai/deployments/Gpt/chat/completions", None), ("gpt/models/chat/completions", None), ("gpt/openai/deployments/gpt-5.4/chat/completions", "gpt-5.4"), ("gpt/openai/deployments/other-group/chat/completions", "other-group"), ("openai/deployments/victim/gpt/chat/completions", "victim"), + ("gpt/openai/deployments/GPT-5.4/chat/completions", "GPT-5.4"), ], ) def test_foreign_azure_deployment_names_a_segment_outside_the_group(endpoint, expected): assert foreign_azure_deployment(endpoint, "gpt", lambda: frozenset({"gpt-5.4-mini"})) == expected +def test_foreign_azure_deployment_skips_the_router_when_the_segment_is_the_group_itself(): + def served_models(): + raise AssertionError("the router must not be consulted for the group's own name") + + assert foreign_azure_deployment("gpt/openai/deployments/Gpt/chat/completions", "gpt", served_models) is None + + @pytest.mark.parametrize( "endpoint, expected", [ From 5ea447e295109b4079d1b41a70f82f1d88024e6a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:19:23 -0700 Subject: [PATCH 043/107] fix(azure): keep the model group segment exact and casefold only deployment names in the relay guard --- litellm/llms/azure/passthrough/transformation.py | 5 ++--- .../passthrough/test_azure_passthrough_transformation.py | 4 ++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 69767bb29a3..1b5a4083ebe 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -82,11 +82,10 @@ def foreign_azure_deployment( if match is None: return None deployment: Final = match.group(1) - folded: Final = deployment.casefold() - if folded == model_group.casefold(): + if deployment == model_group: return None served: Final = frozenset(name.casefold() for name in served_models()) - return None if folded in served else deployment + return None if deployment.casefold() in served else deployment def without_api_version(api_base: str) -> str: diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 798ef1f5a4f..c7e86616ee2 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -457,7 +457,7 @@ def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_da ("openai/deployments/gpt/chat/completions", None), ("gpt/openai/deployments/gpt-5.4-mini/chat/completions", None), ("gpt/openai/deployments/GPT-5.4-MINI/chat/completions", None), - ("gpt/openai/deployments/Gpt/chat/completions", None), + ("gpt/openai/deployments/Gpt/chat/completions", "Gpt"), ("gpt/models/chat/completions", None), ("gpt/openai/deployments/gpt-5.4/chat/completions", "gpt-5.4"), ("gpt/openai/deployments/other-group/chat/completions", "other-group"), @@ -473,7 +473,7 @@ def test_foreign_azure_deployment_skips_the_router_when_the_segment_is_the_group def served_models(): raise AssertionError("the router must not be consulted for the group's own name") - assert foreign_azure_deployment("gpt/openai/deployments/Gpt/chat/completions", "gpt", served_models) is None + assert foreign_azure_deployment("gpt/openai/deployments/gpt/chat/completions", "gpt", served_models) is None @pytest.mark.parametrize( From c6a0381595b3085937e514c7af3f54c2ceb77598 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:38:01 -0700 Subject: [PATCH 044/107] fix(token_counter): bound concurrent HuggingFace encodes so a burst of large counts cannot exhaust memory --- litellm/constants.py | 1 + litellm/litellm_core_utils/token_counter.py | 6 +++- .../litellm_core_utils/test_token_counter.py | 31 +++++++++++++++++++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index a9112c90a54..05d442dd070 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -397,6 +397,7 @@ TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range( minimum=1, maximum=1_000_000_000, ) +TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES: Final = 4 MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512)) MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512)) OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index e3af64fd3bf..8624e79faea 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -3,6 +3,7 @@ import base64 import io import struct +import threading from collections.abc import Callable, Iterable, Mapping, Sequence from typing import Final, Literal, cast @@ -22,6 +23,7 @@ from litellm.constants import ( MAX_TILE_HEIGHT, MAX_TILE_WIDTH, TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS, + TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES, TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.default_encoding import encoding as default_encoding @@ -320,6 +322,7 @@ Type for a function that counts tokens in a string. """ EXTRAPOLATION_SAMPLES: Final = 16 +_HF_ENCODE_SLOTS: Final = threading.BoundedSemaphore(TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES) def _get_tiktoken_count_function( @@ -587,7 +590,8 @@ def _get_exact_count_function( tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"] def count_tokens(text: str) -> int: - return len(tokenizer.encode_batch_fast([text])[0]) + with _HF_ENCODE_SLOTS: + return len(tokenizer.encode_batch_fast([text])[0]) return count_tokens elif tokenizer_json["type"] == "openai_tokenizer": diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index e3dd5feb18d..655be5146cd 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -2,8 +2,10 @@ # This tests litellm.token_counter.token_counter() function import asyncio import importlib +import threading import time import traceback +from concurrent.futures import ThreadPoolExecutor from typing import Final from unittest.mock import MagicMock @@ -16,6 +18,7 @@ import litellm from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants +from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.token_counter import ( _get_exact_count_function, @@ -162,6 +165,34 @@ def test_count_at_or_below_the_cap_is_exact(): assert count_exactly.call_args_list == [(("a" * 5_000,),)] +class _SlowEncoder: + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self.in_flight = 0 + self.peak_in_flight = 0 + + def encode_batch_fast(self, texts: list[str]) -> list[list[int]]: + with self._lock: + self.in_flight += 1 + self.peak_in_flight = max(self.peak_in_flight, self.in_flight) + time.sleep(0.1) + with self._lock: + self.in_flight -= 1 + return [[0] * len(text) for text in texts] + + +def test_huggingface_counts_run_at_most_the_configured_number_at_once(): + encoder: Final = _SlowEncoder() + count: Final = _get_exact_count_function(None, {"type": "huggingface_tokenizer", "tokenizer": encoder}) + burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES + + with ThreadPoolExecutor(max_workers=burst) as pool: + counts: Final = tuple(pool.map(count, ["abc"] * burst)) + + assert counts == (3,) * burst + assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES + + def test_token_counter_applies_the_default_cap(): max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] From 84a99dfc3188fdefbf9c39abc87b7621d126328e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 18:24:51 -0700 Subject: [PATCH 045/107] fix(azure_ai): surface rejected MAI image params as 400 and drop comments --- .../image_generation/mai_transformation.py | 60 +++++++++---------- .../test_mai_image_generation.py | 31 ++++++---- 2 files changed, 48 insertions(+), 43 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 3421cc0a4a1..2d9d3246b4c 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.exceptions import UnsupportedParamsError from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -21,16 +22,7 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_WIDTH = 1024 DEFAULT_HEIGHT = 1024 - # The MAI endpoint produces exactly one image per request. Its documented - # body is model/prompt/width/height (plus `image` for edits) — there is no - # count field, and `n` (or the native `sampleCount`) is accepted and - # ignored, so a request for more silently comes back with one. MAX_IMAGES_PER_REQUEST: Final = 1 - - # Provider-side bounds on the generated image. Both are enforced by the - # MAI endpoint, which 400s with "'width' must be at least 768 pixels." - # Only `size` is checked against them: `width`/`height` pass through - # unmapped, which keeps a future model with different bounds reachable. MIN_DIMENSION_PX: Final = 768 MAX_TOTAL_PX: Final = 1024 * 1024 @@ -158,25 +150,27 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if k in supported_params: if k == "size" and v: - self._map_size_param(v, optional_params) - elif k == "n" and v is not None and v > self.MAX_IMAGES_PER_REQUEST: + self._map_size_param(v, optional_params, model) + elif k == "n" and v is not None and int(v) > self.MAX_IMAGES_PER_REQUEST: if not drop_params: - raise ValueError( + raise self._unsupported( + model, f"n={v} is not supported for model {model}. The Azure AI MAI image " f"endpoint returns exactly {self.MAX_IMAGES_PER_REQUEST} image per " "request and ignores any count, so a larger value would silently " "return fewer images than requested. Send one request per image, or " - "set drop_params=True to drop n." + "set drop_params=True to drop n.", ) else: optional_params[k] = v elif k in ("width", "height"): optional_params[k] = v elif not drop_params: - raise ValueError( + raise self._unsupported( + model, f"Parameter {k} is not supported for model {model}. " f"Supported parameters are {supported_params} and width/height. " - f"Set drop_params=True to drop unsupported parameters." + f"Set drop_params=True to drop unsupported parameters.", ) if "width" not in optional_params: @@ -187,7 +181,11 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): optional_params.pop("size", None) return optional_params - def _map_size_param(self, size: str, optional_params: dict) -> None: + @staticmethod + def _unsupported(model: str, message: str) -> UnsupportedParamsError: + return UnsupportedParamsError(message=message, llm_provider="azure_ai", model=model) + + def _map_size_param(self, size: str, optional_params: dict, model: str) -> None: size_mapping: Final = { "1024x1024": (1024, 1024), "1792x1024": (1792, 1024), @@ -202,34 +200,32 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): try: width, height = map(int, size.lower().split("x")) except ValueError: - raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').") + raise self._unsupported( + model, f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')." + ) else: - raise ValueError( + raise self._unsupported( + model, f"Unsupported size value: '{size}'. " - f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." + f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string.", ) - self._validate_dimensions(size=size, width=width, height=height) + self._validate_dimensions(model=model, size=size, width=width, height=height) optional_params["width"] = width optional_params["height"] = height - def _validate_dimensions(self, size: str, width: int, height: int) -> None: - """Reject a `size` the MAI endpoint would 400 on. - - Several OpenAI-standard sizes are outside MAI's bounds: 512x512 and - 256x256 fall under the per-side minimum, and 1792x1024 / 1024x1792 - exceed the total pixel budget. Checking here turns an opaque provider - 400 into an error that names the constraint. - """ + def _validate_dimensions(self, model: str, size: str, width: int, height: int) -> None: if width < self.MIN_DIMENSION_PX or height < self.MIN_DIMENSION_PX: - raise ValueError( + raise self._unsupported( + model, f"Unsupported size value: '{size}'. Azure AI MAI image models require width and " - f"height of at least {self.MIN_DIMENSION_PX} pixels." + f"height of at least {self.MIN_DIMENSION_PX} pixels.", ) if width * height > self.MAX_TOTAL_PX: - raise ValueError( + raise self._unsupported( + model, f"Unsupported size value: '{size}'. Azure AI MAI image models accept at most " - f"{self.MAX_TOTAL_PX} total pixels ({width}x{height} is {width * height})." + f"{self.MAX_TOTAL_PX} total pixels ({width}x{height} is {width * height}).", ) def transform_image_generation_response( diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py index 6cbc79fff2b..29aa5967c9a 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -4,6 +4,7 @@ import httpx import pytest import litellm +from litellm.exceptions import UnsupportedParamsError from litellm.llms.azure.azure import AzureChatCompletion from litellm.llms.azure.image_generation import get_azure_image_generation_config from litellm.llms.azure.image_generation.http_utils import ( @@ -180,7 +181,7 @@ class TestAzureMAIImageGeneration: def test_map_openai_params_unsupported_size_raises(self): config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(ValueError, match="Unsupported size value: 'auto'"): + with pytest.raises(UnsupportedParamsError, match="Unsupported size value: 'auto'"): config.map_openai_params( non_default_params={"size": "auto"}, optional_params={}, @@ -190,7 +191,7 @@ class TestAzureMAIImageGeneration: def test_map_openai_params_invalid_custom_size_raises(self): config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"): + with pytest.raises(UnsupportedParamsError, match="Invalid size format: '1024xabc'"): config.map_openai_params( non_default_params={"size": "1024xabc"}, optional_params={}, @@ -200,9 +201,8 @@ class TestAzureMAIImageGeneration: @pytest.mark.parametrize("size", ["512x512", "256x256", "700x1400"]) def test_map_openai_params_size_below_minimum_dimension_raises(self, size): - """MAI requires >= 768px per side; the OpenAI size table offered smaller ones.""" config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(ValueError, match="at least 768 pixels"): + with pytest.raises(UnsupportedParamsError, match="at least 768 pixels"): config.map_openai_params( non_default_params={"size": size}, optional_params={}, @@ -212,9 +212,8 @@ class TestAzureMAIImageGeneration: @pytest.mark.parametrize("size", ["1792x1024", "1024x1792"]) def test_map_openai_params_size_over_total_pixel_budget_raises(self, size): - """MAI caps total pixels at 1024*1024, so both landscape/portrait sizes 400 upstream.""" config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(ValueError, match="at most 1048576 total pixels"): + with pytest.raises(UnsupportedParamsError, match="at most 1048576 total pixels"): config.map_openai_params( non_default_params={"size": size}, optional_params={}, @@ -223,7 +222,6 @@ class TestAzureMAIImageGeneration: ) def test_map_openai_params_explicit_width_height_not_range_checked(self): - """width/height pass through unmapped, so a future model's bounds stay reachable.""" config = AzureFoundryMAIImageGenerationConfig() optional_params = config.map_openai_params( non_default_params={"width": 1792, "height": 1024}, @@ -234,11 +232,10 @@ class TestAzureMAIImageGeneration: assert optional_params["width"] == 1792 assert optional_params["height"] == 1024 - @pytest.mark.parametrize("n", [2, 4]) + @pytest.mark.parametrize("n", [2, 4, "2"]) def test_map_openai_params_multi_image_n_raises(self, n): - """The MAI endpoint returns one image and ignores any count, so n>1 must not pass silently.""" config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(ValueError, match="returns exactly 1 image per request"): + with pytest.raises(UnsupportedParamsError, match="returns exactly 1 image per request"): config.map_openai_params( non_default_params={"n": n}, optional_params={}, @@ -266,9 +263,21 @@ class TestAzureMAIImageGeneration: ) assert optional_params["n"] == 1 + @pytest.mark.parametrize("params", [{"n": 2}, {"size": "512x512"}, {"size": "1792x1024"}]) + def test_image_generation_rejected_params_surface_as_400(self, params): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.image_generation( + model="azure_ai/MAI-Image-2.5", + prompt="A photograph of a red fox", + api_key="test-key", + api_base="https://my-resource.services.ai.azure.com", + **params, + ) + assert exc_info.value.status_code == 400 + def test_map_openai_params_unsupported_param_raises(self): config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(ValueError, match="Parameter quality is not supported"): + with pytest.raises(UnsupportedParamsError, match="Parameter quality is not supported"): config.map_openai_params( non_default_params={"quality": "hd"}, optional_params={}, From 38a3de764154a600441a37c746d78c7ca89d57bf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 18:46:18 -0700 Subject: [PATCH 046/107] fix(azure_ai): use the pixel cap the live MAI endpoint enforces --- .../image_generation/mai_transformation.py | 2 +- .../test_mai_image_generation.py | 23 ++++++++++++++++++- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 2d9d3246b4c..7f7a1044689 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -24,7 +24,7 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): MAX_IMAGES_PER_REQUEST: Final = 1 MIN_DIMENSION_PX: Final = 768 - MAX_TOTAL_PX: Final = 1024 * 1024 + MAX_TOTAL_PX: Final = 1_056_768 @staticmethod def get_mai_image_generation_url( diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py index 29aa5967c9a..a0bc3fb00b7 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -213,7 +213,7 @@ class TestAzureMAIImageGeneration: @pytest.mark.parametrize("size", ["1792x1024", "1024x1792"]) def test_map_openai_params_size_over_total_pixel_budget_raises(self, size): config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(UnsupportedParamsError, match="at most 1048576 total pixels"): + with pytest.raises(UnsupportedParamsError, match="at most 1056768 total pixels"): config.map_openai_params( non_default_params={"size": size}, optional_params={}, @@ -221,6 +221,27 @@ class TestAzureMAIImageGeneration: drop_params=True, ) + @pytest.mark.parametrize("size", ["1032x1024", "1376x768"]) + def test_map_openai_params_size_at_live_pixel_cap_passes_through(self, size): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"size": size}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + assert optional_params["width"] * optional_params["height"] == 1_056_768 + + def test_map_openai_params_size_one_pixel_over_live_cap_raises(self): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(UnsupportedParamsError, match="at most 1056768 total pixels"): + config.map_openai_params( + non_default_params={"size": "1033x1024"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + def test_map_openai_params_explicit_width_height_not_range_checked(self): config = AzureFoundryMAIImageGenerationConfig() optional_params = config.map_openai_params( From 8e3052ff650460d8ba226381446c6a9b95e87cd8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:20:57 -0700 Subject: [PATCH 047/107] fix(azure_ai): honour global drop_params for image params MAI cannot serve get_optional_params_image_gen only forwarded the per-call drop_params flag to provider configs, so litellm_settings drop_params: true never dropped the n the MAI generations endpoint ignores. The MAI edits config also advertised and forwarded size, which that endpoint ignores. Non-numeric and non-positive n now surface as a 400 instead of a 500 or a pass-through. --- .../azure_ai/image_edit/mai_transformation.py | 62 +---------------- .../image_generation/mai_transformation.py | 10 ++- litellm/utils.py | 2 +- .../test_mai_image_edit_transformation.py | 67 ++++++++++--------- .../test_mai_image_generation.py | 38 ++++++++++- 5 files changed, 83 insertions(+), 96 deletions(-) diff --git a/litellm/llms/azure_ai/image_edit/mai_transformation.py b/litellm/llms/azure_ai/image_edit/mai_transformation.py index e639c20292b..55b179e9591 100644 --- a/litellm/llms/azure_ai/image_edit/mai_transformation.py +++ b/litellm/llms/azure_ai/image_edit/mai_transformation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final import httpx from httpx._types import RequestFiles @@ -13,7 +13,6 @@ from litellm.llms.azure_ai.image_generation.mai_transformation import ( from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str -from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.llms.openai import FileTypes from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ImageResponse @@ -26,65 +25,8 @@ if TYPE_CHECKING: class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig): """Azure AI Foundry MAI image editing (e.g. MAI-Image-2.5).""" - DEFAULT_SIZE = "1024x1024" - def get_supported_openai_params(self, model: str) -> list: - return ["prompt", "image", "model", "n", "size"] - - def map_openai_params( - self, - image_edit_optional_params: ImageEditOptionalRequestParams, - model: str, - drop_params: bool, - ) -> dict: - optional_params: Final[dict[str, Any]] = {} - supported_params: Final = self.get_supported_openai_params(model) - - for key, value in dict(image_edit_optional_params).items(): - if value is None or key in optional_params: - continue - - if key in supported_params: - if key == "size" and value: - size_param = cast(str, value) - self._validate_size_param(size_param) - optional_params[key] = size_param - else: - optional_params[key] = value - elif not drop_params: - raise ValueError( - f"Parameter {key} is not supported for model {model}. " - f"Supported parameters are {supported_params}. " - f"Set drop_params=True to drop unsupported parameters." - ) - - if "size" not in optional_params: - optional_params["size"] = self.DEFAULT_SIZE - - return optional_params - - def _validate_size_param(self, size: str) -> None: - known_sizes: Final = { - "1024x1024", - "1792x1024", - "1024x1792", - "512x512", - "256x256", - } - - if size in known_sizes: - return - - if "x" in size: - try: - tuple(map(int, size.lower().split("x", 1))) - return - except ValueError: - raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').") - - raise ValueError( - f"Unsupported size value: '{size}'. Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." - ) + return ["prompt", "image", "model", "n"] def validate_environment( self, diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 7f7a1044689..67b1a8bcab3 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -151,7 +151,7 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if k in supported_params: if k == "size" and v: self._map_size_param(v, optional_params, model) - elif k == "n" and v is not None and int(v) > self.MAX_IMAGES_PER_REQUEST: + elif k == "n" and v is not None and self._image_count(v, model) != self.MAX_IMAGES_PER_REQUEST: if not drop_params: raise self._unsupported( model, @@ -185,6 +185,14 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): def _unsupported(model: str, message: str) -> UnsupportedParamsError: return UnsupportedParamsError(message=message, llm_provider="azure_ai", model=model) + def _image_count(self, n: object, model: str) -> int: + if isinstance(n, int): + return n + try: + return int(str(n)) + except ValueError: + raise self._unsupported(model, f"n={n!r} is not a whole number of images for model {model}.") + def _map_size_param(self, size: str, optional_params: dict, model: str) -> None: size_mapping: Final = { "1024x1024": (1024, 1024), diff --git a/litellm/utils.py b/litellm/utils.py index bc2f4a86f12..63ad4bdd945 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3396,7 +3396,7 @@ def get_optional_params_image_gen( non_default_params=non_default_params, optional_params=optional_params, model=model or "", - drop_params=drop_params if drop_params is not None else False, + drop_params=litellm.drop_params is True or drop_params is True, ) elif ( custom_llm_provider == "openai" diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py index 284a912d9a4..75e046825a3 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py @@ -6,6 +6,7 @@ import pytest import litellm +from litellm.images.utils import ImageEditRequestUtils from litellm.llms.azure_ai.image_edit import ( AzureFoundryMAIImageEditConfig, get_azure_ai_image_edit_config, @@ -70,44 +71,48 @@ class TestAzureMAIImageEdit: assert "/mai/v1/images/edits" in url assert "api-version=preview" in url - def test_map_openai_params_keeps_size(self): - config = AzureFoundryMAIImageEditConfig() - optional_params = config.map_openai_params( - image_edit_optional_params={"size": "1792x1024", "n": 1}, + def test_get_optional_params_image_edit_size_raises_400(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError, match="size") as exc_info: + ImageEditRequestUtils.get_optional_params_image_edit( + model="MAI-Image-2.5", + image_edit_provider_config=AzureFoundryMAIImageEditConfig(), + image_edit_optional_params={"size": "1024x1024", "n": 1}, + ) + assert exc_info.value.status_code == 400 + + def test_get_optional_params_image_edit_size_dropped_with_drop_params(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + optional_params = ImageEditRequestUtils.get_optional_params_image_edit( model="MAI-Image-2.5", + image_edit_provider_config=AzureFoundryMAIImageEditConfig(), + image_edit_optional_params={"size": "1024x1024", "n": 1}, drop_params=True, ) - assert optional_params["size"] == "1792x1024" + assert "size" not in optional_params assert optional_params["n"] == 1 - assert "width" not in optional_params - assert "height" not in optional_params - def test_map_openai_params_defaults_size(self): - config = AzureFoundryMAIImageEditConfig() - optional_params = config.map_openai_params( - image_edit_optional_params={}, + def test_get_optional_params_image_edit_without_size_forwards_nothing_extra(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + optional_params = ImageEditRequestUtils.get_optional_params_image_edit( model="MAI-Image-2.5", - drop_params=True, + image_edit_provider_config=AzureFoundryMAIImageEditConfig(), + image_edit_optional_params={}, ) - assert optional_params["size"] == "1024x1024" + assert optional_params == {} - def test_map_openai_params_unsupported_size_raises(self): - config = AzureFoundryMAIImageEditConfig() - with pytest.raises(ValueError, match="Unsupported size value: 'auto'"): - config.map_openai_params( - image_edit_optional_params={"size": "auto"}, - model="MAI-Image-2.5", - drop_params=True, - ) - - def test_map_openai_params_invalid_size_format_raises(self): - config = AzureFoundryMAIImageEditConfig() - with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"): - config.map_openai_params( - image_edit_optional_params={"size": "1024xabc"}, - model="MAI-Image-2.5", - drop_params=True, + def test_image_edit_size_surfaces_as_400(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.image_edit( + model="azure_ai/MAI-Image-2.5", + image=io.BytesIO(b"fake-image-bytes"), + prompt="Turn this into a studio product shot", + size="1024x1024", + api_key="test-key", + api_base="https://my-resource.services.ai.azure.com", ) + assert exc_info.value.status_code == 400 def test_transform_image_edit_request_uses_image_field(self): config = AzureFoundryMAIImageEditConfig() @@ -117,14 +122,14 @@ class TestAzureMAIImageEdit: model="MAI-Image-2.5", prompt="Turn this into a studio product shot", image=image_bytes, - image_edit_optional_request_params={"size": "1024x1024", "n": 1}, + image_edit_optional_request_params={"n": 1}, litellm_params={}, headers={}, ) assert data["model"] == "MAI-Image-2.5" assert data["prompt"] == "Turn this into a studio product shot" - assert data["size"] == "1024x1024" + assert "size" not in data assert data["n"] == 1 assert len(files) == 1 assert files[0][0] == "image" diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py index a0bc3fb00b7..b8632a7ca5c 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -253,8 +253,8 @@ class TestAzureMAIImageGeneration: assert optional_params["width"] == 1792 assert optional_params["height"] == 1024 - @pytest.mark.parametrize("n", [2, 4, "2"]) - def test_map_openai_params_multi_image_n_raises(self, n): + @pytest.mark.parametrize("n", [2, 4, "2", 0, -1]) + def test_map_openai_params_n_other_than_one_raises(self, n): config = AzureFoundryMAIImageGenerationConfig() with pytest.raises(UnsupportedParamsError, match="returns exactly 1 image per request"): config.map_openai_params( @@ -264,6 +264,38 @@ class TestAzureMAIImageGeneration: drop_params=False, ) + def test_map_openai_params_non_numeric_n_raises_400(self): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(UnsupportedParamsError, match="not a whole number of images") as exc_info: + config.map_openai_params( + non_default_params={"n": "abc"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + assert exc_info.value.status_code == 400 + + def test_get_optional_params_image_gen_global_drop_params_drops_multi_image_n(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", True) + optional_params = get_optional_params_image_gen( + model="MAI-Image-2.5", + n=4, + custom_llm_provider="azure_ai", + provider_config=AzureFoundryMAIImageGenerationConfig(), + ) + assert "n" not in optional_params + assert optional_params["width"] == 1024 + + def test_get_optional_params_image_gen_without_any_drop_params_still_raises(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(UnsupportedParamsError, match="returns exactly 1 image per request"): + get_optional_params_image_gen( + model="MAI-Image-2.5", + n=4, + custom_llm_provider="azure_ai", + provider_config=AzureFoundryMAIImageGenerationConfig(), + ) + def test_map_openai_params_multi_image_n_dropped_with_drop_params(self): config = AzureFoundryMAIImageGenerationConfig() optional_params = config.map_openai_params( @@ -284,7 +316,7 @@ class TestAzureMAIImageGeneration: ) assert optional_params["n"] == 1 - @pytest.mark.parametrize("params", [{"n": 2}, {"size": "512x512"}, {"size": "1792x1024"}]) + @pytest.mark.parametrize("params", [{"n": 2}, {"n": "abc"}, {"size": "512x512"}, {"size": "1792x1024"}]) def test_image_generation_rejected_params_surface_as_400(self, params): with pytest.raises(litellm.BadRequestError) as exc_info: litellm.image_generation( From b5d40bef2b830b6d5ccbb3f3fae74eb10968b9cd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:25:04 -0700 Subject: [PATCH 048/107] fix(token_counter): bound offloaded token counts with a dedicated capacity limiter Replace the thread semaphore around HuggingFace encodes with an anyio CapacityLimiter applied at every offloaded count site through offload_token_count, so waiting counts no longer hold slots in the shared 40-thread pool and inline counts on the event loop never block on the bound. Rename TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES to TOKEN_COUNTER_MAX_CONCURRENT_COUNTS --- litellm/constants.py | 2 +- litellm/litellm_core_utils/token_counter.py | 21 ++++++++++----- .../hooks/parallel_request_limiter_v3.py | 4 +-- litellm/proxy/proxy_server.py | 6 ++--- litellm/proxy/utils.py | 4 +-- litellm/router.py | 5 ++-- .../complexity_router/complexity_router.py | 6 ++--- .../io_token_rate_limit_check.py | 4 +-- .../prompt_caching_deployment_check.py | 6 ++--- .../litellm_core_utils/test_token_counter.py | 26 +++++++++++++------ 10 files changed, 52 insertions(+), 32 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 05d442dd070..0bd7def539d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -397,7 +397,7 @@ TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range( minimum=1, maximum=1_000_000_000, ) -TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES: Final = 4 +TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = 4 MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512)) MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512)) OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 8624e79faea..e71ce5dea9c 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -3,13 +3,14 @@ import base64 import io import struct -import threading -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from typing import Final, Literal, cast +import anyio import httpx import tiktoken from tokenizers import Tokenizer +from typing_extensions import ParamSpec, TypeVar import litellm from litellm import verbose_logger @@ -23,9 +24,10 @@ from litellm.constants import ( MAX_TILE_HEIGHT, MAX_TILE_WIDTH, TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS, - TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES, + TOKEN_COUNTER_MAX_CONCURRENT_COUNTS, TOKEN_COUNTER_MAX_EXACT_CHARS, ) +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.default_encoding import encoding as default_encoding from litellm.litellm_core_utils.url_utils import safe_get from litellm.llms.custom_httpx.http_handler import _get_httpx_client @@ -322,7 +324,15 @@ Type for a function that counts tokens in a string. """ EXTRAPOLATION_SAMPLES: Final = 16 -_HF_ENCODE_SLOTS: Final = threading.BoundedSemaphore(TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES) +T_ParamSpec: Final = ParamSpec("T_ParamSpec") +T_Retval = TypeVar("T_Retval") +_COUNT_OFFLOAD_LIMITER: Final = anyio.CapacityLimiter(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) + + +def offload_token_count( + function: Callable[T_ParamSpec, T_Retval], +) -> Callable[T_ParamSpec, Awaitable[T_Retval]]: + return asyncify(function, limiter=_COUNT_OFFLOAD_LIMITER) def _get_tiktoken_count_function( @@ -590,8 +600,7 @@ def _get_exact_count_function( tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"] def count_tokens(text: str) -> int: - with _HF_ENCODE_SLOTS: - return len(tokenizer.encode_batch_fast([text])[0]) + return len(tokenizer.encode_batch_fast([text])[0]) return count_tokens elif tokenizer_json["type"] == "openai_tokenizer": diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 74266801e50..117be770953 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -28,10 +28,10 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( ESTIMATED_OUTPUT_TOKENS_FIELD, @@ -3304,7 +3304,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit=min_configured_otpm_limit, call_type=call_type, ) - raw_estimated_input_tokens: Final = await asyncify(self._estimate_precise_input_tokens)( + raw_estimated_input_tokens: Final = await offload_token_count(self._estimate_precise_input_tokens)( data=data, model=requested_model, call_type=call_type ) estimated_input_tokens: Final = max(raw_estimated_input_tokens, 1) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 77354e9885f..0e3e5c251c8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -67,6 +67,7 @@ from litellm.litellm_core_utils.litellm_logging import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import ( UI_TEAM_ID, CallbackDelete, @@ -274,7 +275,6 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) -from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, @@ -12713,7 +12713,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)( + _tokenizer_used: Final = await offload_token_count(litellm.utils._select_tokenizer)( model=model_to_use, custom_tokenizer=custom_tokenizer ) @@ -12728,7 +12728,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None ) - total_tokens: Final = await asyncify(litellm.token_counter)( + total_tokens: Final = await offload_token_count(litellm.token_counter)( model=model_to_use, text=prompt, messages=counted_messages, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1014c2b9b6a..59c09081890 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,7 +91,6 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert -from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, independent_snapshot, @@ -100,6 +99,7 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms import load_guardrail_translation_mappings from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( @@ -2570,7 +2570,7 @@ class ProxyLogging: original_exception=original_exception, ) - request_data.update(await asyncify(_failure_fields_to_lift)(request_data)) + request_data.update(await offload_token_count(_failure_fields_to_lift)(request_data)) # Remove before callbacks iterate — not serialisable request_data.pop("litellm_logging_obj", None) diff --git a/litellm/router.py b/litellm/router.py index 95cabfad4bd..73abd97d80d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -55,7 +55,7 @@ from litellm.constants import ( SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.asyncify import asyncify, run_async_function +from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, coerce_token_limit, @@ -86,6 +86,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_credentials_in_payload, mask_sensitive_structure, ) +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, vector_store_request_metadata, @@ -11958,7 +11959,7 @@ class Router: try: if not self._pre_call_checks_need_token_count(model, healthy_deployments): return None - return await asyncify(self._count_pre_call_check_tokens)( + return await offload_token_count(self._count_pre_call_check_tokens)( messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter request_kwargs=request_kwargs, diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 7d4497fb6f7..cc1bb54255b 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2486,14 +2486,14 @@ class ComplexityRouter(CustomLogger): """Real-tokenizer count of the resolved messages plus the out-of-band carriers, off the event loop; None when counting fails, and the gate then leaves the placement alone.""" import litellm - from litellm.litellm_core_utils.asyncify import asyncify + from litellm.litellm_core_utils.token_counter import offload_token_count out_of_band: Final = self._out_of_band_request_text(request_kwargs) try: - counted: Final = await asyncify(litellm.token_counter)( + counted: Final = await offload_token_count(litellm.token_counter)( messages=cast(list, resolved_messages) # cast-ok: token_counter only iterates the sequence ) - return counted + (await asyncify(litellm.token_counter)(text=out_of_band) if out_of_band else 0) + return counted + (await offload_token_count(litellm.token_counter)(text=out_of_band) if out_of_band else 0) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.debug("ComplexityRouter: context-window token count failed. Got - %s", e) return None diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index dca5a2be845..fbd3e18e357 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -21,7 +21,7 @@ import litellm from litellm import token_counter from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.types.router import RouterCacheEnum, RouterErrors from litellm.utils import get_utc_datetime @@ -467,7 +467,7 @@ async def async_io_token_pre_call_check( request_kwargs: Final = get_io_token_rate_limit_request_kwargs() _model: Final = (deployment.get("litellm_params") or {}).get("model") or "" - estimated_input: Final = await asyncify(_estimate_input_tokens)(request_kwargs, model=_model) + estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model) max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment) dt: Final = get_utc_datetime() diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 9ea528b9700..0589e290b47 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -14,7 +14,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, ) from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt @@ -62,7 +62,7 @@ class PromptCachingDeploymentCheck(CustomLogger): if request_kwargs is not None and request_kwargs.get("_target_order") is not None: return healthy_deployments - if messages is not None and await asyncify(is_prompt_caching_valid_prompt)( + if messages is not None and await offload_token_count(is_prompt_caching_valid_prompt)( messages=messages, model=model, min_token_count=_get_min_token_count_for_deployments(healthy_deployments), @@ -140,7 +140,7 @@ class PromptCachingDeploymentCheck(CustomLogger): return ## PROMPT CACHING - cache model id, if prompt caching valid prompt + provider - if await asyncify(is_prompt_caching_valid_prompt)( + if await offload_token_count(is_prompt_caching_valid_prompt)( model=model, messages=cast(list[AllMessageValues], messages), ): diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 655be5146cd..7cc7c6c763c 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -5,10 +5,10 @@ import importlib import threading import time import traceback -from concurrent.futures import ThreadPoolExecutor from typing import Final from unittest.mock import MagicMock +import anyio.to_thread import pytest import tiktoken @@ -18,12 +18,13 @@ import litellm from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants -from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES +from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.token_counter import ( _get_exact_count_function, _get_extrapolating_count_function, _get_tiktoken_count_function, + offload_token_count, ) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from tests.large_text import text @@ -181,16 +182,25 @@ class _SlowEncoder: return [[0] * len(text) for text in texts] -def test_huggingface_counts_run_at_most_the_configured_number_at_once(): +@pytest.mark.asyncio +async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool(): encoder: Final = _SlowEncoder() count: Final = _get_exact_count_function(None, {"type": "huggingface_tokenizer", "tokenizer": encoder}) - burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES + shared_pool: Final = anyio.to_thread.current_default_thread_limiter() + burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS - with ThreadPoolExecutor(max_workers=burst) as pool: - counts: Final = tuple(pool.map(count, ["abc"] * burst)) + async def shared_pool_borrowed_until_done(counting: asyncio.Future[list[int]]) -> tuple[int, ...]: + if counting.done(): + return () + await asyncio.sleep(0.01) + return (shared_pool.borrowed_tokens, *await shared_pool_borrowed_until_done(counting)) - assert counts == (3,) * burst - assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES + counting: Final = asyncio.ensure_future(asyncio.gather(*(offload_token_count(count)("abc") for _ in range(burst)))) + borrowed: Final = await shared_pool_borrowed_until_done(counting) + + assert await counting == [3] * burst + assert len(borrowed) > 1 and max(borrowed) == 0 + assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS def test_token_counter_applies_the_default_cap(): From 2bca7ff673f07812c544ce42a165dda931b8a730 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:27:14 -0700 Subject: [PATCH 049/107] fix(anthropic): replay OpenAI encrypted reasoning byte for byte behind /v1/messages Reasoning items the Responses API returned for a /v1/messages turn were rebuilt from their summary text on every replay, so the prompt the model saw changed between turns and the prompt cache never matched. The bridge now asks for reasoning.encrypted_content, carries it in the thinking signature (or as a redacted_thinking block when there is no summary), and replays it verbatim as the reasoning item's encrypted_content. Anthropic replay paths drop those tagged blocks so a cross-model resume never forwards OpenAI bytes to Anthropic --- .../transformation.py | 6 +- .../prompt_templates/common_utils.py | 80 +++++++-- .../prompt_templates/factory.py | 10 +- litellm/llms/anthropic/common_utils.py | 13 +- .../responses_adapters/streaming_iterator.py | 66 ++++++- .../responses_adapters/transformation.py | 84 +++++---- ...ore_utils_prompt_templates_common_utils.py | 64 +++++++ ...llm_core_utils_prompt_templates_factory.py | 12 +- ...t_responses_adapters_streaming_iterator.py | 74 +++++++- .../test_responses_adapters_transformation.py | 168 +++++++++++++++++- .../anthropic/test_anthropic_common_utils.py | 24 +++ 11 files changed, 535 insertions(+), 66 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 4fe069b0b7d..135f34afe46 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -25,7 +25,7 @@ import litellm from litellm import ModelResponse from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( - responses_reasoning_item_from_thinking_blocks, + responses_reasoning_items_from_thinking_blocks, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.bridges.completion_transformation import ( @@ -129,8 +129,8 @@ def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]: return stored raw_blocks: Final = msg.get("thinking_blocks") or () blocks: Final = cast("Iterable[ChatCompletionThinkingBlock]", raw_blocks) # cast-ok: untyped client json - from_thinking: Final = responses_reasoning_item_from_thinking_blocks(blocks) - return [] if from_thinking is None else [dict(from_thinking)] # mutable-ok: API message payload + replayed: Final = responses_reasoning_items_from_thinking_blocks(blocks) + return [dict(item) for item in replayed] # mutable-ok: API message payload def _build_reasoning_item( diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 00b80839dde..918cb982773 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1823,14 +1823,11 @@ def _extract_reasoning_content(message: dict) -> tuple[str | None, str | None]: return None, message_content -def _readable_thinking_text( - block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, -) -> str: +def _readable_thinking_text(block: Mapping[str, object]) -> str: """The text a chat model can read back, empty for redacted blocks and malformed ones.""" if block.get("type") != "thinking": return "" - thinking: Final = cast(ChatCompletionThinkingBlock, block).get("thinking") # cast-ok: narrowed by the type tag - return str(thinking or "") + return str(block.get("thinking") or "") def reasoning_content_from_thinking_blocks( @@ -1843,24 +1840,83 @@ def reasoning_content_from_thinking_blocks( return "\n".join(text for block in thinking_blocks if (text := _readable_thinking_text(block))) -def responses_reasoning_item_from_thinking_blocks( - thinking_blocks: Iterable[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], -) -> ChatCompletionReasoningItem | None: - """Build a Responses API `reasoning` input item from Anthropic thinking blocks. +ENCRYPTED_REASONING_SIGNATURE_PREFIX: Final = "litellm_encrypted_reasoning:" - The item carries no `id`: the Responses API rejects an empty one and 404s on any id it - did not mint itself, while an item without an id is always accepted. + +def encrypted_reasoning_signature(encrypted_content: str) -> str: + """The opaque value a Responses API reasoning item's `encrypted_content` travels in. + + Anthropic clients echo a thinking block's `signature` and a redacted block's `data` + back verbatim, so either field can carry the encrypted reasoning across turns; the + prefix tells the two apart from a signature Anthropic minted. """ + return f"{ENCRYPTED_REASONING_SIGNATURE_PREFIX}{encrypted_content}" + + +def encrypted_content_from_signature(signature: object) -> str | None: + if not isinstance(signature, str) or not signature.startswith(ENCRYPTED_REASONING_SIGNATURE_PREFIX): + return None + return signature.removeprefix(ENCRYPTED_REASONING_SIGNATURE_PREFIX) or None + + +def _encrypted_content_of_block(block: Mapping[str, object]) -> str | None: + match block.get("type"): + case "thinking": + return encrypted_content_from_signature(block.get("signature")) + case "redacted_thinking": + return encrypted_content_from_signature(block.get("data")) + case _: + return None + + +def is_encrypted_reasoning_block(block: object) -> bool: + """A thinking or redacted_thinking block carrying Responses API encrypted reasoning. + + Only the Responses API that minted the content can read it back, so an Anthropic + backend has to drop such a block rather than fail signature verification on it. + """ + if not isinstance(block, Mapping): + return False + mapping: Final = cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance + return _encrypted_content_of_block(mapping) is not None + + +def _reasoning_replay_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str: + index, block = indexed_block + return f"encrypted:{index}" if is_encrypted_reasoning_block(block) else "summary" + + +def _reasoning_item_from_block_group(group: tuple[Mapping[str, object], ...]) -> ChatCompletionReasoningItem | None: summary: Final[list[ChatCompletionReasoningSummaryTextBlock]] = [ # mutable-ok: API message payload ChatCompletionReasoningSummaryTextBlock(type="summary_text", text=text) - for block in thinking_blocks + for block in group if (text := _readable_thinking_text(block)) ] + encrypted_content: Final = _encrypted_content_of_block(group[0]) + if encrypted_content is not None: + return ChatCompletionReasoningItem(type="reasoning", summary=summary, encrypted_content=encrypted_content) if not summary: return None return ChatCompletionReasoningItem(type="reasoning", summary=summary) +def responses_reasoning_items_from_thinking_blocks( + thinking_blocks: Iterable[Mapping[str, object]], +) -> tuple[ChatCompletionReasoningItem, ...]: + """Build Responses API `reasoning` input items from Anthropic thinking blocks. + + A block carrying encrypted reasoning replays the item it came from byte for byte; + a run of plain thinking blocks collapses into one summary-only item. No item carries + an `id`: the Responses API 404s on any id it did not mint itself and rejects an empty + one, while an item without an id is always accepted. + """ + return tuple( + item + for _, group in groupby(enumerate(thinking_blocks), key=_reasoning_replay_group_key) + if (item := _reasoning_item_from_block_group(tuple(block for _, block in group))) is not None + ) + + def _parse_content_for_reasoning( message_text: str | None, ) -> tuple[str | None, str | None]: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 56c1d605700..ece619e3883 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -46,6 +46,7 @@ from litellm.types.utils import GenericImageParsingChunk from .common_utils import ( convert_content_list_to_str, infer_content_type_from_url_and_content, + is_encrypted_reasoning_block, is_non_content_values_set, parse_tool_call_arguments, ) @@ -2299,13 +2300,16 @@ def sanitize_messages_for_tool_calling( def _is_unsignable_thinking_block(block: object) -> bool: - """A `thinking` block that Anthropic cannot accept on input. + """A thinking block that Anthropic cannot accept on input. Anthropic verifies the thinking signature cryptographically, so a block whose signature is null, empty, or missing (e.g. from an open-source reasoning model) - is rejected with a 400 and must be dropped rather than blanked or repaired. - `redacted_thinking` blocks carry no signature and are always kept. + is rejected with a 400 and must be dropped rather than blanked or repaired, and + so is a block whose signature or data carries another provider's encrypted + reasoning. A `redacted_thinking` block Anthropic minted is always kept. """ + if is_encrypted_reasoning_block(block): + return True if not isinstance(block, dict) or block.get("type") != "thinking": return False signature: Final = block.get("signature") diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 2b57883cc13..f46227239d8 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -21,6 +21,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, + is_encrypted_reasoning_block, ) from litellm.litellm_core_utils.prompt_templates.factory import ( THOUGHT_SIGNATURE_SEPARATOR, @@ -1235,8 +1236,10 @@ def strip_empty_content_blocks_from_anthropic_messages( on the unified ``/v1/messages`` path. ``/v1/chat/completions`` already handles this in ``anthropic_messages_pt``; this helper provides the equivalent guarantee for the native Anthropic Messages path. - ``redacted_thinking`` blocks are never touched: they carry opaque - ``data`` instead of thinking text. + A thinking or ``redacted_thinking`` block whose signature or data carries + another provider's encrypted reasoning (a turn served by the Responses API + bridge) is dropped too, since Anthropic cannot verify it; every other + ``redacted_thinking`` block is left alone. Messages whose content is a list and becomes empty after stripping are omitted, matching :func:`strip_thinking_blocks_from_anthropic_messages`. @@ -1249,7 +1252,11 @@ def strip_empty_content_blocks_from_anthropic_messages( out.append(m) continue content = m["content"] - filtered = [b for b in content if not _is_empty_text_block(b) and not is_empty_thinking_block(b)] + filtered = [ # mutable-ok: rebuilt message content list + b + for b in content + if not _is_empty_text_block(b) and not is_empty_thinking_block(b) and not is_encrypted_reasoning_block(b) + ] if len(filtered) == len(content): out.append(m) elif filtered: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 2e0a6a9df8f..f753e87fee3 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -9,13 +9,19 @@ from typing import TYPE_CHECKING, Any, Final from litellm import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, +) from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( refusal_stop_details, responses_output_refusal_text, ) from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage -from .transformation import LiteLLMAnthropicToResponsesAPIAdapter +from .transformation import ( + REASONING_SUMMARY_PART_SEPARATOR, + LiteLLMAnthropicToResponsesAPIAdapter, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject @@ -29,9 +35,10 @@ class AnthropicResponsesStreamWrapper: response.created -> message_start response.output_item.added -> content_block_start (if message/function_call) response.output_text.delta -> content_block_delta (text_delta) + response.reasoning_summary_part.added -> content_block_delta (thinking_delta separator) response.reasoning_summary_text.delta -> content_block_delta (thinking_delta) response.function_call_arguments.delta -> content_block_delta (input_json_delta) - response.output_item.done -> content_block_stop + response.output_item.done -> content_block_delta (signature_delta) + content_block_stop response.completed -> message_delta + message_stop """ @@ -94,6 +101,38 @@ class AnthropicResponsesStreamWrapper: ) return block_idx + @staticmethod + def _field(source: object, name: str) -> object: + return source.get(name) if isinstance(source, dict) else getattr(source, name, None) + + def _close_reasoning_item(self, item: object, item_id: str | None) -> None: + block_idx: Final = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index + encrypted_content: Final = self._field(item, "encrypted_content") + signature: Final = ( + encrypted_reasoning_signature(encrypted_content) + if isinstance(encrypted_content, str) and encrypted_content + else None + ) + if block_idx < 0 and signature is None: + return + if block_idx < 0: + redacted_idx: Final = self._open_block( + item_id, + {"type": "redacted_thinking", "data": signature}, # mutable-ok: API message payload + ) + stop: Final = {"type": "content_block_stop", "index": redacted_idx} # mutable-ok: API message payload + self._chunk_queue.append(stop) + return + if signature is not None: + self._chunk_queue.append( + { # mutable-ok: API message payload + "type": "content_block_delta", + "index": block_idx, + "delta": {"type": "signature_delta", "signature": signature}, # mutable-ok: API message payload + } + ) + self._chunk_queue.append({"type": "content_block_stop", "index": block_idx}) # mutable-ok: API message payload + def _process_event(self, event: object) -> None: """Convert one Responses API event into zero or more Anthropic chunks queued for emission.""" event_type = getattr(event, "type", None) @@ -175,6 +214,26 @@ class AnthropicResponsesStreamWrapper: ) return + if event_type == "response.reasoning_summary_part.added": + part_item_id: Final = self._field(event, "item_id") + summary_index: Final = self._field(event, "summary_index") + part_block_idx: Final = ( + self._item_id_to_block_index.get(part_item_id, -1) if isinstance(part_item_id, str) else -1 + ) + if part_block_idx < 0 or not isinstance(summary_index, int) or summary_index == 0: + return + self._chunk_queue.append( + { # mutable-ok: API message payload + "type": "content_block_delta", + "index": part_block_idx, + "delta": { # mutable-ok: API message payload + "type": "thinking_delta", + "thinking": REASONING_SUMMARY_PART_SEPARATOR, + }, + } + ) + return + # ---- reasoning summary text delta ---- if event_type == "response.reasoning_summary_text.delta": item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) @@ -220,6 +279,9 @@ class AnthropicResponsesStreamWrapper: item_id = ( getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item else None ) + if self._field(item, "type") == "reasoning": + self._close_reasoning_item(item, item_id) + return block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index if block_idx < 0: return diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index f1daf2be42a..07419db443c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -13,7 +13,8 @@ from typing import Any, Final, cast from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, - responses_reasoning_item_from_thinking_blocks, + encrypted_reasoning_signature, + responses_reasoning_items_from_thinking_blocks, with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.reasoning_effort_utils import ( @@ -33,6 +34,7 @@ from litellm.types.llms.anthropic import ( AnthropicFinishReason, AnthropicMessagesRequest, AnthropicMessagesToolChoice, + AnthropicResponseContentBlockRedactedThinking, AnthropicResponseContentBlockText, AnthropicResponseContentBlockThinking, AnthropicResponseContentBlockToolUse, @@ -43,11 +45,13 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicUsage, ) from litellm.types.llms.openai import ( - ChatCompletionThinkingBlock, ResponseAPIUsage, ResponsesAPIResponse, ) +REASONING_SUMMARY_PART_SEPARATOR: Final = "\n\n" +RESPONSES_INCLUDE_ENCRYPTED_REASONING: Final = "reasoning.encrypted_content" + class LiteLLMAnthropicToResponsesAPIAdapter: """ @@ -163,49 +167,55 @@ class LiteLLMAnthropicToResponsesAPIAdapter: return str(getattr(part, "text", None) or "") @classmethod - def _thinking_blocks_from_reasoning_item( + def _thinking_block_from_reasoning_item( cls, summary: Iterable[object], - ) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload - """Anthropic thinking blocks for one Responses reasoning item. + encrypted_content: object, + ) -> dict[str, Any] | None: # mutable-ok: API message payload + """The one Anthropic block for a Responses reasoning item. - The signature stays empty: only Anthropic can sign a thinking block, and a stand-in - value would be replayed as a real one and rejected by every backend that verifies it. + The item's encrypted reasoning rides the block's opaque field (`signature`, or + `data` when there is no summary text) so the client echoes it back and the next + turn replays the very item OpenAI produced; without it the signature stays empty, + since only Anthropic can sign a thinking block. """ - return tuple( - AnthropicResponseContentBlockThinking( - type="thinking", - thinking=text, - signature=None, - ).model_dump() - for part in summary - if (text := cls._summary_part_text(part)) + text: Final = REASONING_SUMMARY_PART_SEPARATOR.join( + part_text for part in summary if (part_text := cls._summary_part_text(part)) ) + if not isinstance(encrypted_content, str) or not encrypted_content: + if not text: + return None + return AnthropicResponseContentBlockThinking(type="thinking", thinking=text, signature=None).model_dump() + signature: Final = encrypted_reasoning_signature(encrypted_content) + if not text: + return AnthropicResponseContentBlockRedactedThinking(type="redacted_thinking", data=signature).model_dump() + return AnthropicResponseContentBlockThinking(type="thinking", thinking=text, signature=signature).model_dump() @staticmethod def _assistant_block_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str: """Group a run of consecutive thinking blocks together; keep every other block alone.""" index, block = indexed_block - return "thinking" if block.get("type") == "thinking" else f"block:{index}" + return "thinking" if block.get("type") in ("thinking", "redacted_thinking") else f"block:{index}" @classmethod - def _assistant_group_to_input_item( + def _assistant_group_to_input_items( cls, group: tuple[Mapping[str, object], ...] - ) -> dict[str, Any] | None: # mutable-ok: API message payload + ) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload first: Final = group[0] btype: Final = first.get("type") - if btype == "thinking": - blocks: Final = cast(tuple[ChatCompletionThinkingBlock, ...], group) # cast-ok: untrusted client payload - reasoning_item: Final = responses_reasoning_item_from_thinking_blocks(blocks) - return None if reasoning_item is None else dict(reasoning_item) # mutable-ok: API message payload + if btype in ("thinking", "redacted_thinking"): + replayed: Final = responses_reasoning_items_from_thinking_blocks(group) + return tuple(dict(item) for item in replayed) # mutable-ok: API message payload if btype == "tool_use": - return { # mutable-ok: API message payload - "type": "function_call", - "call_id": first.get("id", ""), - "name": first.get("name", ""), - "arguments": json.dumps(first.get("input", {})), # mutable-ok: API message payload - } - return None + return ( + { # mutable-ok: API message payload + "type": "function_call", + "call_id": first.get("id", ""), + "name": first.get("name", ""), + "arguments": json.dumps(first.get("input", {})), # mutable-ok: API message payload + }, + ) + return () def translate_messages_to_responses_input( self, @@ -362,7 +372,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: input_items.extend( item for _, group in groupby(enumerate(blocks), key=self._assistant_block_group_key) - if (item := self._assistant_group_to_input_item(tuple(block for _, block in group))) is not None + for item in self._assistant_group_to_input_items(tuple(block for _, block in group)) ) asst_parts: list[dict[str, Any]] = [ # mutable-ok: API message payload {"type": "output_text", "text": block.get("text", "")} # mutable-ok: API message payload @@ -572,6 +582,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: ) if reasoning: responses_kwargs["reasoning"] = reasoning + responses_kwargs["include"] = [RESPONSES_INCLUDE_ENCRYPTED_REASONING] # mutable-ok: json list # output_format / output_config.format -> text format # output_format: {"type": "json_schema", "schema": {...}} @@ -634,7 +645,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: for item in response.output: if isinstance(item, ResponseReasoningItem): - content.extend(self._thinking_blocks_from_reasoning_item(item.summary)) + reasoning_block = self._thinking_block_from_reasoning_item(item.summary, item.encrypted_content) + if reasoning_block is not None: + content.append(reasoning_block) elif isinstance(item, ResponseOutputMessage): for part in item.content: @@ -684,11 +697,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter: ).model_dump() ) elif item_type == "reasoning": - content.extend( - self._thinking_blocks_from_reasoning_item( - cast(Iterable[object], item.get("summary") or ()), # cast-ok: untyped provider json - ) + reasoning_block = self._thinking_block_from_reasoning_item( + cast(Iterable[object], item.get("summary") or ()), # cast-ok: untyped provider json + item.get("encrypted_content"), ) + if reasoning_block is not None: + content.append(reasoning_block) elif item_type == "function_call": try: input_data = json.loads(item.get("arguments", "{}")) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index c037f928593..9df124576a0 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -10,10 +10,14 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, add_system_prompt_to_messages, + encrypted_content_from_signature, + encrypted_reasoning_signature, get_file_ids_from_messages, get_format_from_file_id, handle_any_messages_to_chat_completion_str_messages_conversion, hoist_images_from_tool_messages, + is_encrypted_reasoning_block, + responses_reasoning_items_from_thinking_blocks, split_concatenated_json_objects, update_messages_with_model_file_ids, ) @@ -1554,3 +1558,63 @@ class TestRequestContainsImageContent: for _ in range(50): nested = {"type": "tool_result", "content": [nested]} assert request_contains_image_content([{"role": "user", "content": [nested]}]) is False + + +class TestEncryptedReasoningReplay: + """Regression for https://github.com/BerriAI/litellm/issues/40288.""" + + def test_signature_round_trips_the_encrypted_content(self): + assert encrypted_content_from_signature(encrypted_reasoning_signature("gAAAA_bytes")) == "gAAAA_bytes" + + @pytest.mark.parametrize("signature", [None, "", "ErcBCkgIValidAnthropicSignature", "litellm_encrypted_reasoning:", 7]) + def test_anything_else_is_not_encrypted_content(self, signature): + assert encrypted_content_from_signature(signature) is None + + def test_encrypted_thinking_block_replays_its_own_item(self): + items = responses_reasoning_items_from_thinking_blocks( + [{"type": "thinking", "thinking": "Plan.", "signature": encrypted_reasoning_signature("gAAAA_1")}] + ) + assert items == ( + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "Plan."}], "encrypted_content": "gAAAA_1"}, + ) + + def test_encrypted_redacted_block_replays_with_an_empty_summary(self): + items = responses_reasoning_items_from_thinking_blocks( + [{"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_1")}] + ) + assert items == ({"type": "reasoning", "summary": [], "encrypted_content": "gAAAA_1"},) + + def test_plain_blocks_collapse_into_one_summary_item_around_encrypted_ones(self): + items = responses_reasoning_items_from_thinking_blocks( + [ + {"type": "thinking", "thinking": "A.", "signature": None}, + {"type": "thinking", "thinking": "B.", "signature": ""}, + {"type": "thinking", "thinking": "C.", "signature": encrypted_reasoning_signature("gAAAA_c")}, + {"type": "redacted_thinking", "data": "anthropic-minted-opaque-data"}, + {"type": "thinking", "thinking": "D."}, + ] + ) + assert items == ( + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "A."}, {"type": "summary_text", "text": "B."}]}, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "C."}], "encrypted_content": "gAAAA_c"}, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "D."}]}, + ) + assert all("id" not in item for item in items) + + def test_blocks_without_text_or_encrypted_content_produce_nothing(self): + assert responses_reasoning_items_from_thinking_blocks([{"type": "thinking", "thinking": ""}]) == () + assert responses_reasoning_items_from_thinking_blocks([]) == () + + @pytest.mark.parametrize( + ("block", "expected"), + [ + ({"type": "thinking", "thinking": "x", "signature": encrypted_reasoning_signature("g")}, True), + ({"type": "redacted_thinking", "data": encrypted_reasoning_signature("g")}, True), + ({"type": "thinking", "thinking": "x", "signature": "ErcBCkgIValid"}, False), + ({"type": "redacted_thinking", "data": "EmwKAhgBEgy"}, False), + ({"type": "text", "text": encrypted_reasoning_signature("g")}, False), + ("not a block", False), + ], + ) + def test_is_encrypted_reasoning_block(self, block, expected): + assert is_encrypted_reasoning_block(block) is expected diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index dd2d45f00c6..66d10fd1407 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -191,8 +191,16 @@ def test_bedrock_converse_assistant_with_empty_thinking_block_and_tool_calls(): {"type": "thinking", "thinking": "oss reasoning", "signature": None}, {"type": "thinking", "thinking": "oss reasoning", "signature": ""}, {"type": "thinking", "thinking": "oss reasoning"}, + {"type": "thinking", "thinking": "openai reasoning", "signature": "litellm_encrypted_reasoning:gAAAA"}, + {"type": "redacted_thinking", "data": "litellm_encrypted_reasoning:gAAAA"}, + ], + ids=[ + "null_signature", + "empty_signature", + "missing_signature", + "encrypted_reasoning_signature", + "encrypted_reasoning_redacted_data", ], - ids=["null_signature", "empty_signature", "missing_signature"], ) def test_anthropic_messages_pt_drops_unsignable_thinking_block(thinking_block): """Open-source reasoning models (DeepSeek-R1, Qwen, etc.) emit thinking blocks @@ -219,7 +227,7 @@ def test_anthropic_messages_pt_drops_unsignable_thinking_block(thinking_block): assistant = next(m for m in result if m["role"] == "assistant") content = assistant["content"] assert all( - block.get("type") != "thinking" for block in content + block.get("type") not in ("thinking", "redacted_thinking") for block in content ), f"unsignable thinking block must be dropped, got {content!r}" assert any( block.get("type") == "text" and block.get("text") == "2+2 equals 4." diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index d9df5df426a..bfe2d6b7cea 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -10,6 +10,9 @@ from types import SimpleNamespace sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, +) from litellm.llms.anthropic.experimental_pass_through.responses_adapters.streaming_iterator import ( AnthropicResponsesStreamWrapper, ) @@ -114,7 +117,7 @@ class TestReasoningItemWithoutSummaryText: """ @staticmethod - def _gpt_turn(reasoning_summary_deltas: list) -> list: + def _gpt_turn(reasoning_summary_deltas: list, encrypted_content: str | None = None) -> list: return [ {"type": "response.created"}, {"type": "response.output_item.added", "item": {"type": "reasoning", "id": "rs_1"}}, @@ -122,7 +125,10 @@ class TestReasoningItemWithoutSummaryText: {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": delta} for delta in reasoning_summary_deltas ), - {"type": "response.output_item.done", "item": {"type": "reasoning", "id": "rs_1"}}, + { + "type": "response.output_item.done", + "item": {"type": "reasoning", "id": "rs_1", "encrypted_content": encrypted_content}, + }, {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}, {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hello"}, {"type": "response.output_item.done", "item": {"type": "message", "id": "msg_1"}}, @@ -171,6 +177,70 @@ class TestReasoningItemWithoutSummaryText: assert not [c for c in chunks if c.get("delta", {}).get("type") == "signature_delta"] +_ENCRYPTED_REASONING = "gAAAAABp_encrypted_reasoning_bytes_only_openai_can_read" + + +class TestEncryptedReasoningIsStreamedForReplay: + """Regression for https://github.com/BerriAI/litellm/issues/40288. + + The client echoes a thinking block's signature (or a redacted block's data) back on the + next turn, so the item's ``encrypted_content`` has to reach it through one of those. + """ + + def test_encrypted_content_is_streamed_as_the_signature_before_the_block_closes(self): + chunks = _drain_async( + TestReasoningItemWithoutSummaryText._gpt_turn( + reasoning_summary_deltas=["Weighing options"], encrypted_content=_ENCRYPTED_REASONING + ) + ) + + assert [(c["type"], c.get("index"), c.get("delta", {}).get("type")) for c in chunks[1:5]] == [ + ("content_block_start", 0, None), + ("content_block_delta", 0, "thinking_delta"), + ("content_block_delta", 0, "signature_delta"), + ("content_block_stop", 0, None), + ] + assert chunks[3]["delta"]["signature"] == encrypted_reasoning_signature(_ENCRYPTED_REASONING) + + def test_reasoning_without_summary_streams_a_redacted_thinking_block(self): + chunks = _drain_async( + TestReasoningItemWithoutSummaryText._gpt_turn( + reasoning_summary_deltas=[], encrypted_content=_ENCRYPTED_REASONING + ) + ) + + assert [(c["type"], c.get("index")) for c in chunks[1:]] == [ + ("content_block_start", 0), + ("content_block_stop", 0), + ("content_block_start", 1), + ("content_block_delta", 1), + ("content_block_stop", 1), + ] + assert chunks[1]["content_block"] == { + "type": "redacted_thinking", + "data": encrypted_reasoning_signature(_ENCRYPTED_REASONING), + } + + def test_summary_parts_are_separated_inside_the_one_thinking_block(self): + """Two summary parts read as two paragraphs, not as one run-on sentence.""" + events = [ + {"type": "response.created"}, + {"type": "response.output_item.added", "item": {"type": "reasoning", "id": "rs_1"}}, + {"type": "response.reasoning_summary_part.added", "item_id": "rs_1", "summary_index": 0}, + {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": "First."}, + {"type": "response.reasoning_summary_part.added", "item_id": "rs_1", "summary_index": 1}, + {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": "Second."}, + {"type": "response.output_item.done", "item": {"type": "reasoning", "id": "rs_1"}}, + ] + chunks = _process_all(events) + + thinking = "".join( + c["delta"]["thinking"] for c in chunks if c.get("delta", {}).get("type") == "thinking_delta" + ) + assert thinking == "First.\n\nSecond." + assert [c["type"] for c in chunks].count("content_block_start") == 1 + + class TestToolUseBlockClosedExactlyOnce: """Regression for https://github.com/BerriAI/litellm/issues/37273. 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 9f8414afa38..79eaf6b60f2 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 @@ -19,6 +19,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, + encrypted_reasoning_signature, ) from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, @@ -566,6 +567,66 @@ class TestTranslateMessagesToResponsesInput: result = _translate_messages(messages) assert "id" not in result[0] + def test_thinking_block_with_encrypted_signature_replays_the_encrypted_content(self): + """Regression for https://github.com/BerriAI/litellm/issues/40288 (inbound fault site).""" + messages = [ + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "Private reasoning.", + "signature": encrypted_reasoning_signature("gAAAA_turn_one"), + } + ], + } + ] + result = _translate_messages(messages) + assert result == [ + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Private reasoning."}], + "encrypted_content": "gAAAA_turn_one", + } + ] + + def test_redacted_thinking_with_encrypted_data_replays_the_encrypted_content(self): + messages = [ + { + "role": "assistant", + "content": [{"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_turn_one")}], + } + ] + result = _translate_messages(messages) + assert result == [{"type": "reasoning", "summary": [], "encrypted_content": "gAAAA_turn_one"}] + + def test_each_encrypted_thinking_block_stays_its_own_reasoning_item(self): + """Two upstream items must not be merged into one, or the encrypted content of one is lost.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "First.", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "thinking", "thinking": "Second.", "signature": encrypted_reasoning_signature("gAAAA_2")}, + ], + } + ] + result = _translate_messages(messages) + assert [item["encrypted_content"] for item in result] == ["gAAAA_1", "gAAAA_2"] + + def test_anthropic_signed_thinking_block_replays_as_a_summary_only_item(self): + """A real Anthropic signature is opaque here, so it never masquerades as encrypted content.""" + messages = [ + { + "role": "assistant", + "content": [{"type": "thinking", "thinking": "Private reasoning.", "signature": "ErcBCkgIValid"}], + } + ] + result = _translate_messages(messages) + assert result == [ + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "Private reasoning."}]} + ] + def test_consecutive_thinking_blocks_become_one_reasoning_item(self): """Summary parts of one upstream reasoning item are regrouped into that item.""" messages = [ @@ -1101,6 +1162,13 @@ class TestTranslateRequestBroaderCoverage: req = _make_request(thinking={"type": "disabled"}) kwargs = _ADAPTER.translate_request(req) assert "reasoning" not in kwargs + assert "include" not in kwargs + + def test_thinking_asks_for_the_encrypted_reasoning(self): + """The documented way to get reasoning that survives store=false is to ask for it.""" + req = _make_request(thinking={"type": "enabled", "budget_tokens": 12000}) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["include"] == ["reasoning.encrypted_content"] def test_metadata_user_id_mapped_to_user(self): req = _make_request(metadata={"user_id": "user-42"}) @@ -1229,7 +1297,9 @@ def _make_function_call_item(call_id: str, name: str, arguments: str) -> MagicMo return item -def _make_reasoning_item(summaries: List[str], item_id: str = "rs_test_1") -> MagicMock: +def _make_reasoning_item( + summaries: List[str], item_id: str = "rs_test_1", encrypted_content: str | None = None +) -> MagicMock: """Build a mock ResponseReasoningItem.""" from openai.types.responses import ResponseReasoningItem # type: ignore[import] @@ -1242,9 +1312,13 @@ def _make_reasoning_item(summaries: List[str], item_id: str = "rs_test_1") -> Ma item = MagicMock(spec=ResponseReasoningItem) item.id = item_id item.summary = summary_mocks + item.encrypted_content = encrypted_content return item +_ENCRYPTED_REASONING = "gAAAAABp_encrypted_reasoning_bytes_only_openai_can_read" + + class TestTranslateResponse: """Responses API -> AnthropicMessagesResponse conversion.""" @@ -1369,7 +1443,81 @@ class TestTranslateResponse: reasoning = _make_reasoning_item(["Part one.", "Part two."], item_id="rs_abc123") response = _make_mock_response(output=[reasoning]) result: Any = _ADAPTER.translate_response(response) - assert [block["signature"] for block in result["content"]] == [None, None] + assert [block["signature"] for block in result["content"]] == [None] + assert "rs_abc123" not in json.dumps(result["content"]) + + def test_summary_parts_join_into_one_thinking_block(self): + """One reasoning item is one block, so its signature is echoed back exactly once.""" + reasoning = _make_reasoning_item(["Part one.", "Part two."]) + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert [block["thinking"] for block in result["content"]] == ["Part one.\n\nPart two."] + + def test_encrypted_content_rides_the_thinking_signature(self): + """Regression for https://github.com/BerriAI/litellm/issues/40288 (outbound fault site).""" + reasoning = _make_reasoning_item(["Part one."], item_id="rs_abc123", encrypted_content=_ENCRYPTED_REASONING) + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert result["content"] == [ + { + "type": "thinking", + "thinking": "Part one.", + "signature": encrypted_reasoning_signature(_ENCRYPTED_REASONING), + } + ] + + def test_reasoning_without_summary_becomes_redacted_thinking(self): + """With summaries off the encrypted reasoning still has to reach the client to be replayed.""" + reasoning = _make_reasoning_item([], encrypted_content=_ENCRYPTED_REASONING) + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert result["content"] == [ + {"type": "redacted_thinking", "data": encrypted_reasoning_signature(_ENCRYPTED_REASONING)} + ] + + def test_dict_reasoning_item_carries_its_encrypted_content(self): + response = _make_mock_response( + output=[ + { + "type": "reasoning", + "id": "rs_dict_1", + "encrypted_content": _ENCRYPTED_REASONING, + "summary": [{"type": "summary_text", "text": "Weighing the options."}], + } + ] + ) + result: Any = _ADAPTER.translate_response(response) + assert result["content"][0]["signature"] == encrypted_reasoning_signature(_ENCRYPTED_REASONING) + + def test_reasoning_item_round_trip_is_byte_stable(self): + """Regression for https://github.com/BerriAI/litellm/issues/40288. + + The reasoning item the next turn replays must be the one OpenAI produced, with its + encrypted reasoning intact, and identical on every later turn so the prompt cache + prefix keeps matching. + """ + reasoning = _make_reasoning_item(["Part one.", "Part two."], encrypted_content=_ENCRYPTED_REASONING) + turn: Any = _ADAPTER.translate_response(_make_mock_response(output=[reasoning])) + history = [{"role": "assistant", "content": turn["content"]}] + + replayed_items = [_translate_messages(history) for _ in range(2)] + + assert replayed_items[0] == replayed_items[1] + assert replayed_items[0] == [ + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Part one.\n\nPart two."}], + "encrypted_content": _ENCRYPTED_REASONING, + } + ] + + def test_redacted_reasoning_round_trip_replays_the_encrypted_content(self): + reasoning = _make_reasoning_item([], encrypted_content=_ENCRYPTED_REASONING) + turn: Any = _ADAPTER.translate_response(_make_mock_response(output=[reasoning])) + + replayed = _translate_messages([{"role": "assistant", "content": turn["content"]}]) + + assert replayed == [{"type": "reasoning", "summary": [], "encrypted_content": _ENCRYPTED_REASONING}] def test_dict_reasoning_item_becomes_thinking_block(self): """A reasoning item arriving as a plain dict is kept, not dropped.""" @@ -1385,14 +1533,26 @@ class TestTranslateResponse: result: Any = _ADAPTER.translate_response(response) assert result["content"] == [{"type": "thinking", "thinking": "Weighing the options.", "signature": None}] - def test_thinking_blocks_are_dropped_when_replayed_to_anthropic(self): + @pytest.mark.parametrize( + ("summaries", "encrypted_content"), + [ + (["Part one."], None), + (["Part one."], _ENCRYPTED_REASONING), + ([], _ENCRYPTED_REASONING), + ], + ids=["unsigned_thinking", "encrypted_thinking", "encrypted_redacted_thinking"], + ) + def test_thinking_blocks_are_dropped_when_replayed_to_anthropic(self, summaries, encrypted_content): """Replaying this turn to an Anthropic model must not send a signature it cannot verify.""" from litellm.litellm_core_utils.prompt_templates.factory import ( _drop_unsignable_thinking_blocks, ) - response = _make_mock_response(output=[_make_reasoning_item(["Part one."], item_id="rs_abc123")]) + response = _make_mock_response( + output=[_make_reasoning_item(summaries, item_id="rs_abc123", encrypted_content=encrypted_content)] + ) result: Any = _ADAPTER.translate_response(response) + assert len(result["content"]) == 1 assert _drop_unsignable_thinking_blocks(result["content"]) == [] def test_usage_mapped_correctly(self): 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 ae620fdd6dc..336cbf57e29 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -1541,6 +1541,30 @@ class TestAnthropicThinkingSignatureSelfHeal: out = strip_empty_content_blocks_from_anthropic_messages(msgs) assert [b["type"] for b in out[0]["content"]] == ["thinking"] + def test_strip_drops_encrypted_reasoning_blocks_from_the_responses_bridge(self): + """A session resumed on an Anthropic model replays reasoning only OpenAI can verify.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, + ) + from litellm.llms.anthropic.common_utils import ( + strip_empty_content_blocks_from_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, + {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"}, + {"type": "text", "text": "The answer."}, + ], + } + ] + out = strip_empty_content_blocks_from_anthropic_messages(msgs) + assert [b["type"] for b in out[0]["content"]] == ["redacted_thinking", "text"] + assert len(msgs[0]["content"]) == 4 + def test_strip_empty_text_blocks_treats_null_text_as_empty(self): from litellm.llms.anthropic.common_utils import ( strip_empty_content_blocks_from_anthropic_messages, From dd707da81c8e7ec9611da406c569b61b4e5385bc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:28:58 -0700 Subject: [PATCH 050/107] fix(bedrock): forward raw reasoning_effort to DeepSeek V3, drop it only for R1 --- .../bedrock/chat/converse_transformation.py | 52 ++++++----- .../chat/test_converse_transformation.py | 89 +++++++++++++++---- 2 files changed, 103 insertions(+), 38 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 6fe69b012ca..a9f34cdd683 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -422,13 +422,13 @@ class AmazonConverseConfig(BaseConfig): """ Handle the reasoning_effort parameter based on the model type. - - GPT-OSS models: passed through unchanged via additionalModelRequestFields. + - GPT-OSS and DeepSeek V3 models: passed through unchanged via additionalModelRequestFields. - OpenAI GPT-5.x models: mapped to ``reasoning.effort`` via additionalModelRequestFields. - Nova 2 models: transformed to reasoningConfig. - Anthropic models: mapped to ``thinking`` (and ``output_config.effort`` on adaptive Claude 4.6 / 4.7). """ - if "gpt-oss" in model: + if "gpt-oss" in model or "deepseek" in model: optional_params["reasoning_effort"] = reasoning_effort elif "openai.gpt-5" in model: reasoning: Final[BedrockConverseGptReasoningEffortBlock] = {"effort": reasoning_effort} @@ -509,14 +509,19 @@ class AmazonConverseConfig(BaseConfig): ) thinking["budget_tokens"] = BEDROCK_MIN_THINKING_BUDGET_TOKENS - def _model_accepts_anthropic_thinking_param(self, model: str, base_model: str) -> bool: - """Whether the model accepts the Anthropic-shaped ``thinking`` / ``reasoning_effort`` request field. + def _is_deepseek_model(self, model: str, base_model: str) -> bool: + return "deepseek" in model or "deepseek" in base_model - The Converse mapping serializes ``thinking`` into ``additionalModelRequestFields`` in Anthropic's - shape. Only Anthropic Claude reasoning models accept that field; DeepSeek models reason natively - and reject it (a 400 when it leaks through), even though they advertise ``supports_reasoning``. + def _is_deepseek_r1_model(self, model: str, base_model: str) -> bool: + return "deepseek.r1" in model or "deepseek.r1" in base_model + + def _model_accepts_anthropic_thinking_param(self, model: str, base_model: str) -> bool: + """Whether the model accepts the Anthropic-shaped ``thinking`` request field. + + Only Claude reasoning models accept it. DeepSeek advertises ``supports_reasoning`` but reasons + natively: R1 returns a 400 when the field is sent and V3 silently ignores it. """ - if "deepseek" in model or "deepseek" in base_model: + if self._is_deepseek_model(model=model, base_model=base_model): return False return ( "claude-3-7" in model @@ -526,17 +531,13 @@ class AmazonConverseConfig(BaseConfig): or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider) ) - def _model_reasons_natively_and_rejects_request_param(self, model: str, base_model: str) -> bool: - """Whether the model reasons natively and rejects any reasoning request field on Converse. + def _model_rejects_reasoning_effort_param(self, model: str, base_model: str) -> bool: + """Whether the model returns a 400 for every ``reasoning_effort`` shape on Converse. - The Converse mapping serializes ``thinking`` into ``additionalModelRequestFields`` in Anthropic's - shape and ``reasoning_effort`` into a provider-specific shape. DeepSeek reasons on its own and - returns a 400 when either field is sent, even though it advertises ``supports_reasoning``, so both - must be dropped for it. Every other model either accepts one of those shapes (Claude ``thinking``, - gpt-oss / Nova 2 ``reasoning_effort``) or is an opaque ARN we can't introspect, so we leave those - untouched rather than silently degrading reasoning. + DeepSeek R1 always reasons and rejects any reasoning request field. DeepSeek V3 accepts a raw + ``reasoning_effort`` like gpt-oss does, and every other model maps it to a shape it accepts. """ - return "deepseek" in model or "deepseek" in base_model + return self._is_deepseek_r1_model(model=model, base_model=base_model) def get_supported_openai_params(self, model: str) -> list[str]: from litellm.utils import supports_function_calling @@ -595,6 +596,9 @@ class AmazonConverseConfig(BaseConfig): if "gpt-oss" in model or "openai.gpt-5" in model or "openai.gpt-5" in base_model: supported_params.append("reasoning_effort") + elif self._is_deepseek_model(model=model, base_model=base_model): + if not self._is_deepseek_r1_model(model=model, base_model=base_model): + supported_params.append("reasoning_effort") elif self._is_nova_2_model(model): # Nova 2 models support reasoning_effort (transformed to reasoningConfig) # These models use a different reasoning structure than Anthropic's thinking parameter @@ -891,8 +895,10 @@ class AmazonConverseConfig(BaseConfig): drop_params: bool, ) -> dict: is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) - drop_reasoning_request_param: Final = self._model_reasons_natively_and_rejects_request_param( - model=model, base_model=BedrockModelInfo.get_base_model(model) + base_model: Final = BedrockModelInfo.get_base_model(model) + drop_thinking_param: Final = self._is_deepseek_model(model=model, base_model=base_model) + drop_reasoning_effort_param: Final = self._model_rejects_reasoning_effort_param( + model=model, base_model=base_model ) for param, value in non_default_params.items(): @@ -942,9 +948,9 @@ class AmazonConverseConfig(BaseConfig): optional_params["_parallel_tool_use_config"] = { "tool_choice": {"type": "auto", "disable_parallel_tool_use": not value} } - if param == "thinking" and drop_reasoning_request_param: + if param == "thinking" and drop_thinking_param: verbose_logger.debug( - "Dropping unsupported `thinking` param for Bedrock model=%s; it reasons natively and rejects it.", + "Dropping unsupported `thinking` param for Bedrock model=%s; it reasons natively.", model, ) elif param == "thinking" and "openai.gpt-5" not in model: @@ -973,9 +979,9 @@ class AmazonConverseConfig(BaseConfig): AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( model=model, optional_params=optional_params, custom_llm_provider="bedrock" ) - elif param == "reasoning_effort" and isinstance(value, str) and drop_reasoning_request_param: + elif param == "reasoning_effort" and isinstance(value, str) and drop_reasoning_effort_param: verbose_logger.debug( - "Dropping unsupported `reasoning_effort` param for Bedrock model=%s; it reasons natively and rejects it.", + "Dropping unsupported `reasoning_effort` param for Bedrock model=%s; it always reasons and rejects it.", model, ) elif param == "reasoning_effort" and isinstance(value, str): diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 2fa32cd2037..e336cb8a14a 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -873,31 +873,55 @@ def test_get_supported_openai_params(): ], ) def test_bedrock_deepseek_does_not_advertise_thinking(model): - """DeepSeek reasons natively on Bedrock and rejects the Anthropic-shaped - `thinking`/`reasoning_effort` field, so it must not be advertised as supported - (otherwise it leaks into additionalModelRequestFields and Bedrock 400s).""" + """DeepSeek reasons natively on Bedrock and does not take the Anthropic-shaped `thinking` + field (R1 400s on it, V3 ignores it), so it must not be advertised as supported.""" config = AmazonConverseConfig() supported_params = config.get_supported_openai_params(model=model) assert "thinking" not in supported_params - assert "reasoning_effort" not in supported_params + assert "output_config" not in supported_params -def test_bedrock_deepseek_r1_thinking_raises_without_drop_params(): - """Passing `thinking` to Bedrock DeepSeek R1 must fail client-side with a clear - UnsupportedParamsError instead of leaking through and hitting a Bedrock 400.""" +@pytest.mark.parametrize("model", ["bedrock/us.deepseek.r1-v1:0", "bedrock/converse/us.deepseek.r1-v1:0"]) +def test_bedrock_deepseek_r1_does_not_advertise_reasoning_effort(model): + """DeepSeek R1 always reasons and returns a 400 for any reasoning_effort shape.""" + config = AmazonConverseConfig() + assert "reasoning_effort" not in config.get_supported_openai_params(model=model) + + +@pytest.mark.parametrize("model", ["bedrock/deepseek.v3-v1:0", "bedrock/deepseek.v3.2", "bedrock/us.deepseek.v3.2"]) +def test_bedrock_deepseek_v3_advertises_reasoning_effort(model): + """DeepSeek V3 on Bedrock accepts a raw reasoning_effort in additionalModelRequestFields.""" + config = AmazonConverseConfig() + assert "reasoning_effort" in config.get_supported_openai_params(model=model) + + +@pytest.mark.parametrize("model", ["us.deepseek.r1-v1:0", "deepseek.v3.2"]) +def test_bedrock_deepseek_thinking_raises_without_drop_params(model): + """Passing `thinking` to Bedrock DeepSeek must fail client-side with a clear + UnsupportedParamsError instead of leaking through to Bedrock.""" with pytest.raises(litellm.UnsupportedParamsError): litellm.utils.get_optional_params( - model="us.deepseek.r1-v1:0", + model=model, custom_llm_provider="bedrock", thinking={"type": "enabled", "budget_tokens": 1024}, ) -def test_bedrock_deepseek_r1_thinking_dropped_does_not_leak_into_request(): +def test_bedrock_deepseek_r1_reasoning_effort_raises_without_drop_params(): + with pytest.raises(litellm.UnsupportedParamsError): + litellm.utils.get_optional_params( + model="us.deepseek.r1-v1:0", + custom_llm_provider="bedrock", + reasoning_effort="high", + ) + + +@pytest.mark.parametrize("model", ["us.deepseek.r1-v1:0", "deepseek.v3.2"]) +def test_bedrock_deepseek_thinking_dropped_does_not_leak_into_request(model): """With drop_params, `thinking` is dropped rather than forwarded into - additionalModelRequestFields for Bedrock DeepSeek R1.""" + additionalModelRequestFields for Bedrock DeepSeek.""" optional_params = litellm.utils.get_optional_params( - model="us.deepseek.r1-v1:0", + model=model, custom_llm_provider="bedrock", thinking={"type": "enabled", "budget_tokens": 1024}, drop_params=True, @@ -906,20 +930,20 @@ def test_bedrock_deepseek_r1_thinking_dropped_does_not_leak_into_request(): config = AmazonConverseConfig() request = config._transform_request( - model="bedrock/converse/us.deepseek.r1-v1:0", + model=f"bedrock/converse/{model}", messages=[{"role": "user", "content": "Say hi in one word."}], optional_params=optional_params, litellm_params={}, headers={}, ) - assert "thinking" not in request.get("additionalModelRequestFields", {}) + assert "thinking" not in (request.get("additionalModelRequestFields") or {}) @pytest.mark.parametrize("param", ["thinking", "reasoning_effort"]) def test_bedrock_deepseek_r1_reasoning_params_not_forwarded_by_map(param): """Even when map_openai_params is called directly (bypassing the supported-params - gate), DeepSeek R1 must not forward the Anthropic-shaped thinking/reasoning_effort - into additionalModelRequestFields, since Bedrock rejects it with a 400.""" + gate), DeepSeek R1 must not forward thinking/reasoning_effort into + additionalModelRequestFields, since Bedrock rejects both with a 400.""" config = AmazonConverseConfig() model = "bedrock/converse/us.deepseek.r1-v1:0" value = {"type": "enabled", "budget_tokens": 1024} if param == "thinking" else "high" @@ -943,6 +967,41 @@ def test_bedrock_deepseek_r1_reasoning_params_not_forwarded_by_map(param): assert request.get("additionalModelRequestFields") is None +def test_bedrock_deepseek_v3_reasoning_effort_forwarded_raw(): + """DeepSeek V3 takes reasoning_effort verbatim in additionalModelRequestFields, never + converted into the Anthropic `thinking` block that Claude models get.""" + config = AmazonConverseConfig() + model = "bedrock/deepseek.v3.2" + optional_params = config.map_openai_params( + non_default_params={"reasoning_effort": "high", "max_tokens": 100}, + optional_params={}, + model=model, + drop_params=False, + ) + assert "thinking" not in optional_params + + request = config._transform_request( + model=model, + messages=[{"role": "user", "content": "Say hi in one word."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + assert request["additionalModelRequestFields"] == {"reasoning_effort": "high"} + + +def test_bedrock_deepseek_v3_thinking_dropped_by_map(): + config = AmazonConverseConfig() + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 1024}, "max_tokens": 100}, + optional_params={}, + model="bedrock/deepseek.v3.2", + drop_params=False, + ) + assert "thinking" not in optional_params + assert "reasoning_effort" not in optional_params + + @pytest.mark.parametrize( "model, param, value, kept_key", [ From 751431dc3f782b1aa93908533dbc9d9e0062b400 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:41:38 -0700 Subject: [PATCH 051/107] fix(anthropic): strip bridge reasoning in the native messages transform, not the empty-block pass --- litellm/llms/anthropic/common_utils.py | 36 ++++++++--- .../messages/transformation.py | 3 +- ..._anthropic_messages_encrypted_reasoning.py | 50 ++++++++++++++++ .../anthropic/test_anthropic_common_utils.py | 59 +++++++++++++++---- 4 files changed, 126 insertions(+), 22 deletions(-) create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index f46227239d8..cb24f4153f9 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1202,6 +1202,30 @@ def strip_thinking_blocks_from_anthropic_messages(messages: list[Any]) -> list[A return out +def _without_encrypted_reasoning_blocks(message: dict) -> dict | None: # mutable-ok: Anthropic message payload shape + content: Final = message.get("content") + if not isinstance(content, list): + return message + kept: Final = [b for b in content if not is_encrypted_reasoning_block(b)] # mutable-ok: API message payload + if len(kept) == len(content): + return message + if not kept: + return None + return {**message, "content": kept} # mutable-ok: API message payload + + +def strip_encrypted_reasoning_blocks_from_anthropic_messages( + messages: Sequence[dict], # mutable-ok: Anthropic message payload shape +) -> list[dict]: # mutable-ok: AnthropicMessagesRequest.messages is typed list[dict] + """ + Drop thinking / redacted_thinking blocks that carry another provider's encrypted + reasoning (a turn the Responses API bridge served) before the request reaches + Anthropic, which cannot verify them. Anthropic's own signed blocks are kept. + """ + stripped: Final = (_without_encrypted_reasoning_blocks(m) for m in messages) + return [m for m in stripped if m is not None] # mutable-ok: API message payload + + def strip_thinking_blocks_from_anthropic_messages_request_dict( data: dict[str, Any], ) -> None: @@ -1236,10 +1260,8 @@ def strip_empty_content_blocks_from_anthropic_messages( on the unified ``/v1/messages`` path. ``/v1/chat/completions`` already handles this in ``anthropic_messages_pt``; this helper provides the equivalent guarantee for the native Anthropic Messages path. - A thinking or ``redacted_thinking`` block whose signature or data carries - another provider's encrypted reasoning (a turn served by the Responses API - bridge) is dropped too, since Anthropic cannot verify it; every other - ``redacted_thinking`` block is left alone. + ``redacted_thinking`` blocks are never touched: they carry opaque + ``data`` instead of thinking text. Messages whose content is a list and becomes empty after stripping are omitted, matching :func:`strip_thinking_blocks_from_anthropic_messages`. @@ -1252,11 +1274,7 @@ def strip_empty_content_blocks_from_anthropic_messages( out.append(m) continue content = m["content"] - filtered = [ # mutable-ok: rebuilt message content list - b - for b in content - if not _is_empty_text_block(b) and not is_empty_thinking_block(b) and not is_encrypted_reasoning_block(b) - ] + filtered = [b for b in content if not _is_empty_text_block(b) and not is_empty_thinking_block(b)] if len(filtered) == len(content): out.append(m) elif filtered: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 8267da157ad..9f9346fad4d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -25,6 +25,7 @@ from ...common_utils import ( AnthropicModelInfo, optionally_handle_anthropic_oauth, strip_advisor_blocks_from_messages, + strip_encrypted_reasoning_blocks_from_anthropic_messages, ) DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" @@ -613,7 +614,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): messages = strip_advisor_blocks_from_messages(messages) anthropic_messages_request: Final[AnthropicMessagesRequest] = AnthropicMessagesRequest( - messages=messages, + messages=strip_encrypted_reasoning_blocks_from_anthropic_messages(messages), max_tokens=max_tokens, model=model, **anthropic_messages_optional_request_params, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py new file mode 100644 index 00000000000..c64e9d392e5 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py @@ -0,0 +1,50 @@ +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, +) +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, +) + + +def _transform(messages): + return AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-sonnet-4-5", + messages=messages, + anthropic_messages_optional_request_params={"max_tokens": 1024}, + litellm_params={}, + headers={}, + ) + + +def test_reasoning_replayed_from_the_responses_bridge_never_reaches_anthropic(): + """Claude Code resumed on a Claude model echoes the thinking blocks a gpt turn produced.""" + messages = [ + {"role": "user", "content": "Solve it."}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, + {"type": "text", "text": "The answer."}, + ], + }, + {"role": "user", "content": "And the next one?"}, + ] + request = _transform(messages) + assert request["messages"][1]["content"] == [{"type": "text", "text": "The answer."}] + assert len(messages[1]["content"]) == 3 + + +def test_anthropic_signed_thinking_blocks_are_forwarded_untouched(): + messages = [ + {"role": "user", "content": "Solve it."}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": "EqQBCkYIAxgCIkA_anthropic_signed"}, + {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"}, + {"type": "text", "text": "The answer."}, + ], + }, + ] + assert _transform(messages)["messages"] == messages 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 336cbf57e29..97d22274b1c 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -42,12 +42,15 @@ FAKE_AUTH_TOKEN = "sk-ant-aut01-fake-auth-token-for-testing-123456789" def test_is_claude_code_one_shot_subagent_request(messages, system, expected): from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request - assert is_claude_code_one_shot_subagent_request( - messages=messages, - system=system, - tools=None, - user_agent="claude-cli/2.1.263 (external, cli)", - ) is expected + assert ( + is_claude_code_one_shot_subagent_request( + messages=messages, + system=system, + tools=None, + user_agent="claude-cli/2.1.263 (external, cli)", + ) + is expected + ) class TestOptionallyHandleAnthropicOAuth: @@ -1541,8 +1544,8 @@ class TestAnthropicThinkingSignatureSelfHeal: out = strip_empty_content_blocks_from_anthropic_messages(msgs) assert [b["type"] for b in out[0]["content"]] == ["thinking"] - def test_strip_drops_encrypted_reasoning_blocks_from_the_responses_bridge(self): - """A session resumed on an Anthropic model replays reasoning only OpenAI can verify.""" + def test_strip_keeps_encrypted_reasoning_blocks_for_the_responses_bridge(self): + """The /v1/messages handler runs this before dispatch, so the bridge must still see the replay.""" from litellm.litellm_core_utils.prompt_templates.common_utils import ( encrypted_reasoning_signature, ) @@ -1556,14 +1559,46 @@ class TestAnthropicThinkingSignatureSelfHeal: "content": [ {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, - {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"}, {"type": "text", "text": "The answer."}, ], } ] - out = strip_empty_content_blocks_from_anthropic_messages(msgs) - assert [b["type"] for b in out[0]["content"]] == ["redacted_thinking", "text"] - assert len(msgs[0]["content"]) == 4 + assert strip_empty_content_blocks_from_anthropic_messages(msgs) == msgs + + def test_strip_encrypted_reasoning_drops_only_the_bridge_tagged_blocks(self): + """A session resumed on an Anthropic model replays reasoning only OpenAI can verify.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, + ) + from litellm.llms.anthropic.common_utils import ( + strip_encrypted_reasoning_blocks_from_anthropic_messages, + ) + + msgs = [ + {"role": "user", "content": "Solve it."}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_3")}, + {"type": "thinking", "thinking": "native", "signature": "EqQBCkYIAxgCIkA_anthropic_signed"}, + {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"}, + {"type": "text", "text": "The answer."}, + ], + }, + ] + out = strip_encrypted_reasoning_blocks_from_anthropic_messages(msgs) + assert [m["role"] for m in out] == ["user", "assistant"] + assert [b["type"] for b in out[1]["content"]] == ["thinking", "redacted_thinking", "text"] + assert out[1]["content"][0]["signature"] == "EqQBCkYIAxgCIkA_anthropic_signed" + assert len(msgs[1]["content"]) == 2 + assert len(msgs[2]["content"]) == 4 def test_strip_empty_text_blocks_treats_null_text_as_empty(self): from litellm.llms.anthropic.common_utils import ( From 1bcb587d336c8a141809ca0fe165d1b305f1122f Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:32:18 -0700 Subject: [PATCH 052/107] fix(mcp): challenge gateway-owned per-server authentication --- .../mcp_server/auth/user_api_key_auth_mcp.py | 15 +---- .../mcp_server/discoverable_endpoints.py | 11 ++-- .../types/mcp_server/mcp_server_manager.py | 18 +++++- .../auth/test_user_api_key_auth_mcp.py | 38 ++++++----- .../mcp_server/test_discoverable_endpoints.py | 63 +++++++++---------- 5 files changed, 75 insertions(+), 70 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 4bac65125b4..7e6de474b0b 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -179,16 +179,7 @@ def _gateway_dcr_challenge_target( mcp_servers: list[str] | None, client_ip: str | None, ) -> str | None: - """The single path-named server this request targets, iff it resolves to a - gateway-managed oauth2 server — the one per-server shape the gateway's own keyless - DCR flow serves end to end, so the 401 challenge may advertise the per-server - protected-resource metadata (whose ``authorization_servers`` names the gateway). - - Multi-server CSV paths, header/path mismatches, unknown names, and every - client-forwarded or delegated mode return ``None``: those cells keep their existing - challenge (or absence of one), and a challenge is never emitted for a name the - public discovery routes would 404, so this reveals exactly the server set the - per-server protected-resource metadata already reveals.""" + """Resolve a single path target whose sign-in metadata advertises the gateway.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) @@ -217,7 +208,7 @@ def _is_gateway_dcr_challenge_scope( the caller is not a cold-start DCR client), on the scopes the gateway's keyless flow serves: the aggregate ``/mcp`` endpoint, an ``x-mcp-servers``-scoped request (the resource the client configured is still ``/mcp``), or a per-server path whose - single target is a gateway-managed oauth2 server. Every other named target keeps + single target advertises gateway-owned sign-in. Every other named target keeps its existing behavior, failing closed to the original admission error.""" if not _is_litellm_auth_admission_error(exc): return False @@ -236,7 +227,7 @@ def _gateway_dcr_challenge( ) -> HTTPException: """The RFC 9728 challenge pointing the client at the protected-resource metadata matching the scope it requested: the per-server document (same URL spelling the - request arrived on) when the single target is a gateway-managed oauth2 server, + request arrived on) when the single target advertises gateway-owned sign-in, else the gateway's aggregate document. Either way the client discovers the gateway as its authorization server and starts the same sign-in flow. diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index cab4b6c161a..f18acb7d88d 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2310,8 +2310,7 @@ async def _build_oauth_protected_resource_response( it. Only the legacy ``is_oauth_passthrough`` opt-in rewrites ``resource`` to the gateway's own URL so clients present the bearer token back to the gateway. - An explicitly named gateway-managed oauth2 server (interactive with - gateway-vaulted per-user tokens, or M2M) advertises the gateway's own + An explicitly named server with gateway-owned sign-in advertises the gateway's own authorization server (``{base}/mcp``): a keyless DCR client that configured the per-server URL completes the same sign-in flow the aggregate ``/mcp`` endpoint supports and is admitted with a gateway session bearer. The per-server relay @@ -2401,11 +2400,6 @@ async def _build_oauth_protected_resource_response( if obo_response is not None: return obo_response - # An OBO server with no configured issuer falls through to the gateway default so discovery still - # returns metadata; every other non-oauth2 named server 404s to avoid enumeration. - if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange: - _raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth-protected resource") - if explicitly_named and mcp_server is not None and mcp_server.advertises_gateway_authorization_server: return { "authorization_servers": [f"{request_base_url}/mcp"], @@ -2413,6 +2407,9 @@ async def _build_oauth_protected_resource_response( "scopes_supported": (mcp_server.scopes if mcp_server.scopes else []), } + if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange: + _raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth-protected resource") + return { "authorization_servers": [ (f"{request_base_url}/{mcp_server_name}" if mcp_server_name else f"{request_base_url}") diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 84ffd50eea1..ab2e8a70754 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -250,7 +250,23 @@ class MCPServer(BaseModel): @property def advertises_gateway_authorization_server(self) -> bool: """Whether named discovery should advertise the aggregate gateway authorization server.""" - return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay + if self.auth_type == MCPAuth.oauth2: + return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay + if self.auth_type not in ( + None, + MCPAuth.none, + MCPAuth.api_key, + MCPAuth.bearer_token, + MCPAuth.basic, + MCPAuth.authorization, + MCPAuth.token, + MCPAuth.aws_sigv4, + ): + return False + return not any( + header.lower() in ("authorization", "x-api-key", "api-key", "apikey") + for header in (self.extra_headers or ()) + ) @property def is_true_passthrough(self) -> bool: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 20f4719e4bf..e6c8d4ee039 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -7193,14 +7193,13 @@ class TestAggregateGatewayDcrChallenge: www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] assert www_authenticate == f"Bearer {self._EXPECTED_RESOURCE_METADATA}" - async def test_per_server_challenge_for_gateway_managed_oauth2(self): - """Anonymous request to a per-server path whose single target is a gateway-managed - oauth2 server: 401 plus the RFC 9728 challenge advertising the PER-SERVER - protected-resource metadata in the same URL spelling the request used, so a keyless - DCR client configured with either per-server spelling discovers the gateway as the - authorization server (LIT-4864). Covers interactive and M2M, which the gateway can - both serve end to end.""" - from litellm.types.mcp import MCPAuth + @pytest.mark.parametrize( + "auth_type", + (None, "none", "api_key", "bearer_token", "basic", "aws_sigv4", "authorization", "token", "oauth2"), + ) + @pytest.mark.parametrize("bearer_presented", (False, True)) + async def test_per_server_challenge_for_gateway_owned_auth(self, auth_type, bearer_presented): + """Gateway admission challenges are independent of upstream authentication.""" from litellm.types.mcp_server.mcp_server_manager import MCPServer server = MCPServer( @@ -7209,7 +7208,7 @@ class TestAggregateGatewayDcrChallenge: server_name="github", url="https://upstream.example/mcp", transport="http", - auth_type=MCPAuth.oauth2, + auth_type=auth_type, ) for path, expected_metadata_path in ( ("/mcp/github", "/.well-known/oauth-protected-resource/mcp/github"), @@ -7223,10 +7222,16 @@ class TestAggregateGatewayDcrChallenge: ): mock_mgr.get_mcp_server_by_name.return_value = server with pytest.raises(HTTPException) as exc_info: - await MCPRequestHandler.process_mcp_request(self._scope(path=path)) + await MCPRequestHandler.process_mcp_request( + self._scope( + path=path, + extra_headers=((b"authorization", b"Bearer invalid-key"),) if bearer_presented else (), + ) + ) assert exc_info.value.status_code == 401 www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] - assert www_authenticate == f'Bearer resource_metadata="http://testserver{expected_metadata_path}"' + error = 'error="invalid_token", ' if bearer_presented else "" + assert www_authenticate == f'Bearer {error}resource_metadata="http://testserver{expected_metadata_path}"' async def test_per_server_challenge_keeps_spelling_under_server_root_path(self): """On a sub-path deployment the challenge must still advertise the spelling the client @@ -7303,10 +7308,7 @@ class TestAggregateGatewayDcrChallenge: ) def test_challenge_target_excludes_every_non_gateway_managed_mode(self): - """Unit pin of the challenge-target owner: only a resolved gateway-managed oauth2 - target (interactive or M2M) yields a per-server challenge; delegate-auth oauth2 - (whose keyless flow is upstream PKCE via the relay), every client-forwarded auth - type, OBO, api_key, unknown names, and CSV paths yield None (LIT-4864).""" + """Gateway challenges exclude unresolved, delegated, and client-forwarded targets.""" from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( _gateway_dcr_challenge_target, ) @@ -7333,7 +7335,11 @@ class TestAggregateGatewayDcrChallenge: (_server(MCPAuth.true_passthrough), None), (_server(MCPAuth.oauth_delegate), None), (_server(MCPAuth.oauth_delegate, dcr_bridge=True), None), - (_server(MCPAuth.api_key), None), + (_server(MCPAuth.api_key), "srv"), + (_server(MCPAuth.none, extra_headers=["Authorization"]), None), + (_server(None, extra_headers=["X-API-Key"]), None), + (_server(MCPAuth.none, extra_headers=["Authorization"], oauth_passthrough=True), None), + (_server(MCPAuth.oauth2_id_jag), None), (None, None), ] for resolved, expected in cases: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 763200c3709..a7a65bfe466 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7651,12 +7651,6 @@ async def test_token_endpoint_client_secret_basic_without_secret_returns_400(): assert exc_info.value.status_code == 400 -# ------------------------------------------------------------------- -# Non-oauth2 (auth_type=none, access-group gated) servers must not be -# driven through the gateway OAuth authorize/token/register/discovery -# flow, and must not be advertised as OAuth-protected in discovery docs. -# ------------------------------------------------------------------- - def _access_group_none_server(server_name="access_group_server"): """A non-oauth2, access-group gated MCP server: no client_id, no OAuth.""" @@ -7794,35 +7788,38 @@ async def test_register_client_rejects_non_oauth2_server(): @pytest.mark.asyncio -async def test_oauth_protected_resource_404_for_non_oauth2_server(): - """Discovery must not advertise a none-auth server as an OAuth-protected resource.""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _build_oauth_protected_resource_response, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - except ImportError: - pytest.skip("MCP discoverable endpoints not available") +@pytest.mark.parametrize( + "auth_type", (None, "none", "api_key", "bearer_token", "basic", "aws_sigv4", "authorization", "token") +) +@pytest.mark.parametrize("use_standard_pattern", (False, True)) +async def test_oauth_protected_resource_for_gateway_owned_auth(auth_type, use_standard_pattern): + from starlette.requests import Request + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_oauth_protected_resource_response, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + server = _access_group_none_server().model_copy(update={"auth_type": auth_type}) + request = Request( + {"type": "http", "scheme": "https", "path": "/", "headers": [(b"host", b"litellm.example.com")]} + ) global_mcp_server_manager.registry.clear() - server = _access_group_none_server() global_mcp_server_manager.registry[server.server_id] = server - - mock_request = MagicMock() - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - try: - with pytest.raises(HTTPException) as exc_info: - await _build_oauth_protected_resource_response( - request=mock_request, - mcp_server_name="access_group_server", - use_standard_pattern=False, - ) - assert exc_info.value.status_code == 404 - assert "not an OAuth-protected resource" in str(exc_info.value.detail) + response = await _build_oauth_protected_resource_response( + request=request, + mcp_server_name="access_group_server", + use_standard_pattern=use_standard_pattern, + ) + resource_path = "/mcp/access_group_server" if use_standard_pattern else "/access_group_server/mcp" + assert response == { + "resource": f"https://litellm.example.com{resource_path}", + "authorization_servers": ["https://litellm.example.com/mcp"], + "scopes_supported": [], + } finally: global_mcp_server_manager.registry.clear() @@ -7914,9 +7911,7 @@ async def test_oauth_protected_resource_passthrough_none_auth_not_404(): @pytest.mark.asyncio async def test_oauth_protected_resource_404_for_unknown_server_name(): - """A discovery request for an unknown server name returns the same 404 as a non-oauth2 - server (not a 200 metadata doc with broken URLs), so the well-known paths cannot be used - to enumerate non-OAuth server names.""" + """Unknown server names must not produce metadata advertising nonexistent resources.""" try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _build_oauth_protected_resource_response, From feeda1a36ce9f4155e443bbb264950ac0139f7a0 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:32:06 -0700 Subject: [PATCH 053/107] fix(mcp): report resolved upstream authentication in debug headers --- .../_experimental/mcp_server/mcp_debug.py | 161 ++++++++++++------ .../mcp_server/mcp_server_manager.py | 34 +++- .../outbound_credentials/resolver.py | 31 ++++ .../mcp_server/outbound_credentials/types.py | 26 ++- .../proxy/_experimental/mcp_server/server.py | 14 +- .../outbound_credentials/test_resolver.py | 66 +++++++ .../mcp_server/test_mcp_debug.py | 161 ++++++++++-------- .../mcp_server/test_mcp_server_manager.py | 106 ++++++++++++ 8 files changed, 470 insertions(+), 129 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_debug.py b/litellm/proxy/_experimental/mcp_server/mcp_debug.py index 6aaa00ae415..b25ed27becc 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_debug.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_debug.py @@ -22,7 +22,16 @@ Response headers returned (all values are masked for safety): x-mcp-debug-auth-resolution Which auth priority was used for the outbound MCP call: ``per-request-header``, ``m2m-client-credentials``, ``static-token``, - ``oauth2-passthrough``, or ``no-auth``. + ``oauth2-passthrough``, ``stored-user-token``, ``token-exchange``, + ``id-jag``, ``aws-sigv4``, ``extra-headers``, or ``no-auth``. + ``unresolved`` means no outcome was available before the first response + frame; ``multiple`` means several servers resolved credentials; + ``not-applicable`` covers stdio; ``resolution-failed`` is a resolver error. + + x-mcp-debug-auth-resolutions + For multiple servers, a JSON map of server IDs to resolution labels. + At most 32 entries are included; x-mcp-debug-auth-resolutions-truncated + is true when additional servers were omitted. No credentials are included. x-mcp-debug-outbound-url The upstream MCP server URL that will receive the request. @@ -58,10 +67,16 @@ header is free for OAuth2 discovery:: Symptom: ``x-mcp-debug-oauth2-token`` shows ``(none)`` and ``x-mcp-debug-auth-resolution`` shows ``no-auth``. -This means the client didn't go through the OAuth2 flow. Check that: -1. The ``Authorization`` header is NOT set as a static header in the client config. -2. The ``.well-known/oauth-protected-resource`` endpoint returns valid metadata. -3. The MCP server in LiteLLM config has ``auth_type: oauth2``. +``no-auth`` means the resolved upstream client carries no authentication. +An absent inbound OAuth2 token does not imply the user skipped OAuth: the gateway +can retrieve a stored per-user token, reported as ``stored-user-token``. +``unresolved`` is used when a stream starts before credential resolution, or a +request (such as initialization or a cached tool listing) resolves no credential. +Debug reporting does not fetch credentials or delay a streaming frame to resolve them. +``extra-headers`` identifies supplied headers that won over the resolver or were +the only headers supplied; their values are never inspected to guess a scheme. +``per-request-header`` denotes a legacy credential override, including a BYOK +credential supplied by the gateway; it does not imply a caller-supplied token. **Common issue: M2M token used instead of user token** @@ -69,8 +84,8 @@ Symptom: ``x-mcp-debug-auth-resolution`` shows ``m2m-client-credentials``. This means the server has ``client_id``/``client_secret``/``token_url`` configured and LiteLLM is fetching a machine-to-machine token instead of -using the per-user OAuth2 token. If you want per-user tokens, remove the -client credentials from the server config. +using the per-user OAuth2 token. For gateway-stored per-user tokens, +configure ``oauth2_flow: authorization_code``. Usage from Claude Code:: @@ -85,14 +100,16 @@ Usage with curl:: http://localhost:4000/mcp/atlassian_mcp """ -from typing import TYPE_CHECKING, Final +import json +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final +from starlette.requests import HTTPConnection from starlette.types import Message, Send from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker - -if TYPE_CHECKING: - from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution # Header the client sends to opt into debug mode MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug" @@ -101,6 +118,83 @@ MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug" _RESPONSE_HEADER_PREFIX: Final = "x-mcp-debug" +MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: Final = "litellm.mcp.auth_diagnostics" + + +def record_auth_resolution(server_id: str, source: AuthResolution) -> None: + from mcp.server.lowlevel.server import request_ctx + + context: Final[object] = request_ctx.get(None) + request: Final[object] = getattr(context, "request", None) + if isinstance(request, HTTPConnection): + diagnostics: Final[object] = request.scope.get(MCP_AUTH_DIAGNOSTICS_SCOPE_KEY) + if isinstance(diagnostics, MCPAuthDiagnostics): + diagnostics.record(server_id, source) + + +class MCPAuthDiagnostics: + def __init__(self) -> None: + self._outcomes: tuple[tuple[str, AuthResolution], ...] = () + + def record(self, server_id: str, resolution: AuthResolution) -> None: + self._outcomes = tuple(item for item in self._outcomes if item[0] != server_id) + ((server_id, resolution),) + + def resolution(self) -> str: + match self._outcomes: + case (): + return AuthResolution.unresolved.value + case ((_, source),): + return source.value + case _: + return AuthResolution.multiple.value + + def headers(self) -> Mapping[str, str]: + if len(self._outcomes) <= 1: + return MappingProxyType({"x-mcp-debug-auth-resolution": self.resolution()}) + return MappingProxyType( + { + "x-mcp-debug-auth-resolution": AuthResolution.multiple.value, + "x-mcp-debug-auth-resolutions": json.dumps( + { + server_id: source.value for server_id, source in self._outcomes[:32] + }, # mutable-ok: JSON encoder requires a concrete dict + separators=(",", ":"), + ensure_ascii=True, + ), + **( + MappingProxyType({"x-mcp-debug-auth-resolutions-truncated": "true"}) + if len(self._outcomes) > 32 + else MappingProxyType({}) + ), + } + ) + + +class _DiagnosticSend: + def __init__(self, send: Send, headers: Mapping[str, str], resolution: Callable[[], Mapping[str, str]]) -> None: + self._send = send + self._headers = headers + self._resolution = resolution + self._start: Message | None = None + + async def __call__(self, message: Message) -> None: + if message["type"] == "http.response.start": + self._start = message + return + if self._start is not None: + start: Final = self._start + self._start = None + headers: Final = MappingProxyType({**self._headers, **self._resolution()}) + await self._send( + { # mutable-ok: ASGI send consumes a mutable message mapping + **start, + "headers": tuple(start.get("headers", ())) + + tuple((key.encode(), value.encode()) for key, value in headers.items()), + } + ) + await self._send(message) + + class MCPDebug: """ Static helper class for MCP OAuth2 debug headers. @@ -144,37 +238,6 @@ class MCPDebug: return val.strip().lower() in ("true", "1", "yes") return False - @staticmethod - def resolve_auth_resolution( - server: "MCPServer", - mcp_auth_header: str | None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None, - oauth2_headers: dict[str, str] | None, - ) -> str: - """ - Determine which auth priority will be used for the outbound MCP call. - - Returns one of: ``per-request-header``, ``m2m-client-credentials``, - ``static-token``, ``oauth2-passthrough``, or ``no-auth``. - """ - from litellm.types.mcp import MCPAuth - - has_server_specific: Final = bool( - mcp_server_auth_headers - and ( - mcp_server_auth_headers.get(server.alias or "") or mcp_server_auth_headers.get(server.server_name or "") - ) - ) - if has_server_specific or mcp_auth_header: - return "per-request-header" - if server.has_client_credentials: - return "m2m-client-credentials" - if server.authentication_token: - return "static-token" - if oauth2_headers and server.auth_type == MCPAuth.oauth2: - return "oauth2-passthrough" - return "no-auth" - @staticmethod def build_debug_headers( *, @@ -244,12 +307,17 @@ class MCPDebug: return debug @staticmethod - def wrap_send_with_debug_headers(send: Send, debug_headers: dict[str, str]) -> Send: + def wrap_send_with_debug_headers( + send: Send, debug_headers: Mapping[str, str], resolution: Callable[[], Mapping[str, str]] | None = None + ) -> Send: """ Return a new ASGI ``send`` callable that injects *debug_headers* into the ``http.response.start`` message. """ + if resolution is not None: + return _DiagnosticSend(send, debug_headers, resolution) + async def _send_with_debug(message: Message) -> None: if message["type"] == "http.response.start": headers: Final = list(message.get("headers", [])) @@ -266,8 +334,6 @@ class MCPDebug: raw_headers: dict[str, str] | None, scope: dict, mcp_servers: list[str] | None, - mcp_auth_header: str | None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None, oauth2_headers: dict[str, str] | None, client_ip: str | None, ) -> dict[str, str]: @@ -288,16 +354,13 @@ class MCPDebug: server_url: str | None = None server_auth_type: str | None = None - auth_resolution = "no-auth" + auth_resolution: Final = AuthResolution.unresolved.value for server_name in mcp_servers or []: server = global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=client_ip) if server: server_url = server.url server_auth_type = server.auth_type - auth_resolution = MCPDebug.resolve_auth_resolution( - server, mcp_auth_header, mcp_server_auth_headers, oauth2_headers - ) break scope_headers: Final = MCPRequestHandler._safe_get_headers_from_scope(scope) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d7f238f142c..116bcc715d4 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -80,6 +80,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( raise_classified_list_failure, upstream_auth_challenge, ) +from litellm.proxy._experimental.mcp_server.mcp_debug import record_auth_resolution from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( MCPPerUserTokenCache, mcp_per_user_token_cache, @@ -108,12 +109,14 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import ( LazyPerUserOAuthTokenStore, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_provider import ( build_token_exchanger, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( DEFAULT_CREDENTIAL_HEADER, AuthorizationCodeConfig, + AuthResolution, ClientCredentialsConfig, CredError, IdJagConfig, @@ -3832,13 +3835,21 @@ class MCPServerManager: (authorization_code's browser-OAuth 401, token_exchange's RFC 9728 challenge) or maps any other ``CredError`` onto its public HTTP status; it never returns an error as a value. """ - match await provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): - case Ok(auth): + match await resolve_credentials_with_source(provider, to_subject(user_api_key_auth, subject_token), spec): + case Ok(credential): + auth: Final = credential.auth # NoOpAuth has no header_name and so never conflicts. header_name: Final[str | None] = getattr(auth, "header_name", None) if header_name is None or not extra_headers: + source: Final = ( + AuthResolution.extra_headers + if credential.source == AuthResolution.no_auth and extra_headers + else credential.source + ) + record_auth_resolution(server.server_id, source) return auth, extra_headers if not has_header(extra_headers, header_name): + record_auth_resolution(server.server_id, credential.source) return auth, extra_headers if isinstance( spec.config, @@ -3853,11 +3864,14 @@ class MCPServerManager: # one-shot 401 refetch is lost with it). Drop only the header the resolved # credential is about to occupy, so a static credential the operator aimed at a # DIFFERENT header still reaches upstream. + record_auth_resolution(server.server_id, credential.source) return auth, without_header(extra_headers, header_name) # Other modes: an Authorization already supplied via extra_headers (a forwarded caller # header or static_headers) is intentional and wins; v1 applies those last. + record_auth_resolution(server.server_id, AuthResolution.extra_headers) return None, extra_headers case Error(err): + record_auth_resolution(server.server_id, AuthResolution.failed) if err.tag == "unauthorized" and isinstance(spec.config, AuthorizationCodeConfig): # authorization_code's missing per-user token -> the per-server browser-OAuth # challenge, built here where the full MCPServer is in hand. @@ -3960,6 +3974,7 @@ class MCPServerManager: Returns: Configured MCP client instance. """ + record_auth_resolution(server.server_id, AuthResolution.unresolved) resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) transport: Final = resolved_server.transport or MCPTransport.sse spec = None if transport == MCPTransport.stdio else _to_server_spec_fail_closed(resolved_server) @@ -4032,6 +4047,7 @@ class MCPServerManager: env=resolved_env, ) + record_auth_resolution(server.server_id, AuthResolution.not_applicable) return MCPClient( server_url="", # Not used for stdio transport_type=transport, @@ -4086,6 +4102,20 @@ class MCPServerManager: aws_session_name=resolved_server.aws_session_name, ) + legacy_source: Final = ( + AuthResolution.aws_sigv4 + if aws_auth is not None + else AuthResolution.extra_headers + if extra_headers and has_header(extra_headers, auth_header_name or "Authorization") + else AuthResolution.per_request_header + if mcp_auth_header + else AuthResolution.static_token + if auth_value + else AuthResolution.extra_headers + if extra_headers + else AuthResolution.no_auth + ) + record_auth_resolution(server.server_id, legacy_source) return MCPClient( server_url=server_url, transport_type=transport, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 8328aae01ab..85c7f68719d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -65,6 +65,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, + AuthResolution, AuthSpecKind, AwsSigV4Config, Byok, @@ -76,6 +77,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( NoneConfig, PassthroughConfig, PrivateKeyJwtAuth, + ResolvedCredential, ServerSpec, SharedKey, Subject, @@ -448,3 +450,32 @@ def _client_auth_fingerprint(client_auth: ClientAuth) -> str: def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]: return Error(CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet")) + + +async def resolve_credentials_with_source( + provider: UpstreamCredentialProvider, subject: Subject, server: ServerSpec +) -> Result[ResolvedCredential, CredError]: + match await provider.resolve_credentials(subject, server): + case Error(err): + return Error(err) + case Ok(auth): + if isinstance(auth, NoOpAuth): + return Ok(ResolvedCredential(auth, AuthResolution.no_auth)) + match server.config: + case NoneConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.no_auth)) + case ApiKeyConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.static_token)) + case PassthroughConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.oauth2_passthrough)) + case ClientCredentialsConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.client_credentials)) + case TokenExchangeConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.token_exchange)) + case IdJagConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.id_jag)) + case AuthorizationCodeConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.stored_user_token)) + case AwsSigV4Config(): + return Ok(ResolvedCredential(auth, AuthResolution.aws_sigv4)) + assert_never(server.config) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 632dc57dcf6..d186724fd9f 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -26,10 +26,11 @@ union (see `result.py`), not `expression.Result`. from __future__ import annotations from collections.abc import Mapping -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum from typing import Annotated, Final, Literal +import httpx from expression import case, tag, tagged_union from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator from typing_extensions import assert_never @@ -46,6 +47,29 @@ from litellm.types.mcp import ( ) +class AuthResolution(str, Enum): + no_auth = "no-auth" + stored_user_token = "stored-user-token" + static_token = "static-token" + per_request_header = "per-request-header" + oauth2_passthrough = "oauth2-passthrough" + client_credentials = "m2m-client-credentials" + token_exchange = "token-exchange" + id_jag = "id-jag" + aws_sigv4 = "aws-sigv4" + extra_headers = "extra-headers" + not_applicable = "not-applicable" + unresolved = "unresolved" + failed = "resolution-failed" + multiple = "multiple" + + +@dataclass(frozen=True, slots=True) +class ResolvedCredential: + auth: httpx.Auth = field(repr=False) + source: AuthResolution + + class AuthSpecKind(str, Enum): """The server's statically-declared upstream-auth mode — derived from its `config`. diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 9a52da0cab1..b53c55ef5fa 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -49,7 +49,11 @@ from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_gateway_server_name, _mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # server-owned request mode ) -from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug +from litellm.proxy._experimental.mcp_server.mcp_debug import ( + MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, + MCPAuthDiagnostics, + MCPDebug, +) from litellm.proxy._experimental.mcp_server.oauth_utils import ( _redact_mcp_resource_url, get_passthrough_www_authenticate, @@ -4472,13 +4476,13 @@ if MCP_AVAILABLE: raw_headers=raw_headers, scope=dict(scope), mcp_servers=mcp_servers, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, client_ip=_client_ip, ) - if _debug_headers: - send = MCPDebug.wrap_send_with_debug_headers(send, _debug_headers) + diagnostics: Final = MCPAuthDiagnostics() if _debug_headers else None + if diagnostics is not None: + scope[MCP_AUTH_DIAGNOSTICS_SCOPE_KEY] = diagnostics + send = MCPDebug.wrap_send_with_debug_headers(send, _debug_headers, diagnostics.headers) # Ensure session managers are initialized if not _SESSION_MANAGERS_INITIALIZED: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 6b3098d9e60..2eab053ab73 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -1203,3 +1203,69 @@ async def test_passthrough_ignores_the_carrier_and_keeps_the_callers_slot(): assert isinstance(result, Ok) headers, _ = await _emitted_async(result.ok) assert headers["Authorization"] == "caller-token" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("config", "subject", "expected_source", "expected_header"), + [ + (NoneConfig(), _SUBJECT, "no-auth", None), + (PassthroughConfig(), _SUBJECT, "no-auth", None), + (PassthroughConfig(), _with_inbound("Bearer caller-token"), "oauth2-passthrough", "Bearer caller-token"), + (ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-key"))), _SUBJECT, "static-token", "Bearer static-key"), + (AuthorizationCodeConfig(), Subject(tenant_id="", subject_id="alice"), "stored-user-token", "Bearer stored-alice"), + ], +) +async def test_resolved_source_matches_the_credential_sent_upstream(config, subject, expected_source, expected_header): + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source + + store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="stored-alice")}) + provider = UpstreamCredentialProvider(oauth_token_store=store) + result = await resolve_credentials_with_source(provider, subject, _spec(config)) + assert isinstance(result, Ok) + assert result.ok.source.value == expected_source + assert _emitted(result.ok.auth).get("Authorization") == expected_header + assert "stored-alice" not in repr(result.ok) + assert "static-key" not in repr(result.ok) + + +@pytest.mark.asyncio +async def test_resolved_source_preserves_missing_user_token_error(): + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source + + result = await resolve_credentials_with_source(UpstreamCredentialProvider(), _SUBJECT, _spec(AuthorizationCodeConfig())) + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + + +@pytest.mark.asyncio +async def test_minted_token_sources_match_egress_and_do_not_fetch_twice(): + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source + + source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at"))) + m2m = await resolve_credentials_with_source(UpstreamCredentialProvider(client_credentials_source=source), _SUBJECT, _spec(_M2M)) + assert isinstance(m2m, Ok) + headers, _ = await _emitted_async(m2m.ok.auth) + assert headers["Authorization"] == "Bearer m2m-at" + assert m2m.ok.source.value == "m2m-client-credentials" + assert source.gets == ["s"] + + exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at"))) + exchanged = await resolve_credentials_with_source(UpstreamCredentialProvider(token_exchanger=exchanger), _with_inbound("subject"), _spec(_OBO)) + assert isinstance(exchanged, Ok) + assert _emitted(exchanged.ok.auth)["Authorization"] == "Bearer exchanged-at" + assert exchanged.ok.source.value == "token-exchange" + assert len(exchanger.calls) == 1 + + +@pytest.mark.asyncio +async def test_id_jag_source_describes_final_token_after_both_exchanges(): + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source + + endpoint = _FakeTokenEndpoint(_two_leg_ok("resource-token")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + result = await resolve_credentials_with_source(provider, _with_inbound("identity-token"), _spec(_id_jag_config())) + assert isinstance(result, Ok) + assert result.ok.source.value == "id-jag" + assert _emitted(result.ok.auth)["Authorization"] == "Bearer resource-token" + assert len(endpoint.calls) == 2 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py index d299239f68e..5039209c351 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py @@ -3,11 +3,17 @@ Tests for MCPDebug — MCP OAuth2 debug response headers. """ import asyncio -from unittest.mock import MagicMock +from typing import Final + +import pytest +from starlette.types import Message + +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution from litellm.proxy._experimental.mcp_server.mcp_debug import ( MCP_DEBUG_REQUEST_HEADER, MCPDebug, + MCPAuthDiagnostics, ) @@ -166,77 +172,6 @@ class TestBuildDebugHeaders: assert set(headers.keys()) == expected_keys -class TestResolveAuthResolution: - def _make_server(self, **kwargs): - server = MagicMock() - server.alias = kwargs.get("alias", "test") - server.server_name = kwargs.get("server_name", "test") - server.has_client_credentials = kwargs.get("has_client_credentials", False) - server.authentication_token = kwargs.get("authentication_token", None) - server.auth_type = kwargs.get("auth_type", None) - return server - - def test_per_request_header(self): - server = self._make_server() - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header="Bearer xxx", - mcp_server_auth_headers=None, - oauth2_headers=None, - ) - assert result == "per-request-header" - - def test_server_specific_header(self): - server = self._make_server(alias="atlas") - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers={"atlas": {"Authorization": "Bearer xxx"}}, - oauth2_headers=None, - ) - assert result == "per-request-header" - - def test_m2m(self): - server = self._make_server(has_client_credentials=True) - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers=None, - oauth2_headers=None, - ) - assert result == "m2m-client-credentials" - - def test_static_token(self): - server = self._make_server(authentication_token="static-tok") - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers=None, - oauth2_headers=None, - ) - assert result == "static-token" - - def test_oauth2_passthrough(self): - server = self._make_server(auth_type="oauth2") - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers=None, - oauth2_headers={"Authorization": "Bearer eyJ..."}, - ) - assert result == "oauth2-passthrough" - - def test_no_auth(self): - server = self._make_server() - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers=None, - oauth2_headers=None, - ) - assert result == "no-auth" - - class TestWrapSendWithDebugHeaders: def test_injects_headers(self): captured = [] @@ -269,3 +204,85 @@ class TestWrapSendWithDebugHeaders: asyncio.run(wrapped(body_msg)) assert captured[0] == body_msg + + +@pytest.mark.asyncio +@pytest.mark.parametrize("source", tuple(AuthResolution)) +async def test_debug_uses_resolution_recorded_after_response_start(source: AuthResolution) -> None: + captured: Final[list[Message]] = [] + diagnostics: Final = MCPAuthDiagnostics() + + async def send(message: Message) -> None: + captured.append(message) + + wrapped: Final = MCPDebug.wrap_send_with_debug_headers(send, {}, diagnostics.headers) + await wrapped({"type": "http.response.start", "status": 200, "headers": []}) + assert captured == [] + diagnostics.record("s1", source) + body: Final[Message] = {"type": "http.response.body", "body": b"data: pong\n\n", "more_body": True} + await wrapped(body) + assert dict(captured[0]["headers"])[b"x-mcp-debug-auth-resolution"] == source.value.encode() + assert captured[1] == body + + +@pytest.mark.asyncio +async def test_early_stream_frame_reports_unresolved_without_waiting() -> None: + captured: Final[list[Message]] = [] + diagnostics: Final = MCPAuthDiagnostics() + + async def send(message: Message) -> None: + captured.append(message) + + wrapped: Final = MCPDebug.wrap_send_with_debug_headers(send, {}, diagnostics.headers) + await wrapped({"type": "http.response.start", "status": 200, "headers": []}) + await wrapped({"type": "http.response.body", "body": b": ping\n\n", "more_body": True}) + diagnostics.record("s1", AuthResolution.stored_user_token) + await wrapped({"type": "http.response.body", "body": b"data: pong\n\n", "more_body": False}) + assert len(captured) == 3 + assert dict(captured[0]["headers"])[b"x-mcp-debug-auth-resolution"] == b"unresolved" + + +def test_diagnostics_keep_requests_separate_and_do_not_collapse_multiple_servers() -> None: + alice: Final = MCPAuthDiagnostics() + bob: Final = MCPAuthDiagnostics() + alice.record("s1", AuthResolution.stored_user_token) + assert bob.resolution() == "unresolved" + alice.record("s1", AuthResolution.token_exchange) + assert alice.resolution() == "token-exchange" + alice.record("s2", AuthResolution.static_token) + assert alice.resolution() == "multiple" + assert alice.headers()["x-mcp-debug-auth-resolutions"] == '{"s1":"token-exchange","s2":"static-token"}' + + +@pytest.mark.asyncio +async def test_concurrent_mcp_messages_record_on_their_own_http_scope() -> None: + from unittest.mock import MagicMock + + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.mcp_debug import ( + MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, + record_auth_resolution, + ) + + session: Final = MagicMock() + first: Final = MCPAuthDiagnostics() + second: Final = MCPAuthDiagnostics() + + async def record(diagnostics: MCPAuthDiagnostics, source: AuthResolution) -> None: + context: Final = RequestContext( + request_id=1, meta=None, session=session, lifespan_context=None, + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + ) + token: Final = request_ctx.set(context) + try: + await asyncio.sleep(0) + record_auth_resolution("same-server", source) + finally: + request_ctx.reset(token) + + await asyncio.gather(record(first, AuthResolution.stored_user_token), record(second, AuthResolution.per_request_header)) + assert first.resolution() == "stored-user-token" + assert second.resolution() == "per-request-header" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 46fef83092d..99447acd371 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -12567,3 +12567,109 @@ async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch, with pytest.raises(HTTPException) as exc_info: await call assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("config", "extra_headers", "expected_source", "expected_authorization"), + [ + ("stored", None, "stored-user-token", "Bearer stored-token"), + ("stored", {"aUtHoRiZaTiOn": "Bearer injected"}, "stored-user-token", "Bearer stored-token"), + ("static", None, "static-token", "Bearer static-token"), + ("static", {"authorization": "Bearer injected"}, "extra-headers", "Bearer injected"), + ("none", None, "no-auth", None), + ("none", {"Authorization": "Bearer injected"}, "extra-headers", "Bearer injected"), + ], +) +async def test_debug_resolution_matches_final_header_conflict_winner( + config, extra_headers, expected_source, expected_authorization +): + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + from starlette.requests import Request + from pydantic import SecretStr + + from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import MCPAuthenticatedUser + from litellm.proxy._experimental.mcp_server.mcp_debug import MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, MCPAuthDiagnostics + from litellm.proxy._experimental.mcp_server.outbound_credentials import ( + ApiKeyConfig, AuthorizationCodeConfig, NoneConfig, ServerSpec, SharedKey, UpstreamCredentialProvider, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + class Store: + def __init__(self): + self.calls = 0 + + async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: + self.calls += 1 + return OAuthToken(access_token="stored-token") if user_id == "alice" else None + + store = Store() + context = MCPAuthenticatedUser(UserAPIKeyAuth(user_id="alice")) + diagnostics = MCPAuthDiagnostics() + token = request_ctx.set(RequestContext( + request_id=1, meta=None, session=MagicMock(), lifespan_context=None, + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + )) + selected = { + "stored": AuthorizationCodeConfig(), + "static": ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-token"))), + "none": NoneConfig(), + }[config] + try: + auth, remaining = await MCPServerManager()._resolve_v2_auth( + server=MCPServer( + server_id="s", name="s", transport="http", url="https://up.example/mcp", + static_headers={"Authorization": "Bearer configured"}, + ), + spec=ServerSpec(server_id="s", resource="https://up.example/mcp", config=selected), + provider=UpstreamCredentialProvider(oauth_token_store=store), + subject_token=None, + user_api_key_auth=context.user_api_key_auth, + extra_headers=extra_headers, + ) + request = httpx.Request("GET", "https://up.example/mcp", headers=remaining) + if auth is not None: + next(auth.auth_flow(request)) + assert diagnostics.resolution() == expected_source + assert request.headers.get("Authorization") == expected_authorization + assert store.calls == (1 if config == "stored" else 0) + finally: + request_ctx.reset(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["http", "stdio"]) +async def test_debug_reports_legacy_signing_and_non_http_transport(transport): + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.mcp_debug import MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, MCPAuthDiagnostics + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + diagnostics = MCPAuthDiagnostics() + token = request_ctx.set(RequestContext( + request_id=1, meta=None, session=MagicMock(), lifespan_context=None, + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + )) + try: + server = MCPServer( + server_id="signed", name="signed", transport=transport, + url="https://up.example/mcp", auth_type="aws_sigv4", + aws_access_key_id="AKIDEXAMPLE", aws_secret_access_key="test-signing-secret", + aws_region_name="us-east-1", aws_service_name="execute-api", + command="python", args=["-c", "pass"], + ) + client = await MCPServerManager()._create_mcp_client(server) + if transport == "stdio": + assert diagnostics.resolution() == "not-applicable" + else: + assert diagnostics.resolution() == "aws-sigv4" + request = httpx.Request("POST", "https://up.example/mcp", content=b"{}") + next(client._aws_auth.auth_flow(request)) + assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256 ") + assert "Credential=AKIDEXAMPLE/" in request.headers["Authorization"] + finally: + request_ctx.reset(token) From f49ebc93eaf108d2312779346485ce6808c26184 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:11:43 -0700 Subject: [PATCH 054/107] fix(anthropic): request encrypted reasoning only where the Responses provider returns it The bridge now asks for reasoning.encrypted_content whenever the provider's Responses config lists include, independent of the client's thinking block, and leaves it out for providers such as Perplexity that reject the param. Bridge-tagged blocks are stripped on the chat adapter path too, so a mid session model switch to Gemini or Bedrock no longer forwards them as real signatures, a bare prefix counts as bridge-tagged, and non-mapping messages pass through the strip untouched. --- .../prompt_templates/common_utils.py | 18 ++++++--- litellm/llms/anthropic/common_utils.py | 2 + .../adapters/transformation.py | 4 +- .../responses_adapters/handler.py | 17 ++++++++- .../responses_adapters/transformation.py | 9 ++++- ...ore_utils_prompt_templates_common_utils.py | 3 ++ ...al_pass_through_adapters_transformation.py | 38 +++++++++++++++++++ .../test_responses_adapters_handler.py | 23 +++++++++++ .../test_responses_adapters_transformation.py | 12 +++++- .../anthropic/test_anthropic_common_utils.py | 9 +++++ 10 files changed, 126 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 918cb982773..75a3aefa98b 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1853,22 +1853,30 @@ def encrypted_reasoning_signature(encrypted_content: str) -> str: return f"{ENCRYPTED_REASONING_SIGNATURE_PREFIX}{encrypted_content}" +def _carries_encrypted_reasoning(signature: object) -> bool: + return isinstance(signature, str) and signature.startswith(ENCRYPTED_REASONING_SIGNATURE_PREFIX) + + def encrypted_content_from_signature(signature: object) -> str | None: - if not isinstance(signature, str) or not signature.startswith(ENCRYPTED_REASONING_SIGNATURE_PREFIX): + if not isinstance(signature, str) or not _carries_encrypted_reasoning(signature): return None return signature.removeprefix(ENCRYPTED_REASONING_SIGNATURE_PREFIX) or None -def _encrypted_content_of_block(block: Mapping[str, object]) -> str | None: +def _encrypted_reasoning_field(block: Mapping[str, object]) -> object: match block.get("type"): case "thinking": - return encrypted_content_from_signature(block.get("signature")) + return block.get("signature") case "redacted_thinking": - return encrypted_content_from_signature(block.get("data")) + return block.get("data") case _: return None +def _encrypted_content_of_block(block: Mapping[str, object]) -> str | None: + return encrypted_content_from_signature(_encrypted_reasoning_field(block)) + + def is_encrypted_reasoning_block(block: object) -> bool: """A thinking or redacted_thinking block carrying Responses API encrypted reasoning. @@ -1878,7 +1886,7 @@ def is_encrypted_reasoning_block(block: object) -> bool: if not isinstance(block, Mapping): return False mapping: Final = cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance - return _encrypted_content_of_block(mapping) is not None + return _carries_encrypted_reasoning(_encrypted_reasoning_field(mapping)) def _reasoning_replay_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str: diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index cb24f4153f9..ca520116606 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1203,6 +1203,8 @@ def strip_thinking_blocks_from_anthropic_messages(messages: list[Any]) -> list[A def _without_encrypted_reasoning_blocks(message: dict) -> dict | None: # mutable-ok: Anthropic message payload shape + if not isinstance(message, Mapping): + return message content: Final = message.get("content") if not isinstance(content, list): return message diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index db890662132..aaffc84b872 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -113,6 +113,7 @@ from litellm.litellm_core_utils.reasoning_effort_utils import ( from litellm.llms.anthropic.common_utils import ( is_empty_unsigned_thinking_block, normalize_anthropic_tool_use_id, + strip_encrypted_reasoning_blocks_from_anthropic_messages, ) from litellm.llms.anthropic.experimental_pass_through.context_management import ( PolyfillResult, @@ -417,7 +418,8 @@ class LiteLLMAnthropicMessagesAdapter: model: str | None = None, ) -> list: new_messages: Final[list[AllMessageValues]] = [] - for m in messages: + replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) + for m in replayable_messages: user_message: ChatCompletionUserMessage | None = None tool_message_list: list[ChatCompletionToolMessage] = [] new_user_content_list: list[ChatCompletionTextObject | ChatCompletionImageObject] = [] diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index 0445c23ed8c..24825cc87e0 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -19,6 +19,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.utils import ProviderConfigManager from ..utils import litellm_logging_obj_from_kwargs, local_model_name from .streaming_iterator import AnthropicResponsesStreamWrapper @@ -34,6 +35,15 @@ def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str, return extra_kwargs or {} +def _provider_returns_encrypted_reasoning(model: str, custom_llm_provider: object) -> bool: + provider: Final = ( + custom_llm_provider if isinstance(custom_llm_provider, str) else litellm.get_llm_provider(model=model)[1] + ) + provider_model: Final = local_model_name(model, provider) + responses_config: Final = ProviderConfigManager.get_provider_responses_api_config(provider, provider_model) + return responses_config is not None and "include" in responses_config.get_supported_openai_params(provider_model) + + def _build_responses_kwargs( *, max_tokens: int, @@ -85,8 +95,13 @@ def _build_responses_kwargs( request_data["output_format"] = output_format anthropic_request: Final = AnthropicMessagesRequest(**request_data) - responses_kwargs: Final = _ADAPTER.translate_request(anthropic_request) forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs) + responses_kwargs: Final = _ADAPTER.translate_request( + anthropic_request, + include_encrypted_reasoning=_provider_returns_encrypted_reasoning( + model, forwarded_kwargs.get("custom_llm_provider") + ), + ) # Normalize reasoning effort based on model capabilities # (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported) diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 07419db443c..1fdb0318bab 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -505,10 +505,16 @@ class LiteLLMAnthropicToResponsesAPIAdapter: def translate_request( self, anthropic_request: AnthropicMessagesRequest, + include_encrypted_reasoning: bool = True, ) -> dict[str, Any]: """ Translate a full Anthropic /v1/messages request dict to litellm.responses() / litellm.aresponses() kwargs. + + ``include_encrypted_reasoning`` asks the provider for ``reasoning.encrypted_content`` + on every call, so a reasoning model's items can be replayed intact next turn even + when the client sent no ``thinking`` block; pass False for a provider whose + Responses API rejects ``include``. """ model: Final[str] = anthropic_request["model"] messages_list: Final = cast( @@ -538,6 +544,8 @@ class LiteLLMAnthropicToResponsesAPIAdapter: "model": model, "input": input_items, } + if include_encrypted_reasoning: + responses_kwargs["include"] = [RESPONSES_INCLUDE_ENCRYPTED_REASONING] # mutable-ok: API request payload if system and not developer_parts: if isinstance(system, str): @@ -582,7 +590,6 @@ class LiteLLMAnthropicToResponsesAPIAdapter: ) if reasoning: responses_kwargs["reasoning"] = reasoning - responses_kwargs["include"] = [RESPONSES_INCLUDE_ENCRYPTED_REASONING] # mutable-ok: json list # output_format / output_config.format -> text format # output_format: {"type": "json_schema", "schema": {...}} diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 9df124576a0..ecb07d8c937 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -7,6 +7,7 @@ import pytest from litellm.litellm_core_utils.prompt_templates.common_utils import ( + ENCRYPTED_REASONING_SIGNATURE_PREFIX, TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, add_system_prompt_to_messages, @@ -1610,6 +1611,8 @@ class TestEncryptedReasoningReplay: [ ({"type": "thinking", "thinking": "x", "signature": encrypted_reasoning_signature("g")}, True), ({"type": "redacted_thinking", "data": encrypted_reasoning_signature("g")}, True), + ({"type": "thinking", "thinking": "x", "signature": ENCRYPTED_REASONING_SIGNATURE_PREFIX}, True), + ({"type": "redacted_thinking", "data": ENCRYPTED_REASONING_SIGNATURE_PREFIX}, True), ({"type": "thinking", "thinking": "x", "signature": "ErcBCkgIValid"}, False), ({"type": "redacted_thinking", "data": "EmwKAhgBEgy"}, False), ({"type": "text", "text": encrypted_reasoning_signature("g")}, False), 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 30465ca25ba..279ed7a8b33 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 @@ -9,6 +9,7 @@ import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_PLACEHOLDER, + encrypted_reasoning_signature, ) from litellm.litellm_core_utils.prompt_templates.factory import ( THOUGHT_SIGNATURE_SEPARATOR, @@ -407,6 +408,43 @@ def test_translate_anthropic_messages_to_openai_thinking_blocks(): assert result[1]["tool_calls"][0]["id"] == "toolu_01234" +def test_translate_anthropic_messages_to_openai_drops_bridge_encrypted_reasoning_blocks(): + """A session that moves from an OpenAI reasoning model to a chat provider replays reasoning only OpenAI can read. + + Gemini rejects the whole request when such a block reaches it as a thought_signature, so the + adapter drops those blocks and keeps the provider-signed ones. + """ + + anthropic_messages = [ + AnthropicMessagesUserMessageParam( + role="user", + content=[{"type": "text", "text": "Who drinks water?"}], + ), + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, + {"type": "text", "text": "The Norwegian."}, + ], + ), + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + {"type": "thinking", "thinking": "native", "signature": "EqQBCkYIAxgCIkA_signed"}, + {"type": "text", "text": "Still the Norwegian."}, + ], + ), + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=anthropic_messages) + + assert [m["role"] for m in result] == ["user", "assistant", "assistant"] + assert not result[1].get("thinking_blocks") + assert result[1]["content"] == "The Norwegian." + assert [b["signature"] for b in result[2]["thinking_blocks"]] == ["EqQBCkYIAxgCIkA_signed"] + + def test_translate_anthropic_messages_to_openai_sets_reasoning_content(): """Reasoning-aware chat providers read reasoning_content, so thinking text must land there. diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py index 3383813245a..2cbfaa17d23 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py @@ -54,6 +54,29 @@ def test_build_responses_kwargs_prefers_explicit_prompt_cache_key_over_derived() assert responses_kwargs["prompt_cache_key"] == "explicit-key" +def test_build_responses_kwargs_asks_openai_for_encrypted_reasoning_without_thinking(): + responses_kwargs = _build_responses_kwargs( + max_tokens=1024, + messages=MESSAGES, + model="openai/gpt-5.6-luna", + extra_kwargs={"custom_llm_provider": "openai"}, + ) + assert responses_kwargs["include"] == ["reasoning.encrypted_content"] + assert "reasoning" not in responses_kwargs + + +def test_build_responses_kwargs_skips_include_for_a_responses_provider_that_rejects_it(): + responses_kwargs = _build_responses_kwargs( + max_tokens=1024, + messages=MESSAGES, + model="perplexity/sonar", + thinking={"type": "enabled", "budget_tokens": 4096}, + extra_kwargs={"custom_llm_provider": "perplexity"}, + ) + assert "include" not in responses_kwargs + assert "reasoning" in responses_kwargs + + def test_build_responses_kwargs_without_metadata_sets_no_prompt_cache_key(): responses_kwargs = _build_responses_kwargs( max_tokens=1024, 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 79eaf6b60f2..9b7b15ace61 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 @@ -1162,7 +1162,6 @@ class TestTranslateRequestBroaderCoverage: req = _make_request(thinking={"type": "disabled"}) kwargs = _ADAPTER.translate_request(req) assert "reasoning" not in kwargs - assert "include" not in kwargs def test_thinking_asks_for_the_encrypted_reasoning(self): """The documented way to get reasoning that survives store=false is to ask for it.""" @@ -1170,6 +1169,17 @@ class TestTranslateRequestBroaderCoverage: kwargs = _ADAPTER.translate_request(req) assert kwargs["include"] == ["reasoning.encrypted_content"] + def test_encrypted_reasoning_is_asked_for_without_a_thinking_block(self): + """A reasoning model reasons whether or not the client sent `thinking`, so the replay needs it either way.""" + kwargs = _ADAPTER.translate_request(_make_request()) + assert kwargs["include"] == ["reasoning.encrypted_content"] + + def test_encrypted_reasoning_is_not_asked_for_when_the_provider_rejects_include(self): + req = _make_request(thinking={"type": "enabled", "budget_tokens": 12000}) + kwargs = _ADAPTER.translate_request(req, include_encrypted_reasoning=False) + assert kwargs["reasoning"] == {"effort": "high"} + assert "include" not in kwargs + def test_metadata_user_id_mapped_to_user(self): req = _make_request(metadata={"user_id": "user-42"}) kwargs = _ADAPTER.translate_request(req) 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 97d22274b1c..91652c89092 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -1600,6 +1600,15 @@ class TestAnthropicThinkingSignatureSelfHeal: assert len(msgs[1]["content"]) == 2 assert len(msgs[2]["content"]) == 4 + def test_strip_encrypted_reasoning_leaves_malformed_messages_for_the_provider_to_reject(self): + """A bare string in messages must reach Anthropic as a 400, not die in the stripper as a 500.""" + from litellm.llms.anthropic.common_utils import ( + strip_encrypted_reasoning_blocks_from_anthropic_messages, + ) + + msgs = ["hi", {"role": "user", "content": "hello"}] + assert strip_encrypted_reasoning_blocks_from_anthropic_messages(msgs) == msgs + def test_strip_empty_text_blocks_treats_null_text_as_empty(self): from litellm.llms.anthropic.common_utils import ( strip_empty_content_blocks_from_anthropic_messages, From babf18f37f383b99e697344019482ff97dfea1c7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:20:56 -0700 Subject: [PATCH 055/107] test(token_counter): type the project ITPM lag test's request body --- tests/test_litellm/proxy/hooks/test_tpm_concurrent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 398abf60b54..bdaca9ffc2d 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -3682,7 +3682,7 @@ async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loo handler, _cache = rate_limiter stash = get_or_create_request_stash() warm_tokenizer("claude-fable-5") - data: Dict[str, Any] = { + data: dict[str, object] = { "model": "claude-fable-5", "messages": [{"role": "user", "content": text * 100}], } From cce1d2087b9ae997ec406d6d38857642f1461844 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:22:22 -0700 Subject: [PATCH 056/107] fix(responses): echo a named tool_choice in the Responses API shape on the chat-completions bridge A streamed /v1/responses request with tool_choice {"type": "function", "name": ...} that reaches a chat-completions-only deployment failed with HTTP 500 before the first byte: the synthetic response.created and response.in_progress events copied the chat-shaped tool_choice into ResponsesAPIResponse, whose ToolChoice type expects the flat Responses API shape. The non-streamed path echoed "auto" regardless of the request. Both paths now normalize the request's tool_choice through the existing chat transform and map it back to the Responses API vocabulary, validated by a TypeAdapter(ToolChoice), so a named function is echoed as {"type": "function", "name": ...} and a missing tool_choice is echoed as "auto". Fixes #33689 --- .../streaming_iterator.py | 11 ++-- .../transformation.py | 18 ++++++- .../test_litellm_completion_responses.py | 52 ++++++++++++++++++ .../test_streaming_iterator_transformation.py | 53 +++++++++++++++++++ 4 files changed, 126 insertions(+), 8 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index d7f8cd8f8bd..660dd8f0c92 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -437,14 +437,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): response_created_event_data["temperature"] = self.responses_api_request["temperature"] if "text" in self.responses_api_request: response_created_event_data["text"] = self.responses_api_request["text"] - if "tool_choice" in self.responses_api_request: - # Transform tool_choice from dict format (e.g., {"type": "auto"}) to string format - response_created_event_data["tool_choice"] = ( - LiteLLMCompletionResponsesConfig._transform_tool_choice(self.responses_api_request["tool_choice"]) - or "auto" + response_created_event_data["tool_choice"] = ( + LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + self.responses_api_request.get("tool_choice") ) - else: - response_created_event_data["tool_choice"] = "auto" + ) if "tools" in self.responses_api_request: response_created_event_data["tools"] = self.responses_api_request["tools"] else: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b2d1a69e0d8..878f493b58e 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -27,6 +27,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( ) from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_create_params import ResponseInputParam +from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam from openai.types.responses.tool_param import FunctionToolParam from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -68,6 +69,7 @@ from litellm.types.llms.openai import ( ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, ResponsesAPIStatus, + ToolChoice, ValidChatCompletionMessageContentTypes, ValidChatCompletionMessageContentTypesLiteral, ) @@ -126,6 +128,7 @@ _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) _TEXT_ADAPTER: Final = TypeAdapter(str) +_RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice) @runtime_checkable @@ -267,6 +270,17 @@ class LiteLLMCompletionResponsesConfig: # Return as-is for unknown formats return tool_choice + @staticmethod + def _transform_tool_choice_for_responses_api_response(tool_choice: object) -> ToolChoice: + normalized: Final = LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice) + match normalized: + case None: + return "auto" + case {"type": "function", "function": {"name": str(function_name)}}: + return ToolChoiceFunctionParam(type="function", name=function_name) + case _: + return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(normalized) + @staticmethod def _should_drop_derived_web_search_options(model: str, custom_llm_provider: str | None) -> bool: """ @@ -2263,7 +2277,9 @@ class LiteLLMCompletionResponsesConfig: ), parallel_tool_calls=getattr(chat_completion_response, "parallel_tool_calls", False), temperature=getattr(chat_completion_response, "temperature", 0), - tool_choice=getattr(chat_completion_response, "tool_choice", "auto"), + tool_choice=LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + responses_api_request.get("tool_choice") + ), tools=getattr(chat_completion_response, "tools", []), top_p=getattr(chat_completion_response, "top_p", None), max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None), diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 2068f10ea2d..2aab660b5b7 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1421,6 +1421,58 @@ class TestToolChoiceTransformation: ) assert result == "required" + @pytest.mark.parametrize( + "request_tool_choice,expected", + [ + ({"type": "function", "name": "run_command"}, {"type": "function", "name": "run_command"}), + ({"type": "function", "function": {"name": "run_command"}}, {"type": "function", "name": "run_command"}), + ({"type": "custom", "name": "ApplyPatch"}, {"type": "function", "name": "ApplyPatch"}), + ({"type": "tool"}, "required"), + ({"type": "auto"}, "auto"), + ("required", "required"), + ("none", "none"), + (None, "auto"), + ], + ) + def test_transform_tool_choice_for_responses_api_response(self, request_tool_choice, expected): + result = LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + request_tool_choice + ) + assert result == expected + + def test_non_streamed_response_echoes_named_tool_choice_in_responses_api_shape(self): + chat_completion_response = ModelResponse( + id="chatcmpl-named-tool-choice", + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion", + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_pwd", + type="function", + function=Function(name="run_command", arguments='{"command":"pwd"}'), + ) + ], + ), + ) + ], + ) + + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Run the command pwd.", + responses_api_request={"tool_choice": {"type": "function", "name": "run_command"}}, + chat_completion_response=chat_completion_response, + ) + + assert responses_api_response.tool_choice == {"type": "function", "name": "run_command"} + class TestContentTypeTransformation: """Test content type transformation from Responses API to Chat Completion format""" diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 719d51c11e3..d4b565f82a1 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -628,3 +628,56 @@ def test_streamed_anthropic_tool_call_events_correlate_on_normalized_item_id(): assert item_dones[0].item.call_id == "toolu_01AbCdEf" for evt in deltas + dones: assert evt.item_id == added[0].item.id + + +def _tool_call_chunk(finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_pwd", + "type": "function", + "function": {"name": "run_command", "arguments": '{"command":"pwd"}'}, + "index": 0, + } + ], + ), + finish_reason=finish_reason, + ) + ], + ) + + +def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape(): + iterator = LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper([_tool_call_chunk(finish_reason="tool_calls")]), + request_input="Run the command pwd.", + responses_api_request={ + "tools": [{"type": "function", "name": "run_command", "parameters": {"type": "object"}}], + "tool_choice": {"type": "function", "name": "run_command"}, + }, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + events = list(iterator) + + response_events = [event for event in events if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES] + assert [event.type for event in response_events] == [ + "response.created", + "response.in_progress", + "response.completed", + ] + assert response_events[0].response.tool_choice == {"type": "function", "name": "run_command"} + assert response_events[1].response.tool_choice == {"type": "function", "name": "run_command"} + assert any(getattr(event, "type", None) == "response.output_item.done" for event in events) From 3020a13e24f56a4068055eaadce039cd527f7faa Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:34:12 -0700 Subject: [PATCH 057/107] test(e2e): every spend row a virtual key writes joins its token across all write paths One aliased key owned by a user with an email drives chat, queued chat, messages, responses, embeddings, the Gemini passthrough, a batch file upload, and a batch create against a live proxy. Each row must carry api_key equal to the key's LiteLLM_VerificationToken.token and the alias in metadata, and /spend/logs?api_key= and /user/daily/activity must report the key with its alias and email. Health-check rows must keep the literal service-account key, and the batch cost row for a completed marker batch must join the key that created it. A re-hashed api_key (the v1.99.0 regression fixed by #39568 and #39572) now fails the Buildkite e2e stage naming the write path Resolves MAT-180 --- tests/e2e/CLAUDE.md | 5 +- .../coverage_registry/quota_management.yaml | 4 + tests/e2e/models.py | 2 + .../spend_tracking/conftest.py | 1 + .../spend_tracking/spend_e2e_client.py | 284 +++++++++++++- .../test_key_attribution_e2e.py | 349 ++++++++++++++++++ 6 files changed, 642 insertions(+), 3 deletions(-) create mode 100644 tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 89c04208d65..b92d84b8b25 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -181,13 +181,14 @@ quota_management... | team_multi_window | fallback | spend_counter chat_completions | stream | messages_bridge | embeddings | cache_hit | key_rollup | concurrent_burst | tags | end_user - | per_model | failure | spend_calculate | pagination + | per_model | failure | spend_calculate | pagination | key_attribution assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm | blocks_then_resets | resets_windows_independently | alerts_without_blocking | isolates_per_model | isolates_per_member | isolates_per_group | enforced_across_keys | routes_to_fallback | reseed_matches_db | reports_spend | logs_cost | zero_cost | matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows - | writes_failure_row | returns_cost | keeps_total + | writes_failure_row | returns_cost | keeps_total | joins_key | reports_alias_and_email + | health_rows_keep_service_account | batch_cost_joins_key e.g. quota_management.ratelimit.rpm.blocks_over_limit exercised_on=[chat_completions, messages] quota_management.budget.key.blocks_over_limit exercised_on=[chat_completions] ``` diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index d0afcaca848..7013b59b83d 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -58,3 +58,7 @@ - {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"} - {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"} - {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e stack does not turn it on yet"} +- {id: quota_management.spend_tracking.key_attribution.joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [joins_key], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Every spend row a virtual key writes across chat, queued chat, messages, responses, embeddings, the Gemini passthrough, file upload, and batch create carries api_key equal to the key's token hash and the key alias, the join the usage APIs depend on; a re-hashed token shows up as an unattributed key-hash-* row (#39568, #39572)"} +- {id: quota_management.spend_tracking.key_attribution.reports_alias_and_email, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [reports_alias_and_email], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native], source: "proxy/management_endpoints/internal_user_endpoints.py", rationale: "/spend/logs?api_key= returns every one of the key's rows with its alias and /user/daily/activity aggregates them under the key's token with key_alias and user_email; /spend/logs carries no email field, so the email is asserted on daily activity only"} +- {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"} +- {id: quota_management.spend_tracking.key_attribution.batch_cost_joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [batch_cost_joins_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The _batch_cost row the batch cost poller writes once a batch completes carries the submitting key's token hash; it rides a cross-run marker batch whose metadata records the expected hash, since completion can lag by up to the polling interval. Cold start with no completed marker is a documented vacuous pass"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 62810e6cfd9..faf8557498b 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -84,6 +84,7 @@ class KeyGenerateBody(BaseModel): class KeyGenerateResponse(BaseModel): key: str + token: str | None = None key_alias: str | None = None models: list[str] = [] max_budget: float | None = None @@ -672,6 +673,7 @@ class GuardrailRunRecord(BaseModel): class SpendLogMetadata(BaseModel): + user_api_key_alias: str | None = None applied_guardrails: list[str] | None = None guardrail_information: list[GuardrailRunRecord] | None = None diff --git a/tests/e2e/quota_management/spend_tracking/conftest.py b/tests/e2e/quota_management/spend_tracking/conftest.py index 0597c9af400..9c8ffd18144 100644 --- a/tests/e2e/quota_management/spend_tracking/conftest.py +++ b/tests/e2e/quota_management/spend_tracking/conftest.py @@ -36,6 +36,7 @@ DRIVER_MODELS: tuple[tuple[str, str, str], ...] = ( ("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "ANTHROPIC_API_KEY"), ("openai-text-embedding-3-small", "openai/text-embedding-3-small", "OPENAI_API_KEY"), ("openai-responses-codex", "openai/gpt-5.3-codex", "OPENAI_API_KEY"), + ("openai-gpt-4o-mini", "openai/gpt-4o-mini", "OPENAI_API_KEY"), ) diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 056799b8499..42ebfdb98af 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -15,9 +15,14 @@ import time from collections.abc import Callable from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from typing import Final + +from pydantic import BaseModel, Field from e2e_config import unique_marker from e2e_http import ( + FileUploadForm, + Headers, NoBody, ProbeResult, Result, @@ -35,6 +40,8 @@ from models import ( DateRangeParams, EmbedBody, EmbedResponse, + KeyGenerateBody, + KeyGenerateResponse, OpenAPISchema, SpendCalculateBody, SpendCalculateResponse, @@ -43,13 +50,24 @@ from models import ( SpendLogsPageParams, SpendTagsResponse, TagSpend, + UserDeleteBody, + UserDeleteResponse, + UserNewBody, + UserNewResponse, + UserRole, ) -from proxy_client import ProxyClient +from proxy_client import Converged, ProxyClient, await_converged __all__ = [ + "BatchCreateBody", + "BatchObject", + "DailyActivityKeyBreakdown", + "FileObject", "ProbeResult", + "ResponseIdentity", "SpendClient", "SpendLogRow", + "StreamingResponse", "build_client", "is_ok", "unique_marker", @@ -57,6 +75,112 @@ __all__ = [ ] +class GeminiApiKeyHeaders(Headers): + x_goog_api_key: str = Field(serialization_alias="x-goog-api-key") + content_type: str = Field(default="application/json", serialization_alias="Content-Type") + + +class GeminiPart(BaseModel): + text: str + + +class GeminiContent(BaseModel): + parts: list[GeminiPart] + + +class GeminiGenerationConfig(BaseModel): + maxOutputTokens: int + + +class GeminiGenerateBody(BaseModel): + contents: list[GeminiContent] + generationConfig: GeminiGenerationConfig + + +class ResponsesBody(BaseModel): + model: str + input: str + cache: dict[str, bool] | None = {"no-cache": True} + + +class QueuedChatBody(ChatBody): + priority: int = 0 + + +class ResponseIdentity(BaseModel): + id: str | None = None + + +class HealthParams(BaseModel): + model: str + + +class ModelQuery(BaseModel): + model: str + + +class FileObject(BaseModel): + id: str + + +class BatchCreateBody(BaseModel): + input_file_id: str + endpoint: str = "/v1/chat/completions" + completion_window: str = "24h" + model: str + metadata: dict[str, str] + + +class BatchObject(BaseModel): + id: str + status: str + metadata: dict[str, str] | None = None + created_at: int | None = None + + +class BatchList(BaseModel): + data: list[BatchObject] = [] + + +class BatchListQuery(BaseModel): + model: str + limit: int + + +class DailyActivityParams(BaseModel): + start_date: str + end_date: str + api_key: str + + +class DailyActivityKeyMetadata(BaseModel): + key_alias: str | None = None + team_id: str | None = None + user_email: str | None = None + + +class DailyActivityKeyMetrics(BaseModel): + api_requests: int = 0 + + +class DailyActivityKeyBreakdown(BaseModel): + metrics: DailyActivityKeyMetrics + metadata: DailyActivityKeyMetadata + + +class DailyActivityBreakdown(BaseModel): + api_keys: dict[str, DailyActivityKeyBreakdown] = {} + + +class DailyActivityRow(BaseModel): + date: str + breakdown: DailyActivityBreakdown + + +class DailyActivityResponse(BaseModel): + results: list[DailyActivityRow] = [] + + def _chat_body( model: str, content: str, @@ -207,6 +331,164 @@ class SpendClient: def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult: return self.proxy.transport.probe(path, params=params) + def create_user(self, *, email: str, role: UserRole, user_id: str) -> str: + return unwrap( + self.proxy.transport.post( + "/user/new", + headers=self.proxy.transport.master, + json=UserNewBody(user_email=email, user_role=role, user_id=user_id), + response_type=UserNewResponse, + ) + ).user_id + + def delete_user(self, user_id: str) -> None: + _ = self.proxy.transport.post( + "/user/delete", + headers=self.proxy.transport.master, + json=UserDeleteBody(user_ids=[user_id]), + response_type=UserDeleteResponse, + ) + + def generate_key_record(self, body: KeyGenerateBody) -> KeyGenerateResponse: + return unwrap( + self.proxy.transport.post( + "/key/generate", + headers=self.proxy.transport.master, + json=body, + response_type=KeyGenerateResponse, + ) + ) + + def send_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse: + return self.proxy.transport.send( + "/chat/completions", + headers=self.proxy.transport.bearer(key), + json=_chat_body(model, content, max_tokens=max_tokens), + ) + + def send_queued_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse: + return self.proxy.transport.send( + "/queue/chat/completions", + headers=self.proxy.transport.bearer(key), + json=QueuedChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + ), + ) + + def send_messages(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse: + return self.proxy.transport.send( + "/v1/messages", + headers=self.proxy.transport.bearer(key), + json=AnthropicMessagesBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + ), + ) + + def send_responses(self, key: str, model: str, content: str) -> StreamingResponse: + return self.proxy.transport.send( + "/v1/responses", + headers=self.proxy.transport.bearer(key), + json=ResponsesBody(model=model, input=content), + ) + + def send_embed(self, key: str, model: str, content: str) -> StreamingResponse: + return self.proxy.transport.send( + "/embeddings", + headers=self.proxy.transport.bearer(key), + json=EmbedBody(model=model, input=content), + ) + + def send_gemini_generate(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse: + return self.proxy.transport.send( + f"/gemini/v1beta/models/{model}:generateContent", + headers=GeminiApiKeyHeaders(x_goog_api_key=key), + json=GeminiGenerateBody( + contents=[GeminiContent(parts=[GeminiPart(text=content)])], + generationConfig=GeminiGenerationConfig(maxOutputTokens=max_tokens), + ), + ) + + def upload_batch_file(self, key: str, model: str, content: bytes) -> FileObject: + return unwrap( + self.proxy.transport.upload( + "/v1/files", + headers=self.proxy.transport.bearer(key), + form=FileUploadForm(purpose="batch"), + filename="key_attribution.jsonl", + content=content, + params=ModelQuery(model=model), + response_type=FileObject, + ) + ) + + def create_batch(self, key: str, body: BatchCreateBody) -> BatchObject: + return unwrap( + self.proxy.transport.post( + "/v1/batches", + headers=self.proxy.transport.bearer(key), + json=body, + response_type=BatchObject, + ) + ) + + def list_batches(self, key: str, model: str, *, limit: int) -> list[BatchObject]: + return unwrap( + self.proxy.transport.get( + "/v1/batches", + headers=self.proxy.transport.bearer(key), + params=BatchListQuery(model=model, limit=limit), + response_type=BatchList, + ) + ).data + + def retrieve_batch(self, key: str, batch_id: str) -> BatchObject: + return unwrap( + self.proxy.transport.get( + f"/v1/batches/{batch_id}", + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=BatchObject, + ) + ) + + def health(self, model: str) -> ProbeResult: + return self.proxy.transport.probe("/health", params=HealthParams(model=model)) + + def daily_activity_for_key(self, token: str, *, start: datetime, end: datetime) -> DailyActivityKeyBreakdown | None: + response: Final = unwrap( + self.proxy.transport.get( + "/user/daily/activity", + headers=self.proxy.transport.master, + params=DailyActivityParams( + start_date=start.strftime("%Y-%m-%d"), + end_date=end.strftime("%Y-%m-%d"), + api_key=token, + ), + response_type=DailyActivityResponse, + ) + ) + return next( + (row.breakdown.api_keys[token] for row in response.results if token in row.breakdown.api_keys), + None, + ) + + def poll_daily_activity_for_key( + self, token: str, *, start: datetime, end: datetime, min_requests: int + ) -> DailyActivityKeyBreakdown | None: + outcome: Final = await_converged( + lambda: self.daily_activity_for_key(token, start=start, end=end), + converged=lambda found: found is not None and found.metrics.api_requests >= min_requests, + timeout=self.proxy.poll_timeout, + interval=self.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + def openapi(self) -> OpenAPISchema: return unwrap( self.proxy.transport.get( diff --git a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py new file mode 100644 index 00000000000..986033dbbde --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py @@ -0,0 +1,349 @@ +"""Every spend row a live proxy writes joins its virtual key (MAT-180). + +One virtual key with an alias, owned by a user with an email, drives every spend +write path a key can reach: /chat/completions, /queue/chat/completions, +/v1/messages, /v1/responses, /embeddings, the Gemini native passthrough, a batch +input file upload, and a batch create. Each row those calls write must carry +`api_key` equal to the key's LiteLLM_VerificationToken.token (the sha256 hash +/key/generate returns as `token`), which is the join /spend/logs?api_key= and +/user/daily/activity rely on to report key_alias and user_email. A row keyed by a +re-hashed token (v1.99.0's regression, #39568 and #39572) shows up as a +key-hash-* row with no alias and no email in the customer's usage exports. + +The health-check service account writes rows too; those must stay keyed by the +literal service-account name, never by a hash of it. The batch cost row is +written by the CheckBatchCost poller once the batch completes, up to an hour +later, so it rides a cross-run baton like the batches suite: each run submits a +one-line marker batch whose metadata records the token it expects on the cost +row, and asserts on the newest completed marker from any run (a cold start with +no completed marker is a documented vacuous pass, never a skip). + +/spend/logs carries no email field, so the email assertion lives on +/user/daily/activity alone; /spend/logs is held to the alias in metadata. +""" + +import time +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Final, Iterator + +import pytest + +from models import ChatMessage, KeyGenerateBody +from proxy_client import Converged, await_converged +from spend_e2e_client import ( + BatchCreateBody, + BatchObject, + DailyActivityKeyBreakdown, + ResponseIdentity, + SpendClient, + SpendLogRow, + StreamingResponse, + unique_marker, +) +from pydantic import BaseModel + +pytestmark = pytest.mark.e2e + +CHAT_MODEL: Final = "gemini-2.5-flash" +MESSAGES_MODEL: Final = "claude-haiku-4-5" +RESPONSES_MODEL: Final = "openai-responses-codex" +EMBED_MODEL: Final = "openai-text-embedding-3-small" +BATCH_MODEL: Final = "openai-gpt-4o-mini" +BATCH_BACKEND_MODEL: Final = "gpt-4o-mini" +HEALTH_SERVICE_ACCOUNT: Final = "litellm-internal-health-check" +BATON_MARKER_KEY: Final = "litellm_e2e_suite" +BATON_MARKER_VALUE: Final = "key-attribution-baton" +BATON_EXPECTED_TOKEN_KEY: Final = "expected_api_key" +BATON_POLL_SECONDS: Final = 300.0 +BATON_POLL_INTERVAL_SECONDS: Final = 10.0 +BATON_LIST_LIMIT: Final = 100 +MAX_TOKENS: Final = 8 +WRITE_PATHS: Final = ( + "chat_completions", + "queue_chat_completions", + "messages", + "responses", + "embeddings", + "gemini_passthrough", + "batch_file_upload", + "batch_create", +) + + +class BatchLineBody(BaseModel): + model: str + messages: list[ChatMessage] + max_tokens: int + + +class BatchLine(BaseModel): + custom_id: str + method: str = "POST" + url: str = "/v1/chat/completions" + body: BatchLineBody + + +@dataclass(frozen=True, slots=True) +class AttributedKey: + key: str + token: str + alias: str + email: str + user_id: str + + +@dataclass(frozen=True, slots=True) +class WritePath: + name: str + request_id: str + + +@dataclass(frozen=True, slots=True) +class DrivenKey: + identity: AttributedKey + paths: tuple[WritePath, ...] + started_at: datetime + + +def _body_id(name: str, sent: StreamingResponse) -> WritePath: + assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}" + response_id: Final = ResponseIdentity.model_validate_json(sent.body).id + assert response_id, f"{name} answered without a response id: {sent.body[:300]}" + return WritePath(name=name, request_id=response_id) + + +def _call_id(name: str, sent: StreamingResponse) -> WritePath: + assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}" + assert sent.call_id, f"{name} answered without an x-litellm-call-id header" + return WritePath(name=name, request_id=sent.call_id) + + +def _batch_jsonl(marker: str) -> bytes: + line: Final = BatchLine( + custom_id=marker, + body=BatchLineBody( + model=BATCH_BACKEND_MODEL, + messages=[ChatMessage(role="user", content=f"Reply with the word ok. {marker}")], + max_tokens=MAX_TOKENS, + ), + ) + return f"{line.model_dump_json()}\n".encode() + + +def _drive_batch(client: SpendClient, identity: AttributedKey, marker: str) -> tuple[WritePath, WritePath]: + uploaded: Final = client.upload_batch_file(identity.key, BATCH_MODEL, _batch_jsonl(marker)) + created: Final = client.create_batch( + identity.key, + BatchCreateBody( + input_file_id=uploaded.id, + model=BATCH_MODEL, + metadata={ + BATON_MARKER_KEY: BATON_MARKER_VALUE, + BATON_EXPECTED_TOKEN_KEY: identity.token, + "run": marker, + }, + ), + ) + return ( + WritePath(name="batch_file_upload", request_id=uploaded.id), + WritePath(name="batch_create", request_id=created.id), + ) + + +def _drive_every_write_path(client: SpendClient, identity: AttributedKey) -> tuple[WritePath, ...]: + marker: Final = unique_marker() + prompt: Final = f"Reply with the word ok. {marker}" + key: Final = identity.key + return ( + _body_id("chat_completions", client.send_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), + _body_id("queue_chat_completions", client.send_queued_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), + _body_id("messages", client.send_messages(key, MESSAGES_MODEL, prompt, max_tokens=MAX_TOKENS)), + _body_id("responses", client.send_responses(key, RESPONSES_MODEL, prompt)), + _call_id("embeddings", client.send_embed(key, EMBED_MODEL, prompt)), + _call_id("gemini_passthrough", client.send_gemini_generate(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), + *_drive_batch(client, identity, marker), + ) + + +def _completed_baton(listed: list[BatchObject]) -> BatchObject | None: + return max( + ( + batch + for batch in listed + if batch.status == "completed" and (batch.metadata or {}).get(BATON_MARKER_KEY) == BATON_MARKER_VALUE + ), + key=lambda batch: batch.created_at or 0, + default=None, + ) + + +def _newest_completed_baton(client: SpendClient, key: str) -> BatchObject | None: + outcome: Final = await_converged( + lambda: _completed_baton(client.list_batches(key, BATCH_MODEL, limit=BATON_LIST_LIMIT)), + converged=lambda completed: completed is not None, + timeout=BATON_POLL_SECONDS, + interval=BATON_POLL_INTERVAL_SECONDS, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else None + + +def _health_rows_between(client: SpendClient, started_at: datetime) -> list[SpendLogRow]: + return [ + row + for row in client.proxy.spend_logs_window( + start=started_at - timedelta(minutes=1), end=datetime.now(timezone.utc) + timedelta(minutes=1) + ) + if HEALTH_SERVICE_ACCOUNT in (row.request_tags or []) + ] + + +def _health_rows_since(client: SpendClient, started_at: datetime) -> list[SpendLogRow]: + outcome: Final = await_converged( + lambda: _health_rows_between(client, started_at), + converged=lambda rows: bool(rows), + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + + +class TestKeyAttribution: + @pytest.fixture(scope="class") + def driven(self, client: SpendClient) -> Iterator[DrivenKey]: + marker: Final = unique_marker() + user_id: Final = client.create_user( + email=f"key-attribution-{marker}@example.com", + role="proxy_admin", + user_id=f"key-attribution-{marker}", + ) + record: Final = client.generate_key_record( + KeyGenerateBody(models=[], user_id=user_id, key_alias=f"key-attribution-{marker}") + ) + assert record.token, "/key/generate answered without the key's token hash" + assert record.key_alias, "/key/generate dropped the key alias" + identity: Final = AttributedKey( + key=record.key, + token=record.token, + alias=record.key_alias, + email=f"key-attribution-{marker}@example.com", + user_id=user_id, + ) + started_at: Final = datetime.now(timezone.utc) + try: + yield DrivenKey( + identity=identity, + paths=_drive_every_write_path(client, identity), + started_at=started_at, + ) + finally: + client.proxy.delete_key(identity.key) + client.delete_user(identity.user_id) + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.joins_key", + exercised_on=["chat_completions", "messages", "responses", "embeddings", "batches", "files", "google_native"], + ) + def test_every_write_path_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None: + assert tuple(path.name for path in driven.paths) == WRITE_PATHS + found: Final = tuple( + (path, client.proxy.poll_logs_for_request_id(path.request_id)) for path in driven.paths + ) + unwritten: Final = [path.name for path, rows in found if not rows] + assert not unwritten, f"write paths that produced no spend row within the poll window: {unwritten}" + unjoined: Final = [ + (path.name, row.call_type, row.api_key) + for path, rows in found + for row in rows + if row.api_key != driven.identity.token + ] + assert not unjoined, ( + "spend rows whose api_key does not join LiteLLM_VerificationToken.token " + f"{driven.identity.token}: {unjoined}" + ) + unaliased: Final = [ + (path.name, row.call_type, row.metadata.user_api_key_alias if row.metadata else None) + for path, rows in found + for row in rows + if row.metadata is None or row.metadata.user_api_key_alias != driven.identity.alias + ] + assert not unaliased, f"spend rows written without key alias {driven.identity.alias!r}: {unaliased}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.reports_alias_and_email", + exercised_on=["chat_completions", "messages", "responses", "embeddings", "batches", "files", "google_native"], + ) + def test_spend_logs_by_key_return_every_row_with_the_alias(self, client: SpendClient, driven: DrivenKey) -> None: + expected_ids: Final = frozenset(path.request_id for path in driven.paths) + rows: Final = client.poll_logs_for_key( + driven.identity.key, + min_rows=len(driven.paths), + predicate=lambda found: expected_ids <= frozenset(row.request_id or "" for row in found), + ) + missing: Final = expected_ids - frozenset(row.request_id or "" for row in rows) + assert not missing, ( + f"/spend/logs?api_key= does not return {len(missing)} of {len(expected_ids)} rows for the key: " + f"{sorted(path.name for path in driven.paths if path.request_id in missing)}" + ) + aliases: Final = frozenset(row.metadata.user_api_key_alias if row.metadata else None for row in rows) + assert aliases == {driven.identity.alias}, f"/spend/logs rows carry aliases {sorted(map(str, aliases))}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.reports_alias_and_email", + exercised_on=["chat_completions", "messages", "responses", "embeddings", "batches", "files", "google_native"], + ) + def test_user_daily_activity_reports_alias_and_email(self, client: SpendClient, driven: DrivenKey) -> None: + breakdown: Final[DailyActivityKeyBreakdown | None] = client.poll_daily_activity_for_key( + driven.identity.token, + start=driven.started_at - timedelta(days=1), + end=datetime.now(timezone.utc) + timedelta(days=1), + min_requests=len(driven.paths), + ) + assert breakdown is not None, ( + f"/user/daily/activity?api_key={driven.identity.token} has no api_keys breakdown: " + "the key's rows did not aggregate under its token" + ) + assert breakdown.metrics.api_requests >= len(driven.paths), ( + f"/user/daily/activity counts {breakdown.metrics.api_requests} requests for the key, " + f"expected at least {len(driven.paths)}" + ) + assert breakdown.metadata.key_alias == driven.identity.alias, f"key_alias={breakdown.metadata.key_alias!r}" + assert breakdown.metadata.user_email == driven.identity.email, f"user_email={breakdown.metadata.user_email!r}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.health_rows_keep_service_account", + exercised_on=["chat_completions"], + ) + def test_health_check_rows_keep_the_service_account_key(self, client: SpendClient) -> None: + started_at: Final = datetime.now(timezone.utc) + probe: Final = client.health(CHAT_MODEL) + assert probe.healthy, f"/health?model={CHAT_MODEL} answered {probe.status_code}: {probe.body[:300]}" + rows: Final = _health_rows_since(client, started_at) + assert rows, f"/health?model={CHAT_MODEL} wrote no {HEALTH_SERVICE_ACCOUNT}-tagged spend row" + rehashed: Final = [(row.request_id, row.api_key) for row in rows if row.api_key != HEALTH_SERVICE_ACCOUNT] + assert not rehashed, f"health-check rows keyed by something other than {HEALTH_SERVICE_ACCOUNT!r}: {rehashed}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.batch_cost_joins_key", + exercised_on=["batches"], + ) + def test_completed_batch_cost_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None: + completed: Final = _newest_completed_baton(client, driven.identity.key) + if completed is None: + return + expected_token: Final = (completed.metadata or {}).get(BATON_EXPECTED_TOKEN_KEY) + assert expected_token, f"marker batch {completed.id} lost its expected token metadata: {completed.metadata}" + fetched: Final = client.retrieve_batch(driven.identity.key, completed.id) + assert fetched.status == "completed", f"listed-completed marker retrieved as {fetched.status!r}" + rows: Final = client.proxy.poll_logs_for_request_id( + f"{fetched.id}_batch_cost", + predicate=lambda found: any((row.spend or 0) > 0 for row in found), + ) + priced: Final = [row for row in rows if (row.spend or 0) > 0] + assert priced, f"completed batch {fetched.id} has no positive-cost spend row under {fetched.id}_batch_cost" + unjoined: Final = [(row.call_type, row.api_key) for row in priced if row.api_key != expected_token] + assert not unjoined, f"batch cost rows whose api_key does not join the key's token {expected_token}: {unjoined}" From f5a410585af9fe2c06f0e5fa07ceddb2e0be46e0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:44:29 -0700 Subject: [PATCH 058/107] fix(responses): keep a custom tool_choice type in the bridged echo and type the new tests --- .../transformation.py | 23 +++++++++++++------ .../test_litellm_completion_responses.py | 17 +++++++++----- .../test_streaming_iterator_transformation.py | 9 ++++---- 3 files changed, 32 insertions(+), 17 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 878f493b58e..6b9931b7df2 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -27,9 +27,10 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( ) from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_create_params import ResponseInputParam +from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam from openai.types.responses.tool_param import FunctionToolParam -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger @@ -272,13 +273,21 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_tool_choice_for_responses_api_response(tool_choice: object) -> ToolChoice: - normalized: Final = LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice) - match normalized: - case None: - return "auto" - case {"type": "function", "function": {"name": str(function_name)}}: + if tool_choice is None: + return "auto" + try: + return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(tool_choice) + except ValidationError: + return LiteLLMCompletionResponsesConfig._chat_tool_choice_as_responses_api_tool_choice(tool_choice) + + @staticmethod + def _chat_tool_choice_as_responses_api_tool_choice(tool_choice: object) -> ToolChoice: + match tool_choice, LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice): + case {"type": "custom"}, {"function": {"name": str(custom_name)}}: + return ToolChoiceCustomParam(type="custom", name=custom_name) + case _, {"type": "function", "function": {"name": str(function_name)}}: return ToolChoiceFunctionParam(type="function", name=function_name) - case _: + case _, normalized: return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(normalized) @staticmethod diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 2aab660b5b7..8f2937cb252 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1,4 +1,5 @@ import json +from typing import Final import pytest @@ -1426,7 +1427,9 @@ class TestToolChoiceTransformation: [ ({"type": "function", "name": "run_command"}, {"type": "function", "name": "run_command"}), ({"type": "function", "function": {"name": "run_command"}}, {"type": "function", "name": "run_command"}), - ({"type": "custom", "name": "ApplyPatch"}, {"type": "function", "name": "ApplyPatch"}), + ({"type": "custom", "name": "ApplyPatch"}, {"type": "custom", "name": "ApplyPatch"}), + ({"type": "custom", "custom": {"name": "ApplyPatch"}}, {"type": "custom", "name": "ApplyPatch"}), + ({"type": "function"}, "required"), ({"type": "tool"}, "required"), ({"type": "auto"}, "auto"), ("required", "required"), @@ -1434,14 +1437,16 @@ class TestToolChoiceTransformation: (None, "auto"), ], ) - def test_transform_tool_choice_for_responses_api_response(self, request_tool_choice, expected): - result = LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + def test_transform_tool_choice_for_responses_api_response( + self, request_tool_choice: object, expected: str | dict[str, str] + ) -> None: + result: Final = LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( request_tool_choice ) assert result == expected - def test_non_streamed_response_echoes_named_tool_choice_in_responses_api_shape(self): - chat_completion_response = ModelResponse( + def test_non_streamed_response_echoes_named_tool_choice_in_responses_api_shape(self) -> None: + chat_completion_response: Final = ModelResponse( id="chatcmpl-named-tool-choice", created=1748575031, model="claude-haiku-4-5", @@ -1465,7 +1470,7 @@ class TestToolChoiceTransformation: ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + responses_api_response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( request_input="Run the command pwd.", responses_api_request={"tool_choice": {"type": "function", "name": "run_command"}}, chat_completion_response=chat_completion_response, diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index d4b565f82a1..29061db97dc 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -11,6 +11,7 @@ spend tracking stores, so a follow-up previous_response_id still finds the conve """ import json +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -657,8 +658,8 @@ def _tool_call_chunk(finish_reason: str | None = None) -> ModelResponseStream: ) -def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape(): - iterator = LiteLLMCompletionStreamingIterator( +def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape() -> None: + iterator: Final = LiteLLMCompletionStreamingIterator( model="claude-haiku-4-5", litellm_custom_stream_wrapper=_FakeStreamWrapper([_tool_call_chunk(finish_reason="tool_calls")]), request_input="Run the command pwd.", @@ -670,9 +671,9 @@ def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape(): litellm_metadata={}, ) - events = list(iterator) + events: Final = list(iterator) - response_events = [event for event in events if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES] + response_events: Final = [event for event in events if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES] assert [event.type for event in response_events] == [ "response.created", "response.in_progress", From 07a0ca107459c0aa73da789a82e87688bf8945b9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:54:11 -0700 Subject: [PATCH 059/107] fix(e2e): bill the batch cost baton through a raw-id retrieve and drive the callback replay writer The completed marker batch is now retrieved by its raw provider id with this run's key, so the proxy prices it inline and the {provider_batch_id}_batch_cost row must join that key's token and alias. The CheckBatchCost poller only bills batches it created in the same database, which a stack booted fresh per run never holds for a completed marker, so the old unified-id assertion had no row to find. Every run also replays one callback log through POST /v1/rust_control_plane/logs, the third spend writer, and asserts its row joins the key like the eight request paths --- .../coverage_registry/quota_management.yaml | 6 +- .../spend_tracking/spend_e2e_client.py | 57 +++++- .../test_key_attribution_e2e.py | 178 +++++++++++++----- 3 files changed, 192 insertions(+), 49 deletions(-) diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 7013b59b83d..d3966eed48f 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -58,7 +58,7 @@ - {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"} - {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"} - {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e stack does not turn it on yet"} -- {id: quota_management.spend_tracking.key_attribution.joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [joins_key], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Every spend row a virtual key writes across chat, queued chat, messages, responses, embeddings, the Gemini passthrough, file upload, and batch create carries api_key equal to the key's token hash and the key alias, the join the usage APIs depend on; a re-hashed token shows up as an unattributed key-hash-* row (#39568, #39572)"} -- {id: quota_management.spend_tracking.key_attribution.reports_alias_and_email, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [reports_alias_and_email], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native], source: "proxy/management_endpoints/internal_user_endpoints.py", rationale: "/spend/logs?api_key= returns every one of the key's rows with its alias and /user/daily/activity aggregates them under the key's token with key_alias and user_email; /spend/logs carries no email field, so the email is asserted on daily activity only"} +- {id: quota_management.spend_tracking.key_attribution.joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [joins_key], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Every spend row a virtual key writes across chat, queued chat, messages, responses, embeddings, the Gemini passthrough, file upload, batch create, and a replayed callback log carries api_key equal to the key's token hash and the key alias, the join the usage APIs depend on; a re-hashed token shows up as an unattributed key-hash-* row (#39568, #39572)"} +- {id: quota_management.spend_tracking.key_attribution.reports_alias_and_email, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [reports_alias_and_email], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/management_endpoints/internal_user_endpoints.py", rationale: "/spend/logs?api_key= returns every one of the key's rows with its alias and /user/daily/activity aggregates them under the key's token with key_alias and user_email; /spend/logs carries no email field, so the email is asserted on daily activity only"} - {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"} -- {id: quota_management.spend_tracking.key_attribution.batch_cost_joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [batch_cost_joins_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The _batch_cost row the batch cost poller writes once a batch completes carries the submitting key's token hash; it rides a cross-run marker batch whose metadata records the expected hash, since completion can lag by up to the polling interval. Cold start with no completed marker is a documented vacuous pass"} +- {id: quota_management.spend_tracking.key_attribution.batch_cost_joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [batch_cost_joins_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "Retrieving a completed batch by its raw provider id prices it inline against the retrieving key, so the {provider_batch_id}_batch_cost row must carry that key's token hash and alias; the target is the newest completed cross-run marker batch this proxy has not billed yet, since completion lags by up to 24h, and a cold start with no completed marker is a documented vacuous pass. The CheckBatchCost poller's own unified-id row needs the batch create row in the same database, which a stack booted fresh per run never holds for a completed marker"} diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 42ebfdb98af..18b714d8aeb 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -17,8 +17,6 @@ from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Final -from pydantic import BaseModel, Field - from e2e_config import unique_marker from e2e_http import ( FileUploadForm, @@ -57,9 +55,12 @@ from models import ( UserRole, ) from proxy_client import Converged, ProxyClient, await_converged +from pydantic import BaseModel, Field __all__ = [ "BatchCreateBody", + "CallbackLogMetadata", + "CallbackLogPayload", "BatchObject", "DailyActivityKeyBreakdown", "FileObject", @@ -147,6 +148,44 @@ class BatchListQuery(BaseModel): limit: int +class ProviderQuery(BaseModel): + provider: str + + +class CallbackLogMetadata(BaseModel): + user_api_key_hash: str + user_api_key_alias: str + user_api_key_user_id: str + + +class CallbackLogPayload(BaseModel): + id: str + litellm_call_id: str + model: str + call_type: str = "acompletion" + start_time: float = Field(serialization_alias="startTime") + end_time: float = Field(serialization_alias="endTime") + response_cost: float + prompt_tokens: int + completion_tokens: int + total_tokens: int + metadata: CallbackLogMetadata + + +class CallbackLogRecord(BaseModel): + status: str = "success" + standard_logging_payload: CallbackLogPayload + + +class CallbackLogsRequest(BaseModel): + records: list[CallbackLogRecord] + + +class CallbackLogsResponse(BaseModel): + processed: int + failed: int + + class DailyActivityParams(BaseModel): start_date: str end_date: str @@ -445,16 +484,26 @@ class SpendClient: ) ).data - def retrieve_batch(self, key: str, batch_id: str) -> BatchObject: + def retrieve_batch(self, key: str, batch_id: str, *, provider: str) -> BatchObject: return unwrap( self.proxy.transport.get( f"/v1/batches/{batch_id}", headers=self.proxy.transport.bearer(key), - params=NoBody(), + params=ProviderQuery(provider=provider), response_type=BatchObject, ) ) + def replay_callback_log(self, key: str, payload: CallbackLogPayload) -> CallbackLogsResponse: + return unwrap( + self.proxy.transport.post( + "/v1/rust_control_plane/logs", + headers=self.proxy.transport.bearer(key), + json=CallbackLogsRequest(records=[CallbackLogRecord(standard_logging_payload=payload)]), + response_type=CallbackLogsResponse, + ) + ) + def health(self, model: str) -> ProbeResult: return self.proxy.transport.probe("/health", params=HealthParams(model=model)) diff --git a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py index 986033dbbde..cac4fd7fb3c 100644 --- a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py @@ -3,7 +3,8 @@ One virtual key with an alias, owned by a user with an email, drives every spend write path a key can reach: /chat/completions, /queue/chat/completions, /v1/messages, /v1/responses, /embeddings, the Gemini native passthrough, a batch -input file upload, and a batch create. Each row those calls write must carry +input file upload, a batch create, and a replayed callback log (POST +/v1/rust_control_plane/logs, the writer an external gateway feeds). Each row those calls write must carry `api_key` equal to the key's LiteLLM_VerificationToken.token (the sha256 hash /key/generate returns as `token`), which is the join /spend/logs?api_key= and /user/daily/activity rely on to report key_alias and user_email. A row keyed by a @@ -11,29 +12,40 @@ re-hashed token (v1.99.0's regression, #39568 and #39572) shows up as a key-hash-* row with no alias and no email in the customer's usage exports. The health-check service account writes rows too; those must stay keyed by the -literal service-account name, never by a hash of it. The batch cost row is -written by the CheckBatchCost poller once the batch completes, up to an hour -later, so it rides a cross-run baton like the batches suite: each run submits a -one-line marker batch whose metadata records the token it expects on the cost -row, and asserts on the newest completed marker from any run (a cold start with -no completed marker is a documented vacuous pass, never a skip). +literal service-account name, never by a hash of it. A batch's cost row lands +only once the batch completes, up to 24h later, so the batch cost path rides a +cross-run baton like the batches suite: each run submits a one-line marker batch +and never cancels it, and the newest completed marker from any run that this +proxy has not billed yet is retrieved by its raw provider id with this run's +key. That retrieve is the writer under test: a raw id is never poller-owned, so +the proxy prices it inline against the retrieving key, and the +{provider_batch_id}_batch_cost row must join this run's token with its alias. +The CheckBatchCost poller's own row (the unified id, billed against the +submitting key) needs the batch's create row in the same database, which a stack +booted fresh per run never holds for a completed marker, so that writer is out +of this test's reach. A cold start with no completed marker is a documented +vacuous pass, never a skip. /spend/logs carries no email field, so the email assertion lives on /user/daily/activity alone; /spend/logs is held to the alias in metadata. """ +import base64 import time +from collections.abc import Iterator from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Final, Iterator +from typing import Final import pytest - -from models import ChatMessage, KeyGenerateBody +from models import ChatMessage, KeyGenerateBody, SpendLogsParams from proxy_client import Converged, await_converged +from pydantic import BaseModel from spend_e2e_client import ( BatchCreateBody, BatchObject, + CallbackLogMetadata, + CallbackLogPayload, DailyActivityKeyBreakdown, ResponseIdentity, SpendClient, @@ -41,7 +53,6 @@ from spend_e2e_client import ( StreamingResponse, unique_marker, ) -from pydantic import BaseModel pytestmark = pytest.mark.e2e @@ -51,14 +62,17 @@ RESPONSES_MODEL: Final = "openai-responses-codex" EMBED_MODEL: Final = "openai-text-embedding-3-small" BATCH_MODEL: Final = "openai-gpt-4o-mini" BATCH_BACKEND_MODEL: Final = "gpt-4o-mini" +BATCH_PROVIDER: Final = "openai" HEALTH_SERVICE_ACCOUNT: Final = "litellm-internal-health-check" BATON_MARKER_KEY: Final = "litellm_e2e_suite" BATON_MARKER_VALUE: Final = "key-attribution-baton" -BATON_EXPECTED_TOKEN_KEY: Final = "expected_api_key" -BATON_POLL_SECONDS: Final = 300.0 +BATON_POLL_SECONDS: Final = 30.0 BATON_POLL_INTERVAL_SECONDS: Final = 10.0 BATON_LIST_LIMIT: Final = 100 MAX_TOKENS: Final = 8 +REPLAY_RESPONSE_COST: Final = 0.0001 +REPLAY_PROMPT_TOKENS: Final = 5 +REPLAY_COMPLETION_TOKENS: Final = 1 WRITE_PATHS: Final = ( "chat_completions", "queue_chat_completions", @@ -68,6 +82,7 @@ WRITE_PATHS: Final = ( "gemini_passthrough", "batch_file_upload", "batch_create", + "callback_replay", ) @@ -138,11 +153,7 @@ def _drive_batch(client: SpendClient, identity: AttributedKey, marker: str) -> t BatchCreateBody( input_file_id=uploaded.id, model=BATCH_MODEL, - metadata={ - BATON_MARKER_KEY: BATON_MARKER_VALUE, - BATON_EXPECTED_TOKEN_KEY: identity.token, - "run": marker, - }, + metadata={BATON_MARKER_KEY: BATON_MARKER_VALUE, "run": marker}, ), ) return ( @@ -151,6 +162,32 @@ def _drive_batch(client: SpendClient, identity: AttributedKey, marker: str) -> t ) +def _drive_callback_replay(client: SpendClient, identity: AttributedKey, marker: str) -> WritePath: + request_id: Final = f"callback-replay-{marker}" + finished_at: Final = time.time() + replayed: Final = client.replay_callback_log( + identity.key, + CallbackLogPayload( + id=request_id, + litellm_call_id=request_id, + model=CHAT_MODEL, + start_time=finished_at - 1, + end_time=finished_at, + response_cost=REPLAY_RESPONSE_COST, + prompt_tokens=REPLAY_PROMPT_TOKENS, + completion_tokens=REPLAY_COMPLETION_TOKENS, + total_tokens=REPLAY_PROMPT_TOKENS + REPLAY_COMPLETION_TOKENS, + metadata=CallbackLogMetadata( + user_api_key_hash=identity.token, + user_api_key_alias=identity.alias, + user_api_key_user_id=identity.user_id, + ), + ), + ) + assert replayed.processed == 1 and replayed.failed == 0, f"callback replay rejected the payload: {replayed}" + return WritePath(name="callback_replay", request_id=request_id) + + def _drive_every_write_path(client: SpendClient, identity: AttributedKey) -> tuple[WritePath, ...]: marker: Final = unique_marker() prompt: Final = f"Reply with the word ok. {marker}" @@ -163,31 +200,52 @@ def _drive_every_write_path(client: SpendClient, identity: AttributedKey) -> tup _call_id("embeddings", client.send_embed(key, EMBED_MODEL, prompt)), _call_id("gemini_passthrough", client.send_gemini_generate(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), *_drive_batch(client, identity, marker), + _drive_callback_replay(client, identity, marker), ) -def _completed_baton(listed: list[BatchObject]) -> BatchObject | None: - return max( - ( - batch - for batch in listed - if batch.status == "completed" and (batch.metadata or {}).get(BATON_MARKER_KEY) == BATON_MARKER_VALUE - ), - key=lambda batch: batch.created_at or 0, - default=None, +def _completed_batons(listed: list[BatchObject]) -> tuple[BatchObject, ...]: + return tuple( + sorted( + ( + batch + for batch in listed + if batch.status == "completed" and (batch.metadata or {}).get(BATON_MARKER_KEY) == BATON_MARKER_VALUE + ), + key=lambda batch: batch.created_at or 0, + reverse=True, + ) ) -def _newest_completed_baton(client: SpendClient, key: str) -> BatchObject | None: +def _provider_batch_id(unified_batch_id: str) -> str: + encoded: Final = unified_batch_id.removeprefix("batch_") + decoded: Final = base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)).decode() + return decoded.removeprefix("litellm:").split(";", 1)[0] + + +def _batch_cost_request_id(unified_batch_id: str) -> str: + return f"{_provider_batch_id(unified_batch_id)}_batch_cost" + + +def _newest_unbilled_completed_baton(client: SpendClient, key: str) -> BatchObject | None: outcome: Final = await_converged( - lambda: _completed_baton(client.list_batches(key, BATCH_MODEL, limit=BATON_LIST_LIMIT)), - converged=lambda completed: completed is not None, + lambda: _completed_batons(client.list_batches(key, BATCH_MODEL, limit=BATON_LIST_LIMIT)), + converged=lambda completed: bool(completed), timeout=BATON_POLL_SECONDS, interval=BATON_POLL_INTERVAL_SECONDS, now=time.monotonic, sleep=time.sleep, ) - return outcome.result if isinstance(outcome, Converged) else None + completed: Final = outcome.result if isinstance(outcome, Converged) else () + return next( + ( + batch + for batch in completed + if not client.proxy.spend_logs(SpendLogsParams(request_id=_batch_cost_request_id(batch.id))) + ), + None, + ) def _health_rows_between(client: SpendClient, started_at: datetime) -> list[SpendLogRow]: @@ -246,7 +304,16 @@ class TestKeyAttribution: @pytest.mark.covers( "quota_management.spend_tracking.key_attribution.joins_key", - exercised_on=["chat_completions", "messages", "responses", "embeddings", "batches", "files", "google_native"], + exercised_on=[ + "chat_completions", + "messages", + "responses", + "embeddings", + "batches", + "files", + "google_native", + "rust_control_plane", + ], ) def test_every_write_path_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None: assert tuple(path.name for path in driven.paths) == WRITE_PATHS @@ -275,7 +342,16 @@ class TestKeyAttribution: @pytest.mark.covers( "quota_management.spend_tracking.key_attribution.reports_alias_and_email", - exercised_on=["chat_completions", "messages", "responses", "embeddings", "batches", "files", "google_native"], + exercised_on=[ + "chat_completions", + "messages", + "responses", + "embeddings", + "batches", + "files", + "google_native", + "rust_control_plane", + ], ) def test_spend_logs_by_key_return_every_row_with_the_alias(self, client: SpendClient, driven: DrivenKey) -> None: expected_ids: Final = frozenset(path.request_id for path in driven.paths) @@ -294,7 +370,16 @@ class TestKeyAttribution: @pytest.mark.covers( "quota_management.spend_tracking.key_attribution.reports_alias_and_email", - exercised_on=["chat_completions", "messages", "responses", "embeddings", "batches", "files", "google_native"], + exercised_on=[ + "chat_completions", + "messages", + "responses", + "embeddings", + "batches", + "files", + "google_native", + "rust_control_plane", + ], ) def test_user_daily_activity_reports_alias_and_email(self, client: SpendClient, driven: DrivenKey) -> None: breakdown: Final[DailyActivityKeyBreakdown | None] = client.poll_daily_activity_for_key( @@ -332,18 +417,27 @@ class TestKeyAttribution: exercised_on=["batches"], ) def test_completed_batch_cost_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None: - completed: Final = _newest_completed_baton(client, driven.identity.key) + completed: Final = _newest_unbilled_completed_baton(client, driven.identity.key) if completed is None: return - expected_token: Final = (completed.metadata or {}).get(BATON_EXPECTED_TOKEN_KEY) - assert expected_token, f"marker batch {completed.id} lost its expected token metadata: {completed.metadata}" - fetched: Final = client.retrieve_batch(driven.identity.key, completed.id) + provider_batch_id: Final = _provider_batch_id(completed.id) + fetched: Final = client.retrieve_batch(driven.identity.key, provider_batch_id, provider=BATCH_PROVIDER) assert fetched.status == "completed", f"listed-completed marker retrieved as {fetched.status!r}" + cost_request_id: Final = _batch_cost_request_id(completed.id) rows: Final = client.proxy.poll_logs_for_request_id( - f"{fetched.id}_batch_cost", + cost_request_id, predicate=lambda found: any((row.spend or 0) > 0 for row in found), ) priced: Final = [row for row in rows if (row.spend or 0) > 0] - assert priced, f"completed batch {fetched.id} has no positive-cost spend row under {fetched.id}_batch_cost" - unjoined: Final = [(row.call_type, row.api_key) for row in priced if row.api_key != expected_token] - assert not unjoined, f"batch cost rows whose api_key does not join the key's token {expected_token}: {unjoined}" + assert priced, f"retrieving completed batch {provider_batch_id} wrote no positive-cost row under {cost_request_id}" + unjoined: Final = [ + (row.call_type, row.api_key, row.metadata.user_api_key_alias if row.metadata else None) + for row in priced + if row.api_key != driven.identity.token + or row.metadata is None + or row.metadata.user_api_key_alias != driven.identity.alias + ] + assert not unjoined, ( + f"batch cost rows that do not join the retrieving key's token {driven.identity.token} " + f"with alias {driven.identity.alias!r}: {unjoined}" + ) From 96e46ba8ff43cdbf24d53691685c15b533bfa746 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:06:55 -0700 Subject: [PATCH 060/107] fix(responses): stamp streamed usage cost when the provider usage arrives as a dict Perplexity's Responses payload fails ResponsesAPIResponse validation on truncation "" and is kept as an unvalidated model, so its usage stays a plain dict and _stamp_responses_usage_cost raised AttributeError on every streamed completion once reasoning made the cost non-zero. Validate the dict into ResponseAPIUsage before stamping, keeping a provider-reported cost when it carries one. Resolves LIT-7391 --- litellm/responses/streaming_iterator.py | 17 ++++++- .../responses/test_streaming_iterator.py | 51 +++++++++++++++++++ 2 files changed, 66 insertions(+), 2 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 9f9016c5a7f..658a0e8ad5d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runti import httpx from openai._streaming import SSEDecoder +from pydantic import ValidationError from typing_extensions import TypeIs import litellm @@ -544,7 +545,7 @@ class BaseResponsesAPIStreamingIterator: def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None)) if usage_obj is None: return try: @@ -1293,14 +1294,26 @@ def _add_text_like_part_events( ) +def _usage_as_model(usage: object) -> ResponseAPIUsage | None: + if isinstance(usage, ResponseAPIUsage): + return usage + if not isinstance(usage, dict): + return None + try: + return ResponseAPIUsage.model_validate(usage) + except ValidationError: + return None + + def _stamp_responses_usage_cost( response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None ) -> None: if response_obj is None or logging_obj is None: return - usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None)) if usage_obj is None: return + response_obj.usage = usage_obj # rebind-ok: the stamped cost has to ride on the response the client receives if isinstance(getattr(usage_obj, "cost", None), (int, float)): return try: diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index c226c0b4d09..60ce08666b6 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -368,6 +368,57 @@ def test_stamp_responses_usage_cost_keeps_provider_reported_cost(): logging_obj._response_cost_calculator.assert_not_called() +def _unvalidated_response_with_dict_usage(usage: dict) -> ResponsesAPIResponse: + return ResponsesAPIResponse.model_construct( + id="resp_lit7391", + created_at=int(datetime(2025, 1, 1).timestamp()), + status="completed", + model="perplexity/deepseek-v4-flash-0731", + object="response", + output=[], + truncation="", + usage=usage, + ) + + +def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + from litellm.types.llms.openai import ResponseAPIUsage + + response = _unvalidated_response_with_dict_usage( + { + "input_tokens": 29, + "output_tokens": 120, + "output_tokens_details": {"reasoning_tokens": 117}, + "total_tokens": 149, + "cost": {"currency": "USD", "input_cost": 0, "output_cost": 3e-05, "total_cost": 3e-05}, + } + ) + logging_obj = Mock(spec=LiteLLMLoggingObj) + + _stamp_responses_usage_cost(response, logging_obj) + + assert isinstance(response.usage, ResponseAPIUsage) + assert response.usage.cost == pytest.approx(3e-05) + assert response.usage.output_tokens_details.reasoning_tokens == 117 + logging_obj._response_cost_calculator.assert_not_called() + + +def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + from litellm.types.llms.openai import ResponseAPIUsage + + response = _unvalidated_response_with_dict_usage({"input_tokens": 29, "output_tokens": 120, "total_tokens": 149}) + logging_obj = Mock(spec=LiteLLMLoggingObj) + logging_obj._response_cost_calculator.return_value = 0.000704 + + _stamp_responses_usage_cost(response, logging_obj) + + assert isinstance(response.usage, ResponseAPIUsage) + assert response.usage.cost == pytest.approx(0.000704) + logging_obj._response_cost_calculator.assert_called_once_with(result=response) + + def test_stamp_responses_usage_cost_survives_calculator_failure(): from litellm.responses.streaming_iterator import _stamp_responses_usage_cost From 7e137958e93f91c21f954fabf878132a7534e311 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:12:40 -0700 Subject: [PATCH 061/107] fix(cli): let the apiKeyHelper supply Claude Code's key under lite claude `lite login --config-claude` writes an apiKeyHelper into ~/.claude/settings.json, and `lite claude` then also exported ANTHROPIC_AUTH_TOKEN, so Claude Code opened with its "Both ANTHROPIC_AUTH_TOKEN and apiKeyHelper set" banner and, because the env token wins, never ran the helper that was meant to refresh the key. When the launch key is the stored login key and settings.json carries exactly the helper lite wrote for this base URL, `lite claude` now leaves the token out of the env (dropping an inherited one) so Claude Code asks the helper. An explicit --api-key or LITELLM_PROXY_API_KEY, a helper for another proxy, a hand-written helper, or no settings file keep the previous behavior. --- litellm/proxy/client/cli/commands/agents.py | 29 ++++- .../client/cli/commands/claude_settings.py | 14 +++ .../proxy/client/cli/test_agents.py | 110 ++++++++++++++++++ .../proxy/client/cli/test_claude_settings.py | 43 +++++++ 4 files changed, 193 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 7a0ae9dc955..8903be8ca5c 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -12,6 +12,7 @@ import requests from pydantic import BaseModel, TypeAdapter, ValidationError from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login +from .claude_settings import CLAUDE_SETTINGS_PATH, lite_api_key_helper_configured from .cmd_quoting import quote_for_cmd from .pi import ( LITELLM_PROXY_API_KEY_ENV, @@ -83,6 +84,8 @@ def build_agent_env( base_url: str, api_key: str, profiles: frozenset[str], + *, + export_anthropic_token: bool = True, ) -> dict[str, str]: """Return a copy of base_env wired to route the agent through the proxy. @@ -97,12 +100,19 @@ def build_agent_env( proxy's /v1/models; likewise left alone when already set. pi ignores both base URL variables and instead resolves $LITELLM_PROXY_API_KEY from its synced models.json provider entry. + + With export_anthropic_token=False the bearer is left out (and any inherited + one dropped) so Claude Code asks its configured apiKeyHelper instead; Claude + Code prefers ANTHROPIC_AUTH_TOKEN over the helper and warns when both are set. """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") if PROFILE_ANTHROPIC in profiles: env[ANTHROPIC_BASE_URL_ENV] = root - env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key + if export_anthropic_token: + env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key + else: + env.pop(ANTHROPIC_AUTH_TOKEN_ENV, None) env.pop(ANTHROPIC_API_KEY_ENV, None) if ENABLE_TOOL_SEARCH_ENV not in env: env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE @@ -460,6 +470,7 @@ def run_agent( launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off, reattach_terminal: Callable[[], None] | None = None, preparers: Mapping[str, _Preparer] = MappingProxyType(_PREPARERS), + export_anthropic_token: bool = True, ) -> None: """Validate, wire the environment, and hand off to the agent. @@ -491,7 +502,9 @@ def run_agent( env: Final = MappingProxyType( { - **build_agent_env(env_before_sync, base_url, api_key, profiles), + **build_agent_env( + env_before_sync, base_url, api_key, profiles, export_anthropic_token=export_anthropic_token + ), **(_NO_EXTRA_ENV if isinstance(synced, ModelSyncSkipped) else synced), } ) @@ -529,14 +542,23 @@ def resolve_api_key(ctx: click.Context) -> str: _SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy." +def _helper_supplies_token(ctx_obj: CliContextObj, base_url: str, profiles: frozenset[str]) -> bool: + if PROFILE_ANTHROPIC not in profiles or not ctx_obj.get("api_key_from_token_file"): + return False + return lite_api_key_helper_configured(base_url, CLAUDE_SETTINGS_PATH) + + def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"] started_interactive: Final = _is_interactive() api_key: Final = resolve_api_key(ctx) - display_name, _ = agent_profile(binary) + display_name, profiles = agent_profile(binary) + helper_supplies_token: Final = _helper_supplies_token(ctx_obj, base_url, profiles) click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}") + if helper_supplies_token: + click.echo(f"litellm: {display_name} reads its key from the apiKeyHelper in {CLAUDE_SETTINGS_PATH}") try: run_agent( @@ -545,6 +567,7 @@ def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify [binary, *args], skip_verify=skip_verify, reattach_terminal=(_restore_controlling_terminal if started_interactive else None), + export_anthropic_token=not helper_supplies_token, ) except AgentRunError as e: raise click.ClickException(str(e)) diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index 5e3ce95f088..bc58c2ed815 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -121,6 +121,19 @@ def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str: return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token")) +def lite_api_key_helper_configured(base_url: str, settings_path: Path) -> bool: + """Whether settings_path already carries the apiKeyHelper `lite login --config-claude` writes for base_url. + + Only an exact match counts: a helper for another proxy, a hand-written one, or + settings that cannot be read leave the caller on the env-token path. + """ + try: + configured_helper: Final = load_json_or_empty(settings_path).get(API_KEY_HELPER_KEY) + return configured_helper == resolve_api_key_helper(base_url.rstrip("/")) + except ClaudeSettingsError: + return False + + def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None: """Persistently point Claude Code at base_url, preserving every unrelated setting. @@ -169,6 +182,7 @@ __all__ = ( "SETTINGS_FILE_OWNERS", "ClaudeSettingsError", "SettingsFileOwner", + "lite_api_key_helper_configured", "load_json_or_empty", "merge_claude_settings", "resolve_api_key_helper", diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index a8a6659fe9a..894c2406622 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -28,6 +28,7 @@ from litellm.proxy.client.cli.commands.agents import ( ) AGENTS_MODULE = "litellm.proxy.client.cli.commands.agents" +CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" def _agent_command(name): @@ -163,6 +164,27 @@ class TestBuildAgentEnv: assert env["PATH"] == "/usr/bin" assert base == {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"} + def test_anthropic_profile_leaves_the_bearer_to_the_api_key_helper(self): + env = build_agent_env( + {"ANTHROPIC_AUTH_TOKEN": "stale-token", "ANTHROPIC_API_KEY": "real-key"}, + "http://localhost:4000/", + "sk-key", + frozenset({"anthropic"}), + export_anthropic_token=False, + ) + assert "ANTHROPIC_AUTH_TOKEN" not in env + assert "ANTHROPIC_API_KEY" not in env + assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert env["ENABLE_TOOL_SEARCH"] == "true" + assert env["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + + def test_helper_mode_still_exports_the_openai_key(self): + env = build_agent_env( + {}, "http://localhost:4000", "sk-key", frozenset({"anthropic", "openai"}), export_anthropic_token=False + ) + assert "ANTHROPIC_AUTH_TOKEN" not in env + assert env["OPENAI_API_KEY"] == "sk-key" + class TestAgentLaunchArgs: def test_claude_and_opencode_get_no_extra_args(self): @@ -509,6 +531,25 @@ class TestRunAgent: assert "ANTHROPIC_API_KEY" not in env assert "OPENAI_BASE_URL" not in env + def test_helper_supplied_token_never_reaches_the_launch_env(self): + calls = {} + verified = [] + + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + base_env={"PATH": "/usr/bin", "ANTHROPIC_AUTH_TOKEN": "stale-token"}, + which=lambda name: "/usr/local/bin/claude", + verify=lambda base_url, api_key: verified.append(api_key), + launcher=lambda p, a, e: calls.update(env=dict(e)), + export_anthropic_token=False, + ) + + assert verified == ["sk-key"] + assert "ANTHROPIC_AUTH_TOKEN" not in calls["env"] + assert calls["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + def test_codex_gets_openai_env(self): calls = {} run_agent( @@ -1052,6 +1093,75 @@ class TestAgentCommands: in result.output ) + def _invoke_claude_with_settings(self, tmp_path, settings, obj): + settings_path = tmp_path / "settings.json" + if settings is not None: + settings_path.write_text(json.dumps(settings)) + captured = {} + with ( + patch(f"{AGENTS_MODULE}.CLAUDE_SETTINGS_PATH", settings_path), + patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value="/usr/local/bin/lite"), + patch(f"{AGENTS_MODULE}.run_agent", side_effect=lambda b, k, c, **kw: captured.update(kw)), + ): + result = self.runner.invoke(_agent_command("claude"), [], obj=obj) + assert result.exit_code == 0, result.output + return captured, result.output + + def test_stored_login_with_a_matching_helper_leaves_the_token_to_the_helper(self, tmp_path): + captured, output = self._invoke_claude_with_settings( + tmp_path, + {"apiKeyHelper": "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"}, + {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, + ) + + assert captured["export_anthropic_token"] is False + assert "reads its key from the apiKeyHelper" in output + + def test_explicit_key_is_exported_even_when_a_helper_matches(self, tmp_path): + captured, output = self._invoke_claude_with_settings( + tmp_path, + {"apiKeyHelper": "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"}, + {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": False}, + ) + + assert captured["export_anthropic_token"] is True + assert "apiKeyHelper" not in output + + def test_helper_for_another_proxy_keeps_the_env_token(self, tmp_path): + captured, _ = self._invoke_claude_with_settings( + tmp_path, + {"apiKeyHelper": "/usr/local/bin/lite --base-url https://other.example.com auth print-token"}, + {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, + ) + + assert captured["export_anthropic_token"] is True + + def test_no_claude_settings_keeps_the_env_token(self, tmp_path): + captured, _ = self._invoke_claude_with_settings( + tmp_path, + None, + {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, + ) + + assert captured["export_anthropic_token"] is True + + def test_codex_never_consults_claude_settings(self, tmp_path): + settings_path = tmp_path / "settings.json" + captured = {} + with ( + patch(f"{AGENTS_MODULE}.CLAUDE_SETTINGS_PATH", settings_path), + patch(f"{AGENTS_MODULE}.lite_api_key_helper_configured", side_effect=AssertionError("consulted")), + patch(f"{AGENTS_MODULE}.run_agent", side_effect=lambda b, k, c, **kw: captured.update(kw)), + ): + result = self.runner.invoke( + _agent_command("codex"), + [], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, + ) + + assert result.exit_code == 0, result.output + assert captured["export_anthropic_token"] is True + def test_codex_shows_friendly_name(self): captured = {} with patch( diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index e5f2a9d95bd..8a2bfea3bd7 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -15,6 +15,7 @@ from litellm.proxy.client.cli.commands.claude_settings import ( SETTINGS_FILE_OWNERS, ClaudeSettingsError, SettingsFileOwner, + lite_api_key_helper_configured, resolve_api_key_helper, write_claude_settings, ) @@ -357,3 +358,45 @@ class TestDoesNotDestroyUserOwnedStructure: write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) assert json.loads(settings_path.read_text())["env"] == "not-an-object" + + +class TestLiteApiKeyHelperConfigured: + def _settings(self, tmp_path, payload): + settings_path = tmp_path / "settings.json" + settings_path.write_text(payload) + return settings_path + + def test_recognises_the_helper_lite_login_wrote_for_this_proxy(self, tmp_path, lite_on_path): + settings_path = tmp_path / "settings.json" + write_claude_settings("https://proxy.example.com/", settings_path, ()) + + assert lite_api_key_helper_configured("https://proxy.example.com/", settings_path) is True + assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is True + + def test_a_helper_for_another_proxy_does_not_count(self, tmp_path, lite_on_path): + settings_path = tmp_path / "settings.json" + write_claude_settings("https://other.example.com", settings_path, ()) + + assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False + + def test_a_hand_written_helper_does_not_count(self, tmp_path, lite_on_path): + settings_path = self._settings(tmp_path, json.dumps({"apiKeyHelper": "cat ~/.my-proxy-key"})) + + assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False + + def test_missing_or_helperless_settings_do_not_count(self, tmp_path, lite_on_path): + assert lite_api_key_helper_configured("https://proxy.example.com", tmp_path / "absent.json") is False + helperless = json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://proxy.example.com"}}) + settings_path = self._settings(tmp_path, helperless) + assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False + + def test_unreadable_settings_fall_back_to_false(self, tmp_path, lite_on_path): + settings_path = self._settings(tmp_path, "{not json") + + assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False + + def test_lite_missing_from_path_falls_back_to_false(self, tmp_path): + helper = "/usr/local/bin/lite --base-url https://proxy.example.com auth print-token" + settings_path = self._settings(tmp_path, json.dumps({"apiKeyHelper": helper})) + with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=None): + assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False From e27a549aa919f82e590cf04513e72bf5a75f4a7d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:14:00 -0700 Subject: [PATCH 062/107] test(bedrock): type the probe credential overrides --- .../count_tokens/test_bedrock_count_tokens_handler.py | 6 +++++- .../test_litellm/llms/custom_httpx/test_llm_http_handler.py | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py index 31e65bb6e27..3622ce7f212 100644 --- a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -3,6 +3,7 @@ from unittest.mock import AsyncMock import httpx import pytest +from botocore.credentials import RefreshableCredentials from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -14,7 +15,10 @@ class _ProbedCountTokensHandler(BedrockCountTokensHandler): super().__init__() self._probe = probe - def get_credentials(self, **kwargs): + def get_credentials( + self, + **kwargs: object, # kwargs-ok: mirrors the base resolver's keyword contract, which the probe ignores + ) -> RefreshableCredentials: return self._probe.credentials() 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 b5eb98c4ce6..1140090f9cf 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 @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, Mock, patch import httpx import pytest +from botocore.credentials import RefreshableCredentials import litellm from litellm._logging import verbose_logger @@ -822,7 +823,10 @@ class _ProbedBedrockMessagesConfig(AmazonAnthropicClaudeMessagesConfig): super().__init__() self._probe = probe - def get_credentials(self, **kwargs): + def get_credentials( + self, + **kwargs: object, # kwargs-ok: mirrors the base resolver's keyword contract, which the probe ignores + ) -> RefreshableCredentials: return self._probe.credentials() From 654afb547779eef5f7c575d4a84d24ea3f1163b0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:14:07 -0700 Subject: [PATCH 063/107] fix(e2e): assert the batch cost join on an in-run failed batch instead of a cross-run baton The batch list is served from LiteLLM_ManagedObjectTable whenever the managed files hook is loaded, and the Buildkite e2e stacks bundle a fresh Postgres per build, so a prior run's marker batch is never listed and the baton could only ever pass vacuously. Each run now creates a batch OpenAI fails at validation within seconds, retrieves it by its raw provider id with the same key until it is failed, and asserts the {provider_batch_id}_batch_cost row that retrieve writes joins the key's token hash and alias --- .../coverage_registry/quota_management.yaml | 2 +- .../test_key_attribution_e2e.py | 124 +++++++----------- 2 files changed, 45 insertions(+), 81 deletions(-) diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index d3966eed48f..c92fc0b43bd 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -61,4 +61,4 @@ - {id: quota_management.spend_tracking.key_attribution.joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [joins_key], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Every spend row a virtual key writes across chat, queued chat, messages, responses, embeddings, the Gemini passthrough, file upload, batch create, and a replayed callback log carries api_key equal to the key's token hash and the key alias, the join the usage APIs depend on; a re-hashed token shows up as an unattributed key-hash-* row (#39568, #39572)"} - {id: quota_management.spend_tracking.key_attribution.reports_alias_and_email, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [reports_alias_and_email], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/management_endpoints/internal_user_endpoints.py", rationale: "/spend/logs?api_key= returns every one of the key's rows with its alias and /user/daily/activity aggregates them under the key's token with key_alias and user_email; /spend/logs carries no email field, so the email is asserted on daily activity only"} - {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"} -- {id: quota_management.spend_tracking.key_attribution.batch_cost_joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [batch_cost_joins_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "Retrieving a completed batch by its raw provider id prices it inline against the retrieving key, so the {provider_batch_id}_batch_cost row must carry that key's token hash and alias; the target is the newest completed cross-run marker batch this proxy has not billed yet, since completion lags by up to 24h, and a cold start with no completed marker is a documented vacuous pass. The CheckBatchCost poller's own unified-id row needs the batch create row in the same database, which a stack booted fresh per run never holds for a completed marker"} +- {id: quota_management.spend_tracking.key_attribution.batch_cost_joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [batch_cost_joins_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "The retrieve that first sees a batch in a terminal state writes its {provider_batch_id}_batch_cost row, so the batch each run creates is one OpenAI fails at validation within seconds and the test retrieves it by its raw provider id with the same key until it is failed; that raw-id retrieve is never owned by the CheckBatchCost poller, so it prices the batch inline against the retrieving key and the row must carry that key's token hash and alias. A completed batch with a positive cost is out of one run's reach: OpenAI's completion window is 24h and a stack booted fresh per run lists no earlier run's batches"} diff --git a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py index cac4fd7fb3c..1d7b5e128f5 100644 --- a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py @@ -12,19 +12,17 @@ re-hashed token (v1.99.0's regression, #39568 and #39572) shows up as a key-hash-* row with no alias and no email in the customer's usage exports. The health-check service account writes rows too; those must stay keyed by the -literal service-account name, never by a hash of it. A batch's cost row lands -only once the batch completes, up to 24h later, so the batch cost path rides a -cross-run baton like the batches suite: each run submits a one-line marker batch -and never cancels it, and the newest completed marker from any run that this -proxy has not billed yet is retrieved by its raw provider id with this run's -key. That retrieve is the writer under test: a raw id is never poller-owned, so -the proxy prices it inline against the retrieving key, and the -{provider_batch_id}_batch_cost row must join this run's token with its alias. -The CheckBatchCost poller's own row (the unified id, billed against the -submitting key) needs the batch's create row in the same database, which a stack -booted fresh per run never holds for a completed marker, so that writer is out -of this test's reach. A cold start with no completed marker is a documented -vacuous pass, never a skip. +literal service-account name, never by a hash of it. A batch's cost row is +written by the retrieve that first sees the batch in a terminal state, so the +batch the run creates is one OpenAI fails at validation within seconds (its one +line targets /v1/embeddings under a /v1/chat/completions batch), and the test +retrieves it by its raw provider id with the same key until it is failed. A raw +id is never owned by the CheckBatchCost poller, so that retrieve prices the batch +inline against the retrieving key and its {provider_batch_id}_batch_cost row +must join the key's token with its alias. A completed batch with a positive +cost is out of a single run's reach (OpenAI's completion window is 24h, and a +stack booted fresh per run lists no earlier run's batches), so the poller's own +row is not asserted here. /spend/logs carries no email field, so the email assertion lives on /user/daily/activity alone; /spend/logs is held to the alias in metadata. @@ -38,7 +36,7 @@ from datetime import datetime, timedelta, timezone from typing import Final import pytest -from models import ChatMessage, KeyGenerateBody, SpendLogsParams +from models import KeyGenerateBody from proxy_client import Converged, await_converged from pydantic import BaseModel from spend_e2e_client import ( @@ -64,11 +62,9 @@ BATCH_MODEL: Final = "openai-gpt-4o-mini" BATCH_BACKEND_MODEL: Final = "gpt-4o-mini" BATCH_PROVIDER: Final = "openai" HEALTH_SERVICE_ACCOUNT: Final = "litellm-internal-health-check" -BATON_MARKER_KEY: Final = "litellm_e2e_suite" -BATON_MARKER_VALUE: Final = "key-attribution-baton" -BATON_POLL_SECONDS: Final = 30.0 -BATON_POLL_INTERVAL_SECONDS: Final = 10.0 -BATON_LIST_LIMIT: Final = 100 +BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "expired"}) +FAILED_BATCH_POLL_SECONDS: Final = 120.0 +FAILED_BATCH_POLL_INTERVAL_SECONDS: Final = 5.0 MAX_TOKENS: Final = 8 REPLAY_RESPONSE_COST: Final = 0.0001 REPLAY_PROMPT_TOKENS: Final = 5 @@ -86,17 +82,16 @@ WRITE_PATHS: Final = ( ) -class BatchLineBody(BaseModel): +class EmbeddingLineBody(BaseModel): model: str - messages: list[ChatMessage] - max_tokens: int + input: str -class BatchLine(BaseModel): +class EmbeddingLine(BaseModel): custom_id: str method: str = "POST" - url: str = "/v1/chat/completions" - body: BatchLineBody + url: str = "/v1/embeddings" + body: EmbeddingLineBody @dataclass(frozen=True, slots=True) @@ -134,26 +129,19 @@ def _call_id(name: str, sent: StreamingResponse) -> WritePath: return WritePath(name=name, request_id=sent.call_id) -def _batch_jsonl(marker: str) -> bytes: - line: Final = BatchLine( - custom_id=marker, - body=BatchLineBody( - model=BATCH_BACKEND_MODEL, - messages=[ChatMessage(role="user", content=f"Reply with the word ok. {marker}")], - max_tokens=MAX_TOKENS, - ), - ) +def _endpoint_mismatched_jsonl(marker: str) -> bytes: + line: Final = EmbeddingLine(custom_id=marker, body=EmbeddingLineBody(model=BATCH_BACKEND_MODEL, input=marker)) return f"{line.model_dump_json()}\n".encode() def _drive_batch(client: SpendClient, identity: AttributedKey, marker: str) -> tuple[WritePath, WritePath]: - uploaded: Final = client.upload_batch_file(identity.key, BATCH_MODEL, _batch_jsonl(marker)) + uploaded: Final = client.upload_batch_file(identity.key, BATCH_MODEL, _endpoint_mismatched_jsonl(marker)) created: Final = client.create_batch( identity.key, BatchCreateBody( input_file_id=uploaded.id, model=BATCH_MODEL, - metadata={BATON_MARKER_KEY: BATON_MARKER_VALUE, "run": marker}, + metadata={"run": marker}, ), ) return ( @@ -204,48 +192,26 @@ def _drive_every_write_path(client: SpendClient, identity: AttributedKey) -> tup ) -def _completed_batons(listed: list[BatchObject]) -> tuple[BatchObject, ...]: - return tuple( - sorted( - ( - batch - for batch in listed - if batch.status == "completed" and (batch.metadata or {}).get(BATON_MARKER_KEY) == BATON_MARKER_VALUE - ), - key=lambda batch: batch.created_at or 0, - reverse=True, - ) - ) - - def _provider_batch_id(unified_batch_id: str) -> str: encoded: Final = unified_batch_id.removeprefix("batch_") decoded: Final = base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)).decode() return decoded.removeprefix("litellm:").split(";", 1)[0] -def _batch_cost_request_id(unified_batch_id: str) -> str: - return f"{_provider_batch_id(unified_batch_id)}_batch_cost" +def _driven_batch_id(driven: DrivenKey) -> str: + return next(path.request_id for path in driven.paths if path.name == "batch_create") -def _newest_unbilled_completed_baton(client: SpendClient, key: str) -> BatchObject | None: +def _await_terminal_batch(client: SpendClient, key: str, provider_batch_id: str) -> BatchObject: outcome: Final = await_converged( - lambda: _completed_batons(client.list_batches(key, BATCH_MODEL, limit=BATON_LIST_LIMIT)), - converged=lambda completed: bool(completed), - timeout=BATON_POLL_SECONDS, - interval=BATON_POLL_INTERVAL_SECONDS, + lambda: client.retrieve_batch(key, provider_batch_id, provider=BATCH_PROVIDER), + converged=lambda batch: batch.status in BATCH_TERMINAL_STATUSES, + timeout=FAILED_BATCH_POLL_SECONDS, + interval=FAILED_BATCH_POLL_INTERVAL_SECONDS, now=time.monotonic, sleep=time.sleep, ) - completed: Final = outcome.result if isinstance(outcome, Converged) else () - return next( - ( - batch - for batch in completed - if not client.proxy.spend_logs(SpendLogsParams(request_id=_batch_cost_request_id(batch.id))) - ), - None, - ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result def _health_rows_between(client: SpendClient, started_at: datetime) -> list[SpendLogRow]: @@ -416,23 +382,21 @@ class TestKeyAttribution: "quota_management.spend_tracking.key_attribution.batch_cost_joins_key", exercised_on=["batches"], ) - def test_completed_batch_cost_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None: - completed: Final = _newest_unbilled_completed_baton(client, driven.identity.key) - if completed is None: - return - provider_batch_id: Final = _provider_batch_id(completed.id) - fetched: Final = client.retrieve_batch(driven.identity.key, provider_batch_id, provider=BATCH_PROVIDER) - assert fetched.status == "completed", f"listed-completed marker retrieved as {fetched.status!r}" - cost_request_id: Final = _batch_cost_request_id(completed.id) - rows: Final = client.proxy.poll_logs_for_request_id( - cost_request_id, - predicate=lambda found: any((row.spend or 0) > 0 for row in found), + def test_terminal_batch_cost_row_joins_the_retrieving_key(self, client: SpendClient, driven: DrivenKey) -> None: + provider_batch_id: Final = _provider_batch_id(_driven_batch_id(driven)) + fetched: Final = _await_terminal_batch(client, driven.identity.key, provider_batch_id) + assert fetched.status == "failed", ( + f"endpoint-mismatched batch {provider_batch_id} is {fetched.status!r} after " + f"{FAILED_BATCH_POLL_SECONDS:.0f}s, so its terminal cost row cannot be asserted" ) - priced: Final = [row for row in rows if (row.spend or 0) > 0] - assert priced, f"retrieving completed batch {provider_batch_id} wrote no positive-cost row under {cost_request_id}" + cost_request_id: Final = f"{provider_batch_id}_batch_cost" + rows: Final = client.proxy.poll_logs_for_request_id(cost_request_id) + assert rows, f"retrieving failed batch {provider_batch_id} wrote no cost row under {cost_request_id}" + call_types: Final = tuple(sorted({row.call_type or "" for row in rows})) + assert call_types == ("aretrieve_batch",), f"cost rows under {cost_request_id} carry call types {call_types}" unjoined: Final = [ (row.call_type, row.api_key, row.metadata.user_api_key_alias if row.metadata else None) - for row in priced + for row in rows if row.api_key != driven.identity.token or row.metadata is None or row.metadata.user_api_key_alias != driven.identity.alias From c972bbe80fb575fb0a162e79547edbe1db22298d Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:22:41 -0700 Subject: [PATCH 064/107] fix(mcp): send debug headers immediately for GET streams --- .../proxy/_experimental/mcp_server/mcp_debug.py | 8 ++++++-- litellm/proxy/_experimental/mcp_server/server.py | 4 +++- .../_experimental/mcp_server/test_mcp_debug.py | 15 ++++++++++----- 3 files changed, 19 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_debug.py b/litellm/proxy/_experimental/mcp_server/mcp_debug.py index b25ed27becc..ff30ef99ebd 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_debug.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_debug.py @@ -308,14 +308,18 @@ class MCPDebug: @staticmethod def wrap_send_with_debug_headers( - send: Send, debug_headers: Mapping[str, str], resolution: Callable[[], Mapping[str, str]] | None = None + send: Send, + debug_headers: Mapping[str, str], + resolution: Callable[[], Mapping[str, str]] | None = None, + *, + request_method: str | None = None, ) -> Send: """ Return a new ASGI ``send`` callable that injects *debug_headers* into the ``http.response.start`` message. """ - if resolution is not None: + if resolution is not None and request_method == "POST": return _DiagnosticSend(send, debug_headers, resolution) async def _send_with_debug(message: Message) -> None: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index b53c55ef5fa..72d256e90fc 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -4482,7 +4482,9 @@ if MCP_AVAILABLE: diagnostics: Final = MCPAuthDiagnostics() if _debug_headers else None if diagnostics is not None: scope[MCP_AUTH_DIAGNOSTICS_SCOPE_KEY] = diagnostics - send = MCPDebug.wrap_send_with_debug_headers(send, _debug_headers, diagnostics.headers) + send = MCPDebug.wrap_send_with_debug_headers( + send, _debug_headers, diagnostics.headers, request_method=scope.get("method") + ) # Ensure session managers are initialized if not _SESSION_MANAGERS_INITIALIZED: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py index 5039209c351..7d46bb0237a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py @@ -208,20 +208,25 @@ class TestWrapSendWithDebugHeaders: @pytest.mark.asyncio @pytest.mark.parametrize("source", tuple(AuthResolution)) -async def test_debug_uses_resolution_recorded_after_response_start(source: AuthResolution) -> None: +@pytest.mark.parametrize("method", ("GET", "DELETE", "POST")) +async def test_debug_defers_resolution_until_first_frame_only_for_post(source: AuthResolution, method: str) -> None: captured: Final[list[Message]] = [] diagnostics: Final = MCPAuthDiagnostics() async def send(message: Message) -> None: captured.append(message) - wrapped: Final = MCPDebug.wrap_send_with_debug_headers(send, {}, diagnostics.headers) + wrapped: Final = MCPDebug.wrap_send_with_debug_headers( + send, diagnostics.headers(), diagnostics.headers, request_method=method + ) await wrapped({"type": "http.response.start", "status": 200, "headers": []}) - assert captured == [] + assert len(captured) == (0 if method == "POST" else 1) diagnostics.record("s1", source) body: Final[Message] = {"type": "http.response.body", "body": b"data: pong\n\n", "more_body": True} await wrapped(body) - assert dict(captured[0]["headers"])[b"x-mcp-debug-auth-resolution"] == source.value.encode() + assert dict(captured[0]["headers"])[b"x-mcp-debug-auth-resolution"] == ( + source.value.encode() if method == "POST" else b"unresolved" + ) assert captured[1] == body @@ -233,7 +238,7 @@ async def test_early_stream_frame_reports_unresolved_without_waiting() -> None: async def send(message: Message) -> None: captured.append(message) - wrapped: Final = MCPDebug.wrap_send_with_debug_headers(send, {}, diagnostics.headers) + wrapped: Final = MCPDebug.wrap_send_with_debug_headers(send, {}, diagnostics.headers, request_method="POST") await wrapped({"type": "http.response.start", "status": 200, "headers": []}) await wrapped({"type": "http.response.body", "body": b": ping\n\n", "more_body": True}) diagnostics.record("s1", AuthResolution.stored_user_token) From 5db99ca7ba111c19d7b3aedd434abad436108ddf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:23:45 -0700 Subject: [PATCH 065/107] test(e2e): register the poller batch cost join as its own uncovered cell The retrieve writer and the CheckBatchCost poller are two different spend writers, and the in-run failed batch only proves the first. Split the key_attribution batch cell in two: retrieve_batch_cost_joins_retrieving_key, which test_terminal_batch_cost_row_joins_the_retrieving_key claims, and poller_batch_cost_joins_creating_key, which no test claims yet and so shows up on the coverage dashboard as a P1 gap instead of hiding behind the retrieve leg. The rationale records why one run cannot hand the poller a completed batch on a stack that boots a fresh Postgres per build. --- tests/e2e/CLAUDE.md | 3 ++- tests/e2e/coverage_registry/quota_management.yaml | 3 ++- .../spend_tracking/test_key_attribution_e2e.py | 6 ++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index b92d84b8b25..34fbe9d9247 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -188,7 +188,8 @@ quota_management... | routes_to_fallback | reseed_matches_db | reports_spend | logs_cost | zero_cost | matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows | writes_failure_row | returns_cost | keeps_total | joins_key | reports_alias_and_email - | health_rows_keep_service_account | batch_cost_joins_key + | health_rows_keep_service_account | retrieve_batch_cost_joins_retrieving_key + | poller_batch_cost_joins_creating_key e.g. quota_management.ratelimit.rpm.blocks_over_limit exercised_on=[chat_completions, messages] quota_management.budget.key.blocks_over_limit exercised_on=[chat_completions] ``` diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index c92fc0b43bd..ad0914d455b 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -61,4 +61,5 @@ - {id: quota_management.spend_tracking.key_attribution.joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [joins_key], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Every spend row a virtual key writes across chat, queued chat, messages, responses, embeddings, the Gemini passthrough, file upload, batch create, and a replayed callback log carries api_key equal to the key's token hash and the key alias, the join the usage APIs depend on; a re-hashed token shows up as an unattributed key-hash-* row (#39568, #39572)"} - {id: quota_management.spend_tracking.key_attribution.reports_alias_and_email, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [reports_alias_and_email], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/management_endpoints/internal_user_endpoints.py", rationale: "/spend/logs?api_key= returns every one of the key's rows with its alias and /user/daily/activity aggregates them under the key's token with key_alias and user_email; /spend/logs carries no email field, so the email is asserted on daily activity only"} - {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"} -- {id: quota_management.spend_tracking.key_attribution.batch_cost_joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [batch_cost_joins_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "The retrieve that first sees a batch in a terminal state writes its {provider_batch_id}_batch_cost row, so the batch each run creates is one OpenAI fails at validation within seconds and the test retrieves it by its raw provider id with the same key until it is failed; that raw-id retrieve is never owned by the CheckBatchCost poller, so it prices the batch inline against the retrieving key and the row must carry that key's token hash and alias. A completed batch with a positive cost is out of one run's reach: OpenAI's completion window is 24h and a stack booted fresh per run lists no earlier run's batches"} +- {id: quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [retrieve_batch_cost_joins_retrieving_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "The retrieve that first sees a batch in a terminal state prices it inline and writes its {provider_batch_id}_batch_cost row against the retrieving key, so the batch each run creates is one OpenAI fails at validation within seconds and the test retrieves it by its raw provider id with the same key until it is failed; a raw id is never owned by the CheckBatchCost poller, and the row must carry that key's token hash and alias"} +- {id: quota_management.spend_tracking.key_attribution.poller_batch_cost_joins_creating_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [poller_batch_cost_joins_creating_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The CheckBatchCost poller bills a completed, positive-cost batch created through a unified id against the key that created it, a different writer from the inline retrieve. No test claims this cell yet: OpenAI's completion window is 24h and both e2e stacks boot a fresh Postgres per build, so a completed batch is out of one run's reach and the managed list never shows an earlier run's batch; the cell stays visible as a gap until a run can hand a completed batch to the poller"} diff --git a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py index 1d7b5e128f5..4a2c23927c6 100644 --- a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py @@ -283,9 +283,7 @@ class TestKeyAttribution: ) def test_every_write_path_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None: assert tuple(path.name for path in driven.paths) == WRITE_PATHS - found: Final = tuple( - (path, client.proxy.poll_logs_for_request_id(path.request_id)) for path in driven.paths - ) + found: Final = tuple((path, client.proxy.poll_logs_for_request_id(path.request_id)) for path in driven.paths) unwritten: Final = [path.name for path, rows in found if not rows] assert not unwritten, f"write paths that produced no spend row within the poll window: {unwritten}" unjoined: Final = [ @@ -379,7 +377,7 @@ class TestKeyAttribution: assert not rehashed, f"health-check rows keyed by something other than {HEALTH_SERVICE_ACCOUNT!r}: {rehashed}" @pytest.mark.covers( - "quota_management.spend_tracking.key_attribution.batch_cost_joins_key", + "quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key", exercised_on=["batches"], ) def test_terminal_batch_cost_row_joins_the_retrieving_key(self, client: SpendClient, driven: DrivenKey) -> None: From e27a018a6ce1ecf0fdbfa071da8e50d5b9460d96 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:29:52 -0700 Subject: [PATCH 066/107] fix(router): pin bridge-replayed encrypted reasoning to the deployment that minted it The encrypted_content_affinity check only read the pin from the Responses input, which /v1/messages builds after the router has picked a deployment, so a model group spread across OpenAI orgs sent follow-up turns to the wrong org and got invalid_encrypted_content back. The check now also decodes the pin from bridge-tagged thinking and redacted_thinking blocks in the Anthropic messages. The bridge also keeps a deployment's own include list next to reasoning.encrypted_content instead of replacing it. --- .../prompt_templates/common_utils.py | 4 +- .../responses_adapters/handler.py | 10 +++- .../encrypted_content_affinity_check.py | 50 ++++++++++++++++--- .../test_responses_adapters_handler.py | 10 ++++ .../test_encrypted_content_affinity_check.py | 40 +++++++++++++++ 5 files changed, 104 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 75a3aefa98b..ae9b7f6bc4b 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1873,7 +1873,7 @@ def _encrypted_reasoning_field(block: Mapping[str, object]) -> object: return None -def _encrypted_content_of_block(block: Mapping[str, object]) -> str | None: +def encrypted_content_of_block(block: Mapping[str, object]) -> str | None: return encrypted_content_from_signature(_encrypted_reasoning_field(block)) @@ -1900,7 +1900,7 @@ def _reasoning_item_from_block_group(group: tuple[Mapping[str, object], ...]) -> for block in group if (text := _readable_thinking_text(block)) ] - encrypted_content: Final = _encrypted_content_of_block(group[0]) + encrypted_content: Final = encrypted_content_of_block(group[0]) if encrypted_content is not None: return ChatCompletionReasoningItem(type="reasoning", summary=summary, encrypted_content=encrypted_content) if not summary: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index 24825cc87e0..7731c883d9f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -126,7 +126,7 @@ def _build_responses_kwargs( responses_kwargs["stream"] = True # Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.) - excluded: Final = {"anthropic_messages"} + excluded: Final = frozenset(("anthropic_messages",)) for key, value in forwarded_kwargs.items(): if key == "litellm_logging_obj" and value is not None: from litellm.litellm_core_utils.litellm_logging import ( @@ -147,6 +147,14 @@ def _build_responses_kwargs( if explicit_prompt_cache_key is not None: responses_kwargs["prompt_cache_key"] = explicit_prompt_cache_key + deployment_include: Final = forwarded_kwargs.get("include") + bridge_include: Final = responses_kwargs.get("include") + if isinstance(deployment_include, list) and isinstance(bridge_include, list): + responses_kwargs["include"] = [ + *bridge_include, + *(item for item in deployment_include if item not in bridge_include), + ] + return responses_kwargs diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index b623e31ce06..ba94bd86a56 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -37,6 +37,7 @@ Safe to enable globally: """ import time +from collections.abc import Iterator, Mapping from typing import TYPE_CHECKING, Final, Optional, Protocol, cast import httpx @@ -48,6 +49,7 @@ from litellm.exceptions import ( ServiceUnavailableError, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_content_of_block from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.cooldown_cache import CooldownCacheValue from litellm.types.llms.openai import AllMessageValues @@ -138,15 +140,48 @@ class EncryptedContentAffinityCheck(CustomLogger): # If no encoded ID, check if encrypted_content itself is wrapped encrypted_content = item.get("encrypted_content") if encrypted_content and isinstance(encrypted_content, str): - ( - model_id, - _, - ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content) + model_id = EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(encrypted_content) if model_id: return model_id return None + @staticmethod + def _anthropic_content_blocks(messages: object) -> Iterator[Mapping[str, object]]: + if not isinstance(messages, list): + return iter(()) + return ( + cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance + for message in cast(list[object], messages) # cast-ok: narrowed by isinstance + if isinstance(message, Mapping) + for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance + if isinstance(content, list) + for block in cast(list[object], content) # cast-ok: narrowed by isinstance + if isinstance(block, Mapping) + ) + + @staticmethod + def _model_id_from_wrapped_encrypted_content(encrypted_content: str) -> str | None: + model_id, _ = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content) + return model_id or None + + @staticmethod + def _extract_model_id_from_anthropic_messages(messages: object) -> str | None: + return next( + ( + model_id + for block in EncryptedContentAffinityCheck._anthropic_content_blocks(messages) + if (encrypted_content := encrypted_content_of_block(block)) is not None + if ( + model_id := EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content( + encrypted_content + ) + ) + is not None + ), + None, + ) + @staticmethod def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None: for deployment in healthy_deployments: @@ -248,13 +283,14 @@ class EncryptedContentAffinityCheck(CustomLogger): if "litellm_metadata" in request_kwargs: request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] = True - request_input: Final = request_kwargs.get("input") - model_id: Final = self._extract_model_id_from_input(request_input) + model_id: Final = self._extract_model_id_from_input( + request_kwargs.get("input") + ) or self._extract_model_id_from_anthropic_messages(request_kwargs.get("messages")) if not model_id: return typed_healthy_deployments verbose_router_logger.debug( - "EncryptedContentAffinityCheck: decoded model_id=%s from input item IDs", + "EncryptedContentAffinityCheck: decoded model_id=%s from the request's encrypted content markers", model_id, ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py index 2cbfaa17d23..2aa84f05a86 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py @@ -77,6 +77,16 @@ def test_build_responses_kwargs_skips_include_for_a_responses_provider_that_reje assert "reasoning" in responses_kwargs +def test_build_responses_kwargs_keeps_the_deployment_include_next_to_encrypted_reasoning(): + responses_kwargs = _build_responses_kwargs( + max_tokens=1024, + messages=MESSAGES, + model="openai/gpt-5.6-luna", + extra_kwargs={"custom_llm_provider": "openai", "include": ["file_search_call.results"]}, + ) + assert responses_kwargs["include"] == ["reasoning.encrypted_content", "file_search_call.results"] + + def test_build_responses_kwargs_without_metadata_sets_no_prompt_cache_key(): responses_kwargs = _build_responses_kwargs( max_tokens=1024, diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index dac991a41c4..e858f7eb0a8 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -1656,3 +1656,43 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen assert request_kwargs.get("_encrypted_content_affinity_pinned") is True finally: router.discard() + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_pins_anthropic_messages_replayed_through_the_bridge(): + """ + Claude Code behind /v1/messages replays the encrypted reasoning the bridge packed + into a thinking block's signature (or a redacted block's data). The pin has to be + read from those blocks because the bridge builds the Responses `input` only after + the router has picked a deployment. + """ + check = EncryptedContentAffinityCheck() + deployments = [ + {"model_info": {"id": "openai-org-a"}, "litellm_params": {"model": "openai/gpt-5.1"}}, + {"model_info": {"id": "openai-org-b"}, "litellm_params": {"model": "openai/gpt-5.1"}}, + ] + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA_turn_one", "openai-org-b") + request_kwargs = { + "messages": [ + {"role": "user", "content": "Solve the zebra puzzle"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"}, + {"type": "text", "text": "The zebra owner lives in the green house."}, + ], + }, + {"role": "user", "content": "And who drinks water?"}, + ], + } + + pinned = await check.async_filter_deployments( + model="gpt-5.1", + healthy_deployments=deployments, + messages=None, + request_kwargs=request_kwargs, + ) + + assert [d["model_info"]["id"] for d in pinned] == ["openai-org-b"] + assert request_kwargs["_encrypted_content_affinity_pinned"] is True From 7563d94e65591c2ed3b909c196a3eb76e61625fd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:29:55 -0700 Subject: [PATCH 067/107] test(responses): assert the streamed response.completed event echoes the named tool_choice --- .../test_streaming_iterator_transformation.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 29061db97dc..70f0957890d 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -679,6 +679,9 @@ def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape() -> None: "response.in_progress", "response.completed", ] - assert response_events[0].response.tool_choice == {"type": "function", "name": "run_command"} - assert response_events[1].response.tool_choice == {"type": "function", "name": "run_command"} + assert [event.response.tool_choice for event in response_events] == [ + {"type": "function", "name": "run_command"}, + {"type": "function", "name": "run_command"}, + {"type": "function", "name": "run_command"}, + ] assert any(getattr(event, "type", None) == "response.output_item.done" for event in events) From a68fe4e4d9eb8b55efd33be65ef4e1f1f1932c62 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:31:48 -0700 Subject: [PATCH 068/107] fix(responses): keep the client's usage shape when the logging copy cannot re-validate the response --- litellm/responses/streaming_iterator.py | 41 ++++++++------ .../responses/test_streaming_iterator.py | 54 ++++++++++++++++--- 2 files changed, 72 insertions(+), 23 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 658a0e8ad5d..40ff88fc557 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runti import httpx from openai._streaming import SSEDecoder -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from typing_extensions import TypeIs import litellm @@ -439,18 +439,7 @@ class BaseResponsesAPIStreamingIterator: if self._persist_completed_response_before_logging: self._persist_completed_response_to_cache(is_async=is_async) - # Create a copy for logging to avoid modifying the response object that will be returned to the user - # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) - # to chat completion format (prompt_tokens/completion_tokens) for internal logging - # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with - # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) - logging_response = self.completed_response - if self.completed_response is not None and hasattr(self.completed_response, "model_dump"): - try: - logging_response = type(self.completed_response).model_validate(self.completed_response.model_dump()) - except Exception: - # Fallback to original if serialization fails - pass + logging_response: Final[object] = _logging_copy(self.completed_response) self._restore_provider_response_headers(logging_response) end_time: Final = datetime.now() @@ -489,10 +478,10 @@ class BaseResponsesAPIStreamingIterator: def _restore_provider_response_headers(self, logging_response: object) -> None: """Re-apply the provider's response headers to the copy handed to logging callbacks. - ``model_validate(model_dump())`` above drops pydantic private attributes, so the + ``model_validate(model_dump())`` in ``_logging_copy`` drops pydantic private attributes, so the ``_hidden_params`` the provider transform set on the nested response are lost. Returns early - when that copy fell back to the original event, so logging-only state never lands on the - object the caller is iterating. + when the event was not a pydantic model and logging got the original, so logging-only state + never lands on the object the caller is iterating. """ if logging_response is self.completed_response: return @@ -1294,6 +1283,26 @@ def _add_text_like_part_events( ) +def _logging_copy(event: object) -> object: + """Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never + reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the + deepcopy pickle errors of #17192; when a provider payload fails validation (LIT-7391), shallow + copies of the event and its nested response still keep the caller's ``usage`` attribute separate.""" + if not isinstance(event, BaseModel): + return event + try: + return type(event).model_validate(event.model_dump()) + except Exception: + return _detached_shallow_copy(event) + + +def _detached_shallow_copy(event: BaseModel) -> BaseModel: + nested: Final[object] = getattr(event, "response", None) + if isinstance(nested, BaseModel): + return event.model_copy(update={"response": nested.model_copy()}) + return event.model_copy() + + def _usage_as_model(usage: object) -> ResponseAPIUsage | None: if isinstance(usage, ResponseAPIUsage): return usage diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 60ce08666b6..5e0e794d93e 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -18,6 +18,7 @@ from litellm.responses.streaming_iterator import ( SyncResponsesAPIStreamingIterator, ) from litellm.types.llms.openai import ( + ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents, @@ -329,8 +330,6 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead(): def _responses_api_response_with_usage() -> ResponsesAPIResponse: - from litellm.types.llms.openai import ResponseAPIUsage - return ResponsesAPIResponse( id="resp_lit6427", created_at=int(datetime(2025, 1, 1).timestamp()), @@ -383,8 +382,6 @@ def _unvalidated_response_with_dict_usage(usage: dict) -> ResponsesAPIResponse: def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): from litellm.responses.streaming_iterator import _stamp_responses_usage_cost - from litellm.types.llms.openai import ResponseAPIUsage - response = _unvalidated_response_with_dict_usage( { "input_tokens": 29, @@ -406,8 +403,6 @@ def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost(): from litellm.responses.streaming_iterator import _stamp_responses_usage_cost - from litellm.types.llms.openai import ResponseAPIUsage - response = _unvalidated_response_with_dict_usage({"input_tokens": 29, "output_tokens": 120, "total_tokens": 149}) logging_obj = Mock(spec=LiteLLMLoggingObj) logging_obj._response_cost_calculator.return_value = 0.000704 @@ -586,5 +581,50 @@ async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched(): with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")): iterator._log_completed_response(is_async=True) - assert logged == [iterator.completed_response] + assert len(logged) == 1 + assert logged[0] is not iterator.completed_response + assert logged[0].response is not iterator.completed_response.response + assert logged[0].response._hidden_params["headers"]["apim-request-id"] == "azure-correlation-1" assert iterator.completed_response.response._hidden_params == {} + + +def _unvalidated_completed_config() -> Mock: + """Config whose completed event carries a Perplexity-style response that fails validation + (``truncation: ""``) and already holds the stamped ``ResponseAPIUsage``.""" + mock_config = Mock(spec=BaseResponsesAPIConfig) + + def _transform(model, parsed_chunk, logging_obj): + response = _unvalidated_response_with_dict_usage( + ResponseAPIUsage(input_tokens=29, output_tokens=373, total_tokens=402, cost={"total_cost": 0.0001}) + ) + return ResponseCompletedEvent(type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=response) + + mock_config.transform_streaming_response.side_effect = _transform + return mock_config + + +@pytest.mark.asyncio +async def test_streaming_logging_copy_keeps_client_usage_when_response_fails_validation(): + """LIT-7391: the logging copy cannot round-trip a response that fails validation, and logging + rewrites the assembled response's usage to chat shape in place, so the event handed to logging + must never be the one the caller receives.""" + logging_obj = _logging_obj_stub() + logging_obj.stream = True + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator(headers={}, config=_unvalidated_completed_config(), logging_obj=logging_obj) + events = [event async for event in iterator] + + assert len(logged) == 1 + now = datetime.now() + LiteLLMLoggingObj._get_assembled_streaming_response( + logging_obj, logged[0], start_time=now, end_time=now, is_async=True, streaming_chunks=[] + ) + assert logged[0].response.usage["prompt_tokens"] == 29 + + client_usage = events[-1].response.usage + assert isinstance(client_usage, ResponseAPIUsage) + assert client_usage.input_tokens == 29 + assert client_usage.cost == pytest.approx(0.0001) From fb9bb60e80baed40d80ef0cbd7c0482fbcd20e62 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:55:52 -0700 Subject: [PATCH 069/107] test(mcp): cover auth diagnostics through HTTP handler --- .../mcp_server/test_mcp_server.py | 162 ++++++++---------- 1 file changed, 76 insertions(+), 86 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 9b6eaba3177..12805e355f9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -3,6 +3,7 @@ import contextvars import os from datetime import datetime, timedelta from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1866,96 +1867,85 @@ async def test_streamable_http_session_manager_is_stateless(): @pytest.mark.asyncio -async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless(): - """ - Test that routing correctly sends: - - initialize (no mcp-session-id) → stateful manager (so client gets mcp-session-id) - - tools/list (no mcp-session-id) → stateless manager (curl, Inspector) - """ - try: - from litellm.proxy._experimental.mcp_server.server import ( - handle_streamable_http_mcp, - session_manager_stateful, - session_manager_stateless, +@pytest.mark.parametrize("debug", (False, True)) +@pytest.mark.parametrize( + ("method", "request_body", "stateful"), + ( + ("POST", b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', True), + ("POST", b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}', False), + ("GET", b"", False), + ("DELETE", b"", False), + ), +) +async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless( + debug: bool, method: str, request_body: bytes, stateful: bool +) -> None: + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + from starlette.requests import Request + from starlette.types import Message, Receive, Scope, Send + + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._experimental.mcp_server.mcp_debug import record_auth_resolution + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution + + scope: Final[Scope] = {"type": "http", "method": method, "path": "/mcp", "headers": []} + receive: Final = AsyncMock(return_value={"type": "http.request", "body": request_body, "more_body": False}) + send: Final = AsyncMock() + observe_start: Final = AsyncMock() + body: Final[Message] = {"type": "http.response.body", "body": b"data: pong\n\n", "more_body": True} + + async def handle_request(request_scope: Scope, receive: Receive, outgoing: Send) -> None: + await outgoing({"type": "http.response.start", "status": 200, "headers": []}) + await observe_start(send.await_count) + context: Final = RequestContext( + request_id=1, meta=None, session=MagicMock(), lifespan_context=None, request=Request(request_scope) ) - except ImportError: - pytest.skip("MCP server not available") + token: Final = request_ctx.set(context) + try: + record_auth_resolution("s1", AuthResolution.stored_user_token) + finally: + request_ctx.reset(token) + await outgoing(body) - async def make_request(method_body: bytes, path: str = "/mcp/progress_test"): - scope = { - "type": "http", - "method": "POST", - "path": path, - "headers": [ - (b"content-type", b"application/json"), - (b"authorization", b"Bearer test-key"), - ], - } - receive = AsyncMock( - return_value={ - "type": "http.request", - "body": method_body, - "more_body": False, - } - ) - send = AsyncMock() - - stateless_called = [] - stateful_called = [] - - async def stateless_handle(s, r, se): - stateless_called.append(1) - - async def stateful_handle(s, r, se): - stateful_called.append(1) - - with ( - patch( - "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", - new_callable=AsyncMock, - return_value=(MagicMock(), None, ["progress_test"], None, None, None), + stateless_handle: Final = AsyncMock(side_effect=handle_request) + stateful_handle: Final = AsyncMock(side_effect=handle_request) + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=( + UserAPIKeyAuth(user_id="debug-user"), + None, + None, + None, + None, + {"x-litellm-mcp-debug": "true"} if debug else {}, ), - patch( - "litellm.proxy._experimental.mcp_server.server.set_auth_context", - ), - patch( - "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", - True, - ), - patch.object( - session_manager_stateless, - "handle_request", - side_effect=stateless_handle, - ), - patch.object( - session_manager_stateful, - "handle_request", - side_effect=stateful_handle, - ), - patch.object( - session_manager_stateless, - "_server_instances", - {}, - ), - patch.object( - session_manager_stateful, - "_server_instances", - {}, - ), - ): - await handle_streamable_http_mcp(scope, receive, send) + ), + patch("litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=stateless_handle), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), + ), + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) - return bool(stateless_called), bool(stateful_called) - - # initialize → stateful - init_body = b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05"}}' - stateless_called, stateful_called = await make_request(init_body) - assert stateful_called and not stateless_called, "initialize (no session) should route to stateful, not stateless" - - # tools/list → stateless - tools_body = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}' - stateless_called, stateful_called = await make_request(tools_body) - assert stateless_called and not stateful_called, "tools/list (no session) should route to stateless, not stateful" + assert stateful_handle.await_count == (1 if stateful else 0) + assert stateless_handle.await_count == (0 if stateful else 1) + observe_start.assert_awaited_once_with(0 if debug and method == "POST" else 1) + assert send.await_count == 2 + assert send.call_args_list[0].args[0]["status"] == 200 + assert send.call_args_list[1].args[0] == body + headers: Final = dict(send.call_args_list[0].args[0]["headers"]) + if debug: + assert headers[b"x-mcp-debug-auth-resolution"] == (b"stored-user-token" if method == "POST" else b"unresolved") + else: + assert not any(name.startswith(b"x-mcp-debug") for name in headers) @pytest.mark.asyncio From 892d20d86f1470c302b116c3a6048ce5a9dd30f8 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 17:34:24 -0700 Subject: [PATCH 070/107] fix(mcp): refresh tools when editing server connection settings --- .../_components/editToolPreview.test.ts | 87 ++++++++++++++ .../_components/editToolPreview.ts | 53 +++++++++ .../mcp_server_edit.integration.test.tsx | 111 +++++++++++++++++- .../_components/mcp_server_edit.tsx | 47 ++++++-- .../_components/mcp_tool_configuration.tsx | 2 +- 5 files changed, 285 insertions(+), 15 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts new file mode 100644 index 00000000000..9dcfc1a470e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts @@ -0,0 +1,87 @@ +import { describe, expect, it } from "vitest"; +import { getEditToolPreview } from "./editToolPreview"; + +const saved = { + url: "https://example.com/mcp", + transport: "http", + auth_type: "basic", + static_headers: [{ header: "X-Tenant", value: "original" }], +}; + +describe("getEditToolPreview", () => { + it("keeps saved discovery for unchanged settings and unrelated edits", () => { + expect(getEditToolPreview({ ...saved, server_name: "renamed" }, saved)).toEqual({ kind: "saved" }); + }); + + it("previews URL changes with the existing server credential left for server-side inheritance", () => { + expect(getEditToolPreview({ ...saved, url: "https://correct.example/mcp" }, saved)).toEqual({ + kind: "preview", + config: { + url: "https://correct.example/mcp", + transport: "http", + auth_type: "basic", + static_headers: { "X-Tenant": "original" }, + credentials: undefined, + }, + }); + }); + + it("uses edited static headers and only the static auth value", () => { + expect( + getEditToolPreview( + { + ...saved, + static_headers: [{ header: "X-Tenant", value: "corrected" }], + credentials: { auth_value: "user:password", access_token: "old-oauth-token", client_secret: "old-client" }, + }, + saved, + ), + ).toEqual({ + kind: "preview", + config: { + url: saved.url, + transport: "http", + auth_type: "basic", + static_headers: { "X-Tenant": "corrected" }, + credentials: { auth_value: "user:password" }, + }, + }); + }); + + it.each(["", "https://", "file:///tmp/server"])("does not connect to an incomplete or unsupported URL: %s", (url) => { + expect(getEditToolPreview({ ...saved, url }, saved)).toEqual({ kind: "incomplete" }); + }); + + it("waits for a static header value before connecting", () => { + const values = { ...saved, static_headers: [{ header: "X-Tenant", value: "" }] }; + expect(getEditToolPreview(values, saved)).toEqual({ kind: "incomplete" }); + }); + + it("waits for credentials when switching from None to Basic Auth", () => { + expect(getEditToolPreview(saved, { ...saved, auth_type: "none" })).toEqual({ kind: "incomplete" }); + }); + + it("does not forward old credentials when switching to None", () => { + const result = getEditToolPreview( + { ...saved, auth_type: "none", credentials: { auth_value: "old-secret" } }, + saved, + ); + expect(result.kind).toBe("preview"); + if (result.kind === "preview") expect(result.config.credentials).toBeUndefined(); + }); + + it.each(["oauth2", "true_passthrough", "oauth_delegate", "oauth2_token_exchange", "oauth2_id_jag", "aws_sigv4"])( + "preserves the existing discovery path for %s", + (auth_type) => { + expect(getEditToolPreview({ ...saved, auth_type, url: "https://changed.example/mcp" }, saved)).toEqual({ + kind: "saved", + }); + }, + ); + + it("keeps stdio and OpenAPI on their existing discovery path", () => { + for (const transport of ["stdio", "openapi"]) { + expect(getEditToolPreview({ ...saved, transport }, saved)).toEqual({ kind: "saved" }); + } + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts new file mode 100644 index 00000000000..eb46c33c4e7 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts @@ -0,0 +1,53 @@ +import { AUTH_TYPE, TRANSPORT } from "@/components/mcp_tools/types"; +import { AUTH_TYPES_REQUIRING_AUTH_VALUE, reduceStaticHeaders } from "./createServerPayload"; + +const connectionConfig = (values: Readonly>) => { + const credentials = values.credentials; + const authValue = + credentials && typeof credentials === "object" && "auth_value" in credentials ? credentials.auth_value : undefined; + const needsAuthValue = + typeof values.auth_type === "string" && AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(values.auth_type); + return { + url: typeof values.url === "string" ? values.url : "", + transport: typeof values.transport === "string" ? values.transport : "", + auth_type: typeof values.auth_type === "string" ? values.auth_type : "", + static_headers: Object.fromEntries( + Object.entries(reduceStaticHeaders(values.static_headers)).sort(([a], [b]) => a.localeCompare(b)), + ), + credentials: + needsAuthValue && typeof authValue === "string" && authValue.trim() ? { auth_value: authValue } : undefined, + }; +}; + +type EditToolPreview = + | { readonly kind: "saved" } + | { readonly kind: "incomplete" } + | { readonly kind: "preview"; readonly config: ReturnType }; + +export const getEditToolPreview = ( + values: Readonly>, + initialValues: Readonly>, +): EditToolPreview => { + const staticAuth = + values.auth_type === AUTH_TYPE.NONE || + (typeof values.auth_type === "string" && AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(values.auth_type)); + if (!staticAuth || ![TRANSPORT.HTTP, TRANSPORT.SSE].includes(String(values.transport))) { + return { kind: "saved" }; + } + + const config = connectionConfig(values); + if (JSON.stringify(config) === JSON.stringify(connectionConfig(initialValues))) { + return { kind: "saved" }; + } + + const missingNewCredential = + config.auth_type !== initialValues.auth_type && + AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(config.auth_type) && + config.credentials === undefined; + const validUrl = URL.canParse(config.url) && ["http:", "https:"].includes(new URL(config.url).protocol); + const incompleteHeaders = Object.values(config.static_headers).some((value) => !value.trim()); + if (!validUrl || missingNewCredential || incompleteHeaders) { + return { kind: "incomplete" }; + } + return { kind: "preview", config }; +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.integration.test.tsx index a9664191f3f..f3cd37cc580 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.integration.test.tsx @@ -1,6 +1,9 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; -import { render, screen, waitFor, act } from "@testing-library/react"; +import { render, screen, waitFor, act, fireEvent } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; + +import { selectOption } from "./testUtils"; import MCPServerEdit from "./mcp_server_edit"; import * as networking from "@/components/networking"; @@ -27,10 +30,6 @@ vi.mock("./mcp_server_cost_config", () => ({ default: () =>
, })); -vi.mock("./mcp_tool_configuration", () => ({ - default: () =>
, -})); - const BASE: MCPServer = { server_id: "srv_1", server_name: "srv", @@ -369,3 +368,105 @@ describe("mcp_server_edit save payload contract", () => { } }); }); + +describe("MCPServerEdit live tool preview", () => { + beforeEach(() => { + vi.resetAllMocks(); + vi.mocked(networking.listMCPTools).mockResolvedValue({ + tools: [], + error: "connection_error", + message: "Saved credentials rejected", + }); + vi.mocked(networking.testMCPToolsListRequest).mockResolvedValue({ + tools: [ + { name: "echo", description: "Echo the supplied message", inputSchema: { type: "object", properties: {} } }, + ], + }); + }); + + const renderEditor = (server: MCPServer = BASE) => + render( + , + ); + + it("replaces the saved connection failure with tools after correcting Basic Auth without saving", async () => { + renderEditor(); + expect(await screen.findByText("Saved credentials rejected")).toBeInTheDocument(); + await selectOption("Authentication", "Basic Auth"); + fireEvent.change(screen.getByLabelText("Authentication Value"), { target: { value: "preview:correct" } }); + expect(screen.queryByText("Saved credentials rejected")).not.toBeInTheDocument(); + expect(screen.getByText("Loading tools...")).toBeInTheDocument(); + expect(networking.testMCPToolsListRequest).not.toHaveBeenCalled(); + fireEvent.click(await screen.findByRole("button", { name: "Flat List" })); + expect(screen.getByText("echo")).toBeInTheDocument(); + const expectedConfig = { + server_id: BASE.server_id, + url: BASE.url, + auth_type: "basic", + credentials: { auth_value: "preview:correct" }, + }; + expect(networking.testMCPToolsListRequest).toHaveBeenCalledExactlyOnceWith( + "access-token", + expect.objectContaining(expectedConfig), + ); + expect(networking.updateMCPServer).not.toHaveBeenCalled(); + }); + + it("refreshes tools when a static header is corrected", async () => { + renderEditor({ ...BASE, static_headers: { "X-Preview-Key": "wrong" } }); + expect(await screen.findByText("Saved credentials rejected")).toBeInTheDocument(); + fireEvent.change(screen.getByPlaceholderText("Header value"), { target: { value: "correct" } }); + fireEvent.click(await screen.findByRole("button", { name: "Flat List" })); + expect(screen.getByText("echo")).toBeInTheDocument(); + expect(networking.testMCPToolsListRequest).toHaveBeenCalledExactlyOnceWith( + "access-token", + expect.objectContaining({ static_headers: { "X-Preview-Key": "correct" } }), + ); + }); + + it("coalesces URL edits and ignores an older failed preview after the latest preview succeeds", async () => { + const user = userEvent.setup(); + const older = Promise.withResolvers<{ tools: never[]; error: string; message: string }>(); + vi.mocked(networking.testMCPToolsListRequest).mockImplementationOnce(() => older.promise); + renderEditor(); + expect(await screen.findByText("Saved credentials rejected")).toBeInTheDocument(); + fireEvent.change(screen.getByLabelText("MCP Server URL"), { target: { value: "https://first.example/mcp" } }); + await waitFor(() => expect(networking.testMCPToolsListRequest).toHaveBeenCalledTimes(1)); + await user.clear(screen.getByLabelText("MCP Server URL")); + await user.type(screen.getByLabelText("MCP Server URL"), "https://latest.example/mcp"); + expect(networking.testMCPToolsListRequest).toHaveBeenCalledTimes(1); + fireEvent.click(await screen.findByRole("button", { name: "Flat List" })); + expect(screen.getByText("echo")).toBeInTheDocument(); + expect(networking.testMCPToolsListRequest).toHaveBeenCalledTimes(2); + expect(networking.testMCPToolsListRequest).toHaveBeenLastCalledWith( + "access-token", + expect.objectContaining({ url: "https://latest.example/mcp" }), + ); + await act(async () => older.resolve({ tools: [], error: "connection_error", message: "Older request failed" })); + expect(screen.getByText("echo")).toBeInTheDocument(); + expect(screen.queryByText("Older request failed")).not.toBeInTheDocument(); + }); + + it("ignores a saved-record response after editing and restores saved discovery when changes are reverted", async () => { + const saved = Promise.withResolvers<{ tools: never[]; error: string; message: string }>(); + vi.mocked(networking.listMCPTools).mockImplementationOnce(() => saved.promise); + renderEditor(); + await waitFor(() => expect(networking.listMCPTools).toHaveBeenCalledTimes(1)); + fireEvent.change(screen.getByLabelText("MCP Server URL"), { target: { value: "https://correct.example/mcp" } }); + fireEvent.click(await screen.findByRole("button", { name: "Flat List" })); + expect(screen.getByText("echo")).toBeInTheDocument(); + await act(async () => saved.resolve({ tools: [], error: "connection_error", message: "Stale saved response" })); + expect(screen.getByText("echo")).toBeInTheDocument(); + expect(screen.queryByText("Stale saved response")).not.toBeInTheDocument(); + fireEvent.change(screen.getByLabelText("MCP Server URL"), { target: { value: BASE.url } }); + expect(await screen.findByText("Saved credentials rejected")).toBeInTheDocument(); + expect(networking.listMCPTools).toHaveBeenCalledTimes(2); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx index 8793c45371a..5316a22ce27 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx @@ -52,6 +52,7 @@ import EnvVarsSection from "./EnvVarsSection"; import { validateMCPServerUrl, validateMCPServerName, normalizeToolOverrideMap } from "./utils"; import { EditServerFormValues, buildEditServerPayload, editPayloadErrorMessage } from "./editServerPayload"; import { toast } from "@/lib/toast"; +import { getEditToolPreview } from "./editToolPreview"; import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow"; import { MountedFormField, @@ -449,14 +450,27 @@ const MCPServerEdit: React.FC = ({ } }, [mcpServer]); - // Fetch tools when component mounts for a saved server + const toolPreview = getEditToolPreview(allFieldsValue(form), initialValues); + const toolPreviewKey = JSON.stringify(toolPreview); + useEffect(() => { - if (!mcpServer.server_id || mcpServer.server_id.trim() === "") { + let active = true; + setTools([]); + setToolsError(null); + setIsLoadingTools(false); + if (!accessToken || !mcpServer.server_id) return; + if (toolPreview.kind === "incomplete") { + setToolsError("Complete the URL, authentication, and header settings to load tools."); return; } - fetchTools(); + setIsLoadingTools(true); + const timer = setTimeout(() => fetchTools(() => active), toolPreview.kind === "preview" ? 500 : 0); + return () => { + active = false; + clearTimeout(timer); + }; // eslint-disable-next-line react-hooks/exhaustive-deps - }, [mcpServer, accessToken, userID, oauthTokenResponse?.access_token]); + }, [mcpServer, accessToken, userID, oauthTokenResponse?.access_token, toolPreviewKey]); // Invalidate a token authorized in this edit session once any mint-relevant field diverges from the // identity it was minted against (url, auth_type, oauth_flow_type, client creds/scopes, or the @@ -519,6 +533,7 @@ const MCPServerEdit: React.FC = ({ const previewWithStagedInteractiveToken = async ( isPassthrough: boolean, isBrowserHeldTokenMode: boolean, + isCurrent: () => boolean, ): Promise => { const stagedToken = !isPassthrough && !isBrowserHeldTokenMode && getEffectiveAuthType() === AUTH_TYPE.OAUTH2 @@ -550,6 +565,7 @@ const MCPServerEdit: React.FC = ({ registration_url: values.registration_url, }; const toolsResponse = await testMCPToolsListRequest(accessToken, previewConfig, stagedToken); + if (!isCurrent()) return true; if (toolsResponse.tools && !toolsResponse.error) { setTools(toolsResponse.tools); } else { @@ -557,15 +573,16 @@ const MCPServerEdit: React.FC = ({ setToolsError(toolsResponse.message || "Failed to load tools"); } } catch (error) { + if (!isCurrent()) return true; setTools([]); setToolsError(error instanceof Error ? error.message : "Failed to load tools"); } finally { - setIsLoadingTools(false); + if (isCurrent()) setIsLoadingTools(false); } return true; }; - const fetchTools = async () => { + const fetchTools = async (isCurrent: () => boolean) => { if (!accessToken || !mcpServer.server_id) return; // OBO/M2M/static auth is attached server-side from the stored credential, so @@ -574,6 +591,7 @@ const MCPServerEdit: React.FC = ({ // same way the Tools playground does. let customHeaders: Record | undefined; const isPassthrough = + toolPreview.kind === "saved" && getMcpOAuthMode({ auth_type: mcpServer.auth_type, oauth2_flow: mcpServer.oauth2_flow, @@ -581,9 +599,10 @@ const MCPServerEdit: React.FC = ({ }) === "passthrough"; const isBrowserHeldTokenMode = isClientForwardedTokenMode(getEffectiveAuthType()); - if (await previewWithStagedInteractiveToken(isPassthrough, isBrowserHeldTokenMode)) { + if (await previewWithStagedInteractiveToken(isPassthrough, isBrowserHeldTokenMode, isCurrent)) { return; } + if (!isCurrent()) return; if (isPassthrough || isBrowserHeldTokenMode) { const token = oauthTokenResponse?.access_token ?? @@ -591,6 +610,7 @@ const MCPServerEdit: React.FC = ({ ? getToken(mcpServer.server_id, userID)?.access_token ?? null : null); if (!token) { + setIsLoadingTools(false); setTools([]); setToolsError( isBrowserHeldTokenMode @@ -608,7 +628,15 @@ const MCPServerEdit: React.FC = ({ try { // include_disabled_tools: configuring the allowlist needs the full server // catalog, so tools toggled off still render (as unchecked) instead of vanishing. - const toolsResponse = await listMCPTools(accessToken, mcpServer.server_id, customHeaders, true); + const toolsResponse = + toolPreview.kind === "preview" + ? await testMCPToolsListRequest(accessToken, { + ...toolPreview.config, + server_id: mcpServer.server_id, + server_name: mcpServer.server_name || mcpServer.alias, + }) + : await listMCPTools(accessToken, mcpServer.server_id, customHeaders, true); + if (!isCurrent()) return; if (toolsResponse.tools && !toolsResponse.error) { setTools(toolsResponse.tools); @@ -617,10 +645,11 @@ const MCPServerEdit: React.FC = ({ setToolsError(toolsResponse.message || "Failed to load tools"); } } catch (error) { + if (!isCurrent()) return; setTools([]); setToolsError(error instanceof Error ? error.message : "Failed to load tools"); } finally { - setIsLoadingTools(false); + if (isCurrent()) setIsLoadingTools(false); } }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx index 274bdf63e32..a3d62de9941 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx @@ -433,7 +433,7 @@ const MCPToolConfiguration: React.FC = ({ {isLoadingTools && (
-

Loading tools from spec...

+

Loading tools...

)} From b8d7f68aeba2bd4a599895c910df085c2eb530a1 Mon Sep 17 00:00:00 2001 From: kerry-berri Date: Wed, 9 Sep 2026 17:46:12 -0700 Subject: [PATCH 071/107] Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../src/components/llm_calls/chat_completion.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx index b813a35f687..c12a2c3cfb4 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx @@ -245,7 +245,7 @@ export async function makeOpenAIChatCompletionRequest( // Extract cost from usage object if available if (chunkWithUsage.usage.cost !== undefined && chunkWithUsage.usage.cost !== null) { - const parsedCost = parseFloat(chunkWithUsage.usage.cost); + const parsedCost = Number(chunkWithUsage.usage.cost); if (Number.isFinite(parsedCost)) { usageData.cost = parsedCost; } From 3258b732c74a344863e0445bd9e9bd99c248efc8 Mon Sep 17 00:00:00 2001 From: mrinal Date: Thu, 10 Sep 2026 00:46:54 +0000 Subject: [PATCH 072/107] chore(ui): bump smol-toml to fix GHSA-7w5x-hrqm-74c2 osv-scan failure Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/package-lock.json | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 44360e94392..1f459bc50ea 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -11402,9 +11402,9 @@ } }, "node_modules/smol-toml": { - "version": "1.6.1", - "resolved": "https://registry.npmjs.org/smol-toml/-/smol-toml-1.6.1.tgz", - "integrity": "sha512-dWUG8F5sIIARXih1DTaQAX4SsiTXhInKf1buxdY9DIg4ZYPZK5nGM1VRIYmEbDbsHt7USo99xSLFu5Q1IqTmsg==", + "version": "1.8.0", + "resolved": "https://registry.npmjs.org/smol-toml/-/smol-toml-1.8.0.tgz", + "integrity": "sha512-kCZr2V3ch9i00x8zXRhjUNVcjG9ijES5dDudkXvUVCT5QlJNQWElSJdZqyPemffHoLNUYwOcou0Fy+ojN0uHSQ==", "dev": true, "license": "BSD-3-Clause", "engines": { From 245369764e80d113e80ec8269846e5257692b8bc Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 17:49:52 -0700 Subject: [PATCH 073/107] fix(mcp): preserve edited settings in static connection previews --- .../mcp_server/rest_endpoints.py | 11 +++- .../mcp_server/test_rest_endpoints.py | 66 +++++++++++++++++++ 2 files changed, 75 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 5b74caee3a2..a72465f16f4 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1234,7 +1234,7 @@ if MCP_AVAILABLE: return client_id, client_secret, scopes _STAGED_AUTH_VALUE_AUTH_TYPES: Final = frozenset( - (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization) + (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization, MCPAuth.token) ) @dataclass(frozen=True, slots=True) @@ -1318,8 +1318,15 @@ if MCP_AVAILABLE: if _oauth2_flow == "client_credentials" and not request.token_url: _oauth2_flow = None + # Static previews inherit credentials before this step, but must not resolve back to + # the saved record during client creation and discard the edited connection settings. + preview_server_id: Final = ( + "" + if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES or request.auth_type in (None, MCPAuth.none) + else request.server_id or "" + ) server_model: Final = MCPServer( - server_id=request.server_id or "", + server_id=preview_server_id, name=request.alias or request.server_name or "", url=request.url, transport=request.transport, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 3487f634251..771a73c8ad0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -87,6 +87,72 @@ def _route_has_dependency(route, dependency) -> bool: class TestExecuteWithMcpClient: + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("auth_type", "auth_value", "expected_auth"), + ( + (MCPAuth.none, None, {}), + (MCPAuth.basic, "preview:correct", {"Authorization": "Basic cHJldmlldzpjb3JyZWN0"}), + (MCPAuth.basic, None, {"Authorization": "Basic cHJldmlldzpzdG9yZWQ="}), + (MCPAuth.bearer_token, "edited", {"Authorization": "Bearer edited"}), + (MCPAuth.api_key, "edited", {"X-API-Key": "edited"}), + (MCPAuth.token, "edited", {"Authorization": "token edited"}), + (MCPAuth.authorization, "Custom edited", {"Authorization": "Custom edited"}), + ), + ) + async def test_static_preview_uses_edited_connection_instead_of_registered_server( + self, + monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuth, + auth_value: str | None, + expected_auth: dict[str, str], + ) -> None: + from starlette.datastructures import Headers + + from litellm.experimental_mcp_client.client import MCPClient + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy.management_endpoints import mcp_management_endpoints + + saved: Final = MCPServer( + server_id="saved-preview-server", + name="saved", + url="https://stored.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.basic, + authentication_token="preview:stored", + ) + manager: Final = MCPServerManager() + manager.registry = {saved.server_id: saved} + monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) + monkeypatch.setattr(mcp_management_endpoints, "global_mcp_server_manager", manager) + payload: Final = NewMCPServerRequest( + server_id=saved.server_id, + server_name="edited", + url="https://edited.example/mcp", + transport=MCPTransport.sse, + auth_type=auth_type, + credentials={"auth_value": auth_value} if auth_value is not None else None, + static_headers={"X-Preview": "edited"}, + ) + staged: Final = rest_endpoints._stage_server_test(payload, Headers()) + + async def inspect_connection(client: MCPClient) -> dict[str, object]: + return {"url": client.server_url, "transport": client.transport_type, "headers": client._get_auth_headers()} + + result: Final = await rest_endpoints._execute_with_mcp_client( + staged.request, + inspect_connection, + mcp_auth_header=staged.mcp_auth_header, + oauth2_headers=staged.oauth2_headers, + ) + assert result == { + "url": "https://edited.example/mcp", + "transport": MCPTransport.sse, + "headers": {"X-Preview": "edited", **expected_auth}, + } + assert manager.get_mcp_server_by_id(saved.server_id) is saved + assert saved.url == "https://stored.example/mcp" + @pytest.mark.asyncio async def test_redacts_stack_trace(self, monkeypatch): async def fake_create_client(*args, **kwargs): From b554d79914ea0b0a6f12ad17d944d064c3a8f689 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:01:11 -0700 Subject: [PATCH 074/107] fix(cli): resolve Claude Code's settings file through CLAUDE_CONFIG_DIR Claude Code reads settings.json from CLAUDE_CONFIG_DIR when it is set, while lite login --config-claude wrote to ~/.claude/settings.json and lite claude checked that same file for its apiKeyHelper. With the override set, lite could drop ANTHROPIC_AUTH_TOKEN because the helper lives in a file Claude Code never reads, leaving it with no key at all. claude_settings_path(environ) now picks the file the way Claude Code does, and both commands go through it. The CLI test directory gets an autouse conftest that isolates HOME, USERPROFILE, and CLAUDE_CONFIG_DIR per test and fails any test that writes the developer's real settings.json, restoring it first. --- litellm/proxy/client/cli/commands/agents.py | 14 ++++--- litellm/proxy/client/cli/commands/auth.py | 10 +++-- .../client/cli/commands/claude_settings.py | 11 +++++ .../test_litellm/proxy/client/cli/conftest.py | 32 ++++++++++++++ .../proxy/client/cli/test_agents.py | 42 +++++++++++++++---- .../proxy/client/cli/test_auth_commands.py | 7 ++-- .../proxy/client/cli/test_claude_settings.py | 17 ++++++++ .../proxy/client/cli/test_up_commands.py | 7 +--- 8 files changed, 115 insertions(+), 25 deletions(-) create mode 100644 tests/test_litellm/proxy/client/cli/conftest.py diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 8903be8ca5c..ab99754fdc7 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -4,6 +4,7 @@ import subprocess import sys from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from pathlib import Path from types import MappingProxyType from typing import Final, TypeAlias @@ -12,7 +13,7 @@ import requests from pydantic import BaseModel, TypeAdapter, ValidationError from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login -from .claude_settings import CLAUDE_SETTINGS_PATH, lite_api_key_helper_configured +from .claude_settings import claude_settings_path, lite_api_key_helper_configured from .cmd_quoting import quote_for_cmd from .pi import ( LITELLM_PROXY_API_KEY_ENV, @@ -542,10 +543,12 @@ def resolve_api_key(ctx: click.Context) -> str: _SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy." -def _helper_supplies_token(ctx_obj: CliContextObj, base_url: str, profiles: frozenset[str]) -> bool: +def _helper_supplies_token( + ctx_obj: CliContextObj, base_url: str, profiles: frozenset[str], settings_path: Path +) -> bool: if PROFILE_ANTHROPIC not in profiles or not ctx_obj.get("api_key_from_token_file"): return False - return lite_api_key_helper_configured(base_url, CLAUDE_SETTINGS_PATH) + return lite_api_key_helper_configured(base_url, settings_path) def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None: @@ -555,10 +558,11 @@ def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify api_key: Final = resolve_api_key(ctx) display_name, profiles = agent_profile(binary) - helper_supplies_token: Final = _helper_supplies_token(ctx_obj, base_url, profiles) + settings_path: Final = claude_settings_path(os.environ) + helper_supplies_token: Final = _helper_supplies_token(ctx_obj, base_url, profiles, settings_path) click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}") if helper_supplies_token: - click.echo(f"litellm: {display_name} reads its key from the apiKeyHelper in {CLAUDE_SETTINGS_PATH}") + click.echo(f"litellm: {display_name} reads its key from the apiKeyHelper in {settings_path}") try: run_agent( diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 12a288202b6..8e766d94732 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -1,3 +1,4 @@ +import os import sys import time import webbrowser @@ -40,9 +41,9 @@ from litellm.litellm_core_utils.cli_token_utils import ( ) from .claude_settings import ( - CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS, ClaudeSettingsError, + claude_settings_path, write_claude_settings, ) from .pkce_login import ( @@ -778,12 +779,13 @@ def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None: def _configure_claude_code(base_url: str) -> None: - """Point Claude Code at base_url by patching ~/.claude/settings.json.""" + """Point Claude Code at base_url by patching the settings.json it reads.""" + settings_path: Final = claude_settings_path(os.environ) try: - write_claude_settings(base_url, CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS) + write_claude_settings(base_url, settings_path, SETTINGS_FILE_OWNERS) except ClaudeSettingsError as e: raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}") - click.echo(f"\nConfigured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url.rstrip('/')}.") + click.echo(f"\nConfigured Claude Code: {settings_path} now routes through {base_url.rstrip('/')}.") click.echo("Your other Claude Code settings were left untouched. Restart Claude Code to pick this up.") diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index bc58c2ed815..524edcb899a 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -30,6 +30,7 @@ ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DI ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1" CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json" +CLAUDE_CONFIG_DIR_ENV: Final = "CLAUDE_CONFIG_DIR" BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json" @@ -55,6 +56,14 @@ class ClaudeSettingsError(Exception): """Raised for any user-actionable failure while reading or writing Claude Code settings.""" +def claude_settings_path(environ: Mapping[str, str]) -> Path: + """The settings.json Claude Code reads: under CLAUDE_CONFIG_DIR when set, else ~/.claude/settings.json.""" + config_dir: Final = environ.get(CLAUDE_CONFIG_DIR_ENV, "") + if not config_dir: + return CLAUDE_SETTINGS_PATH + return Path(config_dir).expanduser() / "settings.json" + + def load_json_or_empty(path: Path) -> dict[str, JsonValue]: try: content: Final = path.read_bytes() if path.exists() else b"" @@ -173,6 +182,7 @@ __all__ = ( "API_KEY_HELPER_KEY", "AUTOROUTE_BACKUP_PATH", "BACKUP_PATH", + "CLAUDE_CONFIG_DIR_ENV", "CLAUDE_SETTINGS_PATH", "ENABLE_GATEWAY_MODEL_DISCOVERY_KEY", "ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE", @@ -182,6 +192,7 @@ __all__ = ( "SETTINGS_FILE_OWNERS", "ClaudeSettingsError", "SettingsFileOwner", + "claude_settings_path", "lite_api_key_helper_configured", "load_json_or_empty", "merge_claude_settings", diff --git a/tests/test_litellm/proxy/client/cli/conftest.py b/tests/test_litellm/proxy/client/cli/conftest.py new file mode 100644 index 00000000000..50c76d3f125 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/conftest.py @@ -0,0 +1,32 @@ +import os +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest + +REAL_CLAUDE_SETTINGS: Final = Path(os.path.expanduser("~")) / ".claude" / "settings.json" + + +def _current_bytes() -> bytes | None: + return REAL_CLAUDE_SETTINGS.read_bytes() if REAL_CLAUDE_SETTINGS.exists() else None + + +@pytest.fixture(autouse=True) +def isolated_claude_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[Path]: + before: Final = _current_bytes() + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + monkeypatch.setenv("CLAUDE_CONFIG_DIR", str(tmp_path / ".claude")) + yield tmp_path + after: Final = _current_bytes() + if after == before: + return + if before is None: + REAL_CLAUDE_SETTINGS.unlink() + else: + REAL_CLAUDE_SETTINGS.write_bytes(before) + pytest.fail( + f"this test wrote the developer's real {REAL_CLAUDE_SETTINGS}; the original bytes were restored. " + "Resolve the Claude settings path at call time (never Path.home() at import) and point the test at tmp_path" + ) diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index 894c2406622..bebea285edc 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -1093,20 +1093,48 @@ class TestAgentCommands: in result.output ) - def _invoke_claude_with_settings(self, tmp_path, settings, obj): - settings_path = tmp_path / "settings.json" + def _invoke_claude_with_settings(self, tmp_path, settings, obj, *, default_settings=None): + config_dir = tmp_path / "claude-config" + config_dir.mkdir() if settings is not None: - settings_path.write_text(json.dumps(settings)) + (config_dir / "settings.json").write_text(json.dumps(settings)) + default_path = tmp_path / "home-claude" / "settings.json" + default_path.parent.mkdir() + if default_settings is not None: + default_path.write_text(json.dumps(default_settings)) captured = {} with ( - patch(f"{AGENTS_MODULE}.CLAUDE_SETTINGS_PATH", settings_path), + patch(f"{CLAUDE_SETTINGS_MODULE}.CLAUDE_SETTINGS_PATH", default_path), patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value="/usr/local/bin/lite"), patch(f"{AGENTS_MODULE}.run_agent", side_effect=lambda b, k, c, **kw: captured.update(kw)), ): - result = self.runner.invoke(_agent_command("claude"), [], obj=obj) + result = self.runner.invoke( + _agent_command("claude"), [], obj=obj, env={"CLAUDE_CONFIG_DIR": str(config_dir)} + ) assert result.exit_code == 0, result.output return captured, result.output + def test_helper_is_read_from_the_config_dir_claude_code_uses(self, tmp_path): + captured, output = self._invoke_claude_with_settings( + tmp_path, + {"apiKeyHelper": "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"}, + {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, + ) + + assert captured["export_anthropic_token"] is False + assert str(tmp_path / "claude-config" / "settings.json") in output + + def test_helper_only_in_the_default_file_keeps_the_env_token_when_config_dir_points_elsewhere(self, tmp_path): + captured, output = self._invoke_claude_with_settings( + tmp_path, + None, + {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, + default_settings={"apiKeyHelper": "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"}, + ) + + assert captured["export_anthropic_token"] is True + assert "apiKeyHelper" not in output + def test_stored_login_with_a_matching_helper_leaves_the_token_to_the_helper(self, tmp_path): captured, output = self._invoke_claude_with_settings( tmp_path, @@ -1145,11 +1173,9 @@ class TestAgentCommands: assert captured["export_anthropic_token"] is True - def test_codex_never_consults_claude_settings(self, tmp_path): - settings_path = tmp_path / "settings.json" + def test_codex_never_consults_claude_settings(self): captured = {} with ( - patch(f"{AGENTS_MODULE}.CLAUDE_SETTINGS_PATH", settings_path), patch(f"{AGENTS_MODULE}.lite_api_key_helper_configured", side_effect=AssertionError("consulted")), patch(f"{AGENTS_MODULE}.run_agent", side_effect=lambda b, k, c, **kw: captured.update(kw)), ): diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 821323e722c..225396d78f6 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -1416,7 +1416,6 @@ class TestLoginConfigClaude: patch("requests.get", return_value=poll_response), patch("litellm.proxy.client.cli.commands.auth.save_cli_token"), patch("litellm.proxy.client.cli.interface.show_commands"), - patch("litellm.proxy.client.cli.commands.auth.CLAUDE_SETTINGS_PATH", settings_path), patch( "litellm.proxy.client.cli.commands.auth.SETTINGS_FILE_OWNERS", (SettingsFileOwner(backup_path, "lite up", "lite down"),), @@ -1426,7 +1425,9 @@ class TestLoginConfigClaude: return_value="/usr/local/bin/lite", ), ): - result = self.runner.invoke(login, args, obj={"base_url": base_url}) + result = self.runner.invoke( + login, args, obj={"base_url": base_url}, env={"CLAUDE_CONFIG_DIR": str(settings_path.parent)} + ) return result, settings_path, backup_path def test_default_login_does_not_touch_claude_settings(self, tmp_path): @@ -1445,7 +1446,7 @@ class TestLoginConfigClaude: assert written["env"]["ANTHROPIC_BASE_URL"] == "https://test.example.com" assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token" - assert "Configured Claude Code" in result.output + assert f"Configured Claude Code: {settings_path} now routes through https://test.example.com." in result.output def test_flag_preserves_unrelated_settings_on_an_existing_file(self, tmp_path): settings_path = tmp_path / "claude" / "settings.json" diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index 8a2bfea3bd7..0ff92567abc 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -2,6 +2,7 @@ import json import shlex import stat import time +from pathlib import Path from unittest.mock import patch import pytest @@ -12,9 +13,11 @@ from litellm.proxy.client.cli import cli from litellm.proxy.client.cli.commands.claude_settings import ( AUTOROUTE_BACKUP_PATH, BACKUP_PATH, + CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS, ClaudeSettingsError, SettingsFileOwner, + claude_settings_path, lite_api_key_helper_configured, resolve_api_key_helper, write_claude_settings, @@ -360,6 +363,20 @@ class TestDoesNotDestroyUserOwnedStructure: assert json.loads(settings_path.read_text())["env"] == "not-an-object" +class TestClaudeSettingsPath: + def test_defaults_to_the_home_settings_file(self): + assert claude_settings_path({}) == CLAUDE_SETTINGS_PATH + assert claude_settings_path({"CLAUDE_CONFIG_DIR": ""}) == CLAUDE_SETTINGS_PATH + + def test_follows_claude_config_dir_like_claude_code_does(self, tmp_path): + assert claude_settings_path({"CLAUDE_CONFIG_DIR": str(tmp_path)}) == tmp_path / "settings.json" + + def test_expands_a_tilde_in_claude_config_dir(self): + assert claude_settings_path({"CLAUDE_CONFIG_DIR": "~/.claude-work"}) == ( + Path.home() / ".claude-work" / "settings.json" + ) + + class TestLiteApiKeyHelperConfigured: def _settings(self, tmp_path, payload): settings_path = tmp_path / "settings.json" diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index aead1764b0e..8e33e1ed350 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -620,10 +620,7 @@ class TestUpCanInvokeTheRealLoginCommand: ctx.obj = {"base_url": "http://127.0.0.1:9"} ctx.invoke(real_login, pkce=False) - with ( - patch(f"{AUTH_MODULE}.CLAUDE_SETTINGS_PATH", settings_path), - patch(f"{AUTH_MODULE}._start_cli_sso_flow", side_effect=RuntimeError("stop")), - ): - CliRunner().invoke(driver, [], standalone_mode=False) + with patch(f"{AUTH_MODULE}._start_cli_sso_flow", side_effect=RuntimeError("stop")): + CliRunner().invoke(driver, [], standalone_mode=False, env={"CLAUDE_CONFIG_DIR": str(tmp_path)}) assert not settings_path.exists() From 3523f3731cb356bcdf642182ea6b971d65291f1b Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:01:58 -0700 Subject: [PATCH 075/107] fix(mcp): bind inherited preview credentials to the saved origin --- .../mcp_server/rest_endpoints.py | 13 ++++++- .../mcp_server/test_rest_endpoints.py | 39 ++++++++++++++++++- .../_components/editToolPreview.test.ts | 22 ++++++++++- .../_components/editToolPreview.ts | 16 +++++++- .../_components/mcp_server_edit.tsx | 11 ++++-- 5 files changed, 91 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index a72465f16f4..a116c09d844 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1255,7 +1255,18 @@ if MCP_AVAILABLE: MCPRequestHandler, ) - request: Final = _inherit_credentials_from_existing_server(new_mcp_server_request) + saved_server: Final = ( + global_mcp_server_manager.get_mcp_server_by_id(new_mcp_server_request.server_id) + if new_mcp_server_request.server_id + else None + ) + saved_origin: Final = _redact_mcp_resource_url(saved_server.url) if saved_server else None + may_inherit: Final = new_mcp_server_request.auth_type not in _STAGED_AUTH_VALUE_AUTH_TYPES or ( + saved_origin is not None and saved_origin == _redact_mcp_resource_url(new_mcp_server_request.url) + ) + request: Final = ( + _inherit_credentials_from_existing_server(new_mcp_server_request) if may_inherit else new_mcp_server_request + ) mcp_auth_header: Final = ( request.credentials.get("auth_value") if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES and isinstance(request.credentials, dict) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 771a73c8ad0..ee5dabf043f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -128,7 +128,7 @@ class TestExecuteWithMcpClient: payload: Final = NewMCPServerRequest( server_id=saved.server_id, server_name="edited", - url="https://edited.example/mcp", + url="https://stored.example/corrected-mcp", transport=MCPTransport.sse, auth_type=auth_type, credentials={"auth_value": auth_value} if auth_value is not None else None, @@ -146,13 +146,48 @@ class TestExecuteWithMcpClient: oauth2_headers=staged.oauth2_headers, ) assert result == { - "url": "https://edited.example/mcp", + "url": "https://stored.example/corrected-mcp", "transport": MCPTransport.sse, "headers": {"X-Preview": "edited", **expected_auth}, } assert manager.get_mcp_server_by_id(saved.server_id) is saved assert saved.url == "https://stored.example/mcp" + @pytest.mark.parametrize( + "url", ("https://other.example/mcp", "http://stored.example/mcp", "https://stored.example:8443/mcp") + ) + @pytest.mark.parametrize("explicit_credential", (None, "preview:explicit")) + def test_static_preview_does_not_inherit_credentials_across_origins( + self, monkeypatch: pytest.MonkeyPatch, url: str, explicit_credential: str | None + ) -> None: + from starlette.datastructures import Headers + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy.management_endpoints import mcp_management_endpoints + + saved: Final = MCPServer( + server_id="saved-preview-server", + name="saved", + url="https://stored.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.basic, + authentication_token="preview:stored", + ) + manager: Final = MCPServerManager() + manager.registry = {saved.server_id: saved} + monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) + monkeypatch.setattr(mcp_management_endpoints, "global_mcp_server_manager", manager) + payload: Final = NewMCPServerRequest( + server_id=saved.server_id, + url=url, + transport=MCPTransport.http, + auth_type=MCPAuth.basic, + credentials={"auth_value": explicit_credential} if explicit_credential else None, + ) + staged: Final = rest_endpoints._stage_server_test(payload, Headers()) + assert staged.mcp_auth_header == explicit_credential + assert staged.request.credentials == payload.credentials + @pytest.mark.asyncio async def test_redacts_stack_trace(self, monkeypatch): async def fake_create_client(*args, **kwargs): diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts index 9dcfc1a470e..c7147466d43 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts @@ -14,10 +14,10 @@ describe("getEditToolPreview", () => { }); it("previews URL changes with the existing server credential left for server-side inheritance", () => { - expect(getEditToolPreview({ ...saved, url: "https://correct.example/mcp" }, saved)).toEqual({ + expect(getEditToolPreview({ ...saved, url: "https://example.com/corrected-mcp" }, saved)).toEqual({ kind: "preview", config: { - url: "https://correct.example/mcp", + url: "https://example.com/corrected-mcp", transport: "http", auth_type: "basic", static_headers: { "X-Tenant": "original" }, @@ -26,6 +26,24 @@ describe("getEditToolPreview", () => { }); }); + it.each(["https://other.example/mcp", "http://example.com/mcp", "https://example.com:8443/mcp"])( + "requires explicit credentials for a changed origin: %s", + (url) => { + expect(getEditToolPreview({ ...saved, url, static_headers: [] }, saved)).toEqual({ + kind: "incomplete", + message: expect.stringContaining("origin changed"), + }); + const explicit = { ...saved, url, static_headers: [], credentials: { auth_value: "new:secret" } }; + expect(getEditToolPreview(explicit, saved).kind).toBe("preview"); + }, + ); + + it("does not automatically send saved static headers to a new origin", () => { + expect(getEditToolPreview({ ...saved, url: "https://other.example/mcp", auth_type: "none" }, saved).kind).toBe( + "incomplete", + ); + }); + it("uses edited static headers and only the static auth value", () => { expect( getEditToolPreview( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts index eb46c33c4e7..6dea7d7d18e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts @@ -21,7 +21,7 @@ const connectionConfig = (values: Readonly>) => { type EditToolPreview = | { readonly kind: "saved" } - | { readonly kind: "incomplete" } + | { readonly kind: "incomplete"; readonly message?: string } | { readonly kind: "preview"; readonly config: ReturnType }; export const getEditToolPreview = ( @@ -49,5 +49,19 @@ export const getEditToolPreview = ( if (!validUrl || missingNewCredential || incompleteHeaders) { return { kind: "incomplete" }; } + const savedConfig = connectionConfig(initialValues); + const changedOrigin = + !URL.canParse(savedConfig.url) || new URL(config.url).origin !== new URL(savedConfig.url).origin; + const reusesHeader = Object.entries(config.static_headers).some( + ([key, value]) => savedConfig.static_headers[key] === value, + ); + const needsSavedCredential = AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(config.auth_type) && !config.credentials; + if (changedOrigin && (needsSavedCredential || reusesHeader)) { + return { + kind: "incomplete", + message: + "The server origin changed. Enter credentials and replace or remove saved static headers to preview tools.", + }; + } return { kind: "preview", config }; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx index 5316a22ce27..2a37029a2c4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx @@ -454,19 +454,22 @@ const MCPServerEdit: React.FC = ({ const toolPreviewKey = JSON.stringify(toolPreview); useEffect(() => { - let active = true; + const controller = new AbortController(); setTools([]); setToolsError(null); setIsLoadingTools(false); if (!accessToken || !mcpServer.server_id) return; if (toolPreview.kind === "incomplete") { - setToolsError("Complete the URL, authentication, and header settings to load tools."); + setToolsError(toolPreview.message ?? "Complete the URL, authentication, and header settings to load tools."); return; } setIsLoadingTools(true); - const timer = setTimeout(() => fetchTools(() => active), toolPreview.kind === "preview" ? 500 : 0); + const timer = setTimeout( + () => fetchTools(() => !controller.signal.aborted), + toolPreview.kind === "preview" ? 500 : 0, + ); return () => { - active = false; + controller.abort(); clearTimeout(timer); }; // eslint-disable-next-line react-hooks/exhaustive-deps From 493c98f9bffbcf2ea1d5ba30e324427daa9a173b Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:11:20 -0700 Subject: [PATCH 076/107] fix(mcp): normalize default ports in preview origins --- .../_experimental/mcp_server/rest_endpoints.py | 5 ++++- .../mcp_server/test_rest_endpoints.py | 17 ++++++++++++----- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index a116c09d844..0c919ac7657 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1261,8 +1261,11 @@ if MCP_AVAILABLE: else None ) saved_origin: Final = _redact_mcp_resource_url(saved_server.url) if saved_server else None + preview_origin: Final = _redact_mcp_resource_url(new_mcp_server_request.url) may_inherit: Final = new_mcp_server_request.auth_type not in _STAGED_AUTH_VALUE_AUTH_TYPES or ( - saved_origin is not None and saved_origin == _redact_mcp_resource_url(new_mcp_server_request.url) + saved_origin is not None + and preview_origin is not None + and httpx.URL(saved_origin) == httpx.URL(preview_origin) ) request: Final = ( _inherit_credentials_from_existing_server(new_mcp_server_request) if may_inherit else new_mcp_server_request diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index ee5dabf043f..9b0bf1d92f6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -154,11 +154,17 @@ class TestExecuteWithMcpClient: assert saved.url == "https://stored.example/mcp" @pytest.mark.parametrize( - "url", ("https://other.example/mcp", "http://stored.example/mcp", "https://stored.example:8443/mcp") + ("url", "same_origin"), + ( + ("https://other.example/mcp", False), + ("http://stored.example/mcp", False), + ("https://stored.example:8443/mcp", False), + ("https://stored.example:443/mcp", True), + ), ) @pytest.mark.parametrize("explicit_credential", (None, "preview:explicit")) - def test_static_preview_does_not_inherit_credentials_across_origins( - self, monkeypatch: pytest.MonkeyPatch, url: str, explicit_credential: str | None + def test_static_preview_respects_origin_when_inheriting_credentials( + self, monkeypatch: pytest.MonkeyPatch, url: str, same_origin: bool, explicit_credential: str | None ) -> None: from starlette.datastructures import Headers @@ -185,8 +191,9 @@ class TestExecuteWithMcpClient: credentials={"auth_value": explicit_credential} if explicit_credential else None, ) staged: Final = rest_endpoints._stage_server_test(payload, Headers()) - assert staged.mcp_auth_header == explicit_credential - assert staged.request.credentials == payload.credentials + expected: Final = explicit_credential or ("preview:stored" if same_origin else None) + assert staged.mcp_auth_header == expected + assert staged.request.credentials == ({"auth_value": expected} if expected else None) @pytest.mark.asyncio async def test_redacts_stack_trace(self, monkeypatch): From 47eb7f434974604a488e5a11193b3ef43410259c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:17:32 -0700 Subject: [PATCH 077/107] fix(token_counter): give each event loop its own count limiter A single process-wide anyio.CapacityLimiter wakes waiters with the asyncio.Event of whichever loop created it, so a fifth loop in another thread (asyncio.run per request under Celery, gunicorn sync workers, or run_async_function from function_with_fallbacks) waited forever once four counts were in flight, and building it at import time raised AsyncLibraryNotFoundError on anyio below 4.2. The limiter now lives in a RunVar and is created on first use inside the running loop. TOKEN_COUNTER_MAX_CONCURRENT_COUNTS reads from the environment like TOKEN_COUNTER_MAX_EXACT_CHARS, and the proxy's tokenizer lookup runs on the shared thread pool so a Hub download no longer holds a count slot. --- litellm/constants.py | 7 +++- litellm/litellm_core_utils/token_counter.py | 20 +++++++++- litellm/proxy/proxy_server.py | 3 +- .../litellm_core_utils/test_token_counter.py | 38 +++++++++++++++++++ 4 files changed, 64 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 0bd7def539d..d8c83876e87 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -397,7 +397,12 @@ TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range( minimum=1, maximum=1_000_000_000, ) -TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = 4 +TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = get_env_int_in_range( + "TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", + default=4, + minimum=1, + maximum=256, +) MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512)) MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512)) OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index e71ce5dea9c..70562748c6a 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -7,6 +7,7 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from typing import Final, Literal, cast import anyio +import anyio.lowlevel import httpx import tiktoken from tokenizers import Tokenizer @@ -326,13 +327,28 @@ Type for a function that counts tokens in a string. EXTRAPOLATION_SAMPLES: Final = 16 T_ParamSpec: Final = ParamSpec("T_ParamSpec") T_Retval = TypeVar("T_Retval") -_COUNT_OFFLOAD_LIMITER: Final = anyio.CapacityLimiter(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) +_COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter") + + +def _count_offload_limiter_for_this_loop() -> anyio.CapacityLimiter: + existing: Final = _COUNT_OFFLOAD_LIMITER.get(None) + if existing is not None: + return existing + created: Final = anyio.CapacityLimiter(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) + _COUNT_OFFLOAD_LIMITER.set(created) + return created def offload_token_count( function: Callable[T_ParamSpec, T_Retval], ) -> Callable[T_ParamSpec, Awaitable[T_Retval]]: - return asyncify(function, limiter=_COUNT_OFFLOAD_LIMITER) + async def offloaded( + *args: T_ParamSpec.args, + **kwargs: T_ParamSpec.kwargs, # kwargs-ok: ParamSpec keeps the wrapped function's own keyword contract + ) -> T_Retval: + return await asyncify(function, limiter=_count_offload_limiter_for_this_loop())(*args, **kwargs) + + return offloaded def _get_tiktoken_count_function( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0e3e5c251c8..453f57131b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -62,6 +62,7 @@ from litellm.constants import ( LITELLM_UI_SESSION_DURATION, RUNTIME_UPDATABLE_ROUTER_SETTINGS, ) +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, ) @@ -12713,7 +12714,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = await offload_token_count(litellm.utils._select_tokenizer)( + _tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)( model=model_to_use, custom_tokenizer=custom_tokenizer ) diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 7cc7c6c763c..7da04d12569 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -5,6 +5,7 @@ import importlib import threading import time import traceback +from concurrent.futures import Future, wait from typing import Final from unittest.mock import MagicMock @@ -203,6 +204,43 @@ async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool(): assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS +def _count_in_a_fresh_event_loop(text: str, result: Future[int]) -> None: + def slow_count(counted: str) -> int: + time.sleep(0.1) + return len(counted) + + result.set_result(asyncio.run(offload_token_count(slow_count)(text))) + + +def test_offloaded_counts_finish_in_every_event_loop_that_shares_the_process(): + loops: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + results: Final = tuple(Future[int]() for _ in range(loops)) + threads: Final = tuple( + threading.Thread(target=_count_in_a_fresh_event_loop, args=("a" * size, result), daemon=True) + for size, result in enumerate(results, start=1) + ) + for thread in threads: + thread.start() + + _, pending = wait(results, timeout=5) + + assert not pending + assert tuple(result.result() for result in results) == tuple(range(1, loops + 1)) + + +@pytest.mark.parametrize( + ("configured", "expected"), + [("8", 8), ("0", 4), ("not-an-int", 4)], +) +def test_max_concurrent_counts_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): + monkeypatch.setenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", configured) + try: + assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_CONCURRENT_COUNTS == expected + finally: + monkeypatch.delenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS") + importlib.reload(litellm.constants) + + def test_token_counter_applies_the_default_cap(): max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] From 3aa9267df278a4a28399460cc7cbe4708a4bd169 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 18:19:52 -0700 Subject: [PATCH 078/107] fix(ui): drop blank streamed costs instead of reading them as zero Number("") and Number(" ") both return 0, which passed the finite check, so a provider reporting an empty cost got a fabricated $0.000000 metric instead of having the unusable value omitted. Both ingestion sites carried the same inline parsing, so this pulls it into one parseUsageCost helper that keeps finite numbers and non-blank numeric strings and drops everything else, including booleans, arrays and breakdown objects. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01YNw8WvkvCSeTcE5qvergu3 --- .../llm_calls/chat_completion.test.tsx | 6 +++ .../components/llm_calls/chat_completion.tsx | 10 ++--- .../llm_calls/responses_api.test.tsx | 29 ++++++++++++++ .../components/llm_calls/responses_api.tsx | 9 ++--- .../components/llm_calls/usage_cost.test.ts | 39 +++++++++++++++++++ .../src/components/llm_calls/usage_cost.ts | 23 +++++++++++ 6 files changed, 105 insertions(+), 11 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/llm_calls/usage_cost.test.ts create mode 100644 ui/litellm-dashboard/src/components/llm_calls/usage_cost.ts diff --git a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx index 129d51d50a2..71f36bdb6db 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx @@ -472,6 +472,12 @@ describe("chat_completion prompt cache usage", () => { expect(usageData).toEqual(expect.not.objectContaining({ cost: expect.anything() })); }); + + it("omits cost when the provider reports a blank value", async () => { + const usageData = await captureUsage({ cost: " " }); + + expect(usageData).toEqual(expect.not.objectContaining({ cost: expect.anything() })); + }); }); describe("chat_completion response cache", () => { diff --git a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx index c12a2c3cfb4..ffa2877fbd9 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx @@ -5,6 +5,7 @@ import { VectorStoreSearchResponse } from "../chat_ui/types"; import { getProxyBaseUrl } from "@/components/networking"; import { MCPServer, MCPToolset, type MCPEvent } from "@/components/mcp_tools/types"; import { extractPromptCacheTokens } from "@/utils/promptCacheUsage"; +import { parseUsageCost } from "./usage_cost"; const completionAsSingleChunk = (completion: ChatCompletion): ChatCompletionChunk => ({ @@ -243,12 +244,9 @@ export async function makeOpenAIChatCompletionRequest( usageData.reasoningTokens = chunkWithUsage.usage.completion_tokens_details.reasoning_tokens; } - // Extract cost from usage object if available - if (chunkWithUsage.usage.cost !== undefined && chunkWithUsage.usage.cost !== null) { - const parsedCost = Number(chunkWithUsage.usage.cost); - if (Number.isFinite(parsedCost)) { - usageData.cost = parsedCost; - } + const parsedCost = parseUsageCost(chunkWithUsage.usage.cost); + if (parsedCost !== undefined) { + usageData.cost = parsedCost; } onUsageData(usageData); diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx index 7b13f0ac2e3..b55df89d5cc 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx @@ -260,6 +260,35 @@ describe("responses_api", () => { expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ cost: expect.anything() }), ""); }); + it("should omit cost when the proxy reports a blank cost", async () => { + async function* streamWithBlankCost() { + yield { + type: "response.completed", + response: { + id: "resp_blank_cost", + usage: { output_tokens: 12, input_tokens: 12, total_tokens: 24, cost: " " }, + }, + }; + } + mockResponsesCreate.mockResolvedValueOnce(streamWithBlankCost()); + + const onUsageData = vi.fn(); + + await makeOpenAIResponsesRequest( + messages, + mockUpdateTextUI, + "gpt-4", + "test-token", + undefined, + undefined, + undefined, + undefined, + onUsageData, + ); + + expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ cost: expect.anything() }), ""); + }); + it("should replay MCP output items as events for a non-streaming response", async () => { mockResponsesCreate.mockReturnValueOnce( nonStreamingResponse({ diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx index 63085a81b94..ce54c7c6b40 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx @@ -4,6 +4,7 @@ import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { getProxyBaseUrl } from "@/components/networking"; import { toast } from "@/lib/toast"; import { extractPromptCacheTokens } from "@/utils/promptCacheUsage"; +import { parseUsageCost } from "./usage_cost"; import type { MCPEvent } from "@/components/mcp_tools/types"; import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; import { @@ -311,11 +312,9 @@ export async function makeOpenAIResponsesRequest( usageData.reasoningTokens = reasoningTokens; } - if (usage.cost !== undefined && usage.cost !== null) { - const parsedCost = Number(usage.cost); - if (Number.isFinite(parsedCost)) { - usageData.cost = parsedCost; - } + const parsedCost = parseUsageCost(usage.cost); + if (parsedCost !== undefined) { + usageData.cost = parsedCost; } onUsageData(usageData, mcpToolUsed); diff --git a/ui/litellm-dashboard/src/components/llm_calls/usage_cost.test.ts b/ui/litellm-dashboard/src/components/llm_calls/usage_cost.test.ts new file mode 100644 index 00000000000..ee1021c0629 --- /dev/null +++ b/ui/litellm-dashboard/src/components/llm_calls/usage_cost.test.ts @@ -0,0 +1,39 @@ +import { describe, expect, it } from "vitest"; +import { parseUsageCost } from "./usage_cost"; + +describe("parseUsageCost", () => { + it("keeps finite numbers, including zero", () => { + expect(parseUsageCost(0)).toBe(0); + expect(parseUsageCost(0.000063)).toBe(0.000063); + }); + + it("keeps numeric strings", () => { + expect(parseUsageCost("0.00019")).toBe(0.00019); + expect(parseUsageCost(" 0.00019 ")).toBe(0.00019); + }); + + it("drops blank strings instead of fabricating a zero cost", () => { + expect(parseUsageCost("")).toBeUndefined(); + expect(parseUsageCost(" ")).toBeUndefined(); + expect(parseUsageCost("\t\n")).toBeUndefined(); + }); + + it("drops strings with a numeric prefix instead of truncating them", () => { + expect(parseUsageCost("1oops")).toBeUndefined(); + expect(parseUsageCost("0.5 USD")).toBeUndefined(); + }); + + it("drops non-finite numbers", () => { + expect(parseUsageCost(Number.NaN)).toBeUndefined(); + expect(parseUsageCost(Number.POSITIVE_INFINITY)).toBeUndefined(); + }); + + it("drops values that are not numbers or strings", () => { + expect(parseUsageCost(null)).toBeUndefined(); + expect(parseUsageCost(undefined)).toBeUndefined(); + expect(parseUsageCost(true)).toBeUndefined(); + expect(parseUsageCost([])).toBeUndefined(); + expect(parseUsageCost(["0.5"])).toBeUndefined(); + expect(parseUsageCost({ total_cost: 0.5 })).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/llm_calls/usage_cost.ts b/ui/litellm-dashboard/src/components/llm_calls/usage_cost.ts new file mode 100644 index 00000000000..79f56dcb2df --- /dev/null +++ b/ui/litellm-dashboard/src/components/llm_calls/usage_cost.ts @@ -0,0 +1,23 @@ +/** + * Providers and upstream gateways report `usage.cost` unvalidated: it arrives as a number, a numeric + * string, an empty string, or something non-numeric. A cost that does not resolve to a finite number + * must be dropped rather than coerced, because `NaN` survives `JSON.stringify` as `null` and crashes + * the metrics row on the next load. + */ +export function parseUsageCost(rawCost: unknown): number | undefined { + if (typeof rawCost === "number") { + return Number.isFinite(rawCost) ? rawCost : undefined; + } + + if (typeof rawCost !== "string") { + return undefined; + } + + const trimmed = rawCost.trim(); + if (trimmed === "") { + return undefined; + } + + const parsed = Number(trimmed); + return Number.isFinite(parsed) ? parsed : undefined; +} From b3f2de058be47e4613b8e6ad27039af3bb58c7a1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:19:56 -0700 Subject: [PATCH 079/107] fix(cli): only guard the default settings.json with the lite up and autoroute backups on login --config-claude --- litellm/proxy/client/cli/commands/auth.py | 4 +- .../client/cli/commands/claude_settings.py | 6 +++ .../proxy/client/cli/test_auth_commands.py | 37 ++++++++++++++++--- 3 files changed, 40 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 8e766d94732..a0bf9254bbc 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -41,9 +41,9 @@ from litellm.litellm_core_utils.cli_token_utils import ( ) from .claude_settings import ( - SETTINGS_FILE_OWNERS, ClaudeSettingsError, claude_settings_path, + settings_file_owners, write_claude_settings, ) from .pkce_login import ( @@ -782,7 +782,7 @@ def _configure_claude_code(base_url: str) -> None: """Point Claude Code at base_url by patching the settings.json it reads.""" settings_path: Final = claude_settings_path(os.environ) try: - write_claude_settings(base_url, settings_path, SETTINGS_FILE_OWNERS) + write_claude_settings(base_url, settings_path, settings_file_owners(settings_path)) except ClaudeSettingsError as e: raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}") click.echo(f"\nConfigured Claude Code: {settings_path} now routes through {base_url.rstrip('/')}.") diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index 524edcb899a..7e96d2a3a08 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -64,6 +64,11 @@ def claude_settings_path(environ: Mapping[str, str]) -> Path: return Path(config_dir).expanduser() / "settings.json" +def settings_file_owners(settings_path: Path) -> tuple[SettingsFileOwner, ...]: + """The commands whose backups guard settings_path: `lite up` and `lite autoroute up` only ever manage the default file.""" + return SETTINGS_FILE_OWNERS if settings_path == CLAUDE_SETTINGS_PATH else () + + def load_json_or_empty(path: Path) -> dict[str, JsonValue]: try: content: Final = path.read_bytes() if path.exists() else b"" @@ -197,5 +202,6 @@ __all__ = ( "load_json_or_empty", "merge_claude_settings", "resolve_api_key_helper", + "settings_file_owners", "write_claude_settings", ) diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 225396d78f6..9c3c02a5e85 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -1398,9 +1398,10 @@ class TestLoginConfigClaude: def setup_method(self): self.runner = CliRunner() - def _run_login(self, tmp_path, args, base_url="https://test.example.com"): + def _run_login(self, tmp_path, args, base_url="https://test.example.com", *, config_dir_env=None): settings_path = tmp_path / "claude" / "settings.json" backup_path = tmp_path / "claude_settings_backup.json" + env = {"CLAUDE_CONFIG_DIR": str(settings_path.parent)} if config_dir_env is None else config_dir_env poll_response = Mock() poll_response.status_code = 200 poll_response.json.return_value = { @@ -1417,17 +1418,19 @@ class TestLoginConfigClaude: patch("litellm.proxy.client.cli.commands.auth.save_cli_token"), patch("litellm.proxy.client.cli.interface.show_commands"), patch( - "litellm.proxy.client.cli.commands.auth.SETTINGS_FILE_OWNERS", + "litellm.proxy.client.cli.commands.claude_settings.SETTINGS_FILE_OWNERS", (SettingsFileOwner(backup_path, "lite up", "lite down"),), ), + patch( + "litellm.proxy.client.cli.commands.claude_settings.CLAUDE_SETTINGS_PATH", + tmp_path / "default-home" / ".claude" / "settings.json", + ), patch( "litellm.proxy.client.cli.commands.claude_settings.shutil.which", return_value="/usr/local/bin/lite", ), ): - result = self.runner.invoke( - login, args, obj={"base_url": base_url}, env={"CLAUDE_CONFIG_DIR": str(settings_path.parent)} - ) + result = self.runner.invoke(login, args, obj={"base_url": base_url}, env=env) return result, settings_path, backup_path def test_default_login_does_not_touch_claude_settings(self, tmp_path): @@ -1460,6 +1463,30 @@ class TestLoginConfigClaude: assert written["theme"] == "dark" assert written["env"]["KEEP"] == "me" + def test_flag_refuses_while_lite_up_holds_the_default_settings_file(self, tmp_path): + default_settings_path = tmp_path / "default-home" / ".claude" / "settings.json" + (tmp_path / "claude_settings_backup.json").write_text("{}") + + result, _settings_path, _backup_path = self._run_login( + tmp_path, ["--config-claude"], config_dir_env={"CLAUDE_CONFIG_DIR": ""} + ) + + assert result.exit_code != 0 + assert "Login successful!" in result.output + assert "`lite up` is currently managing" in result.output + assert "Run `lite down` first" in result.output + assert not default_settings_path.exists() + + def test_flag_writes_an_alternate_config_dir_even_while_lite_up_holds_the_default_file(self, tmp_path): + (tmp_path / "claude_settings_backup.json").write_text("{}") + + result, settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token" + assert f"Configured Claude Code: {settings_path} now routes through https://test.example.com." in result.output + def test_settings_failure_is_reported_without_claiming_login_failed(self, tmp_path): settings_path = tmp_path / "claude" / "settings.json" settings_path.parent.mkdir(parents=True) From 3169c80252e6e98f6c68ff77869244dd9788b796 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:21:55 -0700 Subject: [PATCH 080/107] fix(mcp): preserve IPv6 hosts when comparing preview origins --- .../mcp_server/rest_endpoints.py | 19 ++++++++++---- .../mcp_server/test_rest_endpoints.py | 25 +++++++++++++------ 2 files changed, 32 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 0c919ac7657..27cb632c843 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1243,6 +1243,17 @@ if MCP_AVAILABLE: mcp_auth_header: str | None oauth2_headers: dict[str, str] | None + def _preview_origin(url: str | None) -> tuple[str, str, int | None] | None: + if not url: + return None + try: + parsed: Final = httpx.URL(url) + except httpx.InvalidURL: + return None + if parsed.scheme not in ("http", "https") or not parsed.host: + return None + return parsed.scheme, parsed.host, parsed.port + def _stage_server_test(new_mcp_server_request: NewMCPServerRequest, headers: Headers) -> _StagedServerTest: """ Resolve the credentials a not-yet-saved server config carries for a preview call. @@ -1260,12 +1271,10 @@ if MCP_AVAILABLE: if new_mcp_server_request.server_id else None ) - saved_origin: Final = _redact_mcp_resource_url(saved_server.url) if saved_server else None - preview_origin: Final = _redact_mcp_resource_url(new_mcp_server_request.url) + saved_origin: Final = _preview_origin(saved_server.url) if saved_server else None + preview_origin: Final = _preview_origin(new_mcp_server_request.url) may_inherit: Final = new_mcp_server_request.auth_type not in _STAGED_AUTH_VALUE_AUTH_TYPES or ( - saved_origin is not None - and preview_origin is not None - and httpx.URL(saved_origin) == httpx.URL(preview_origin) + saved_origin is not None and saved_origin == preview_origin ) request: Final = ( _inherit_credentials_from_existing_server(new_mcp_server_request) if may_inherit else new_mcp_server_request diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 9b0bf1d92f6..16c6aa128d0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -154,17 +154,28 @@ class TestExecuteWithMcpClient: assert saved.url == "https://stored.example/mcp" @pytest.mark.parametrize( - ("url", "same_origin"), + ("saved_url", "url", "same_origin"), ( - ("https://other.example/mcp", False), - ("http://stored.example/mcp", False), - ("https://stored.example:8443/mcp", False), - ("https://stored.example:443/mcp", True), + ("https://stored.example/mcp", "https://other.example/mcp", False), + ("https://stored.example/mcp", "http://stored.example/mcp", False), + ("https://stored.example/mcp", "https://stored.example:8443/mcp", False), + ("https://stored.example/mcp", "https://stored.example:443/mcp", True), + ("http://stored.example/mcp", "http://stored.example:80/edited", True), + ("https://stored.example/mcp", "HTTPS://STORED.EXAMPLE/edited", True), + ("https://[::1]/mcp", "https://[::1]/edited", True), + ("https://[::1]/mcp", "https://[::1]:443/edited", True), + ("https://[::1]/mcp", "https://[::2]/edited", False), + ("https://stored.example/mcp", "https://stored.example:invalid/mcp", False), ), ) @pytest.mark.parametrize("explicit_credential", (None, "preview:explicit")) def test_static_preview_respects_origin_when_inheriting_credentials( - self, monkeypatch: pytest.MonkeyPatch, url: str, same_origin: bool, explicit_credential: str | None + self, + monkeypatch: pytest.MonkeyPatch, + saved_url: str, + url: str, + same_origin: bool, + explicit_credential: str | None, ) -> None: from starlette.datastructures import Headers @@ -174,7 +185,7 @@ class TestExecuteWithMcpClient: saved: Final = MCPServer( server_id="saved-preview-server", name="saved", - url="https://stored.example/mcp", + url=saved_url, transport=MCPTransport.http, auth_type=MCPAuth.basic, authentication_token="preview:stored", From 4adf99557f901fd73525697fdba8c9b2c99d0cc9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:27:12 -0700 Subject: [PATCH 081/107] fix(responses): echo "auto" for a tool_choice the bridge cannot express instead of failing after the provider call --- .../transformation.py | 6 +++-- .../test_litellm_completion_responses.py | 26 +++++++++++++++++++ .../test_streaming_iterator_transformation.py | 20 ++++++++++++++ 3 files changed, 50 insertions(+), 2 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 6b9931b7df2..fca5b0d11cf 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -287,8 +287,10 @@ class LiteLLMCompletionResponsesConfig: return ToolChoiceCustomParam(type="custom", name=custom_name) case _, {"type": "function", "function": {"name": str(function_name)}}: return ToolChoiceFunctionParam(type="function", name=function_name) - case _, normalized: - return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(normalized) + case _, "none" | "auto" | "required" as normalized: + return normalized + case _, _: + return "auto" @staticmethod def _should_drop_derived_web_search_options(model: str, custom_llm_provider: str | None) -> bool: diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 8f2937cb252..46249e50572 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1435,6 +1435,9 @@ class TestToolChoiceTransformation: ("required", "required"), ("none", "none"), (None, "auto"), + ("any", "auto"), + ("run_command", "auto"), + ({"name": "run_command"}, "auto"), ], ) def test_transform_tool_choice_for_responses_api_response( @@ -1478,6 +1481,29 @@ class TestToolChoiceTransformation: assert responses_api_response.tool_choice == {"type": "function", "name": "run_command"} + def test_non_streamed_response_with_unrecognized_tool_choice_echoes_auto(self) -> None: + chat_completion_response: Final = ModelResponse( + id="chatcmpl-unrecognized-tool-choice", + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message(role="assistant", content="/Users/dev"), + ) + ], + ) + + responses_api_response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Run the command pwd.", + responses_api_request={"tool_choice": "any"}, + chat_completion_response=chat_completion_response, + ) + + assert responses_api_response.tool_choice == "auto" + class TestContentTypeTransformation: """Test content type transformation from Responses API to Chat Completion format""" diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 70f0957890d..850ee7ba623 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -685,3 +685,23 @@ def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape() -> None: {"type": "function", "name": "run_command"}, ] assert any(getattr(event, "type", None) == "response.output_item.done" for event in events) + + +def test_streamed_unrecognized_tool_choice_is_echoed_as_auto() -> None: + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper([_tool_call_chunk(finish_reason="tool_calls")]), + request_input="Run the command pwd.", + responses_api_request={ + "tools": [{"type": "function", "name": "run_command", "parameters": {"type": "object"}}], + "tool_choice": "any", + }, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + response_events: Final = [ + event for event in iterator if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES + ] + + assert [event.response.tool_choice for event in response_events] == ["auto", "auto", "auto"] From 4716b46c243b90d833ea08de8bebcc9a20c963d2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:31:08 -0700 Subject: [PATCH 082/107] fix(router): read the affinity pin from the messages argument and strip bridge reasoning on a cross-group route The encrypted_content_affinity check only read the Anthropic history from request_kwargs["messages"], so a caller that passes it through the callback's messages argument alone skipped the pin. Read the argument first and fall back to the kwargs. When the minting deployment is not a candidate of the routed group, the base already strips the Responses input's encrypted reasoning; do the same for the bridge-tagged thinking blocks in Anthropic messages so the routed deployment gets the readable thinking text instead of ciphertext it cannot decrypt. --- .../prompt_templates/common_utils.py | 43 ++++++++++- .../encrypted_content_affinity_check.py | 14 +++- ...ore_utils_prompt_templates_common_utils.py | 60 ++++++++++++++- .../test_encrypted_content_affinity_check.py | 73 +++++++++++++++---- 4 files changed, 166 insertions(+), 24 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index ae9b7f6bc4b..7ca2a8aad6a 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -6,7 +6,7 @@ import io import json import mimetypes import re -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from itertools import groupby from os import PathLike from pathlib import Path @@ -1889,6 +1889,47 @@ def is_encrypted_reasoning_block(block: object) -> bool: return _carries_encrypted_reasoning(_encrypted_reasoning_field(mapping)) +def strip_encrypted_reasoning_from_messages(messages: object) -> None: + """Drop the encrypted reasoning a routed deployment cannot decrypt from Anthropic-shaped + history, keeping the readable thinking text. + + Mutates the content lists in place: the router's fallback snapshot shares these + message objects, so a rebound list would replay the stripped blocks on the fallback hop. + """ + if not isinstance(messages, list): + return + for content in _anthropic_content_lists(cast(list[object], messages)): # cast-ok: untyped client json + _strip_encrypted_reasoning_from_blocks(content) + + +def _anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: + return ( + cast(list[object], content) # cast-ok: narrowed by isinstance + for message in messages + if isinstance(message, Mapping) + for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance + if isinstance(content, list) + ) + + +def _strip_encrypted_reasoning_from_blocks(content: object) -> None: + blocks: Final = cast(list[object], content) # cast-ok: narrowed by the caller's isinstance + stripped: Final = tuple(_without_encrypted_reasoning_block(block) for block in blocks) + blocks[:] = (block for block in stripped if block is not None) # rebind-ok: list shared with fallback snapshot + + +def _without_encrypted_reasoning_block(block: object) -> object | None: + if not is_encrypted_reasoning_block(block): + return block + mapping: Final = cast(Mapping[str, object], block) # cast-ok: narrowed by is_encrypted_reasoning_block + if mapping.get("type") != "thinking" or not mapping.get("thinking"): + return None + kept: Final[dict[str, object]] = { # mutable-ok: thinking block rebuilt without the undecryptable signature + key: value for key, value in mapping.items() if key != "signature" + } + return kept + + def _reasoning_replay_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str: index, block = indexed_block return f"encrypted:{index}" if is_encrypted_reasoning_block(block) else "summary" diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index b2edb13a4bf..cdd70e6baf2 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -48,7 +48,10 @@ from litellm.exceptions import ( ServiceUnavailableError, ) from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_content_of_block +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_content_of_block, + strip_encrypted_reasoning_from_messages, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.cooldown_cache import CooldownCacheValue from litellm.types.llms.openai import AllMessageValues @@ -274,8 +277,9 @@ class EncryptedContentAffinityCheck(CustomLogger): parent_otel_span: Span | None = None, ) -> list[dict]: """ - If the request ``input`` contains litellm-encoded item IDs, decode the - embedded ``model_id`` and pin the request to that deployment. Raises + If the request ``input`` contains litellm-encoded item IDs, or its Anthropic + ``messages`` replay a bridge-tagged thinking block, decode the embedded + ``model_id`` and pin the request to that deployment. Raises ``RateLimitError`` / ``ServiceUnavailableError`` when the originating deployment is a member of the routed model group but currently unavailable and no encryption-boundary peer exists, rather than dispatching a doomed @@ -304,9 +308,10 @@ class EncryptedContentAffinityCheck(CustomLogger): request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] = True request_input: Final = request_kwargs.get("input") + anthropic_messages: Final = messages or request_kwargs.get("messages") model_id: Final = self._extract_model_id_from_input( request_input - ) or self._extract_model_id_from_anthropic_messages(request_kwargs.get("messages")) + ) or self._extract_model_id_from_anthropic_messages(anthropic_messages) if not model_id: return typed_healthy_deployments @@ -363,6 +368,7 @@ class EncryptedContentAffinityCheck(CustomLogger): model, ) ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + strip_encrypted_reasoning_from_messages(anthropic_messages) return typed_healthy_deployments # The origin is a member of the routed group but currently unavailable (cooled down); fail fast diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index ecb07d8c937..71b71d55fa7 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -1,3 +1,4 @@ +import copy import functools import json import os @@ -20,6 +21,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( is_encrypted_reasoning_block, responses_reasoning_items_from_thinking_blocks, split_concatenated_json_objects, + strip_encrypted_reasoning_from_messages, update_messages_with_model_file_ids, ) @@ -1567,7 +1569,9 @@ class TestEncryptedReasoningReplay: def test_signature_round_trips_the_encrypted_content(self): assert encrypted_content_from_signature(encrypted_reasoning_signature("gAAAA_bytes")) == "gAAAA_bytes" - @pytest.mark.parametrize("signature", [None, "", "ErcBCkgIValidAnthropicSignature", "litellm_encrypted_reasoning:", 7]) + @pytest.mark.parametrize( + "signature", [None, "", "ErcBCkgIValidAnthropicSignature", "litellm_encrypted_reasoning:", 7] + ) def test_anything_else_is_not_encrypted_content(self, signature): assert encrypted_content_from_signature(signature) is None @@ -1576,7 +1580,11 @@ class TestEncryptedReasoningReplay: [{"type": "thinking", "thinking": "Plan.", "signature": encrypted_reasoning_signature("gAAAA_1")}] ) assert items == ( - {"type": "reasoning", "summary": [{"type": "summary_text", "text": "Plan."}], "encrypted_content": "gAAAA_1"}, + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Plan."}], + "encrypted_content": "gAAAA_1", + }, ) def test_encrypted_redacted_block_replays_with_an_empty_summary(self): @@ -1596,7 +1604,10 @@ class TestEncryptedReasoningReplay: ] ) assert items == ( - {"type": "reasoning", "summary": [{"type": "summary_text", "text": "A."}, {"type": "summary_text", "text": "B."}]}, + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "A."}, {"type": "summary_text", "text": "B."}], + }, {"type": "reasoning", "summary": [{"type": "summary_text", "text": "C."}], "encrypted_content": "gAAAA_c"}, {"type": "reasoning", "summary": [{"type": "summary_text", "text": "D."}]}, ) @@ -1621,3 +1632,46 @@ class TestEncryptedReasoningReplay: ) def test_is_encrypted_reasoning_block(self, block, expected): assert is_encrypted_reasoning_block(block) is expected + + def test_strip_keeps_the_readable_thinking_and_drops_the_undecryptable_bytes(self): + assistant_content = [ + {"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"}, + {"type": "thinking", "thinking": "packed by the bridge", "signature": encrypted_reasoning_signature("g1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("g2")}, + {"type": "thinking", "thinking": "", "signature": encrypted_reasoning_signature("g3")}, + {"type": "text", "text": "answer"}, + ] + messages = [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": assistant_content}, + {"role": "user", "content": [{"type": "text", "text": "follow-up"}]}, + ] + + strip_encrypted_reasoning_from_messages(messages) + + assert messages[1]["content"] is assistant_content + assert assistant_content == [ + {"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"}, + {"type": "thinking", "thinking": "packed by the bridge"}, + {"type": "text", "text": "answer"}, + ] + assert messages[0] == {"role": "user", "content": "question"} + assert messages[2] == {"role": "user", "content": [{"type": "text", "text": "follow-up"}]} + + @pytest.mark.parametrize( + "messages", + [ + "not a list", + None, + [{"role": "user", "content": None}], + [{"role": "user", "content": "plain string"}], + ["not a message"], + [{"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "ErcBCkgIValid"}]}], + ], + ) + def test_strip_leaves_history_without_bridge_reasoning_untouched(self, messages): + before = copy.deepcopy(messages) + + strip_encrypted_reasoning_from_messages(messages) + + assert messages == before diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 99a650436b8..57c02e43fc1 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -1597,26 +1597,12 @@ async def test_encrypted_content_affinity_pins_anthropic_messages_replayed_throu {"model_info": {"id": "openai-org-a"}, "litellm_params": {"model": "openai/gpt-5.1"}}, {"model_info": {"id": "openai-org-b"}, "litellm_params": {"model": "openai/gpt-5.1"}}, ] - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA_turn_one", "openai-org-b") - request_kwargs = { - "messages": [ - {"role": "user", "content": "Solve the zebra puzzle"}, - { - "role": "assistant", - "content": [ - {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, - {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"}, - {"type": "text", "text": "The zebra owner lives in the green house."}, - ], - }, - {"role": "user", "content": "And who drinks water?"}, - ], - } + request_kwargs = {"model": "gpt-5.1"} pinned = await check.async_filter_deployments( model="gpt-5.1", healthy_deployments=deployments, - messages=None, + messages=_bridge_replayed_anthropic_messages(minted_by="openai-org-b"), request_kwargs=request_kwargs, ) @@ -1624,6 +1610,61 @@ async def test_encrypted_content_affinity_pins_anthropic_messages_replayed_throu assert request_kwargs["_encrypted_content_affinity_pinned"] is True +def _bridge_replayed_anthropic_messages(minted_by: str) -> list: + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA_turn_one", minted_by) + return [ + {"role": "user", "content": "Solve the zebra puzzle"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"}, + { + "type": "thinking", + "thinking": "The bridge packed this one", + "signature": f"litellm_encrypted_reasoning:{wrapped}", + }, + {"type": "text", "text": "The zebra owner lives in the green house."}, + ], + }, + {"role": "user", "content": "And who drinks water?"}, + ] + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_strips_bridge_reasoning_from_messages_routed_to_another_group(): + """ + The /v1/messages twin of the tier-change case: the routed group holds no deployment + of the org that minted the reasoning, so the bridge-tagged blocks are stripped down + to their readable thinking text and the request dispatches to the routed pool. + """ + originating = _make_originating_mock(None, "key-a", model_name="gpt-reasoning-tier") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["openai-org-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + routed_pool = [{"model_info": {"id": "openai-org-b"}, "litellm_params": {"model": "openai/gpt-5-nano"}}] + messages = _bridge_replayed_anthropic_messages(minted_by="openai-org-a") + assistant_content = messages[1]["content"] + request_kwargs = {"model": "gpt-5.1"} + + result = await check.async_filter_deployments( + model="gpt-simple-tier", + healthy_deployments=routed_pool, + messages=messages, + request_kwargs=request_kwargs, + ) + + assert result is routed_pool + assert "_encrypted_content_affinity_pinned" not in request_kwargs + assert messages[1]["content"] is assistant_content + assert assistant_content == [ + {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, + {"type": "thinking", "thinking": "The bridge packed this one"}, + {"type": "text", "text": "The zebra owner lives in the green house."}, + ] + + class TestStripEncryptedReasoningFromInput: def test_keeps_summary_and_drops_encrypted_content_and_id(self): wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a") From a25eccafc952aa654c9d826cfff18e568c32d817 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:32:13 -0700 Subject: [PATCH 083/107] fix(azure_ai): drop the api_base path segments a relay already repeats --- .../azure_ai/passthrough/transformation.py | 18 +++++++++++++++++- ...test_azure_ai_passthrough_transformation.py | 18 ++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index 1689755b042..ed630bc6f2b 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -62,6 +62,22 @@ def foundry_root(api_base: str) -> str: return str(url.copy_with(path="/" + "/".join(root_segments), query=None)).rstrip("/") +def without_repeated_native_prefix(root: str, native_endpoint: str) -> str: + url: Final = httpx.URL(root) + root_segments: Final = tuple(segment for segment in url.path.split("/") if segment) + native_segments: Final = tuple(segment.casefold() for segment in native_endpoint.split("/") if segment) + overlap: Final = next( + ( + length + for length in range(min(len(root_segments), len(native_segments)), 0, -1) + if tuple(segment.casefold() for segment in root_segments[-length:]) == native_segments[:length] + ), + 0, + ) + kept_segments: Final = root_segments[: len(root_segments) - overlap] + return str(url.copy_with(path="/" + "/".join(kept_segments), query=None)).rstrip("/") + + def relay_query_params( request_query_params: Mapping[str, object] | None, deployment_api_version: str | None, @@ -111,7 +127,7 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): raise ValueError("Azure AI api base not found: set `api_base` on the deployment or AZURE_AI_API_BASE") native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params))) - root: Final = foundry_root(base_target_url).removesuffix(f"/{native_endpoint.strip('/')}") + root: Final = without_repeated_native_prefix(foundry_root(base_target_url), native_endpoint) query_params: Final = relay_query_params( request_query_params, api_version_from(litellm_params), base_target_url ) diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 064be954518..1b9036c50c5 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -135,6 +135,24 @@ def test_full_url_api_base_that_already_ends_with_the_native_path_is_not_doubled assert base == "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router" +@pytest.mark.parametrize("relayed_deployment", ["gpt-4o", "GPT-4o"]) +def test_deployment_root_api_base_is_not_repeated_when_the_relay_carries_the_deployment_path(relayed_deployment): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com/openai/deployments/gpt-4o", + api_key="key", + model="gpt-4o", + endpoint=f"aoai-gpt-4o/openai/deployments/{relayed_deployment}/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={"litellm_metadata": {"model_group": "aoai-gpt-4o"}}, + ) + + assert str(url) == ( + f"https://my-resource.openai.azure.com/openai/deployments/{relayed_deployment}/chat/completions" + "?api-version=2024-10-21" + ) + assert base == "https://my-resource.openai.azure.com" + + def test_parse_relay_under_a_models_api_base_targets_the_foundry_root(): url, _ = AzureAIPassthroughConfig().get_complete_url( api_base=f"{FOUNDRY_BASE}/models", From 708381948b0025f6f54e34b9e0e22b8d02070801 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:34:07 -0700 Subject: [PATCH 084/107] fix(cli): compare the resolved settings path when deciding whether the lite up and autoroute backups guard it --- .../proxy/client/cli/commands/claude_settings.py | 2 +- .../proxy/client/cli/test_auth_commands.py | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index 7e96d2a3a08..793728b0350 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -66,7 +66,7 @@ def claude_settings_path(environ: Mapping[str, str]) -> Path: def settings_file_owners(settings_path: Path) -> tuple[SettingsFileOwner, ...]: """The commands whose backups guard settings_path: `lite up` and `lite autoroute up` only ever manage the default file.""" - return SETTINGS_FILE_OWNERS if settings_path == CLAUDE_SETTINGS_PATH else () + return SETTINGS_FILE_OWNERS if settings_path.resolve() == CLAUDE_SETTINGS_PATH.resolve() else () def load_json_or_empty(path: Path) -> dict[str, JsonValue]: diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 9c3c02a5e85..d3e6d8f66ef 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -1477,6 +1477,21 @@ class TestLoginConfigClaude: assert "Run `lite down` first" in result.output assert not default_settings_path.exists() + def test_flag_refuses_while_lite_up_holds_the_default_file_reached_through_a_symlinked_config_dir(self, tmp_path): + default_config_dir = tmp_path / "default-home" / ".claude" + default_config_dir.mkdir(parents=True) + alias = tmp_path / "claude-alias" + alias.symlink_to(default_config_dir, target_is_directory=True) + (tmp_path / "claude_settings_backup.json").write_text("{}") + + result, _settings_path, _backup_path = self._run_login( + tmp_path, ["--config-claude"], config_dir_env={"CLAUDE_CONFIG_DIR": str(alias)} + ) + + assert result.exit_code != 0 + assert "`lite up` is currently managing" in result.output + assert not (default_config_dir / "settings.json").exists() + def test_flag_writes_an_alternate_config_dir_even_while_lite_up_holds_the_default_file(self, tmp_path): (tmp_path / "claude_settings_backup.json").write_text("{}") From 71cf1c7901adae5d7dd244dc98274c4ee109f149 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:34:15 -0700 Subject: [PATCH 085/107] test(e2e): drop the unused batch list client and fail loudly on a user teardown miss The failed-batch redesign left list_batches, BatchList, BatchListQuery and the batch object's metadata and created_at fields with no caller, and they duplicated the batches suite's own client. delete_user discarded its result, so a user that outlived the class fixture went unnoticed; unwrap turns that into a teardown error like delete_key already does. --- .../spend_tracking/spend_e2e_client.py | 33 ++++--------------- 1 file changed, 7 insertions(+), 26 deletions(-) diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 18b714d8aeb..9ac97f57f47 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -135,17 +135,6 @@ class BatchCreateBody(BaseModel): class BatchObject(BaseModel): id: str status: str - metadata: dict[str, str] | None = None - created_at: int | None = None - - -class BatchList(BaseModel): - data: list[BatchObject] = [] - - -class BatchListQuery(BaseModel): - model: str - limit: int class ProviderQuery(BaseModel): @@ -381,11 +370,13 @@ class SpendClient: ).user_id def delete_user(self, user_id: str) -> None: - _ = self.proxy.transport.post( - "/user/delete", - headers=self.proxy.transport.master, - json=UserDeleteBody(user_ids=[user_id]), - response_type=UserDeleteResponse, + _ = unwrap( + self.proxy.transport.post( + "/user/delete", + headers=self.proxy.transport.master, + json=UserDeleteBody(user_ids=[user_id]), + response_type=UserDeleteResponse, + ) ) def generate_key_record(self, body: KeyGenerateBody) -> KeyGenerateResponse: @@ -474,16 +465,6 @@ class SpendClient: ) ) - def list_batches(self, key: str, model: str, *, limit: int) -> list[BatchObject]: - return unwrap( - self.proxy.transport.get( - "/v1/batches", - headers=self.proxy.transport.bearer(key), - params=BatchListQuery(model=model, limit=limit), - response_type=BatchList, - ) - ).data - def retrieve_batch(self, key: str, batch_id: str, *, provider: str) -> BatchObject: return unwrap( self.proxy.transport.get( From d93baa3f2c1ca3c5123ecdee8510ee3c8d827f61 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:52:00 -0700 Subject: [PATCH 086/107] fix(bedrock): sign on a dedicated executor instead of the shared default one asyncio.to_thread puts every Bedrock signing on the loop's default executor, the same pool every provider's async entry point hops through, so signings parked on botocore's refresh lock queued unrelated providers behind Bedrock. run_aws_signing runs them on a 16-thread pool only AWS signing uses --- .../providers/bedrock_agentcore/handler.py | 6 ++--- litellm/constants.py | 1 + litellm/integrations/s3_v2.py | 6 ++--- litellm/integrations/sqs.py | 4 +-- litellm/llms/bedrock/base_aws_llm.py | 20 +++++++++++++- litellm/llms/bedrock/chat/converse_handler.py | 7 +++-- litellm/llms/bedrock/count_tokens/handler.py | 4 +-- litellm/llms/bedrock/embed/embedding.py | 7 +++-- litellm/llms/bedrock/realtime/handler.py | 6 ++--- litellm/llms/custom_httpx/llm_http_handler.py | 4 +-- .../guardrail_hooks/bedrock_guardrails.py | 10 +++---- .../llm_passthrough_endpoints.py | 9 +++---- .../llms/bedrock/test_base_aws_llm.py | 26 +++++++++++++++++++ 13 files changed, 76 insertions(+), 34 deletions(-) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index e01e215b5f0..306a8871b12 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -5,7 +5,6 @@ Sends JSON-RPC envelopes directly to AgentCore endpoints, bypassing the completion bridge that would otherwise strip the envelope. """ -import asyncio import json from collections.abc import AsyncIterator, Mapping from typing import Any, Final @@ -14,6 +13,7 @@ from litellm._logging import verbose_logger from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( BedrockAgentCoreA2ATransformation, ) +from litellm.llms.bedrock.base_aws_llm import run_aws_signing from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -46,7 +46,7 @@ class BedrockAgentCoreA2AHandler: Returns: A2A JSON-RPC response dict from the AgentCore agent """ - url, headers, body = await asyncio.to_thread( + url, headers, body = await run_aws_signing( BedrockAgentCoreA2ATransformation.get_url_and_signed_request, request_id=request_id, params=params, @@ -93,7 +93,7 @@ class BedrockAgentCoreA2AHandler: Yields: A2A streaming response events from the AgentCore agent """ - url, headers, body = await asyncio.to_thread( + url, headers, body = await run_aws_signing( BedrockAgentCoreA2ATransformation.get_url_and_signed_request, request_id=request_id, params=params, diff --git a/litellm/constants.py b/litellm/constants.py index f9389d22dea..060402cc08d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -570,6 +570,7 @@ LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float( LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100) LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000) LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0 +AWS_SIGNING_MAX_THREADS: Final = 16 DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv( "DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield" ) diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index e3297cfd27c..996ff9c75af 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -24,7 +24,7 @@ from litellm.integrations.s3 import ( from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, @@ -366,7 +366,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Sign the request aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers) aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) - await asyncio.to_thread(S3SigV4Auth(credentials, "s3", aws_region_name).add_auth, aws_request) + await run_aws_signing(S3SigV4Auth(credentials, "s3", aws_region_name).add_auth, aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) @@ -597,7 +597,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Sign the request aws_request: Final = AWSRequest(method="GET", url=url, headers=headers) - await asyncio.to_thread(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request) + await run_aws_signing(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 31cc20f5995..d787375ca3c 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -22,7 +22,7 @@ from litellm.constants import ( SQS_SEND_MESSAGE_ACTION, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -295,7 +295,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): data=prepped.body, headers=prepped.headers, ) - await asyncio.to_thread(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request) + await run_aws_signing(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request) signed_headers: Final = dict(aws_request.headers.items()) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index eddd4d6f826..96804a1fa62 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1,12 +1,15 @@ import asyncio import base64 +import contextvars import hashlib import json import os import re import urllib.parse from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor from datetime import datetime +from functools import partial from threading import Lock from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload @@ -17,6 +20,7 @@ from litellm._logging import verbose_logger from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( + AWS_SIGNING_MAX_THREADS, BEDROCK_EMBEDDING_PROVIDERS_LITERAL, BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES, BEDROCK_IAM_CACHE_MAX_ENTRIES, @@ -1697,6 +1701,20 @@ def sign_aws_json_post( _SignParams = ParamSpec("_SignParams") _SignedRequest = TypeVar("_SignedRequest") +AWS_SIGNING_EXECUTOR: Final = ThreadPoolExecutor(max_workers=AWS_SIGNING_MAX_THREADS, thread_name_prefix="aws-signing") + + +async def run_aws_signing( + sign: Callable[_SignParams, _SignedRequest], + /, + *args: _SignParams.args, + **kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped signing signature +) -> _SignedRequest: + context: Final = contextvars.copy_context() + return await asyncio.get_running_loop().run_in_executor( + AWS_SIGNING_EXECUTOR, partial(context.run, sign, *args, **kwargs) + ) + async def sign_request_off_loop_if_aws( provider_config: object, @@ -1706,5 +1724,5 @@ async def sign_request_off_loop_if_aws( **kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature ) -> _SignedRequest: if isinstance(provider_config, SignsRequestsWithAWS): - return await asyncio.to_thread(sign_request, *args, **kwargs) + return await run_aws_signing(sign_request, *args, **kwargs) return sign_request(*args, **kwargs) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 3e8e1226a55..6d48ff3f07c 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -1,4 +1,3 @@ -import asyncio import json from collections.abc import Mapping from types import MappingProxyType @@ -22,7 +21,7 @@ from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper -from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token +from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call @@ -137,7 +136,7 @@ class BedrockConverseLLM(BaseAWSLLM): ) data: Final = json.dumps(request_data) - prepped: Final = await asyncio.to_thread( + prepped: Final = await run_aws_signing( self.get_request_headers, credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", @@ -208,7 +207,7 @@ class BedrockConverseLLM(BaseAWSLLM): ) data: Final = json.dumps(request_data) - prepped: Final = await asyncio.to_thread( + prepped: Final = await run_aws_signing( self.get_request_headers, credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index cae2315e744..d7fb510f057 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -4,13 +4,13 @@ AWS Bedrock CountTokens API handler. Simplified handler leveraging existing LiteLLM Bedrock infrastructure. """ -import asyncio from typing import Any, Final import httpx import litellm from litellm._logging import verbose_logger +from litellm.llms.bedrock.base_aws_llm import run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client @@ -77,7 +77,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): # Extract api_key for bearer token auth if provided api_key: Final = litellm_params.get("api_key", None) headers: Final = {"Content-Type": "application/json"} - signed_headers, signed_body = await asyncio.to_thread( + signed_headers, signed_body = await run_aws_signing( self._sign_request, service_name="bedrock", headers=headers, diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 4254d9b1d8b..be766eaedd0 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -2,7 +2,6 @@ Handles embedding calls to Bedrock's `/invoke` endpoint """ -import asyncio import copy import json import urllib.parse @@ -27,7 +26,7 @@ from litellm.types.llms.bedrock import ( ) from litellm.types.utils import EmbeddingResponse, LlmProviders -from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token +from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing from ..common_utils import BedrockError from .amazon_nova_transformation import AmazonNovaEmbeddingConfig from .amazon_titan_g1_transformation import AmazonTitanG1Config @@ -357,7 +356,7 @@ class BedrockEmbedding(BaseAWSLLM): if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - prepped = await asyncio.to_thread( + prepped = await run_aws_signing( self.get_request_headers, credentials=credentials, aws_region_name=aws_region_name, @@ -638,7 +637,7 @@ class BedrockEmbedding(BaseAWSLLM): credentials=credentials, url=status_url, headers=headers, aws_region_name=aws_region_name ) - prepped: Final = await asyncio.to_thread(sign_status_request) + prepped: Final = await run_aws_signing(sign_status_request) # LOGGING if logging_obj is not None: diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 792e3169a7a..ca2370303f2 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -21,7 +21,7 @@ from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeE from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput -from ..base_aws_llm import BaseAWSLLM +from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig @@ -149,7 +149,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Connecting to %s with model %s", endpoint_uri, model) - credentials: Final = await asyncio.to_thread( + credentials: Final = await run_aws_signing( self.get_credentials, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, @@ -170,7 +170,7 @@ class BedrockRealtime(BaseAWSLLM): "or configure credentials in the environment" ), ) - frozen_credentials: Final = await asyncio.to_thread(credentials.get_frozen_credentials) + frozen_credentials: Final = await run_aws_signing(credentials.get_frozen_credentials) # Initialize Bedrock client with aws_sdk_bedrock_runtime config: Final = Config( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3106810b4b3..13f134fce1d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -77,7 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig -from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws +from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_request_off_loop_if_aws from litellm.llms.custom_httpx.container_handler import raise_for_error_status from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -638,7 +638,7 @@ class BaseLLMHTTPHandler: headers=request_headers, ), ) - return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed)) + return await dispatch_async(*await run_aws_signing(sign_and_log, transformed)) return transform_then_dispatch() diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 0ae2c48a3b4..6a7ac4361b9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -44,7 +44,7 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicM from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token, run_aws_signing from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -917,7 +917,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source, ) return BedrockGuardrailResponse() - credentials, aws_region_name = await asyncio.to_thread( + credentials, aws_region_name = await run_aws_signing( self._load_credentials, bearer_token=bedrock_bearer_token(api_key) ) allow_chunking: Final = not self._content_uses_contextual_grounding(content) @@ -1180,7 +1180,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): **base_request_data, "content": content, } # mutable-ok: outbound JSON request body - prepared_request: Final = await asyncio.to_thread( + prepared_request: Final = await run_aws_signing( self._prepare_request, credentials=credentials, data=bedrock_request_data, @@ -1878,12 +1878,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return BedrockGuardrailResponse() api_key: Final[str | None] = request_data.get("api_key") if request_data else None - credentials, aws_region_name = await asyncio.to_thread( + credentials, aws_region_name = await run_aws_signing( self._load_credentials, bearer_token=bedrock_bearer_token(api_key) ) body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks} - prepared_request: Final = await asyncio.to_thread( + prepared_request: Final = await run_aws_signing( self._prepare_request, credentials=credentials, data=body, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 037f663044f..bf6ee57b4fc 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -8,7 +8,6 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc. from __future__ import annotations -import asyncio import hmac import inspect import json @@ -1131,7 +1130,7 @@ async def bedrock_proxy_route( ) # Add or update query parameters - from litellm.llms.bedrock.base_aws_llm import sign_aws_json_post + from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post from litellm.llms.bedrock.chat import BedrockConverseLLM bedrock_llm: Final = BedrockConverseLLM() @@ -1140,7 +1139,7 @@ async def bedrock_proxy_route( data: Final = await _json_request_body(request) except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) - prepped: Final = await asyncio.to_thread( + prepped: Final = await run_aws_signing( sign_aws_json_post, get_credentials=bedrock_llm.get_credentials, service_name="bedrock", @@ -1236,10 +1235,10 @@ async def comprehend_medical_proxy_route( if "stream" in data: raise HTTPException(status_code=400, detail="'stream' is not a Comprehend Medical request member") - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, sign_aws_json_post + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" - prepped: Final = await asyncio.to_thread( + prepped: Final = await run_aws_signing( sign_aws_json_post, get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name), service_name="comprehendmedical", diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index ca93f58f366..e24ea1ca611 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1,5 +1,6 @@ import asyncio import json +from concurrent.futures import ThreadPoolExecutor import os import threading import time @@ -23,6 +24,7 @@ from litellm.llms.bedrock.base_aws_llm import ( AwsAuthError, BaseAWSLLM, Boto3CredentialsInfo, + run_aws_signing, sign_request_off_loop_if_aws, ) from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe @@ -3247,3 +3249,27 @@ async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credent assert "Authorization" in signed assert probe.served_during_refresh is True + + +async def test_run_aws_signing_leaves_the_default_executor_free_for_other_providers(): + """A signing parked on botocore's refresh lock must not hold a default-executor thread, since every + other provider's async entry point hops through that same executor.""" + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + signing_parked = asyncio.Event() + refresh_done = threading.Event() + + def sign() -> str: + loop.call_soon_threadsafe(signing_parked.set) + refresh_done.wait() + return threading.current_thread().name + + signing = asyncio.create_task(run_aws_signing(sign)) + try: + await asyncio.wait_for(signing_parked.wait(), timeout=5) + other_provider = await asyncio.wait_for(loop.run_in_executor(None, threading.current_thread), timeout=5) + finally: + refresh_done.set() + + assert other_provider.name != await signing + assert (await signing).startswith("aws-signing") From 9c6a385efe1f61cbd9d276ed09427102aec820ff Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:57:37 -0700 Subject: [PATCH 087/107] docs(azure_ai): drop size from the MAI image edit router docstring --- litellm/llms/azure_ai/image_edit/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/image_edit/__init__.py b/litellm/llms/azure_ai/image_edit/__init__.py index 51a23859058..fda9335a5d6 100644 --- a/litellm/llms/azure_ai/image_edit/__init__.py +++ b/litellm/llms/azure_ai/image_edit/__init__.py @@ -23,7 +23,7 @@ def get_azure_ai_image_edit_config(model: str) -> BaseImageEditConfig: """ Get the appropriate image edit config for an Azure AI model. - - MAI models use /mai/v1/images/edits with multipart form data and size + - MAI models use /mai/v1/images/edits with multipart form data - FLUX 2 models use JSON with base64 image - FLUX 1 models use multipart/form-data """ From 8627576c9c37bc475ac35d98ab6f44d1ea8978d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:11:04 -0700 Subject: [PATCH 088/107] fix(azure_ai): keep a deployment named like the first native path segment in the relayed URL --- .../llms/azure_ai/passthrough/transformation.py | 5 +++++ .../test_azure_ai_passthrough_transformation.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index ed630bc6f2b..f2be1d95593 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -62,6 +62,10 @@ def foundry_root(api_base: str) -> str: return str(url.copy_with(path="/" + "/".join(root_segments), query=None)).rstrip("/") +def is_repeated_native_prefix(native_segments: tuple[str, ...], overlap: int) -> bool: + return overlap == len(native_segments) or native_segments[0] == "openai" + + def without_repeated_native_prefix(root: str, native_endpoint: str) -> str: url: Final = httpx.URL(root) root_segments: Final = tuple(segment for segment in url.path.split("/") if segment) @@ -71,6 +75,7 @@ def without_repeated_native_prefix(root: str, native_endpoint: str) -> str: length for length in range(min(len(root_segments), len(native_segments)), 0, -1) if tuple(segment.casefold() for segment in root_segments[-length:]) == native_segments[:length] + and is_repeated_native_prefix(native_segments, length) ), 0, ) diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 1b9036c50c5..c8007acf70f 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -153,6 +153,20 @@ def test_deployment_root_api_base_is_not_repeated_when_the_relay_carries_the_dep assert base == "https://my-resource.openai.azure.com" +def test_deployment_named_like_the_first_native_segment_keeps_its_deployment_root(): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com/openai/deployments/chat", + api_key="key", + model="chat", + endpoint="aoai-chat/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={"litellm_metadata": {"model_group": "aoai-chat"}}, + ) + + assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/chat/chat/completions?api-version=2024-10-21" + assert base == "https://my-resource.openai.azure.com/openai/deployments/chat" + + def test_parse_relay_under_a_models_api_base_targets_the_foundry_root(): url, _ = AzureAIPassthroughConfig().get_complete_url( api_base=f"{FOUNDRY_BASE}/models", From 3dcc09b26ccb0aa781e7bb62007e73e93c2bd9f6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:16:11 -0700 Subject: [PATCH 089/107] fix(bedrock): keep the signing pool to AWS configs on the async-transform path The async-transform path in the shared handler sent every provider's sign_and_log to the AWS pool, so Ollama, Snowflake, and watsonx queued behind Bedrock refreshes there. Only SignsRequestsWithAWS configs take run_aws_signing now, the rest keep the default-executor hop they had. The executor isolation test also runs on its own loop instead of pinning a one-thread default executor on the session-scoped pytest loop --- litellm/llms/custom_httpx/llm_http_handler.py | 9 +++- .../llms/bedrock/test_base_aws_llm.py | 41 +++++++++++-------- .../custom_httpx/test_llm_http_handler.py | 21 ++++++++++ 3 files changed, 51 insertions(+), 20 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 13f134fce1d..e720428847d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -77,7 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig -from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_request_off_loop_if_aws +from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS, run_aws_signing, sign_request_off_loop_if_aws from litellm.llms.custom_httpx.container_handler import raise_for_error_status from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -638,7 +638,12 @@ class BaseLLMHTTPHandler: headers=request_headers, ), ) - return await dispatch_async(*await run_aws_signing(sign_and_log, transformed)) + signed_request: Final = await ( + run_aws_signing(sign_and_log, transformed) + if isinstance(provider_config, SignsRequestsWithAWS) + else asyncio.to_thread(sign_and_log, transformed) + ) + return await dispatch_async(*signed_request) return transform_then_dispatch() diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index e24ea1ca611..2c7e476e9a8 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -3251,25 +3251,30 @@ async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credent assert probe.served_during_refresh is True -async def test_run_aws_signing_leaves_the_default_executor_free_for_other_providers(): +def test_run_aws_signing_leaves_the_default_executor_free_for_other_providers(): """A signing parked on botocore's refresh lock must not hold a default-executor thread, since every - other provider's async entry point hops through that same executor.""" - loop = asyncio.get_running_loop() - loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) - signing_parked = asyncio.Event() - refresh_done = threading.Event() + other provider's async entry point hops through that same executor. The scenario runs on its own loop + so the one-thread default executor it pins never leaks into the session loop.""" - def sign() -> str: - loop.call_soon_threadsafe(signing_parked.set) - refresh_done.wait() - return threading.current_thread().name + async def scenario() -> tuple[str, str]: + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + signing_parked = asyncio.Event() + refresh_done = threading.Event() - signing = asyncio.create_task(run_aws_signing(sign)) - try: - await asyncio.wait_for(signing_parked.wait(), timeout=5) - other_provider = await asyncio.wait_for(loop.run_in_executor(None, threading.current_thread), timeout=5) - finally: - refresh_done.set() + def sign() -> str: + loop.call_soon_threadsafe(signing_parked.set) + refresh_done.wait() + return threading.current_thread().name - assert other_provider.name != await signing - assert (await signing).startswith("aws-signing") + signing = asyncio.create_task(run_aws_signing(sign)) + try: + await asyncio.wait_for(signing_parked.wait(), timeout=5) + other_provider = await asyncio.wait_for(loop.run_in_executor(None, threading.current_thread), timeout=5) + finally: + refresh_done.set() + return other_provider.name, await signing + + other_provider, signing_thread = asyncio.run(scenario()) + assert other_provider != signing_thread + assert signing_thread.startswith("aws-signing") 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 1140090f9cf..cea2d439198 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 @@ -20,6 +20,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( BaseAudioTranscriptionConfig, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( @@ -3431,6 +3432,26 @@ async def test_completion_signs_and_logs_off_the_event_loop_after_the_async_tran assert captured["body"] == {"transformed_by": "async"} assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) assert pre_call_threads and all(thread is not loop_thread for thread in pre_call_threads) + assert not any(thread.name.startswith("aws-signing") for thread in config.sign_threads + pre_call_threads) + + +class _AWSTransformRecordingConfig(SignsRequestsWithAWS, _TransformRecordingConfig): + pass + + +async def test_completion_signs_aws_configs_on_the_aws_signing_pool_after_the_async_transform(): + config = _AWSTransformRecordingConfig(transform_async=True) + pre_call_threads = [] + logging_obj = Mock(dynamic_success_callbacks=None, model_call_details={}) + logging_obj.pre_call.side_effect = lambda **kwargs: pre_call_threads.append(threading.current_thread()) + + pending, captured = _start_async_completion(config, logging_obj) + response = await pending + + assert response.choices[0].message.content == "async" + assert captured["body"] == {"transformed_by": "async"} + assert config.sign_threads and all(thread.name.startswith("aws-signing") for thread in config.sign_threads) + assert pre_call_threads and all(thread.name.startswith("aws-signing") for thread in pre_call_threads) async def test_completion_keeps_sync_transform_request_before_returning_by_default(): From 5a1be5642606eb56945e0dfd91ca02647a592a52 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:27:54 -0700 Subject: [PATCH 090/107] fix(router): drop bridge-tagged reasoning blocks whole when the minting deployment is not in the routed group The cross-group branch of EncryptedContentAffinityCheck removed only the signature from Anthropic-shaped thinking blocks, which left unsigned thinking blocks that Anthropic and Bedrock reject (thinking.signature: Field required). Drop the whole block, the way #40280 drops undecryptable Responses input items, so the routed request carries the conversation text with no reasoning item for those turns --- .../prompt_templates/common_utils.py | 25 +++++++------------ ...ore_utils_prompt_templates_common_utils.py | 4 +-- .../test_encrypted_content_affinity_check.py | 8 +++--- 3 files changed, 16 insertions(+), 21 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 7ca2a8aad6a..d70530534da 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1890,8 +1890,13 @@ def is_encrypted_reasoning_block(block: object) -> bool: def strip_encrypted_reasoning_from_messages(messages: object) -> None: - """Drop the encrypted reasoning a routed deployment cannot decrypt from Anthropic-shaped - history, keeping the readable thinking text. + """Drop the bridge-tagged reasoning blocks a routed deployment cannot decrypt from + Anthropic-shaped history. + + The whole block goes, the way #40280 drops undecryptable Responses ``input`` items: a + provider that did not mint the block rejects it signed (a foreign signature) and unsigned + (a missing signature) alike, so keeping its text as an unsigned thinking block only moves + the 400 from the router to the provider. Mutates the content lists in place: the router's fallback snapshot shares these message objects, so a rebound list would replay the stripped blocks on the fallback hop. @@ -1914,20 +1919,8 @@ def _anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: def _strip_encrypted_reasoning_from_blocks(content: object) -> None: blocks: Final = cast(list[object], content) # cast-ok: narrowed by the caller's isinstance - stripped: Final = tuple(_without_encrypted_reasoning_block(block) for block in blocks) - blocks[:] = (block for block in stripped if block is not None) # rebind-ok: list shared with fallback snapshot - - -def _without_encrypted_reasoning_block(block: object) -> object | None: - if not is_encrypted_reasoning_block(block): - return block - mapping: Final = cast(Mapping[str, object], block) # cast-ok: narrowed by is_encrypted_reasoning_block - if mapping.get("type") != "thinking" or not mapping.get("thinking"): - return None - kept: Final[dict[str, object]] = { # mutable-ok: thinking block rebuilt without the undecryptable signature - key: value for key, value in mapping.items() if key != "signature" - } - return kept + kept: Final = tuple(block for block in blocks if not is_encrypted_reasoning_block(block)) + blocks[:] = kept # rebind-ok: shared with fallback snapshot def _reasoning_replay_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str: diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 71b71d55fa7..7c1445d79c9 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -1633,7 +1633,7 @@ class TestEncryptedReasoningReplay: def test_is_encrypted_reasoning_block(self, block, expected): assert is_encrypted_reasoning_block(block) is expected - def test_strip_keeps_the_readable_thinking_and_drops_the_undecryptable_bytes(self): + def test_strip_drops_every_bridge_tagged_block_and_leaves_no_unsigned_thinking_behind(self): assistant_content = [ {"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"}, {"type": "thinking", "thinking": "packed by the bridge", "signature": encrypted_reasoning_signature("g1")}, @@ -1652,9 +1652,9 @@ class TestEncryptedReasoningReplay: assert messages[1]["content"] is assistant_content assert assistant_content == [ {"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"}, - {"type": "thinking", "thinking": "packed by the bridge"}, {"type": "text", "text": "answer"}, ] + assert all(block["signature"] for block in assistant_content if block["type"] == "thinking") assert messages[0] == {"role": "user", "content": "question"} assert messages[2] == {"role": "user", "content": [{"type": "text", "text": "follow-up"}]} diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 57c02e43fc1..ea8e2eacaa6 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -1635,8 +1635,10 @@ def _bridge_replayed_anthropic_messages(minted_by: str) -> list: async def test_encrypted_content_affinity_strips_bridge_reasoning_from_messages_routed_to_another_group(): """ The /v1/messages twin of the tier-change case: the routed group holds no deployment - of the org that minted the reasoning, so the bridge-tagged blocks are stripped down - to their readable thinking text and the request dispatches to the routed pool. + of the org that minted the reasoning, so the bridge-tagged blocks are dropped whole + and the request dispatches to the routed pool. No unsigned thinking block may be left + behind: Anthropic and Bedrock reject a thinking block with a missing signature the + same way they reject a foreign one. """ originating = _make_originating_mock(None, "key-a", model_name="gpt-reasoning-tier") mock_router = _make_router_mock_with_cooldown( @@ -1660,9 +1662,9 @@ async def test_encrypted_content_affinity_strips_bridge_reasoning_from_messages_ assert messages[1]["content"] is assistant_content assert assistant_content == [ {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, - {"type": "thinking", "thinking": "The bridge packed this one"}, {"type": "text", "text": "The zebra owner lives in the green house."}, ] + assert all(block["signature"] for block in assistant_content if block["type"] == "thinking") class TestStripEncryptedReasoningFromInput: From b64197c2898dc6db980ec5b944ce6189caf3ae84 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:44:36 -0700 Subject: [PATCH 091/107] fix(terraform): preserve undeclared server-side key metadata on update (#40514) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/provider/litellm/client.go | 14 +- terraform/provider/litellm/resource_key.go | 59 +++++++- .../provider/litellm/resource_key_test.go | 133 ++++++++++++++++++ 3 files changed, 202 insertions(+), 4 deletions(-) diff --git a/terraform/provider/litellm/client.go b/terraform/provider/litellm/client.go index 0f825d85d31..0dfbe3dd3c9 100644 --- a/terraform/provider/litellm/client.go +++ b/terraform/provider/litellm/client.go @@ -80,16 +80,24 @@ func (c *Client) UpdateKey(key *Key) (*Key, error) { updateData := map[string]interface{}{ "key": key.Key, "team_id": key.TeamID, - "metadata": key.Metadata, "key_alias": key.KeyAlias, "aliases": key.Aliases, "permissions": key.Permissions, "model_max_budget": key.ModelMaxBudget, - "model_rpm_limit": key.ModelRPMLimit, - "model_tpm_limit": key.ModelTPMLimit, "blocked": key.Blocked, } + // The proxy keeps the stored metadata only when the field is absent, so nil means omit. + if key.Metadata != nil { + updateData["metadata"] = key.Metadata + } + if key.ModelRPMLimit != nil { + updateData["model_rpm_limit"] = key.ModelRPMLimit + } + if key.ModelTPMLimit != nil { + updateData["model_tpm_limit"] = key.ModelTPMLimit + } + // The proxy rejects an empty-string budget_duration with a 400, so only // send it when set. if key.BudgetDuration != "" { diff --git a/terraform/provider/litellm/resource_key.go b/terraform/provider/litellm/resource_key.go index 0d8674f2d4c..a37041fe7e8 100644 --- a/terraform/provider/litellm/resource_key.go +++ b/terraform/provider/litellm/resource_key.go @@ -223,6 +223,7 @@ func resourceKeyRead(ctx context.Context, d *schema.ResourceData, m interface{}) return nil } + key.Metadata = declaredKeyMetadata(key.Metadata, d.Get("metadata").(map[string]interface{})) mapKeyToResourceData(d, key) return nil } @@ -232,15 +233,71 @@ func resourceKeyUpdate(ctx context.Context, d *schema.ResourceData, m interface{ key := &Key{Key: d.Id()} mapResourceDataToKey(d, key) + key.ModelRPMLimit = changedMap(d, "model_rpm_limit") + key.ModelTPMLimit = changedMap(d, "model_tpm_limit") - _, err := c.UpdateKey(key) + metadata, err := plannedKeyMetadata(c, d) if err != nil { return diag.FromErr(fmt.Errorf("error updating key: %s", err)) } + key.Metadata = metadata + + if _, err := c.UpdateKey(key); err != nil { + return diag.FromErr(fmt.Errorf("error updating key: %s", err)) + } return resourceKeyRead(ctx, d, m) } +func changedMap(d *schema.ResourceData, name string) map[string]interface{} { + if !d.HasChange(name) { + return nil + } + return d.Get(name).(map[string]interface{}) +} + +func plannedKeyMetadata(c *Client, d *schema.ResourceData) (map[string]interface{}, error) { + if !d.HasChange("metadata") { + return nil, nil + } + current, err := c.GetKey(d.Id()) + if err != nil { + return nil, err + } + if current == nil { + return nil, fmt.Errorf("key %s no longer exists", d.Id()) + } + oldDeclared, newDeclared := d.GetChange("metadata") + return mergeKeyMetadata(current.Metadata, oldDeclared.(map[string]interface{}), newDeclared.(map[string]interface{})), nil +} + +func declaredKeyMetadata(server, declared map[string]interface{}) map[string]interface{} { + if server == nil { + return nil + } + result := make(map[string]interface{}, len(declared)) + for k := range declared { + if v, ok := server[k]; ok { + result[k] = v + } + } + return result +} + +func mergeKeyMetadata(server, oldDeclared, newDeclared map[string]interface{}) map[string]interface{} { + result := make(map[string]interface{}, len(server)+len(newDeclared)) + for k, v := range server { + result[k] = v + } + for k := range oldDeclared { + delete(result, k) + } + for k, v := range newDeclared { + result[k] = v + } + return result +} + func resourceKeyDelete(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics { c := m.(*Client) diff --git a/terraform/provider/litellm/resource_key_test.go b/terraform/provider/litellm/resource_key_test.go index 91f0061a9ef..d8f3cfe86b3 100644 --- a/terraform/provider/litellm/resource_key_test.go +++ b/terraform/provider/litellm/resource_key_test.go @@ -6,9 +6,11 @@ import ( "io" "net/http" "net/http/httptest" + "reflect" "testing" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" ) func newKeyResourceData(t *testing.T, raw map[string]interface{}) *schema.ResourceData { @@ -254,3 +256,134 @@ func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) { t.Errorf("RPMLimit not parsed: %+v", key.RPMLimit) } } + +// fakeKeyProxy serves /key/info from stored metadata and applies /key/update +// the way the proxy does: an absent "metadata" keeps the stored map, a +// present one replaces it wholesale. +type fakeKeyProxy struct { + metadata map[string]interface{} + updates []map[string]interface{} +} + +func (p *fakeKeyProxy) handler() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/key/info": + json.NewEncoder(w).Encode(map[string]interface{}{ + "key": "hash-1", + "info": map[string]interface{}{"key_alias": "alias-1", "models": []string{"gpt-4o-mini"}, "metadata": p.metadata}, + }) + case "/key/update": + var body map[string]interface{} + json.NewDecoder(r.Body).Decode(&body) + p.updates = append(p.updates, body) + if m, ok := body["metadata"].(map[string]interface{}); ok { + p.metadata = m + } + json.NewEncoder(w).Encode(map[string]interface{}{"key": "hash-1", "metadata": p.metadata}) + default: + http.NotFound(w, r) + } + } +} + +func applyKeyUpdate(t *testing.T, client *Client, stateAttrs map[string]string, config map[string]interface{}) *terraform.InstanceState { + t.Helper() + r := resourceKey() + state := &terraform.InstanceState{ID: "hash-1", Attributes: stateAttrs} + diff, err := r.Diff(context.Background(), state, terraform.NewResourceConfigRaw(config), client) + if err != nil { + t.Fatalf("Diff returned error: %v", err) + } + if diff == nil { + t.Fatalf("expected a non-empty diff between %v and %v", stateAttrs, config) + } + newState, diags := r.Apply(context.Background(), state, diff, client) + if diags.HasError() { + t.Fatalf("Apply returned error: %v", diags) + } + return newState +} + +func TestKeyUpdateWithoutMetadataChangePreservesServerMetadata(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "server_side": "x", "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": float64(5)}}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + newState := applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "max_budget": "10", "metadata.%": "1", "metadata.a": "1"}, + map[string]interface{}{"key_alias": "alias-1", "max_budget": 20, "metadata": map[string]interface{}{"a": "1"}}, + ) + + if len(proxy.updates) != 1 { + t.Fatalf("expected one /key/update call, got %d", len(proxy.updates)) + } + for _, field := range []string{"metadata", "model_rpm_limit", "model_tpm_limit"} { + if _, present := proxy.updates[0][field]; present { + t.Errorf("unchanged %q was sent on /key/update: %v", field, proxy.updates[0][field]) + } + } + if proxy.metadata["server_side"] != "x" { + t.Errorf("server-side metadata lost: %v", proxy.metadata) + } + if got := newState.Attributes["metadata.%"]; got != "1" { + t.Errorf("state metadata should hold only the declared entry, got %v", newState.Attributes) + } + if got := newState.Attributes["metadata.a"]; got != "1" { + t.Errorf("metadata.a = %q, want 1", got) + } +} + +func TestKeyUpdateWithMetadataChangeMergesOverServerMetadata(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "b": "2", "server_side": "x"}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "metadata.%": "2", "metadata.a": "1", "metadata.b": "2"}, + map[string]interface{}{"key_alias": "alias-1", "metadata": map[string]interface{}{"a": "2", "c": "3"}}, + ) + + want := map[string]interface{}{"a": "2", "c": "3", "server_side": "x"} + if !reflect.DeepEqual(proxy.metadata, want) { + t.Errorf("metadata after update = %v, want %v", proxy.metadata, want) + } +} + +func TestKeyUpdateSendsChangedModelLimits(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "model_rpm_limit.%": "1", "model_rpm_limit.gpt-4o-mini": "5"}, + map[string]interface{}{"key_alias": "alias-1", "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 7}}, + ) + + got, ok := proxy.updates[0]["model_rpm_limit"].(map[string]interface{}) + if !ok || got["gpt-4o-mini"] != float64(7) { + t.Errorf("changed model_rpm_limit not sent: %v", proxy.updates[0]) + } +} + +func TestKeyReadKeepsOnlyDeclaredMetadata(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "server_side": "x"}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + d := newKeyResourceData(t, map[string]interface{}{"metadata": map[string]interface{}{"a": "1"}}) + d.SetId("hash-1") + if diags := resourceKeyRead(context.Background(), d, client); diags.HasError() { + t.Fatalf("Read returned error: %v", diags) + } + + want := map[string]interface{}{"a": "1"} + if got := d.Get("metadata"); !reflect.DeepEqual(got, want) { + t.Errorf("metadata in state = %v, want %v", got, want) + } +} From 4ddc5e2c29fe28665888a01b7fb53bd3f6889362 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:44:50 -0700 Subject: [PATCH 092/107] fix(terraform): read team per-model rpm/tpm limits from metadata and clear them on removal (#40439) * fix(terraform): read team per-model rpm/tpm limits from metadata and clear them on removal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(terraform): drop explanatory comments from team per-model limit helper and tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/provider/litellm/resource_team.go | 25 +++++-- .../provider/litellm/resource_team_test.go | 71 +++++++++++++++++++ 2 files changed, 90 insertions(+), 6 deletions(-) diff --git a/terraform/provider/litellm/resource_team.go b/terraform/provider/litellm/resource_team.go index 24c47843cd1..cea0344b72e 100644 --- a/terraform/provider/litellm/resource_team.go +++ b/terraform/provider/litellm/resource_team.go @@ -263,11 +263,11 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { d.Set("team_member_tpm_limit", *teamResp.TeamMemberTPMLimit) } d.Set("team_member_key_duration", GetStringValue(teamResp.TeamMemberKeyDuration, d.Get("team_member_key_duration").(string))) - if teamResp.ModelRPMLimit != nil { - d.Set("model_rpm_limit", teamResp.ModelRPMLimit) + if v := teamModelLimit(teamResp.ModelRPMLimit, teamResp.Metadata, "model_rpm_limit"); v != nil { + d.Set("model_rpm_limit", v) } - if teamResp.ModelTPMLimit != nil { - d.Set("model_tpm_limit", teamResp.ModelTPMLimit) + if v := teamModelLimit(teamResp.ModelTPMLimit, teamResp.Metadata, "model_tpm_limit"); v != nil { + d.Set("model_tpm_limit", v) } if teamResp.AllowedPassthroughRoutes != nil { d.Set("allowed_passthrough_routes", teamResp.AllowedPassthroughRoutes) @@ -364,14 +364,19 @@ func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{} "organization_id", "tpm_limit", "rpm_limit", "max_budget", "budget_duration", "models", "blocked", "team_member_permissions", "model_aliases", "guardrails", "prompts", "team_member_budget", "team_member_budget_duration", "team_member_rpm_limit", - "team_member_tpm_limit", "team_member_key_duration", "model_rpm_limit", - "model_tpm_limit", "allowed_passthrough_routes", + "team_member_tpm_limit", "team_member_key_duration", "allowed_passthrough_routes", } { if v, ok := d.GetOk(key); ok { teamData[key] = v } } + for _, key := range []string{"model_rpm_limit", "model_tpm_limit"} { + if v, ok := d.GetOk(key); ok || d.HasChange(key) { + teamData[key] = v + } + } + if v, ok := d.GetOk("soft_budget"); ok { teamData["soft_budget"] = v } else if d.HasChange("soft_budget") { @@ -404,6 +409,14 @@ func buildTeamMetadata(d *schema.ResourceData) map[string]interface{} { return metadata } +func teamModelLimit(topLevel, metadata map[string]interface{}, key string) map[string]interface{} { + if topLevel != nil { + return topLevel + } + nested, _ := metadata[key].(map[string]interface{}) + return nested +} + func splitTeamMetadata(raw map[string]interface{}) (map[string]string, []string, []string) { metadata := map[string]string{} var tags, alertEmails []string diff --git a/terraform/provider/litellm/resource_team_test.go b/terraform/provider/litellm/resource_team_test.go index 9638378cdfe..ca4114ad24f 100644 --- a/terraform/provider/litellm/resource_team_test.go +++ b/terraform/provider/litellm/resource_team_test.go @@ -250,6 +250,77 @@ func TestTeamReadMapsNewFields(t *testing.T) { } } +func TestTeamReadMapsPerModelLimitsFromMetadata(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{ + "team_id": "team-1", + "team_info": { + "team_id": "team-1", + "team_alias": "eng", + "model_rpm_limit": null, + "model_tpm_limit": null, + "metadata": { + "department": "eng", + "model_rpm_limit": {"gpt-4o-mini": 250}, + "model_tpm_limit": {"gpt-4o-mini": 5000} + } + } + }`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{ + "team_alias": "eng", + "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 100}, + }) + d.SetId("team-1") + + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read returned error: %v", err) + } + if got := d.Get("model_rpm_limit"); !reflect.DeepEqual(got, map[string]interface{}{"gpt-4o-mini": 250}) { + t.Errorf("model_rpm_limit = %v, want server value 250", got) + } + if got := d.Get("model_tpm_limit"); !reflect.DeepEqual(got, map[string]interface{}{"gpt-4o-mini": 5000}) { + t.Errorf("model_tpm_limit = %v, want server value 5000", got) + } + if got := d.Get("metadata"); !reflect.DeepEqual(got, map[string]interface{}{"department": "eng"}) { + t.Errorf("metadata = %v, want per-model limits kept out of the string map", got) + } +} + +func TestTeamUpdateClearsRemovedPerModelLimits(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`) + defer srv.Close() + + res := ResourceLiteLLMTeam() + priorData := schema.TestResourceDataRaw(t, res.Schema, map[string]interface{}{ + "team_alias": "eng", + "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 100}, + "model_tpm_limit": map[string]interface{}{"gpt-4o-mini": 5000}, + }) + priorData.SetId("team-1") + prior := priorData.State() + config := terraform.NewResourceConfigRaw(map[string]interface{}{"team_alias": "eng"}) + diff, err := res.Diff(context.Background(), prior, config, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + d, err := schema.InternalMap(res.Schema).Data(prior, diff) + if err != nil { + t.Fatalf("data failed: %v", err) + } + + if err := resourceLiteLLMTeamUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + for _, k := range []string{"model_rpm_limit", "model_tpm_limit"} { + if got, ok := captured[k]; !ok || !reflect.DeepEqual(got, map[string]interface{}{}) { + t.Errorf("payload %s = %v (present=%v), want explicit empty map", k, got, ok) + } + } +} + // rpm_limit_type / tpm_limit_type are accepted by /team/new but not // /team/update, so create must send them and update must not. func TestTeamLimitTypesSentOnCreateOnly(t *testing.T) { From 18eb6d799880f4fc4f4ea79bdbce6b7839530625 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:45:02 -0700 Subject: [PATCH 093/107] feat(terraform): allow custom team_id on litellm_team (#40459) Add an optional, computed, force-new team_id argument to the litellm_team resource. When set, it is sent to /team/new as the team_id; when omitted the provider keeps generating a UUID. Read mirrors the resource id into state so import and plan stay clean. Resolves LIT-6399 Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/provider/CHANGELOG.md | 1 + terraform/provider/docs/resources/team.md | 16 +++- terraform/provider/litellm/resource_team.go | 17 +++- .../provider/litellm/resource_team_test.go | 82 +++++++++++++++++++ 4 files changed, 113 insertions(+), 3 deletions(-) diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 842bfb4bdb1..de54207ea5a 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -16,6 +16,7 @@ longer signal it. ### Added +- **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement - **jwt_key_mapping**: New `litellm_jwt_key_mapping` resource for the proxy's JWT to virtual key mappings, so JWT clients identified by a claim (`client_id`, `azp`, `sub`) map to virtual keys and inherit their models, budgets and rate limits. Supports `description` and `is_active`, rotating the mapped key in place, and forces replacement when the claim name or value changes - **team**: `soft_budget`, `tags`, and `soft_budget_alerting_emails` attributes on `litellm_team`, matching what `/team/new` and `/team/update` already accept; `soft_budget_alerting_emails` is sent under `metadata`, where the proxy reads it - **user**: New `litellm_user` resource and `litellm_user` / `litellm_users` data sources for managing internal users diff --git a/terraform/provider/docs/resources/team.md b/terraform/provider/docs/resources/team.md index 821d8c1dee3..19575907269 100644 --- a/terraform/provider/docs/resources/team.md +++ b/terraform/provider/docs/resources/team.md @@ -14,6 +14,16 @@ resource "litellm_team" "engineering" { } ``` +### Team with a Custom ID + +```hcl +resource "litellm_team" "platform" { + team_id = "platform-team" + team_alias = "platform" + models = ["gpt-4-proxy"] +} +``` + ### Team with Comprehensive Configuration ```hcl @@ -92,6 +102,8 @@ resource "litellm_team" "model_dependent_team" { The following arguments are supported: +* `team_id` - (Optional) A stable, human-readable ID for the team (for example `platform-team`). If omitted, the provider generates a random UUID. Changing this forces a new team to be created. + * `team_alias` - (Required) A human-readable identifier for the team. * `organization_id` - (Optional) The ID of the organization this team belongs to. @@ -152,7 +164,7 @@ The following arguments are supported: In addition to the arguments above, the following attributes are exported: -* `id` - The unique identifier for the team. +* `id` - The unique identifier for the team, equal to `team_id`. ## Import @@ -162,7 +174,7 @@ Teams can be imported using the team ID: terraform import litellm_team.engineering ``` -Note: The team ID is generated when the team is created and is different from the `team_alias`. +Note: Unless `team_id` is set, the team ID is generated when the team is created and is different from the `team_alias`. ## Note on Team Members diff --git a/terraform/provider/litellm/resource_team.go b/terraform/provider/litellm/resource_team.go index cea0344b72e..bf7d2508077 100644 --- a/terraform/provider/litellm/resource_team.go +++ b/terraform/provider/litellm/resource_team.go @@ -31,6 +31,13 @@ func ResourceLiteLLMTeam() *schema.Resource { }, Schema: map[string]*schema.Schema{ + "team_id": { + Type: schema.TypeString, + Optional: true, + Computed: true, + ForceNew: true, + Description: "Unique ID for the team. Generated by the provider if not provided", + }, "team_alias": { Type: schema.TypeString, Required: true, @@ -162,7 +169,7 @@ func ResourceLiteLLMTeam() *schema.Resource { func resourceLiteLLMTeamCreate(d *schema.ResourceData, m interface{}) error { client := m.(*Client) - teamID := uuid.New().String() + teamID := resolveTeamID(d) teamData := buildTeamData(d, teamID) // Throughput limit types are only accepted by /team/new, not /team/update. @@ -214,6 +221,7 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { teamResp := infoResp.TeamInfo // Update the state with values from the response or fall back to the data passed in during creation + d.Set("team_id", d.Id()) d.Set("team_alias", GetStringValue(teamResp.TeamAlias, d.Get("team_alias").(string))) d.Set("organization_id", GetStringValue(teamResp.OrganizationID, d.Get("organization_id").(string))) @@ -354,6 +362,13 @@ func resourceLiteLLMTeamDelete(d *schema.ResourceData, m interface{}) error { return nil } +func resolveTeamID(d *schema.ResourceData) string { + if v, ok := d.GetOk("team_id"); ok { + return v.(string) + } + return uuid.New().String() +} + func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{} { teamData := map[string]interface{}{ "team_id": teamID, diff --git a/terraform/provider/litellm/resource_team_test.go b/terraform/provider/litellm/resource_team_test.go index ca4114ad24f..35d60401d30 100644 --- a/terraform/provider/litellm/resource_team_test.go +++ b/terraform/provider/litellm/resource_team_test.go @@ -9,6 +9,7 @@ import ( "reflect" "testing" + "github.com/google/uuid" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" ) @@ -85,6 +86,87 @@ func TestTeamCreateSendsSoftBudgetTagsAndAlertEmails(t *testing.T) { } } +func TestTeamCreateSendsConfiguredTeamID(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"platform-team","team_info":{"team_id":"platform-team","team_alias":"platform"},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{ + "team_id": "platform-team", + "team_alias": "platform", + }) + + if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + if got := captured["team_id"]; got != "platform-team" { + t.Fatalf("payload team_id = %v, want platform-team", got) + } + if got := d.Id(); got != "platform-team" { + t.Fatalf("resource id = %q, want platform-team", got) + } + if got := d.Get("team_id"); got != "platform-team" { + t.Fatalf("state team_id = %v, want platform-team", got) + } +} + +func TestTeamCreateGeneratesTeamIDWhenUnset(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"x","team_info":{"team_alias":"eng"},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{"team_alias": "eng"}) + + if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + sent, _ := captured["team_id"].(string) + if _, err := uuid.Parse(sent); err != nil { + t.Fatalf("payload team_id = %q, want a generated UUID: %v", sent, err) + } + if d.Id() != sent || d.Get("team_id") != sent { + t.Fatalf("id = %q, state team_id = %v, want both to equal the sent id %q", d.Id(), d.Get("team_id"), sent) + } +} + +func TestTeamReadSetsTeamIDFromResourceID(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"imported-team","team_info":{"team_id":"imported-team","team_alias":"imported"},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{}) + d.SetId("imported-team") + + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + if got := d.Get("team_id"); got != "imported-team" { + t.Fatalf("team_id = %v, want imported-team", got) + } +} + +func TestTeamIDChangeForcesReplacement(t *testing.T) { + res := ResourceLiteLLMTeam() + priorData := schema.TestResourceDataRaw(t, res.Schema, map[string]interface{}{ + "team_id": "old-team", + "team_alias": "eng", + }) + priorData.SetId("old-team") + config := terraform.NewResourceConfigRaw(map[string]interface{}{ + "team_id": "new-team", + "team_alias": "eng", + }) + diff, err := res.Diff(context.Background(), priorData.State(), config, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + if diff == nil || !diff.RequiresNew() { + t.Fatalf("changing team_id must force replacement, diff = %+v", diff) + } +} + func TestTeamReadMapsTeamInfoEnvelope(t *testing.T) { var captured map[string]interface{} srv := newTeamTestServer(t, &captured, teamInfoWithSoftBudget) From cb9c3a913708845254dd289f8ffa0fb25ec4c1d9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:46:26 -0700 Subject: [PATCH 094/107] fix(terraform): send litellm_key model_max_budget as per-model BudgetConfig objects (#40450) * fix(terraform): send litellm_key model_max_budget as per-model BudgetConfig objects The key resource typed model_max_budget as map(number) and forwarded the bare numbers to /key/generate, which the proxy rejects with a 500 ('int' object is not iterable) because each model entry must be a BudgetConfig object. The attribute is now a JSON string of per-model budget objects, matching litellm_user, litellm_budget and litellm_tag, with a schema version 1 state upgrader that drops the old map value so existing state keeps loading Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(terraform): reject litellm_key model_max_budget JSON that is not per-model budget objects Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(terraform): reject unknown or empty per-model budget fields in litellm_key model_max_budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/provider/CHANGELOG.md | 1 + terraform/provider/README.md | 9 +- terraform/provider/docs/resources/key.md | 11 ++- terraform/provider/litellm/resource_key.go | 91 ++++++++++++++++-- .../provider/litellm/resource_key_test.go | 96 +++++++++++++++++++ .../provider/litellm/resource_key_utils.go | 4 +- 6 files changed, 195 insertions(+), 17 deletions(-) diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index de54207ea5a..798c967c736 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -45,6 +45,7 @@ longer signal it. ### Changed +- **key** (breaking): `model_max_budget` on `litellm_key` is now a JSON string of per-model budget objects (`jsonencode({"gpt-4o-mini" = {budget_limit = 50, time_period = "30d"}})`), matching `litellm_user`, `litellm_budget` and `litellm_tag`. The old `map(number)` form sent bare numbers to `/key/generate`, which the proxy rejects with a 500 (`'int' object is not iterable`), so every key with a non-empty `model_max_budget` failed to apply. Existing state upgrades automatically (schema version 1) and the attribute is refilled from the proxy on the next read; configurations still using the map form must be rewritten - **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying ## [0.4.0] - 2026-08-06 diff --git a/terraform/provider/README.md b/terraform/provider/README.md index 0a6d15c7844..b392fd6279d 100644 --- a/terraform/provider/README.md +++ b/terraform/provider/README.md @@ -103,9 +103,12 @@ resource "litellm_key" "example_key" { permissions = { can_create_keys = "true" } - model_max_budget = { - "gpt-4" = 50.0 - } + model_max_budget = jsonencode({ + "gpt-4" = { + budget_limit = 50.0 + time_period = "30d" + } + }) model_rpm_limit = { "claude-3.5-sonnet" = 30 } diff --git a/terraform/provider/docs/resources/key.md b/terraform/provider/docs/resources/key.md index 5094b77cbec..8b228914649 100644 --- a/terraform/provider/docs/resources/key.md +++ b/terraform/provider/docs/resources/key.md @@ -30,9 +30,12 @@ resource "litellm_key" "example" { permissions = { "can_create_keys" = "true" } - model_max_budget = { - "gpt-4" = 50.0 - } + model_max_budget = jsonencode({ + "gpt-4" = { + budget_limit = 50.0 + time_period = "30d" + } + }) model_rpm_limit = { "gpt-3.5-turbo" = 30 } @@ -81,7 +84,7 @@ The following arguments are supported: * `permissions` - (Optional) Permissions associated with this key. This defines what actions are allowed with this key. -* `model_max_budget` - (Optional) Maximum budget per model. This allows setting different budget limits for each model. +* `model_max_budget` - (Optional) JSON string of per-model budget config, e.g. `jsonencode({"gpt-4" = {budget_limit = 50.0, time_period = "30d"}})`. Each model maps to an object with `budget_limit` (or `max_budget`), `time_period` (or `budget_duration`), `tpm_limit` and `rpm_limit`. * `model_rpm_limit` - (Optional) Requests per minute limit per model. This allows setting different RPM limits for each model. diff --git a/terraform/provider/litellm/resource_key.go b/terraform/provider/litellm/resource_key.go index a37041fe7e8..5d8bc5667a9 100644 --- a/terraform/provider/litellm/resource_key.go +++ b/terraform/provider/litellm/resource_key.go @@ -2,6 +2,7 @@ package litellm import ( "context" + "encoding/json" "fmt" "github.com/hashicorp/go-cty/cty" @@ -10,7 +11,7 @@ import ( ) func resourceKey() *schema.Resource { - return &schema.Resource{ + r := &schema.Resource{ CreateContext: resourceKeyCreate, ReadContext: resourceKeyRead, UpdateContext: resourceKeyUpdate, @@ -18,6 +19,7 @@ func resourceKey() *schema.Resource { Importer: &schema.ResourceImporter{ StateContext: schema.ImportStatePassthroughContext, }, + SchemaVersion: 1, Schema: map[string]*schema.Schema{ "key": { Type: schema.TypeString, @@ -105,9 +107,11 @@ func resourceKey() *schema.Resource { Elem: &schema.Schema{Type: schema.TypeString}, }, "model_max_budget": { - Type: schema.TypeMap, - Optional: true, - Elem: &schema.Schema{Type: schema.TypeFloat, Computed: true}, + Type: schema.TypeString, + Optional: true, + ValidateFunc: validateKeyModelMaxBudget, + DiffSuppressFunc: budgetSuppressEquivalentJSON, + Description: "JSON string of per-model budget config (e.g. '{\"gpt-4o-mini\": {\"budget_limit\": 50, \"time_period\": \"30d\"}}')", }, "model_rpm_limit": { Type: schema.TypeMap, @@ -182,6 +186,79 @@ func resourceKey() *schema.Resource { }, }, } + r.StateUpgraders = []schema.StateUpgrader{{ + Version: 0, + Type: resourceKeyV0Type(r.Schema), + Upgrade: resourceKeyStateUpgradeV0, + }} + return r +} + +// Schema version 0 typed model_max_budget as map(number), which the proxy +// rejects; version 1 stores the per-model BudgetConfig objects as a JSON string. +func resourceKeyV0Type(current map[string]*schema.Schema) cty.Type { + v0 := make(map[string]*schema.Schema, len(current)) + for k, v := range current { + v0[k] = v + } + v0["model_max_budget"] = &schema.Schema{ + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeFloat}, + } + return (&schema.Resource{Schema: v0}).CoreConfigSchema().ImpliedType() +} + +func resourceKeyStateUpgradeV0(_ context.Context, rawState map[string]interface{}, _ interface{}) (map[string]interface{}, error) { + delete(rawState, "model_max_budget") + return rawState, nil +} + +var keyModelBudgetFields = map[string]bool{ + "budget_limit": true, + "max_budget": true, + "time_period": true, + "budget_duration": true, + "tpm_limit": true, + "rpm_limit": true, +} + +func validateKeyModelMaxBudget(v interface{}, k string) ([]string, []error) { + var parsed map[string]json.RawMessage + if err := json.Unmarshal([]byte(v.(string)), &parsed); err != nil || parsed == nil { + return nil, []error{fmt.Errorf("%q must be a JSON object keyed by model name, got %s", k, v)} + } + for model, cfg := range parsed { + var budget map[string]json.RawMessage + if err := json.Unmarshal(cfg, &budget); err != nil || len(budget) == 0 { + return nil, []error{fmt.Errorf("%q[%q] must be a budget object such as {\"budget_limit\": 50, \"time_period\": \"30d\"}, got %s", k, model, cfg)} + } + for field := range budget { + if !keyModelBudgetFields[field] { + return nil, []error{fmt.Errorf("%q[%q] has unknown budget field %q; supported fields are budget_limit, max_budget, time_period, budget_duration, tpm_limit, rpm_limit", k, model, field)} + } + } + } + return nil, nil +} + +func parseKeyModelMaxBudget(raw string) map[string]interface{} { + var parsed map[string]interface{} + if err := json.Unmarshal([]byte(raw), &parsed); err != nil || parsed == nil { + return map[string]interface{}{} + } + return parsed +} + +func keyModelMaxBudgetJSON(modelMaxBudget map[string]interface{}) string { + if len(modelMaxBudget) == 0 { + return "" + } + encoded, err := json.Marshal(modelMaxBudget) + if err != nil { + return "" + } + return string(encoded) } func resourceKeyCreate(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics { @@ -342,7 +419,7 @@ func mapResourceDataToKey(d *schema.ResourceData, key *Key) { key.Aliases = d.Get("aliases").(map[string]interface{}) key.Config = d.Get("config").(map[string]interface{}) key.Permissions = d.Get("permissions").(map[string]interface{}) - key.ModelMaxBudget = d.Get("model_max_budget").(map[string]interface{}) + key.ModelMaxBudget = parseKeyModelMaxBudget(d.Get("model_max_budget").(string)) key.ModelRPMLimit = d.Get("model_rpm_limit").(map[string]interface{}) key.ModelTPMLimit = d.Get("model_tpm_limit").(map[string]interface{}) key.Guardrails = expandStringList(d.Get("guardrails").([]interface{})) @@ -415,9 +492,7 @@ func mapKeyToResourceData(d *schema.ResourceData, key *Key) { if key.Permissions != nil { d.Set("permissions", key.Permissions) } - if key.ModelMaxBudget != nil { - d.Set("model_max_budget", key.ModelMaxBudget) - } + d.Set("model_max_budget", keyModelMaxBudgetJSON(key.ModelMaxBudget)) if key.ModelRPMLimit != nil { d.Set("model_rpm_limit", key.ModelRPMLimit) } diff --git a/terraform/provider/litellm/resource_key_test.go b/terraform/provider/litellm/resource_key_test.go index d8f3cfe86b3..30108856a3a 100644 --- a/terraform/provider/litellm/resource_key_test.go +++ b/terraform/provider/litellm/resource_key_test.go @@ -195,6 +195,102 @@ func TestCreateKeySendsConfigSuppliedKey(t *testing.T) { } } +// The proxy validates each model_max_budget entry as a BudgetConfig object and +// 500s on a bare number, so the JSON string must reach /key/generate as nested +// objects and the proxy's response must map back to equivalent JSON in state. +func TestCreateKeySendsModelMaxBudgetAsBudgetObjects(t *testing.T) { + var captured map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/key/generate" { + body, _ := io.ReadAll(r.Body) + json.Unmarshal(body, &captured) + w.Write([]byte(`{"key": "sk-test", "token_id": "hash-1"}`)) + return + } + w.Write([]byte(`{"key": "hash-1", "info": {"model_max_budget": {"gpt-4o-mini": {"budget_limit": 50, "time_period": "30d", "rpm_limit": 60}}}}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newKeyResourceData(t, map[string]interface{}{ + "model_max_budget": `{"gpt-4o-mini": {"budget_limit": 50, "time_period": "30d"}}`, + }) + + if diags := resourceKeyCreate(context.Background(), d, client); diags.HasError() { + t.Fatalf("create returned error: %v", diags) + } + + budgets, ok := captured["model_max_budget"].(map[string]interface{}) + if !ok { + t.Fatalf("create payload model_max_budget = %v, want object", captured["model_max_budget"]) + } + cfg, ok := budgets["gpt-4o-mini"].(map[string]interface{}) + if !ok { + t.Fatalf("model_max_budget[gpt-4o-mini] = %v, want BudgetConfig object", budgets["gpt-4o-mini"]) + } + if cfg["budget_limit"] != float64(50) || cfg["time_period"] != "30d" { + t.Errorf("BudgetConfig = %v, want budget_limit 50 and time_period 30d", cfg) + } + + var state map[string]interface{} + if err := json.Unmarshal([]byte(d.Get("model_max_budget").(string)), &state); err != nil { + t.Fatalf("state model_max_budget %q is not JSON: %v", d.Get("model_max_budget"), err) + } + if got, _ := state["gpt-4o-mini"].(map[string]interface{}); got["budget_limit"] != float64(50) || got["rpm_limit"] != float64(60) { + t.Errorf("state model_max_budget = %v, want the BudgetConfig read back from /key/info", state) + } +} + +// Schema version 0 stored model_max_budget as map(number); that state cannot +// decode into the version 1 string attribute, so the upgrader must drop it. +func TestKeyStateUpgradeV0DropsMapModelMaxBudget(t *testing.T) { + upgraded, err := resourceKey().StateUpgraders[0].Upgrade(context.Background(), map[string]interface{}{ + "id": "hash-1", + "key_alias": "legacy", + "model_max_budget": map[string]interface{}{"gpt-4o-mini": 50.0}, + }, nil) + if err != nil { + t.Fatalf("upgrade returned error: %v", err) + } + if _, present := upgraded["model_max_budget"]; present { + t.Errorf("upgraded state still carries map model_max_budget: %v", upgraded["model_max_budget"]) + } + if upgraded["key_alias"] != "legacy" { + t.Errorf("upgrade dropped unrelated attribute: %v", upgraded) + } +} + +func TestKeyModelMaxBudgetValidationRequiresBudgetObjects(t *testing.T) { + validate := resourceKey().Schema["model_max_budget"].ValidateFunc + for _, valid := range []string{ + `{}`, + `{"gpt-4o-mini": {"budget_limit": 50, "time_period": "30d"}}`, + `{"gpt-4o-mini": {"max_budget": 50, "rpm_limit": 60}, "gpt-4o": {"budget_duration": "1d", "tpm_limit": 1000}}`, + } { + if _, errs := validate(valid, "model_max_budget"); len(errs) != 0 { + t.Errorf("validate(%s) = %v, want accepted", valid, errs) + } + } + for _, invalid := range []string{ + `null`, + `[]`, + `"gpt-4o-mini"`, + `50`, + `{"gpt-4o-mini": 50}`, + `{"gpt-4o-mini": null}`, + `{"gpt-4o-mini": [50]}`, + `{"gpt-4o-mini": {}}`, + `{"gpt-4o-mini": {"budget_limt": 50}}`, + `{"gpt-4o-mini": {"budget_limit": 50, "max_tokens": 100}}`, + `not json`, + } { + if _, errs := validate(invalid, "model_max_budget"); len(errs) == 0 { + t.Errorf("validate(%s) accepted a value that would send no per-model budget", invalid) + } + } +} + // The proxy 400s on budget_duration: "", so an unset duration must be // omitted from the update payload entirely. func TestUpdateKeyOmitsEmptyBudgetDuration(t *testing.T) { diff --git a/terraform/provider/litellm/resource_key_utils.go b/terraform/provider/litellm/resource_key_utils.go index d426fec05b2..7046ef7097a 100644 --- a/terraform/provider/litellm/resource_key_utils.go +++ b/terraform/provider/litellm/resource_key_utils.go @@ -65,7 +65,7 @@ func buildKeyData(d *schema.ResourceData) map[string]interface{} { keyData["permissions"] = v.(map[string]interface{}) } if v, ok := d.GetOkExists("model_max_budget"); ok { - keyData["model_max_budget"] = v.(map[string]interface{}) + keyData["model_max_budget"] = parseKeyModelMaxBudget(v.(string)) } if v, ok := d.GetOkExists("model_rpm_limit"); ok { keyData["model_rpm_limit"] = v.(map[string]interface{}) @@ -107,7 +107,7 @@ func setKeyResourceData(d *schema.ResourceData, key *Key) error { "aliases": key.Aliases, "config": key.Config, "permissions": key.Permissions, - "model_max_budget": key.ModelMaxBudget, + "model_max_budget": keyModelMaxBudgetJSON(key.ModelMaxBudget), "model_rpm_limit": key.ModelRPMLimit, "model_tpm_limit": key.ModelTPMLimit, "guardrails": key.Guardrails, From fb603fdb6b2894f4bd21f33293c025c1f73e4dc6 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:48:35 -0700 Subject: [PATCH 095/107] fix(terraform): keep prior state when a virtual key update is rejected (#40512) resourceKeyUpdate now calls d.Partial(true) on a failed /key/update so the SDK does not persist the rejected planned values into state, which made the next plan report no changes and silently dropped the update. Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/provider/litellm/resource_key.go | 1 + .../provider/litellm/resource_key_test.go | 41 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/terraform/provider/litellm/resource_key.go b/terraform/provider/litellm/resource_key.go index 5d8bc5667a9..3b14f7ca655 100644 --- a/terraform/provider/litellm/resource_key.go +++ b/terraform/provider/litellm/resource_key.go @@ -315,6 +315,7 @@ func resourceKeyUpdate(ctx context.Context, d *schema.ResourceData, m interface{ metadata, err := plannedKeyMetadata(c, d) if err != nil { + d.Partial(true) return diag.FromErr(fmt.Errorf("error updating key: %s", err)) } key.Metadata = metadata diff --git a/terraform/provider/litellm/resource_key_test.go b/terraform/provider/litellm/resource_key_test.go index 30108856a3a..5e6385735c6 100644 --- a/terraform/provider/litellm/resource_key_test.go +++ b/terraform/provider/litellm/resource_key_test.go @@ -319,6 +319,47 @@ func TestUpdateKeyOmitsEmptyBudgetDuration(t *testing.T) { } } +func TestResourceKeyUpdateFailureKeepsPriorState(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/key/update" { + w.WriteHeader(http.StatusBadRequest) + w.Write([]byte(`{"error":{"message":"Invalid budget_duration 'bad'"}}`)) + return + } + w.Write([]byte(`{"key":"hash-1","info":{"key_alias":"demo","models":["fake-model"]}}`)) + })) + defer srv.Close() + + res := resourceKey() + priorData := newKeyResourceData(t, map[string]interface{}{ + "key_alias": "demo", + "models": []interface{}{"fake-model"}, + }) + priorData.SetId("hash-1") + prior := priorData.State() + config := terraform.NewResourceConfigRaw(map[string]interface{}{ + "key_alias": "demo", + "models": []interface{}{"fake-model"}, + "budget_duration": "bad", + }) + diff, err := res.Diff(context.Background(), prior, config, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + + newState, diags := res.Apply(context.Background(), prior, diff, NewClient(srv.URL, "test-key", true)) + if !diags.HasError() { + t.Fatal("apply succeeded, want the proxy's 400 surfaced as an error") + } + if got, ok := newState.Attributes["budget_duration"]; ok { + t.Errorf("failed update persisted budget_duration=%q into state, want it absent", got) + } + if newState.Attributes["key_alias"] != "demo" { + t.Errorf("prior key_alias lost from state: %v", newState.Attributes) + } +} + // /key/info nests the key's fields under "info"; GetKey must unwrap that // envelope or reads map nothing back into state. func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) { From ea106fd8be7666111623c0ea7480aa4e1634cc82 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 02:57:07 +0000 Subject: [PATCH 096/107] fix(terraform): send changed litellm_key duration on /key/update so expires is recomputed (#40511) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/provider/docs/resources/key.md | 2 +- terraform/provider/litellm/client.go | 3 ++ terraform/provider/litellm/resource_key.go | 8 +++-- .../provider/litellm/resource_key_test.go | 35 +++++++++++++++++++ 4 files changed, 45 insertions(+), 3 deletions(-) diff --git a/terraform/provider/docs/resources/key.md b/terraform/provider/docs/resources/key.md index 8b228914649..0ef0688830f 100644 --- a/terraform/provider/docs/resources/key.md +++ b/terraform/provider/docs/resources/key.md @@ -76,7 +76,7 @@ The following arguments are supported: * `key_alias` - (Optional) Alias for this key. This provides a human-readable identifier for the key. -* `duration` - (Optional) Duration for which this key is valid. This sets an expiration time for the key. +* `duration` - (Optional) How long the key stays valid, e.g. "30d" or "12h". The proxy stores this as an absolute `expires` timestamp. Changing the value resets the expiry to the time of the update plus the new duration; removing it from the configuration leaves the current expiry in place. * `aliases` - (Optional) Map of model aliases. This allows you to create custom names for models when using this key. diff --git a/terraform/provider/litellm/client.go b/terraform/provider/litellm/client.go index 0dfbe3dd3c9..58e52aabefc 100644 --- a/terraform/provider/litellm/client.go +++ b/terraform/provider/litellm/client.go @@ -103,6 +103,9 @@ func (c *Client) UpdateKey(key *Key) (*Key, error) { if key.BudgetDuration != "" { updateData["budget_duration"] = key.BudgetDuration } + if key.Duration != "" { + updateData["duration"] = key.Duration + } // Only add pointer fields if they are explicitly set if key.MaxBudget != nil { diff --git a/terraform/provider/litellm/resource_key.go b/terraform/provider/litellm/resource_key.go index 3b14f7ca655..95bf88c0145 100644 --- a/terraform/provider/litellm/resource_key.go +++ b/terraform/provider/litellm/resource_key.go @@ -88,8 +88,9 @@ func resourceKey() *schema.Resource { Optional: true, }, "duration": { - Type: schema.TypeString, - Optional: true, + Type: schema.TypeString, + Optional: true, + Description: "How long the key stays valid, e.g. \"30d\" or \"12h\". Changing it resets the expiry to the time of the update plus the new duration; removing it leaves the current expiry in place", }, "aliases": { Type: schema.TypeMap, @@ -310,6 +311,9 @@ func resourceKeyUpdate(ctx context.Context, d *schema.ResourceData, m interface{ key := &Key{Key: d.Id()} mapResourceDataToKey(d, key) + if !d.HasChange("duration") { + key.Duration = "" + } key.ModelRPMLimit = changedMap(d, "model_rpm_limit") key.ModelTPMLimit = changedMap(d, "model_tpm_limit") diff --git a/terraform/provider/litellm/resource_key_test.go b/terraform/provider/litellm/resource_key_test.go index 5e6385735c6..722a90ee02d 100644 --- a/terraform/provider/litellm/resource_key_test.go +++ b/terraform/provider/litellm/resource_key_test.go @@ -524,3 +524,38 @@ func TestKeyReadKeepsOnlyDeclaredMetadata(t *testing.T) { t.Errorf("metadata in state = %v, want %v", got, want) } } + +func TestKeyUpdateSendsChangedDuration(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "duration": "30d"}, + map[string]interface{}{"key_alias": "alias-1", "duration": "90d"}, + ) + + if got := proxy.updates[0]["duration"]; got != "90d" { + t.Errorf("update payload duration = %v, want 90d", got) + } +} + +func TestKeyUpdateOmitsUnchangedDuration(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "duration": "30d"}, + map[string]interface{}{"key_alias": "alias-2", "duration": "30d"}, + ) + + if got := proxy.updates[0]["key_alias"]; got != "alias-2" { + t.Fatalf("update payload key_alias = %v, want alias-2", got) + } + if v, present := proxy.updates[0]["duration"]; present { + t.Errorf("update payload unexpectedly contains duration = %v", v) + } +} From 5264f8ed48de6615c903e459aced1dd04137b1b9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 02:57:51 +0000 Subject: [PATCH 097/107] fix(terraform): drop litellm_key from state on 404 instead of failing the plan (#40443) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/provider/litellm/client.go | 20 +++++++++- terraform/provider/litellm/resource_key.go | 2 + .../provider/litellm/resource_key_test.go | 39 +++++++++++++++++++ 3 files changed, 60 insertions(+), 1 deletion(-) diff --git a/terraform/provider/litellm/client.go b/terraform/provider/litellm/client.go index 58e52aabefc..76526a8595c 100644 --- a/terraform/provider/litellm/client.go +++ b/terraform/provider/litellm/client.go @@ -4,6 +4,7 @@ import ( "bytes" "crypto/tls" "encoding/json" + "errors" "fmt" "io" "log" @@ -19,6 +20,20 @@ type Client struct { InsecureSkipVerify bool } +type apiError struct { + StatusCode int + Body string +} + +func (e *apiError) Error() string { + return fmt.Sprintf("API request failed with status code %d: %s", e.StatusCode, e.Body) +} + +func isNotFound(err error) bool { + var apiErr *apiError + return errors.As(err, &apiErr) && apiErr.StatusCode == http.StatusNotFound +} + func NewClient(apiBase, apiKey string, insecureSkipVerify bool) *Client { tr := &http.Transport{ TLSClientConfig: &tls.Config{InsecureSkipVerify: insecureSkipVerify}, @@ -57,6 +72,9 @@ func (c *Client) CreateKey(key *Key) (*Key, error) { func (c *Client) GetKey(keyID string) (*Key, error) { resp, err := c.sendRequest("GET", fmt.Sprintf("/key/info?key=%s", keyID), nil) + if isNotFound(err) { + return nil, nil + } if err != nil { return nil, err } @@ -377,7 +395,7 @@ func (c *Client) sendRequest(method, path string, body interface{}) (map[string] log.Printf("Response body: %s", c.redactSensitiveData(string(bodyBytes))) if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("API request failed with status code %d: %s", resp.StatusCode, string(bodyBytes)) + return nil, &apiError{StatusCode: resp.StatusCode, Body: string(bodyBytes)} } var result map[string]interface{} diff --git a/terraform/provider/litellm/resource_key.go b/terraform/provider/litellm/resource_key.go index 95bf88c0145..ffdc448f5fe 100644 --- a/terraform/provider/litellm/resource_key.go +++ b/terraform/provider/litellm/resource_key.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "log" "github.com/hashicorp/go-cty/cty" "github.com/hashicorp/terraform-plugin-sdk/v2/diag" @@ -297,6 +298,7 @@ func resourceKeyRead(ctx context.Context, d *schema.ResourceData, m interface{}) } if key == nil { + log.Printf("[WARN] Key %s not found, removing from state", d.Id()) d.SetId("") return nil } diff --git a/terraform/provider/litellm/resource_key_test.go b/terraform/provider/litellm/resource_key_test.go index 722a90ee02d..1da8528576d 100644 --- a/terraform/provider/litellm/resource_key_test.go +++ b/terraform/provider/litellm/resource_key_test.go @@ -394,6 +394,45 @@ func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) { } } +func TestResourceKeyReadDropsMissingKeyFromState(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotFound) + w.Write([]byte(`{"error":{"message":"Key not found in database","type":"not_found_error","param":"key","code":"404"}}`)) + })) + defer srv.Close() + + d := newKeyResourceData(t, map[string]interface{}{"key_alias": "stale"}) + d.SetId("deleted-out-of-band") + + diags := resourceKeyRead(context.Background(), d, NewClient(srv.URL, "test-key", true)) + if diags.HasError() { + t.Fatalf("read of a missing key must not error, got: %v", diags) + } + if d.Id() != "" { + t.Errorf("Id = %q, want empty so Terraform plans a recreate", d.Id()) + } +} + +func TestResourceKeyReadStillFailsOnNon404Errors(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error":{"message":"db down"}}`)) + })) + defer srv.Close() + + d := newKeyResourceData(t, map[string]interface{}{"key_alias": "live"}) + d.SetId("still-exists") + + diags := resourceKeyRead(context.Background(), d, NewClient(srv.URL, "test-key", true)) + if !diags.HasError() { + t.Fatal("a 500 from /key/info must surface as an error, not be treated as a deleted key") + } + if d.Id() != "still-exists" { + t.Errorf("Id = %q, want unchanged on a transient error", d.Id()) + } +} + // fakeKeyProxy serves /key/info from stored metadata and applies /key/update // the way the proxy does: an absent "metadata" keeps the stored map, a // present one replaces it wholesale. From 87e2026f0a440279c3a1b32d7351df91d13d2d29 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 03:10:42 +0000 Subject: [PATCH 098/107] fix(terraform): read key fields the proxy stores in metadata back from /key/info (#40513) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/provider/CHANGELOG.md | 1 + terraform/provider/litellm/client.go | 28 ++++++ .../provider/litellm/resource_key_test.go | 88 +++++++++++++++++++ 3 files changed, 117 insertions(+) diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 798c967c736..e5e0a164a83 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -39,6 +39,7 @@ longer signal it. - **team**: Read now decodes the `team_info` envelope `/team/info` actually returns, so team attributes refresh from the proxy instead of always falling back to the prior state - **key**: Read now unwraps the `info` envelope `/key/info` actually returns; previously reads mapped nothing back into state, so drift on a key was never detected +- **key**: Read now picks up `model_rpm_limit`, `model_tpm_limit`, `guardrails`, `tags`, `enforced_params`, `allowed_passthrough_routes`, `rpm_limit_type`, `tpm_limit_type` and `prompts` from `info.metadata`, where the proxy actually stores them; previously they stayed empty in state, so a matching config showed a permanent phantom diff on them and out-of-band changes to them were never detected - **key**: Updates no longer send an empty `budget_duration`, which the proxy rejects with a 400; any update to a key without a configured `budget_duration` previously failed outright - **key**: A config-supplied `key` value (write-only) is now forwarded to `/key/generate`; previously it was silently dropped and the proxy generated a random key instead - **security**: The `litellm_key` data source and `litellm_key_block` resource normalize raw `sk-` keys to their SHA-256 token hash before building request URLs and resource IDs, so plaintext keys no longer land in reverse-proxy access logs, Terraform plan output, or state IDs diff --git a/terraform/provider/litellm/client.go b/terraform/provider/litellm/client.go index 76526a8595c..e68b8a3a80b 100644 --- a/terraform/provider/litellm/client.go +++ b/terraform/provider/litellm/client.go @@ -87,12 +87,40 @@ func (c *Client) GetKey(keyID string) (*Key, error) { info["key"] = k } } + hoistKeyFieldsStoredInMetadata(info) return c.parseKeyResponse(info) } return c.parseKeyResponse(resp) } +var keyFieldsStoredInMetadata = []string{ + "model_rpm_limit", + "model_tpm_limit", + "guardrails", + "tags", + "enforced_params", + "allowed_passthrough_routes", + "rpm_limit_type", + "tpm_limit_type", + "prompts", +} + +func hoistKeyFieldsStoredInMetadata(info map[string]interface{}) { + metadata, ok := info["metadata"].(map[string]interface{}) + if !ok { + return + } + for _, field := range keyFieldsStoredInMetadata { + if existing, present := info[field]; present && existing != nil { + continue + } + if v, present := metadata[field]; present { + info[field] = v + } + } +} + func (c *Client) UpdateKey(key *Key) (*Key, error) { // Create a new map with only the fields that can be updated updateData := map[string]interface{}{ diff --git a/terraform/provider/litellm/resource_key_test.go b/terraform/provider/litellm/resource_key_test.go index 1da8528576d..66291eadcc5 100644 --- a/terraform/provider/litellm/resource_key_test.go +++ b/terraform/provider/litellm/resource_key_test.go @@ -394,6 +394,94 @@ func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) { } } +func TestGetKeyReadsFieldsStoredInMetadata(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "key": "hash-1", + "info": { + "models": ["gpt-4o-mini"], + "metadata": { + "team": "core-infra", + "model_rpm_limit": {"gpt-4o-mini": 7}, + "model_tpm_limit": {"gpt-4o-mini": 10000}, + "guardrails": ["pii-guard"], + "tags": ["prod"], + "enforced_params": ["user"], + "allowed_passthrough_routes": ["/v1/foo"], + "rpm_limit_type": "guaranteed_throughput", + "tpm_limit_type": "dynamic", + "prompts": ["p1"] + } + } + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + key, err := client.GetKey("hash-1") + if err != nil { + t.Fatalf("GetKey returned error: %v", err) + } + if got, ok := key.ModelRPMLimit["gpt-4o-mini"].(float64); !ok || got != 7 { + t.Errorf("ModelRPMLimit = %v, want gpt-4o-mini=7 read from metadata", key.ModelRPMLimit) + } + if got, ok := key.ModelTPMLimit["gpt-4o-mini"].(float64); !ok || got != 10000 { + t.Errorf("ModelTPMLimit = %v, want gpt-4o-mini=10000 read from metadata", key.ModelTPMLimit) + } + if len(key.Guardrails) != 1 || key.Guardrails[0] != "pii-guard" { + t.Errorf("Guardrails = %v, want [pii-guard]", key.Guardrails) + } + if len(key.Tags) != 1 || key.Tags[0] != "prod" { + t.Errorf("Tags = %v, want [prod]", key.Tags) + } + if len(key.EnforcedParams) != 1 || key.EnforcedParams[0] != "user" { + t.Errorf("EnforcedParams = %v, want [user]", key.EnforcedParams) + } + if len(key.AllowedPassthroughRoutes) != 1 || key.AllowedPassthroughRoutes[0] != "/v1/foo" { + t.Errorf("AllowedPassthroughRoutes = %v, want [/v1/foo]", key.AllowedPassthroughRoutes) + } + if key.RPMLimitType != "guaranteed_throughput" || key.TPMLimitType != "dynamic" { + t.Errorf("limit types = %q/%q, want guaranteed_throughput/dynamic", key.RPMLimitType, key.TPMLimitType) + } + if len(key.Prompts) != 1 || key.Prompts[0] != "p1" { + t.Errorf("Prompts = %v, want [p1]", key.Prompts) + } + if key.Metadata["team"] != "core-infra" { + t.Errorf("Metadata = %v, want team=core-infra preserved", key.Metadata) + } +} + +func TestGetKeyPrefersTopLevelOverMetadataCopy(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "key": "hash-1", + "info": { + "tags": ["top-level"], + "guardrails": null, + "metadata": { + "tags": ["from-metadata"], + "guardrails": ["from-metadata"] + } + } + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + key, err := client.GetKey("hash-1") + if err != nil { + t.Fatalf("GetKey returned error: %v", err) + } + if len(key.Tags) != 1 || key.Tags[0] != "top-level" { + t.Errorf("Tags = %v, want [top-level]", key.Tags) + } + if len(key.Guardrails) != 1 || key.Guardrails[0] != "from-metadata" { + t.Errorf("Guardrails = %v, want [from-metadata] (null top-level must not shadow)", key.Guardrails) + } +} + func TestResourceKeyReadDropsMissingKeyFromState(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") From 935611d25d092c5623b44e08d89d7c2e3d565933 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 20:20:07 -0700 Subject: [PATCH 099/107] fix(mcp): respect optional discovery capabilities and quiet unsupported methods --- litellm/experimental_mcp_client/client.py | 53 ++++- .../test_mcp_client.py | 209 ++++++++++++++++++ 2 files changed, 254 insertions(+), 8 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 8c7f4557992..421a3c84d75 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -54,18 +54,22 @@ def missing_streamable_http_client_error() -> ImportError: ) -from mcp.types import CallToolRequestParams as MCPCallToolRequestParams -from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import ( + METHOD_NOT_FOUND, ClientResult, GetPromptRequestParams, GetPromptResult, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, Prompt, ResourceTemplate, ServerNotification, ServerRequest, TextContent, ) +from mcp.types import CallToolRequestParams as MCPCallToolRequestParams +from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import Tool as MCPTool from pydantic import AnyUrl @@ -777,8 +781,19 @@ class MCPClient: """List available prompts from the server.""" verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") - async def _list_prompts_operation(session: ClientSession): - return await session.list_prompts() + async def _list_prompts_operation(session: ClientSession) -> ListPromptsResult: + capabilities: Final = session.get_server_capabilities() + if capabilities is not None and capabilities.prompts is None: + return ListPromptsResult(prompts=[]) + try: + return await session.list_prompts() + except McpError as error: + if error.error.code != METHOD_NOT_FOUND: + raise + verbose_logger.debug( + "MCP client list_prompts is unsupported by %s: %s", self.server_url or "stdio", error + ) + return ListPromptsResult(prompts=[]) try: result: Final = await self.run_with_session(_list_prompts_operation) @@ -854,8 +869,19 @@ class MCPClient: """List available resources from the server.""" verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio") - async def _list_resources_operation(session: ClientSession): - return await session.list_resources() + async def _list_resources_operation(session: ClientSession) -> ListResourcesResult: + capabilities: Final = session.get_server_capabilities() + if capabilities is not None and capabilities.resources is None: + return ListResourcesResult(resources=[]) + try: + return await session.list_resources() + except McpError as error: + if error.error.code != METHOD_NOT_FOUND: + raise + verbose_logger.debug( + "MCP client list_resources is unsupported by %s: %s", self.server_url or "stdio", error + ) + return ListResourcesResult(resources=[]) try: result: Final = await self.run_with_session(_list_resources_operation) @@ -890,8 +916,19 @@ class MCPClient: """List available resource templates from the server.""" verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio") - async def _list_resource_templates_operation(session: ClientSession): - return await session.list_resource_templates() + async def _list_resource_templates_operation(session: ClientSession) -> ListResourceTemplatesResult: + capabilities: Final = session.get_server_capabilities() + if capabilities is not None and capabilities.resources is None: + return ListResourceTemplatesResult(resourceTemplates=[]) + try: + return await session.list_resource_templates() + except McpError as error: + if error.error.code != METHOD_NOT_FOUND: + raise + verbose_logger.debug( + "MCP client list_resource_templates is unsupported by %s: %s", self.server_url or "stdio", error + ) + return ListResourceTemplatesResult(resourceTemplates=[]) try: result: Final = await self.run_with_session(_list_resource_templates_operation) diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index b07e5876e8b..713330a8280 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -12,6 +12,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import anyio import httpx import pytest +import respx from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth from mcp import McpError from mcp.client.streamable_http import streamable_http_client @@ -1686,3 +1687,211 @@ async def test_empty_http_event_stream_uses_the_existing_request_deadline() -> N timeout=3, ) assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +@pytest.mark.parametrize( + "outcome", + ( + "absent", + "other_capability", + "supported", + "method_not_found", + "internal_error", + "unauthorized", + "timeout", + "initialize_not_found", + ), +) +async def test_optional_discovery_capabilities_and_errors( + method: str, outcome: str, caplog: pytest.LogCaptureFixture +) -> None: + import logging + from unittest.mock import Mock + + from mcp.types import JSONRPCRequest + + capability: Final = "prompts" if method == "prompts/list" else "resources" + field: Final = { + "prompts/list": "prompts", + "resources/list": "resources", + "resources/templates/list": "resourceTemplates", + }[method] + advertised: Final = "resources" if capability == "prompts" else "prompts" + entry: Final = { + "prompts/list": {"name": "example"}, + "resources/list": {"name": "example", "uri": "test://example"}, + "resources/templates/list": {"name": "example", "uriTemplate": "test://{name}"}, + }[method] + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = JSONRPCMessage.model_validate_json(request.content).root + if not isinstance(payload, JSONRPCRequest): + return httpx.Response(202) + if outcome == "initialize_not_found": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "error": {"code": -32601, "message": "Initialization rejected"}, + }, + ) + if payload.method == "initialize": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {} + if outcome == "absent" + else {advertised if outcome == "other_capability" else capability: {}}, + "serverInfo": {"name": "discovery", "version": "1"}, + }, + }, + ) + if outcome == "timeout": + raise httpx.ReadTimeout("Optional list timed out", request=request) + if outcome == "unauthorized": + return httpx.Response(401) + if outcome in ("method_not_found", "internal_error", "absent", "other_capability"): + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "error": { + "code": -32603 if outcome == "internal_error" else -32601, + "message": "Optional list rejected", + }, + }, + ) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [entry]}}) + + responder: Final = Mock(side_effect=respond) + caplog.set_level(logging.DEBUG, logger="LiteLLM") + with respx.mock(base_url="https://example.com") as router: + router.route().mock(side_effect=responder) + client: Final = MCPClient(server_url="https://example.com/mcp") + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + result: Final = await operation() + + requests: Final = tuple( + JSONRPCMessage.model_validate_json(call.args[0].content).root + for call in responder.call_args_list + if call.args[0].method == "POST" + ) + assert sum(isinstance(request, JSONRPCRequest) and request.method == method for request in requests) == ( + 0 if outcome in ("absent", "other_capability", "initialize_not_found") else 1 + ) + assert [item.name for item in result] == (["example"] if outcome == "supported" else []) + failures: Final = tuple( + record for record in caplog.records if record.name == "LiteLLM" and record.levelno >= logging.WARNING + ) + if outcome in ("internal_error", "unauthorized", "timeout", "initialize_not_found"): + assert any(record.levelno == logging.ERROR and "failed" in record.message for record in failures) + else: + assert failures == () + if outcome == "method_not_found": + assert any( + record.levelno == logging.DEBUG and "Optional list rejected" in record.message for record in caplog.records + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("supports_first", (True, False)) +async def test_optional_discovery_uses_each_sessions_capabilities(supports_first: bool) -> None: + from unittest.mock import Mock + from mcp.types import JSONRPCRequest + + capabilities: Final = iter(({"resources": {}}, {}) if supports_first else ({}, {"resources": {}})) + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = JSONRPCMessage.model_validate_json(request.content).root + if not isinstance(payload, JSONRPCRequest): + return httpx.Response(202) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": next(capabilities), + "serverInfo": {"name": "changing", "version": "1"}, + } + if payload.method == "initialize" + else {"resources": [{"name": "example", "uri": "test://example"}]} + ) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + + responder: Final = Mock(side_effect=respond) + with respx.mock(base_url="https://example.com") as router: + router.route().mock(side_effect=responder) + client: Final = MCPClient(server_url="https://example.com/mcp") + first: Final = await client.list_resources() + second: Final = await client.list_resources() + + assert [item.name for item in first] == (["example"] if supports_first else []) + assert [item.name for item in second] == ([] if supports_first else ["example"]) + requests: Final = tuple( + JSONRPCMessage.model_validate_json(call.args[0].content).root + for call in responder.call_args_list + if call.args[0].method == "POST" + ) + assert sum(isinstance(request, JSONRPCRequest) and request.method == "resources/list" for request in requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +async def test_optional_discovery_preserves_cancellation(method: str) -> None: + from mcp.types import JSONRPCRequest + + ready: Final = asyncio.Event() + pending: Final = asyncio.Event() + + async def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = JSONRPCMessage.model_validate_json(request.content).root + if not isinstance(payload, JSONRPCRequest): + return httpx.Response(202) + if payload.method == "initialize": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {"resources": {}, "prompts": {}}, + "serverInfo": {"name": "pending", "version": "1"}, + }, + }, + ) + ready.set() + await pending.wait() + return httpx.Response(202) + + with respx.mock(base_url="https://example.com") as router: + router.route().mock(side_effect=respond) + client: Final = MCPClient(server_url="https://example.com/mcp") + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + task: Final = asyncio.create_task(operation()) + try: + await asyncio.wait_for(ready.wait(), timeout=3) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=3) From a0bb021f028efa4b29ac6dca3bdbff67c6f890f0 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 20:41:56 -0700 Subject: [PATCH 100/107] fix(mcp): preserve gateway-owned resource scope through consent --- .../mcp_server/gateway_dcr_flow.py | 12 +++-- .../mcp_server/test_gateway_dcr_flow.py | 46 ++++++++++++++++++- 2 files changed, 52 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 3d94fa345d0..f4889008e94 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -411,7 +411,7 @@ def relative_request_url(request: Request) -> str: def resolve_scoped_resource_server(request: Request, resource: str | None) -> MCPServer | None: - """Resolve an RFC 8707 ``resource`` value to the single gateway-managed oauth2 server it + """Resolve an RFC 8707 ``resource`` value to the single gateway-owned server it names, or ``None`` for every other shape: absent, the aggregate resource, a foreign host, an unparseable value, a multi-server path, an unknown name, or any server mode the keyless gateway flow does not serve (whose protected-resource metadata never directs a @@ -443,7 +443,7 @@ def resolve_scoped_resource_server(request: Request, resource: str | None) -> MC if len(names) != 1: return None server: Final = global_mcp_server_manager.get_mcp_server_by_name(names[0]) - if server is None or not server.is_gateway_managed_oauth2: + if server is None or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server): return None return server @@ -729,11 +729,15 @@ async def _flow_target( server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id) if ( server is None - or not server.is_gateway_managed_oauth2 + or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server) or not await lookup_server_reachability(flow.user_id, server.server_id) ): return "stale", None - state: Final = "m2m" if MCPServerManager.effective_oauth2_flow(server) == "client_credentials" else "interactive" + state: Final = ( + "interactive" + if server.is_gateway_managed_oauth2 and MCPServerManager.effective_oauth2_flow(server) != "client_credentials" + else "m2m" + ) return state, server diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 73a52a8d2e8..8aa4cfc5619 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -833,7 +833,7 @@ async def test_manual_delivery_page_renders_the_url_as_data_never_as_a_shell_com assert 'value="' in body -def _scoped_mcp_server(name="github", **kw): +def _scoped_mcp_server(name="github", auth_type="oauth2", **kw): from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -844,7 +844,7 @@ def _scoped_mcp_server(name="github", **kw): alias=name, url="https://upstream.example/mcp", transport="http", - auth_type=MCPAuth.oauth2, + auth_type=MCPAuth(auth_type) if auth_type is not None else None, **kw, ) @@ -2044,3 +2044,45 @@ async def test_introspect_fails_closed_on_dead_user_and_503s_on_outage(): status, body = await _introspect(minted.token.get_secret_value(), master_key=None) assert (status, body["error"]) == (500, "server_error") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth_type", [None, "none", "api_key", "bearer_token", "basic", "authorization", "token", "aws_sigv4"] +) +@pytest.mark.parametrize("resource", ["https://llm.example.com/mcp/github", "https://llm.example.com/github/mcp"]) +async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(auth_type, resource): + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server(auth_type=auth_type) + vendor = _VendorCredential("absent") + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + response = _scoped_authorize(client_id, resource) + described = await _describe_page(response, scoped_server=server, vendor=vendor) + assert json.loads(described.body) == { + "state": "m2m", + "client_origin": "https://claude.ai", + "server_id": "github-id", + "server_name": "github", + "connected": True, + } + unreachable = await _complete_page(response, scoped_server=server, reachable=_ServerReachability(False)) + assert unreachable.status_code == 400 + cache = DualCache() + completed = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache) + assert completed.status_code == 303 + assert vendor.calls == [] + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + redeemed = await _redeem(code, client_id, cache=cache, resource=resource) + assert redeemed.status_code == 200 + payload = json.loads(redeemed.body) + assert _opened_principal(payload).resource_server_id == "github-id" + renewed = await _redeem( + None, client_id, cache=cache, grant_type="refresh_token", refresh_token=payload["refresh_token"] + ) + assert renewed.status_code == 200 + assert _opened_principal(json.loads(renewed.body)).resource_server_id == "github-id" From 95f6e96ef982d0244cfa8607d5c4fd1f2703867c Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 20:57:33 -0700 Subject: [PATCH 101/107] test(mcp): type auth diagnostics regression parameters --- .../outbound_credentials/test_resolver.py | 11 +++++++---- .../mcp_server/test_mcp_server_manager.py | 13 ++++++++----- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 2eab053ab73..5fab4ceec72 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -19,6 +19,7 @@ from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials import ( ApiKeyConfig, + AuthConfig, AuthorizationCodeConfig, AwsSigV4Config, Byok, @@ -1216,7 +1217,9 @@ async def test_passthrough_ignores_the_carrier_and_keeps_the_callers_slot(): (AuthorizationCodeConfig(), Subject(tenant_id="", subject_id="alice"), "stored-user-token", "Bearer stored-alice"), ], ) -async def test_resolved_source_matches_the_credential_sent_upstream(config, subject, expected_source, expected_header): +async def test_resolved_source_matches_the_credential_sent_upstream( + config: AuthConfig, subject: Subject, expected_source: str, expected_header: str | None +) -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="stored-alice")}) @@ -1230,7 +1233,7 @@ async def test_resolved_source_matches_the_credential_sent_upstream(config, subj @pytest.mark.asyncio -async def test_resolved_source_preserves_missing_user_token_error(): +async def test_resolved_source_preserves_missing_user_token_error() -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source result = await resolve_credentials_with_source(UpstreamCredentialProvider(), _SUBJECT, _spec(AuthorizationCodeConfig())) @@ -1239,7 +1242,7 @@ async def test_resolved_source_preserves_missing_user_token_error(): @pytest.mark.asyncio -async def test_minted_token_sources_match_egress_and_do_not_fetch_twice(): +async def test_minted_token_sources_match_egress_and_do_not_fetch_twice() -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at"))) @@ -1259,7 +1262,7 @@ async def test_minted_token_sources_match_egress_and_do_not_fetch_twice(): @pytest.mark.asyncio -async def test_id_jag_source_describes_final_token_after_both_exchanges(): +async def test_id_jag_source_describes_final_token_after_both_exchanges() -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source endpoint = _FakeTokenEndpoint(_two_leg_ok("resource-token")) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 99447acd371..adf985e9a21 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -5,7 +5,7 @@ import logging import os import sys from datetime import datetime -from typing import Any, Dict, Final, Optional +from typing import Any, Dict, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -12582,8 +12582,11 @@ async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch, ], ) async def test_debug_resolution_matches_final_header_conflict_winner( - config, extra_headers, expected_source, expected_authorization -): + config: Literal["stored", "static", "none"], + extra_headers: dict[str, str] | None, + expected_source: str, + expected_authorization: str | None, +) -> None: from mcp.server.lowlevel.server import request_ctx from mcp.shared.context import RequestContext from starlette.requests import Request @@ -12598,7 +12601,7 @@ async def test_debug_resolution_matches_final_header_conflict_winner( from litellm.types.mcp_server.mcp_server_manager import MCPServer class Store: - def __init__(self): + def __init__(self) -> None: self.calls = 0 async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: @@ -12641,7 +12644,7 @@ async def test_debug_resolution_matches_final_header_conflict_winner( @pytest.mark.asyncio @pytest.mark.parametrize("transport", ["http", "stdio"]) -async def test_debug_reports_legacy_signing_and_non_http_transport(transport): +async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Literal["http", "stdio"]) -> None: from mcp.server.lowlevel.server import request_ctx from mcp.shared.context import RequestContext from starlette.requests import Request From 00b631883da02740ad285ed053d20488eba896a1 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 9 Sep 2026 22:03:16 -0700 Subject: [PATCH 102/107] fix(proxy): surface runtime-registered callbacks in UI Logging page (#38974) * fix(proxy): surface runtime-registered callbacks in /get/config/callbacks Config-file callbacks fire at runtime but never appear in the UI Logging and Alerts page because /get/config/callbacks only reads the DB-merged config. Append runtime-registered callbacks from LoggingCallbackManager as read-only rows, deduplicated against configured rows via alias normalization. UI hides edit/delete/test actions for read-only rows. * fix: filter internal proxy hooks from runtime callbacks, update test - Filter _PROXY*, ShadowEval, ServiceLogging, SkillsInjection, ResponsesID prefixes - Update test to exclude read_only rows from count assertions - Still allows deployment/guardrail callbacks to surface if configured Note: comprehensive internal-hook filtering deferred, live-pr-risk will observe real behavior on running proxy. * fix: guard non-list config callbacks in get_config, use monkeypatch in tests - Line-concat type error: normalize_callback now returns empty list for non-list types (dict/tuple/set) instead of passing through unchanged; prevents TypeError when config values are non-list - Test quality TQ005: replace manual try/finally save-restore of litellm.callbacks with monkeypatch.setattr in test_get_config_callbacks_appends_runtime_only_callbacks and test_get_config_callbacks_redacts_runtime_only_row_secrets_for_view_only_admin - Ruff format: wrap _internal_callback_prefixes tuple and isinstance check across multiple lines to respect 120-char limit - All three new tests pass * fix: rework runtime callback inventory filtering and dedup - Filter internal proxy hooks by name: _PROXY_ prefix plus fixed internal names (cache, _ProxyDBLogger, deployment callbacks, service hooks) - Hide guardrail instances and runtime instances of already configured callbacks via CustomLoggerRegistry class lookup - Sort runtime rows and dedup per mode for stable output - normalize_callback returns tuples for str/None/list config values and empty for any other type - Tests mock get_callbacks_by_type explicitly and pin the exact row set; UI test covers read_only action hiding * fix: list dict-shaped callback config values by their keys Dict-valued success_callback/failure_callback/callbacks settings previously listed their keys as editable rows; keep that behavior instead of dropping them to read-only runtime rows. Adds a pin test for the dict shape. * fix: mark dotted-path callbacks read-only to prevent duplicate display Configured callbacks loaded from dotted Python paths (e.g. custom_callbacks.my_logger) are never matched against runtime instances by name because the registry uses short canonical names (e.g. langsmith, arize). Mark these rows read-only to prevent the UI from attempting delete operations that would fail at the endpoint level anyway. * fix: dedupe dotted-path callbacks by instance module instead of marking them read-only A dotted-path callback loaded from config registers as an object, so it surfaces at runtime under its class name and never matched the configured string, producing a second row. Marking the config row read_only hid the duplicate but also hid delete, which does work for these rows. Match the live instance back to its configured entry by module and drop it from the runtime rows, so the callback stays a single editable row. * test: cover dotted-path dedup across success, failure, and callbacks modes * fix(proxy): filter runtime callback inventory by object identity and label read-only rows in the UI Runtime-only rows were filtered by callback name, which missed initialized CustomLogger instances, router and proxy hook methods, guardrails, and user functions. The inventory now inspects the live callback objects through a public LoggingCallbackManager.get_callback_objects accessor and hides litellm-internal hooks, guardrails, and instances of already configured callbacks. The dashboard shows a Read only label for runtime-only rows instead of an empty action cell Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): keep configured-callback assertions minimal when runtime rows are present Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): hide internal cache string callback from runtime callback inventory Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): hide auto-registered vector store hook from callback inventory Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep YAML OTel-family callbacks listed next to a configured one arize, weave_otel and langfuse_otel all initialize OpenTelemetry subclasses, so hiding runtime callbacks by configured class made one saved OTel callback swallow its YAML siblings. Match runtime instances by their own callback_name and only fall back to class identity for bare OpenTelemetry Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover scalar and null YAML callback keys in callback inventory Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): drop docstrings that restate callback inventory helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep runtime-only s3 and sqs callbacks in UI Logging inventory _is_litellm_internal_callback checked registry membership with the display alias (s3, sqs), which is not a registry key, so runtime-only S3Logger and SQSLogger instances were classified as internal and dropped from /get/config/callbacks. Check the registered name instead and cover both loggers in the internal-exclusion regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../logging_callback_manager.py | 10 + litellm/proxy/proxy_server.py | 94 ++- tests/proxy_unit_tests/test_proxy_server.py | 2 +- .../proxy/proxy_server/test_routes_config.py | 564 ++++++++++++++---- .../LoggingCallbacksTable.test.tsx | 18 + .../LoggingCallbacksTableColumns.tsx | 10 + .../LoggingCallbacks/types.ts | 1 + 7 files changed, 576 insertions(+), 123 deletions(-) diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 9b612993a69..31523c9309d 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -441,7 +441,15 @@ class LoggingCallbackManager: return result + def get_callback_objects(self) -> tuple[tuple[str, CustomLogger | Callable], ...]: + return tuple( + (self._get_callback_string(callback), callback) + for callback in self._get_all_callbacks() + if not isinstance(callback, str) + ) + def _get_callback_string(self, callback: CustomLogger | Callable | str) -> str: + from litellm.integrations.opentelemetry import OpenTelemetry from litellm.litellm_core_utils.custom_logger_registry import ( CustomLoggerRegistry, ) @@ -449,6 +457,8 @@ class LoggingCallbackManager: """Convert a callback to its string representation""" if isinstance(callback, str): return callback + elif isinstance(callback, OpenTelemetry) and callback.callback_name is not None: + return callback.callback_name elif isinstance(callback, CustomLogger): # Try to get the string representation from the registry callback_str: Final = CustomLoggerRegistry.get_callback_str_from_class_type(type(callback)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 30ff47ac9f9..3bbe932f5d0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -268,7 +268,7 @@ from litellm.constants import ( WEEKLY_SPEND_REPORT_JOB_ID, ) from litellm.exceptions import RejectedRequestError -from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.agentic_loop_settings import ( @@ -17655,6 +17655,66 @@ async def delete_callback( ) +def _normalize_callback_alias(callback_name: str) -> str: + callback_aliases: Final = ( + ("opentelemetry", "otel"), + ("s3_v2", "s3"), + ("aws_sqs", "sqs"), + ("custom_callback_api", "generic_api"), + ) + return next( + (canonical_name for alias, canonical_name in callback_aliases if alias == callback_name), + callback_name, + ) + + +def _callback_module_name(callback: CustomLogger | Callable[..., object]) -> str: + if inspect.ismethod(callback): + return callback.__func__.__module__ + if inspect.isfunction(callback): + return callback.__module__ + return type(callback).__module__ + + +def _is_litellm_internal_callback(callback_name: str, callback: CustomLogger | Callable[..., object]) -> bool: + from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry + + module_owner: Final = _callback_module_name(callback).partition(".")[0] + is_registered_integration: Final = callback_name in CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE + return not is_registered_integration and module_owner in ("litellm", "litellm_enterprise") + + +def _is_instance_of_configured_callback( + callback_name: str, callback: CustomLogger | Callable[..., object], configured_classes: tuple[type, ...] +) -> bool: + """Self-naming OTel-family instances (`arize`, `weave_otel`) match by name, so a configured `logfire` (a bare + `OpenTelemetry`) does not hide YAML-configured siblings of the same class.""" + from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry + + class_derived_name: Final = CustomLoggerRegistry.get_callback_str_from_class_type(type(callback)) + return isinstance(callback, configured_classes) and callback_name in (class_derived_name, type(callback).__name__) + + +def _hidden_runtime_callback_names(configured_callback_names: frozenset[str]) -> frozenset[str]: + from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry + + configured_classes: Final = tuple( + CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE[name] + for name in configured_callback_names + if name in CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE + ) + configured_modules: Final = frozenset(name.rsplit(".", 1)[0] for name in configured_callback_names if "." in name) + internal_callback_names: Final = frozenset({"cache", "vector_store_pre_call_hook"}) + return internal_callback_names | frozenset( + callback_name + for callback_name, callback in litellm.logging_callback_manager.get_callback_objects() + if isinstance(callback, CustomGuardrail) + or _is_litellm_internal_callback(callback_name, callback) + or _is_instance_of_configured_callback(callback_name, callback, configured_classes) + or _callback_module_name(callback) in configured_modules + ) + + @router.get( "/get/config/callbacks", tags=["config.yaml"], @@ -17687,10 +17747,10 @@ async def get_config( # Normalize string callbacks to lists def normalize_callback(callback): if isinstance(callback, str): - return [callback] - elif callback is None: - return [] - return callback + return (callback,) + if callback is None: + return () + return tuple(callback) if isinstance(callback, (list, dict)) else () _success_callbacks = normalize_callback(_success_callbacks) _failure_callbacks = normalize_callback(_failure_callbacks) @@ -17721,6 +17781,30 @@ async def get_config( for _callback in _success_and_failure_callbacks: _data_to_return.append(process_callback(_callback, "success_and_failure", environment_variables)) + configured_callback_names: Final = frozenset( + _normalize_callback_alias(callback) + for callback in (_success_callbacks + _failure_callbacks + _success_and_failure_callbacks) + ) + runtime_callbacks_by_type: Final = litellm.logging_callback_manager.get_callbacks_by_type() + hidden_callback_names: Final = _hidden_runtime_callback_names(configured_callback_names) + runtime_callback_rows: Final = tuple( + (_normalize_callback_alias(callback_name), callback_type) + for callback_type, callback_names in ( + ("success", runtime_callbacks_by_type["success"]), + ("failure", runtime_callbacks_by_type["failure"]), + ("success_and_failure", runtime_callbacks_by_type["success_and_failure"]), + ) + for callback_name in callback_names + if callback_name not in hidden_callback_names + ) + runtime_only_rows: Final = sorted( + frozenset(row for row in runtime_callback_rows if row[0] not in configured_callback_names) + ) + _data_to_return.extend( + dict(process_callback(callback_name, callback_type, environment_variables), read_only=True) + for callback_name, callback_type in runtime_only_rows + ) + _data_to_return = _apply_callback_role_gate(_data_to_return, is_full_admin) # Check if slack alerting is on diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 54cce9cdd78..d06eb0426c9 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2920,7 +2920,7 @@ async def test_get_config_callbacks_with_all_types(client_no_auth): assert result["status"] == "success" assert "callbacks" in result - callbacks = result["callbacks"] + callbacks = [cb for cb in result["callbacks"] if not cb.get("read_only", False)] # Verify we have all 5 callbacks (2 success + 1 failure + 2 success_and_failure) assert len(callbacks) == 5 diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index dcb63b8ca82..2d9c1bd8b46 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -13,9 +13,12 @@ Routes covered: from __future__ import annotations +import asyncio import json from unittest.mock import AsyncMock, MagicMock +import pytest + from .conftest import VOLATILE_KEYS, normalize @@ -229,10 +232,7 @@ def test_config_update_no_db_error(client, auth_as, monkeypatch): json={"general_settings": {"alerting": ["slack"]}}, ) assert response.status_code != 200 - assert ( - "db" in str(response.json()).lower() - or "connect" in str(response.json()).lower() - ) + assert "db" in str(response.json()).lower() or "connect" in str(response.json()).lower() # --------------------------------------------------------------------------- @@ -273,9 +273,7 @@ def test_config_field_update_happy_admin(client, auth_as, mock_prisma, monkeypat } -def test_config_field_update_non_admin_rejected( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_update_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin cannot update config fields — returns 400 with not-allowed detail (handler uses 400 for the auth gate, not 403).""" from litellm.proxy import proxy_server as ps @@ -335,9 +333,7 @@ def test_config_field_info_happy_admin(client, auth_as, mock_prisma, monkeypatch monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/field/info", params={"field_name": "max_parallel_requests"} - ) + response = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) assert response.status_code == 200 assert normalize(response.json()) == { "field_name": "max_parallel_requests", @@ -345,9 +341,7 @@ def test_config_field_info_happy_admin(client, auth_as, mock_prisma, monkeypatch } -def test_config_field_info_non_admin_rejected( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_info_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin (INTERNAL_USER) is denied — admin-view gate fires.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -356,9 +350,7 @@ def test_config_field_info_non_admin_rejected( monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.INTERNAL_USER): - response = client.get( - "/config/field/info", params={"field_name": "max_parallel_requests"} - ) + response = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) assert response.status_code == 400 assert "error" in response.json().get("detail", {}) @@ -375,16 +367,12 @@ def test_config_field_info_field_not_in_db(client, auth_as, mock_prisma, monkeyp monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/field/info", params={"field_name": "max_parallel_requests"} - ) + response = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) assert response.status_code == 400 assert "not in DB" in response.json().get("detail", {}).get("error", "") -def test_config_field_info_redacts_nested_secret_for_view_only_admin( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_info_redacts_nested_secret_for_view_only_admin(client, auth_as, mock_prisma, monkeypatch): """A view-only admin reading a structured field must not receive nested credentials. database_args carries aws_web_identity_token (a DynamoDB role-assumption credential); it must come back redacted while non-secret @@ -405,9 +393,7 @@ def test_config_field_info_redacts_nested_secret_for_view_only_admin( monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): - response = client.get( - "/config/field/info", params={"field_name": "database_args"} - ) + response = client.get("/config/field/info", params={"field_name": "database_args"}) assert response.status_code == 200 value = response.json()["field_value"] assert value["aws_web_identity_token"] == "REDACTED" @@ -415,9 +401,7 @@ def test_config_field_info_redacts_nested_secret_for_view_only_admin( assert value["user_table_name"] == "LiteLLM_UserTable" -def test_config_field_info_full_admin_sees_nested_secret( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_info_full_admin_sees_nested_secret(client, auth_as, mock_prisma, monkeypatch): """The redaction must not over-redact for a full PROXY_ADMIN, who needs the real nested value to populate the edit form.""" from litellm.proxy import proxy_server as ps @@ -435,18 +419,14 @@ def test_config_field_info_full_admin_sees_nested_secret( monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/field/info", params={"field_name": "database_args"} - ) + response = client.get("/config/field/info", params={"field_name": "database_args"}) assert response.status_code == 200 value = response.json()["field_value"] assert value["aws_web_identity_token"] == "sk-super-secret-token" assert value["region_name"] == "us-east-1" -def test_config_field_info_redacts_top_level_scalar_for_view_only( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_info_redacts_top_level_scalar_for_view_only(client, auth_as, mock_prisma, monkeypatch): """The top-level scalar branch must also redact for a view-only admin. database_url carries DB credentials and is not caught by the name masker, so it is in the explicit secret set.""" @@ -460,9 +440,7 @@ def test_config_field_info_redacts_top_level_scalar_for_view_only( monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): - response = client.get( - "/config/field/info", params={"field_name": "database_url"} - ) + response = client.get("/config/field/info", params={"field_name": "database_url"}) assert response.status_code == 200 assert response.json()["field_value"] == "REDACTED" @@ -476,17 +454,12 @@ def test_redact_general_setting_value_recurses_list_of_dicts(): {"path": "/foo", "headers": {"Authorization": "Bearer sk-x"}}, {"path": "/bar", "client_secret": "sk-y"}, ] - redacted = ps._redact_general_setting_value( - "some_list_field", value, is_full_admin=False - ) + redacted = ps._redact_general_setting_value("some_list_field", value, is_full_admin=False) assert redacted[0]["headers"]["Authorization"] == "REDACTED" assert redacted[0]["path"] == "/foo" assert redacted[1]["client_secret"] == "REDACTED" assert redacted[1]["path"] == "/bar" - assert ( - ps._redact_general_setting_value("some_list_field", value, is_full_admin=True) - == value - ) + assert ps._redact_general_setting_value("some_list_field", value, is_full_admin=True) == value def test_redact_secret_values_in_obj_fails_closed_at_max_depth(): @@ -504,22 +477,16 @@ def test_redact_secret_values_in_obj_fails_closed_at_max_depth(): for _ in range(ps._REDACT_SECRET_MAX_DEPTH + 2): nested = {"wrap": nested} - out = ps._redact_general_setting_value( - "some_struct_field", nested, is_full_admin=False - ) + out = ps._redact_general_setting_value("some_struct_field", nested, is_full_admin=False) # the secret must not survive anywhere in the returned tree assert "sk-leak-bottom" not in repr(out) # full admin is unaffected by the cap — the value comes back untouched - admin_out = ps._redact_general_setting_value( - "some_struct_field", nested, is_full_admin=True - ) + admin_out = ps._redact_general_setting_value("some_struct_field", nested, is_full_admin=True) assert admin_out is nested -def test_config_list_redacts_pass_through_secret_for_view_only( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_list_redacts_pass_through_secret_for_view_only(client, auth_as, mock_prisma, monkeypatch): """/config/list must not leak pass_through_endpoints upstream credentials to a view-only admin. pass_through_endpoints is a known secret-bearing field, so a non-admin gets it redacted; a full admin still sees it.""" @@ -546,24 +513,16 @@ def test_config_list_redacts_pass_through_secret_for_view_only( ) def _pass_through_value(body): - return next( - entry["field_value"] - for entry in body - if entry["field_name"] == "pass_through_endpoints" - ) + return next(entry["field_value"] for entry in body if entry["field_name"] == "pass_through_endpoints") with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): - view_resp = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + view_resp = client.get("/config/list", params={"config_type": "general_settings"}) assert view_resp.status_code == 200 assert "sk-UPSTREAM-SECRET" not in view_resp.text assert _pass_through_value(view_resp.json()) == "REDACTED" with auth_as(LitellmUserRoles.PROXY_ADMIN): - admin_resp = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + admin_resp = client.get("/config/list", params={"config_type": "general_settings"}) assert admin_resp.status_code == 200 admin_value = _pass_through_value(admin_resp.json()) assert admin_value[0]["headers"]["Authorization"] == "Bearer sk-UPSTREAM-SECRET" @@ -587,9 +546,7 @@ def test_config_list_happy_admin(client, auth_as, mock_prisma, monkeypatch): monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + response = client.get("/config/list", params={"config_type": "general_settings"}) assert response.status_code == 200 body = response.json() assert isinstance(body, list) @@ -695,9 +652,7 @@ def test_config_list_non_admin_rejected(client, auth_as, mock_prisma, monkeypatc monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.INTERNAL_USER): - response = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + response = client.get("/config/list", params={"config_type": "general_settings"}) assert response.status_code == 400 assert "role" in response.json().get("detail", {}).get("error", "").lower() @@ -710,9 +665,7 @@ def test_config_list_no_db_error(client, auth_as, monkeypatch): monkeypatch.setattr(ps, "prisma_client", None) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + response = client.get("/config/list", params={"config_type": "general_settings"}) assert response.status_code == 400 assert "error" in response.json().get("detail", {}) @@ -756,9 +709,7 @@ def test_config_field_delete_happy_admin(client, auth_as, mock_prisma, monkeypat } -def test_config_field_delete_non_admin_rejected( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_delete_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin caller hits the 400 not-allowed branch with role in detail.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -778,9 +729,7 @@ def test_config_field_delete_non_admin_rejected( assert "role" in response.json().get("detail", {}).get("error", "").lower() -def test_config_field_delete_field_not_in_config( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_delete_field_not_in_config(client, auth_as, mock_prisma, monkeypatch): """If there is no general_settings row at all, returns 400 'not in config'.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -825,9 +774,7 @@ def test_config_callback_delete_happy_admin(client, auth_as, mock_prisma, monkey monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.post( - "/config/callback/delete", json={"callback_name": "langfuse"} - ) + response = client.post("/config/callback/delete", json={"callback_name": "langfuse"}) assert response.status_code == 200 # `deleted_at` is an ISO timestamp generated at request time — extend # the volatile set just for this assertion so dict-equality still works. @@ -840,9 +787,7 @@ def test_config_callback_delete_happy_admin(client, auth_as, mock_prisma, monkey } -def test_config_callback_delete_non_admin_rejected( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_callback_delete_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin caller is rejected with 400 not-allowed.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -852,9 +797,7 @@ def test_config_callback_delete_non_admin_rejected( monkeypatch.setattr(ps, "store_model_in_db", True) with auth_as(LitellmUserRoles.INTERNAL_USER): - response = client.post( - "/config/callback/delete", json={"callback_name": "langfuse"} - ) + response = client.post("/config/callback/delete", json={"callback_name": "langfuse"}) assert response.status_code == 400 assert "role" in response.json().get("detail", {}).get("error", "").lower() @@ -869,22 +812,15 @@ def test_config_callback_delete_not_found(client, auth_as, mock_prisma, monkeypa monkeypatch.setattr(ps, "store_model_in_db", True) fake_proxy_config = MagicMock() - fake_proxy_config.get_config = AsyncMock( - return_value={"litellm_settings": {"success_callback": ["slack"]}} - ) + fake_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": {"success_callback": ["slack"]}}) monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.post( - "/config/callback/delete", json={"callback_name": "langfuse"} - ) + response = client.post("/config/callback/delete", json={"callback_name": "langfuse"}) # The handler re-raises HTTPException(404) verbatim (only generic # `Exception` becomes a 500 ProxyException), so pin 404 strictly. assert response.status_code == 404 - assert ( - "langfuse" in str(response.json()).lower() - or "not found" in str(response.json()).lower() - ) + assert "langfuse" in str(response.json()).lower() or "not found" in str(response.json()).lower() # --------------------------------------------------------------------------- @@ -948,10 +884,7 @@ def test_get_config_callbacks_internal_error(client, auth_as, mock_prisma, monke with auth_as(LitellmUserRoles.PROXY_ADMIN): response = client.get("/get/config/callbacks") assert response.status_code >= 400 - assert ( - "boom" in str(response.json()).lower() - or "error" in str(response.json()).lower() - ) + assert "boom" in str(response.json()).lower() or "error" in str(response.json()).lower() _CALLBACK_ENV_FIXTURE = { @@ -985,14 +918,10 @@ def _install_callbacks_config(monkeypatch, mock_prisma): def _callback_variables(body: dict, name: str) -> dict: - return next( - cb["variables"] for cb in body["callbacks"] if cb["name"] == name - ) + return next(cb["variables"] for cb in body["callbacks"] if cb["name"] == name) -def test_get_config_callbacks_redacts_secret_env_vars_for_view_only_admin( - client, auth_as, mock_prisma, monkeypatch -): +def test_get_config_callbacks_redacts_secret_env_vars_for_view_only_admin(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy._types import LitellmUserRoles _install_callbacks_config(monkeypatch, mock_prisma) @@ -1024,9 +953,7 @@ def test_get_config_callbacks_redacts_secret_env_vars_for_view_only_admin( assert otel_vars["OTEL_ENDPOINT"] == _CALLBACK_ENV_FIXTURE["OTEL_ENDPOINT"] -def test_get_config_callbacks_full_admin_still_sees_secret_env_vars( - client, auth_as, mock_prisma, monkeypatch -): +def test_get_config_callbacks_full_admin_still_sees_secret_env_vars(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy._types import LitellmUserRoles _install_callbacks_config(monkeypatch, mock_prisma) @@ -1047,9 +974,7 @@ def test_get_config_callbacks_full_admin_still_sees_secret_env_vars( assert otel_vars["OTEL_HEADERS"] == _CALLBACK_ENV_FIXTURE["OTEL_HEADERS"] -def test_get_config_callbacks_redacts_slack_webhook_urls_for_view_only_admin( - client, auth_as, mock_prisma, monkeypatch -): +def test_get_config_callbacks_redacts_slack_webhook_urls_for_view_only_admin(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -1170,6 +1095,413 @@ def test_get_config_callbacks_redacts_email_alerting_vars_for_view_only_admin( assert admin_email["SMTP_HOST"] == "smtp.resend.com" +def test_get_config_callbacks_appends_runtime_only_callbacks(client, auth_as, mock_prisma, monkeypatch): + """LIT-5281: a YAML callback that the DB callback list replaced in the merged config still runs, so it must + show up as a read_only row next to the editable DB-configured one.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": ["langfuse"]}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + from litellm.integrations.langsmith import LangsmithLogger + from litellm.integrations.opentelemetry import OpenTelemetry + + monkeypatch.setattr(litellm, "success_callback", ["langfuse", LangsmithLogger()]) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", [OpenTelemetry()]) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + + assert [(cb["name"], cb["type"], cb.get("read_only", False)) for cb in response.json()["callbacks"]] == [ + ("langfuse", "success", False), + ("langsmith", "success", True), + ("otel", "success_and_failure", True), + ] + + +def test_get_config_callbacks_accepts_scalar_and_null_yaml_callbacks(client, auth_as, mock_prisma, monkeypatch): + """`success_callback: langfuse` (a YAML scalar) is one configured callback, not eight single-letter ones, and a + `callbacks: null` key contributes nothing.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": "langfuse", "failure_callback": None, "callbacks": None}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + from litellm.integrations.langsmith import LangsmithLogger + + monkeypatch.setattr(litellm, "success_callback", ["langfuse", LangsmithLogger()]) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + + assert [(cb["name"], cb["type"], cb.get("read_only", False)) for cb in response.json()["callbacks"]] == [ + ("langfuse", "success", False), + ("langsmith", "success", True), + ] + + +def test_get_config_callbacks_deduplicates_configured_and_runtime(client, auth_as, mock_prisma, monkeypatch): + """A configured callback shows once as editable, whether the runtime holds its string or an initialized instance + (arize initializes an ArizeLogger, logfire a bare OpenTelemetry that only its class identifies).""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": ["langfuse", "arize", "logfire"]}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + from litellm.integrations.arize.arize import ArizeLogger + from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig + + arize_logger = ArizeLogger(config=OpenTelemetryConfig(exporter="console"), callback_name="arize") + monkeypatch.setattr(litellm, "success_callback", ["langfuse", arize_logger, OpenTelemetry()]) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + + assert [(cb["name"], cb["type"], cb.get("read_only", False)) for cb in response.json()["callbacks"]] == [ + ("langfuse", "success", False), + ("arize", "success", False), + ("logfire", "success", False), + ] + + +def test_get_config_callbacks_keeps_yaml_otel_family_callbacks_next_to_configured_one( + client, auth_as, mock_prisma, monkeypatch +): + """LIT-5281: arize, weave_otel and langfuse_otel all initialize OpenTelemetry subclasses. Saving one of them + from the dashboard replaces the YAML `callbacks` list, so the YAML siblings keep running and must stay listed + under their own names instead of being hidden as duplicates of the configured OTel callback.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"callbacks": ["langfuse_otel"]}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + from litellm.integrations.arize.arize import ArizeLogger + from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger + from litellm.integrations.langsmith import LangsmithLogger + from litellm.integrations.opentelemetry import OpenTelemetryConfig + from litellm.integrations.weave.weave_otel import WeaveOtelLogger + + console_config = OpenTelemetryConfig(exporter="console") + monkeypatch.setattr(litellm, "success_callback", [LangsmithLogger()]) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr( + litellm, + "callbacks", + [ + ArizeLogger(config=console_config, callback_name="arize"), + WeaveOtelLogger(config=console_config), + LangfuseOtelLogger(config=console_config, callback_name="langfuse_otel"), + ], + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + + assert [(cb["name"], cb["type"], cb.get("read_only", False)) for cb in response.json()["callbacks"]] == [ + ("langfuse_otel", "success_and_failure", False), + ("arize", "success_and_failure", True), + ("langsmith", "success", True), + ("weave_otel", "success_and_failure", True), + ] + + +def _dotted_path_test_function(*args, **kwargs): + pass + + +@pytest.mark.parametrize("handler_kind", ["instance", "function"]) +@pytest.mark.parametrize( + "config_key,expected_type", + [ + ("success_callback", "success"), + ("failure_callback", "failure"), + ("callbacks", "success_and_failure"), + ], +) +def test_get_config_callbacks_deduplicates_dotted_path_callback( + client, auth_as, mock_prisma, monkeypatch, config_key, expected_type, handler_kind +): + """A dotted-path callback stays a single editable row instead of duplicating under its class or function name.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + import litellm + from litellm.integrations.custom_logger import CustomLogger + + class _DottedPathTestHandler(CustomLogger): + pass + + dotted_handler = _DottedPathTestHandler() if handler_kind == "instance" else _dotted_path_test_function + dotted_path = f"{__name__}.dotted_handler" + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {config_key: [dotted_path]}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + monkeypatch.setattr(litellm, "callbacks", [dotted_handler]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + + assert response.status_code == 200 + callbacks = response.json()["callbacks"] + assert [(callback["name"], callback["type"], callback.get("read_only", False)) for callback in callbacks] == [ + (dotted_path, expected_type, False) + ] + + +def test_get_config_callbacks_lists_dict_shaped_config_callbacks(client, auth_as, mock_prisma, monkeypatch): + """Dict-shaped success_callback config values list their keys as editable rows.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": {"langsmith": {"batch_size": 1}}}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_callbacks_by_type", + MagicMock(return_value={"success": ["langsmith"], "failure": [], "success_and_failure": []}), + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + + assert response.status_code == 200 + callbacks = response.json()["callbacks"] + assert [(callback["name"], callback.get("read_only", False)) for callback in callbacks] == [("langsmith", False)] + + +def test_get_config_callbacks_excludes_internal_runtime_callbacks(client, auth_as, mock_prisma, monkeypatch): + """Proxy infrastructure callbacks are excluded from callback inventory.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": []}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + + import litellm + from litellm._service_logger import ServiceLogging + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_logger import CustomLogger + from litellm.integrations.langsmith import LangsmithLogger + from litellm.integrations.s3_v2 import S3Logger + from litellm.integrations.sqs import SQSLogger + from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import VectorStorePreCallHook + from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter + from litellm.router import Router + + class _InventoryTestGuardrail(CustomGuardrail): + pass + + class _UserCodeLogger(CustomLogger): + pass + + def user_code_function(*args, **kwargs): + pass + + async def build_aws_loggers() -> tuple[S3Logger, SQSLogger]: + return S3Logger(s3_bucket_name="inventory-bucket"), SQSLogger(sqs_queue_url="https://sqs.example/inventory") + + s3_logger, sqs_logger = asyncio.run(build_aws_loggers()) + router = Router(model_list=[]) + monkeypatch.setattr(litellm, "input_callback", []) + monkeypatch.setattr( + litellm, "success_callback", [LangsmithLogger(), s3_logger, router.sync_deployment_callback_on_success] + ) + monkeypatch.setattr(litellm, "_async_success_callback", [sqs_logger, router.deployment_callback_on_success]) + monkeypatch.setattr(litellm, "failure_callback", [user_code_function]) + monkeypatch.setattr(litellm, "_async_failure_callback", [router.async_deployment_callback_on_failure]) + monkeypatch.setattr( + litellm, + "callbacks", + [ + _PROXY_MaxBudgetLimiter(), + _PROXY_LiteLLMManagedFiles(internal_usage_cache=MagicMock(), prisma_client=MagicMock()), + ServiceLogging(), + VectorStorePreCallHook(), + _InventoryTestGuardrail(guardrail_name="inventory-test-guardrail"), + _UserCodeLogger(), + ], + ) + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + assert "cache" in litellm.success_callback and "cache" in litellm._async_success_callback + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + + assert response.status_code == 200 + assert [ + (callback["name"], callback["type"], callback["read_only"]) for callback in response.json()["callbacks"] + ] == [ + ("_UserCodeLogger", "success_and_failure", True), + ("langsmith", "success", True), + ("s3", "success", True), + ("sqs", "success", True), + ("user_code_function", "failure", True), + ] + + +def test_get_config_callbacks_redacts_runtime_only_row_secrets_for_view_only_admin( + client, auth_as, mock_prisma, monkeypatch +): + """Runtime-only callback rows are subject to the same redaction gate as configured.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": []}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", ["otel"]) + + with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + body = response.json() + + callbacks = body["callbacks"] + otel_cb = next((cb for cb in callbacks if cb["name"] == "otel"), None) + assert otel_cb is not None + assert otel_cb["type"] == "success_and_failure" + assert otel_cb["read_only"] is True + assert otel_cb["variables"]["OTEL_HEADERS"] == "REDACTED" + assert otel_cb["variables"]["OTEL_ENDPOINT"] == _CALLBACK_ENV_FIXTURE["OTEL_ENDPOINT"] + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + admin_response = client.get("/get/config/callbacks") + assert admin_response.status_code == 200 + admin_body = admin_response.json() + admin_otel = next((cb for cb in admin_body["callbacks"] if cb["name"] == "otel"), None) + assert admin_otel is not None + assert admin_otel["variables"]["OTEL_HEADERS"] == _CALLBACK_ENV_FIXTURE["OTEL_HEADERS"] + + # --------------------------------------------------------------------------- # GET /config/yaml # --------------------------------------------------------------------------- @@ -1183,9 +1515,7 @@ def test_config_yaml_returns_demo_payload(client, auth_as): response = client.request("GET", "/config/yaml", json={}) shape = { "status": response.status_code, - "media_type_yaml": response.headers.get("content-type", "").startswith( - "application/json" - ), + "media_type_yaml": response.headers.get("content-type", "").startswith("application/json"), "has_body": len(response.content) > 0, } assert shape == { diff --git a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx index 3ef7dee01b0..a4d76e2cc06 100644 --- a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx @@ -97,6 +97,24 @@ describe("LoggingCallbacksTable", () => { expect(onDelete).toHaveBeenCalledWith(callback); }); + it("shows a read-only label instead of the actions menu for runtime-only callback rows", () => { + render( + , + ); + expect(screen.getByTestId("callback-actions-langfuse-success")).toBeInTheDocument(); + expect(screen.queryByTestId("callback-actions-datadog-success")).not.toBeInTheDocument(); + expect(screen.getAllByText("Read only")).toHaveLength(1); + }); + // Regression: `/get_callbacks` returns the same `name` twice when a // callback is registered for both success and failure (e.g. `generic_api` // → POST to spend-log on both 200 and 4xx/5xx). The UI used to ignore diff --git a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTableColumns.tsx b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTableColumns.tsx index 2263fe03b3d..950ef5de6e4 100644 --- a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTableColumns.tsx @@ -50,6 +50,16 @@ interface CallbackRowActionsProps { } function CallbackRowActions({ callback, onTest, onEdit, onDelete }: CallbackRowActionsProps) { + if (callback.read_only) { + return ( + + Read only + + ); + } return ( Date: Wed, 9 Sep 2026 22:04:37 -0700 Subject: [PATCH 103/107] feat(proxy): expose reversible Claude Code model listing aliases (#40515) Encode complete non-Claude source names and include source_model in the Claude Code listing. Preserve configured route and alias precedence, normalize once before model policy checks, and select CLI models using explicit source identity instead of name stripping or positional joins. Resolves LIT-7360 Claude-Session: https://claude.ai/code/session_01WyqeRhfZGm26zAnHx9P3kq Co-authored-by: Claude Fable 5.1 --- litellm/llms/anthropic/common_utils.py | 27 +++- litellm/proxy/auth/user_api_key_auth.py | 41 ++++++ litellm/proxy/client/cli/README.md | 4 +- .../proxy/client/cli/commands/configure.py | 68 +++++++--- litellm/proxy/client/cli/commands/pi.py | 50 +++++-- .../proxy/common_utils/model_listing_utils.py | 123 +++++++++++++++++- litellm/proxy/proxy_server.py | 40 ++++++ .../anthropic/test_anthropic_common_utils.py | 22 ++++ .../proxy/auth/test_user_api_key_auth.py | 43 ++++++ .../client/cli/test_configure_commands.py | 73 ++++++++++- .../test_litellm/proxy/client/cli/test_pi.py | 38 ++++++ .../common_utils/test_model_listing_utils.py | 109 ++++++++++++++++ .../proxy/proxy_server/test_routes_models.py | 50 +++++++ 13 files changed, 647 insertions(+), 41 deletions(-) create mode 100644 tests/test_litellm/proxy/common_utils/test_model_listing_utils.py diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index ca520116606..87c4ec8938e 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -73,8 +73,13 @@ _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) _CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/") + + def is_claude_code_user_agent(user_agent: str) -> bool: - return user_agent.startswith("claude-cli/") + """Claude Code sends its API calls through the Anthropic SDK as `claude-cli/` and its own + fetches, such as gateway model discovery, as `claude-code/`""" + return user_agent.startswith(_CLAUDE_CODE_USER_AGENT_PREFIXES) def _validated_claude_code_mapping(value: object) -> dict[object, object] | None: @@ -1656,11 +1661,16 @@ def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: def _anthropic_model_entry( - model: ModelInfoResponse, created_at: str, display_names: Mapping[str, str] + model: ModelInfoResponse, created_at: str, display_names: Mapping[str, str], listed_ids: Mapping[str, str] ) -> Mapping[str, object]: + listed_id: Final = listed_ids.get(model["id"]) + source: Final[Mapping[str, object]] = ( + MappingProxyType({"source_model": model["id"]}) if listed_id is not None else MappingProxyType({}) + ) return { # mutable-ok: JSON response body, serialized by the route and never mutated "type": "model", - "id": model["id"], + "id": listed_id or model["id"], + **source, "display_name": display_names.get(model["id"], model["id"]), "created_at": created_at, "max_input_tokens": model.get("max_input_tokens"), @@ -1671,6 +1681,7 @@ def _anthropic_model_entry( def create_anthropic_model_list_response( models: Sequence[ModelInfoResponse], display_names: Mapping[str, str] = MappingProxyType({}), + listed_ids: Mapping[str, str] = MappingProxyType({}), ) -> Mapping[str, object]: """Build the Anthropic-native /v1/models envelope. @@ -1680,17 +1691,19 @@ def create_anthropic_model_list_response( over from the OpenAI-shaped listing, named as the Messages API names them, and are always present because the vendor shape declares them nullable, not optional. display_names maps a listed model id to a configured human-readable name; ids - without an entry fall back to the id itself, matching the vendor behavior + without an entry fall back to the id itself, matching the vendor behavior. + listed_ids maps a model id to the id the caller should see it under (the Claude + Code view); ids without an entry are listed as they are """ created_at: Final = ( datetime.fromtimestamp(DEFAULT_MODEL_CREATED_AT_TIME, tz=timezone.utc).isoformat().replace("+00:00", "Z") ) data: Final = [ # mutable-ok: JSON response body, serialized by the route and never mutated - _anthropic_model_entry(model, created_at, display_names) for model in models + _anthropic_model_entry(model, created_at, display_names, listed_ids) for model in models ] return { # mutable-ok: JSON response body, serialized by the route and never mutated "data": data, "has_more": False, - "first_id": models[0]["id"] if models else None, - "last_id": models[-1]["id"] if models else None, + "first_id": data[0]["id"] if data else None, + "last_id": data[-1]["id"] if data else None, } diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index bfccb703e76..64c0da1c28f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -92,6 +92,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_set_request_parsed_body, populate_request_with_path_params, ) +from litellm.proxy.common_utils.model_listing_utils import claude_code_requested_group from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, @@ -183,6 +184,44 @@ def _get_model_from_request_context( ) +_CLAUDE_MODEL_ROUTES: Final = frozenset( + f"/{prefix}{endpoint}" for prefix in ("", "v1/") for endpoint in ("messages", "chat/completions", "responses") +) +_CLAUDE_MODEL_NORMALIZED: Final = "litellm.claude_model_normalized" + + +async def _normalize_claude_model( + request_data: dict, valid_token: UserAPIKeyAuth, request: Request | None, route: str +) -> None: + from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_config, proxy_logging_obj + + if route not in _CLAUDE_MODEL_ROUTES or llm_router is None: + return + if request is not None and request.scope.get(_CLAUDE_MODEL_NORMALIZED) is True: + return + requested: Final = _get_model_from_request_context(request_data, route, request, llm_router) + if not isinstance(requested, str) or requested != request_data.get("model"): + return + if not requested.startswith("claude-router-") and not requested.lower().endswith("[1m]"): + return + settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=valid_token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + aliases: Final = settings.get("model_group_alias") if isinstance(settings, Mapping) else None + source: Final = claude_code_requested_group( + requested, llm_router, valid_token.team_id, (valid_token.aliases, valid_token.team_model_aliases, aliases) + ) + if request is not None: + request.scope[_CLAUDE_MODEL_NORMALIZED] = True + if source is None: + return + request_data["model"] = source + _safe_set_request_parsed_body(request=request, parsed_body=request_data) + if request is not None: + request._json = request_data + request._body = orjson.dumps(request_data) + + def _get_model_names_for_budget_checks( model: str | list[str] | None, ) -> list[str]: @@ -2768,6 +2807,7 @@ async def _authorize_authenticated_request( """ ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request) + await _normalize_claude_model(request_data, user_api_key_auth_obj, request, route) # Single authorization point. Builder paths MUST NOT call common_checks. # Route through the same exception handler the builder uses so @@ -3131,6 +3171,7 @@ async def _enforce_key_and_fallback_model_access( Key-level model allowlist and client fallbacks (same as standard auth). Not included in common_checks — common_checks enforces team/user/project model access only. """ + await _normalize_claude_model(request_data, valid_token, request, route) config: Final = valid_token.config if config != {}: diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 03d01ff7f66..4045857a237 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -490,7 +490,7 @@ lite codex exec "summarize the repo" Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process. -The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, and older versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. +The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/` wins, and inside the TUI the `/model` picker lists every synced litellm model. @@ -548,7 +548,7 @@ lite --base-url https://your-proxy.example.com configure claude --api-key sk-... claude ``` -With `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) the key is written into `env.ANTHROPIC_AUTH_TOKEN`. Without one, your `lite login` credential is used the way `--config-claude` uses it, through `apiKeyHelper`, so a later `lite login` (or a `--pkce` renewal) picks up on its own and nothing secret lands in the file; a missing or stale login is refreshed first. Either way the command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (the ones whose id contains `claude` or `anthropic`) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key, which has to be on `/v1/models` for the key. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window and sends no thinking parameters for it, so either name the group like a Claude model id or append `[1m]` to opt into the 1M window. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control +With `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) the key is written into `env.ANTHROPIC_AUTH_TOKEN`. Without one, your `lite login` credential is used the way `--config-claude` uses it, through `apiKeyHelper`, so a later `lite login` (or a `--pkce` renewal) picks up on its own and nothing secret lands in the file; a missing or stale login is refreshed first. Either way the command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (under `claude-router-` for a group whose id contains neither `claude` nor `anthropic`, since Claude Code lists only those) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key, which has to be on `/v1/models` for the key. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window (the proxy appends `[1m]` for a group whose configured or known input window reaches 1M) and sends no thinking parameters for it, so name the group like a Claude model id to change that. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control Plain `lite configure`, with no agent named, asks the same things interactively: which agents to wire (Claude Code today) and which of the proxy's models to start on, picked from `/v1/models` with a type-to-filter prompt diff --git a/litellm/proxy/client/cli/commands/configure.py b/litellm/proxy/client/cli/commands/configure.py index 83fbd323924..9d329f8d8f2 100644 --- a/litellm/proxy/client/cli/commands/configure.py +++ b/litellm/proxy/client/cli/commands/configure.py @@ -1,16 +1,23 @@ """`lite configure claude` and `lite unconfigure claude`: persistent Claude Code wiring, undoable.""" import os -import re import sys from collections.abc import Callable, Sequence +from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final import click from InquirerPy import inquirer from InquirerPy.base.control import Choice +from litellm.proxy.common_utils.model_listing_utils import ( + CLAUDE_CODE_CLIENT, + CLAUDE_CODE_PICKER_PATTERN, + GATEWAY_CLIENT_HEADER, +) + from .auth import CliContextObj, context_secret_vault, get_stored_api_key from .claude_settings import ( STARTING_MODEL_ROLE, @@ -30,14 +37,16 @@ from .claude_settings import ( settings_file_owners, unconfigure_claude_settings, ) -from .pi import ListingFailure, PiSyncError, fetch_model_ids +from .pi import ListedModel, ListingFailure, PiSyncError, fetch_model_listing from .up import ensure_fresh_login _LISTED_MODELS_SHOWN: Final = 20 _CLAUDE_TARGET: Final = "claude" _TARGETS: Final = ((_CLAUDE_TARGET, "Claude Code (CLI)"),) _KEEP_DEFAULT_MODEL: Final = "Keep Claude Code's own default" -_CLAUDE_CODE_PICKER_FILTER: Final = re.compile(r"claude|anthropic", re.IGNORECASE) +_CLAUDE_CODE_VIEW: Final = MappingProxyType( + {"anthropic-version": "2023-06-01", GATEWAY_CLIENT_HEADER: CLAUDE_CODE_CLIENT} +) _MODEL_OPTION_HELP: Final = ( f"Proxy model to set as {STARTING_MODEL_ROLE}. Must be listed on /v1/models for the key; without it, " "Claude Code keeps its own default and a pin an earlier configure made is let go of. Nothing pins Claude " @@ -65,7 +74,16 @@ def resolve_credential(ctx: click.Context, api_key: str | None) -> tuple[ClaudeC return ApiKeyHelper(resolve_api_key_helper(base_url)), stored -def _start(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, tuple[str, ...]]: +@dataclass(frozen=True, slots=True) +class _Listing: + models: tuple[ListedModel, ...] + + @property + def ids(self) -> tuple[str, ...]: + return tuple(model.id for model in self.models) + + +def _start(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, _Listing]: """Every configure path begins the same way: the local ownership check first, so a `lite up` session is refused before any login prompt or request, then the credential, then the listing.""" settings_path: Final = claude_settings_path(os.environ) @@ -88,21 +106,28 @@ def _listing_error(base_url: str, error: PiSyncError) -> str: return f"{error.message} The proxy at {base_url} answered, so check that it is a LiteLLM proxy and is healthy." -def _listed_models(base_url: str, key: str) -> tuple[str, ...]: - listed: Final = fetch_model_ids(base_url, key) +def _listed_models(base_url: str, key: str) -> _Listing: + listed: Final = fetch_model_listing(base_url, key, headers=_CLAUDE_CODE_VIEW) if isinstance(listed, PiSyncError): raise click.ClickException(_listing_error(base_url, listed)) - return listed + return _Listing(listed) + + +def _starting_model(model: str, listing: _Listing) -> str | None: + source: Final = next((listed.id for listed in listing.models if listed.source_model == model), None) + return source or next((listed.id for listed in listing.models if listed.id == model), None) def _model_choice(model: str | None) -> ModelChoice: return StartOn(model) if model is not None else UnpinModel() -def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listed: Sequence[str], model: str | None) -> None: +def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listing: _Listing, model: str | None) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"] - if model is not None and model not in listed: + listed: Final = listing.ids + starting: Final = _starting_model(model, listing) if model is not None else None + if model is not None and starting is None: shown: Final = ", ".join(listed[:_LISTED_MODELS_SHOWN]) more: Final = f", and {len(listed) - _LISTED_MODELS_SHOWN} more" if len(listed) > _LISTED_MODELS_SHOWN else "" raise click.ClickException( @@ -113,29 +138,32 @@ def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listed: Sequ configure_claude_settings( base_url, credential, - _model_choice(model), + _model_choice(starting), settings_path, configure_state_path(settings_path), settings_file_owners(settings_path), ) except ClaudeSettingsError as e: raise click.ClickException(str(e)) - in_picker: Final = sum(1 for listed_model in listed if _CLAUDE_CODE_PICKER_FILTER.search(listed_model)) + in_picker: Final = sum(1 for listed_model in listed if CLAUDE_CODE_PICKER_PATTERN.search(listed_model)) click.echo(f"Configured Claude Code: {settings_path} now routes through {base_url}.") + click.echo( "Credential: your virtual key, stored in the file as ANTHROPIC_AUTH_TOKEN." if isinstance(credential, StaticToken) else "Credential: your `lite login`, read through apiKeyHelper on every request, so a later login renews it." ) click.echo( - f"Starting model: {model} ({STARTING_MODEL_ROLE}); switch any time with /model." - if model is not None + f"Starting model: {starting} ({STARTING_MODEL_ROLE}); switch any time with /model." + if starting is not None else "Starting model: not pinned (Claude Code's default, or a model you set yourself); switch with /model, or " "pass --model to start on a proxy model." ) click.echo( - f"/model will list {in_picker} of the proxy's {len(listed)} models (Claude Code shows only ids containing " - "'claude' or 'anthropic')." + f"/model will list all {len(listed)} of the proxy's models." + if in_picker == len(listed) + else f"/model will list {in_picker} of the proxy's {len(listed)} models: Claude Code shows only ids containing " + "'claude' or 'anthropic', and this proxy does not list the rest under such names." ) click.echo("Start `claude` from any terminal. Undo with `lite unconfigure claude`.") if isinstance(credential, StaticToken) and settings_path.is_symlink(): @@ -173,8 +201,10 @@ def interactive_configure( targets: Final = pick_targets() if _CLAUDE_TARGET not in targets: return - credential, listed = _start(ctx, None) - _apply_claude(ctx, credential, listed, pick_model(listed)) + credential, listing = _start(ctx, None) + _apply_claude( + ctx, credential, listing, pick_model(tuple(model.source_model or model.id for model in listing.models)) + ) @click.group(name="configure", invoke_without_command=True) @@ -218,8 +248,8 @@ def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) setting is kept, and what changed is recorded so `lite unconfigure claude` can put it back. Assumes the proxy is already running. """ - credential, listed = _start(ctx, api_key) - _apply_claude(ctx, credential, listed, model) + credential, listing = _start(ctx, api_key) + _apply_claude(ctx, credential, listing, model) @unconfigure_group.command(name="claude") diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index 70c89a853e3..9810e81ae36 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -13,10 +13,11 @@ from dataclasses import dataclass from enum import StrEnum from pathlib import Path from types import MappingProxyType -from typing import Final +from typing import Annotated, Final import requests -from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator +from pydantic.types import StringConstraints PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR" PI_PROVIDER_NAME: Final = "litellm" @@ -51,12 +52,25 @@ class ModelLimits: max_tokens: int | None -class _Model(BaseModel): - id: str +_NonEmptyString = Annotated[str, StringConstraints(min_length=1)] + + +class ListedModel(BaseModel): + model_config = ConfigDict(frozen=True) + + id: _NonEmptyString + source_model: _NonEmptyString | None = None class _ModelList(BaseModel): - data: tuple[_Model, ...] + data: tuple[ListedModel, ...] + + @model_validator(mode="after") + def unique_id_mappings(self) -> "_ModelList": + mappings: Final = frozenset((model.id, model.source_model or model.id) for model in self.data) + if len(frozenset(model.id for model in self.data)) != len(mappings): + raise ValueError("model ids must not map to multiple source models") + return self class _ModelGroup(BaseModel): @@ -69,17 +83,18 @@ class _ModelGroupList(BaseModel): data: tuple[_ModelGroup, ...] -def fetch_model_ids( +def fetch_model_listing( base_url: str, api_key: str, *, get: Callable[..., requests.Response] = requests.get, -) -> tuple[str, ...] | PiSyncError: + headers: Mapping[str, str] = MappingProxyType({}), +) -> tuple[ListedModel, ...] | PiSyncError: url: Final = base_url.rstrip("/") + "/v1/models" try: resp: Final = get( url, - headers={"Authorization": f"Bearer {api_key}"}, # mutable-ok: requests headers require a dict + headers={"Authorization": f"Bearer {api_key}", **headers}, # mutable-ok: requests headers require a dict timeout=10, ) except requests.RequestException as e: @@ -94,10 +109,21 @@ def fetch_model_ids( listing: Final = _ModelList.model_validate(resp.json()) except (ValueError, ValidationError) as e: return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}", kind=ListingFailure.BAD_BODY) - ids: Final = tuple(dict.fromkeys(model.id for model in listing.data)) - if not ids: + models: Final = tuple(dict.fromkeys(listing.data)) + if not models: return PiSyncError("The proxy returned no models for your key.", kind=ListingFailure.EMPTY) - return ids + return models + + +def fetch_model_ids( + base_url: str, + api_key: str, + *, + get: Callable[..., requests.Response] = requests.get, + headers: Mapping[str, str] = MappingProxyType({}), +) -> tuple[str, ...] | PiSyncError: + listed: Final = fetch_model_listing(base_url, api_key, get=get, headers=headers) + return listed if isinstance(listed, PiSyncError) else tuple(dict.fromkeys(model.id for model in listed)) _NO_LIMITS: Final[Mapping[str, ModelLimits]] = MappingProxyType({}) @@ -222,11 +248,13 @@ __all__ = ( "LITELLM_PROXY_API_KEY_ENV", "PI_CONFIG_DIR_ENV", "PI_PROVIDER_NAME", + "ListedModel", "ListingFailure", "ModelLimits", "PiSyncError", "fetch_model_ids", "fetch_model_limits", + "fetch_model_listing", "models_json_path", "provider_block", "sync_models_json", diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 213a697b3dd..4b3ff2711b3 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -10,12 +10,24 @@ legacy internal names with `general_settings.use_team_public_model_name: false`. from __future__ import annotations -from collections.abc import Mapping, Sequence +import re +from collections.abc import Container, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +import litellm + if TYPE_CHECKING: from litellm.router import Router + from litellm.types.proxy.model_listing import ModelInfoResponse + +CLAUDE_CODE_PICKER_PATTERN: Final = re.compile(r"claude|anthropic", re.IGNORECASE) +GATEWAY_CLIENT_HEADER: Final = "x-gateway-client" +CLAUDE_CODE_CLIENT: Final = "claude-code" +_CLAUDE_CODE_ALIAS_PREFIX: Final = "claude-router-" +_ONE_MILLION_SUFFIX: Final = "[1m]" +_ONE_MILLION_TOKENS: Final = 1_000_000 def configured_display_names( @@ -40,6 +52,115 @@ def configured_display_names( ) +def _unmarked(name: str) -> str: + return name[: -len(_ONE_MILLION_SUFFIX)] if name.lower().endswith(_ONE_MILLION_SUFFIX) else name + + +def _compatibility_id(model_id: str) -> str: + return f"{_CLAUDE_CODE_ALIAS_PREFIX}{model_id.encode().hex()}" + + +def _decoded_compatibility_id(view_id: str) -> str | None: + encoded: Final = _unmarked(view_id).removeprefix(_CLAUDE_CODE_ALIAS_PREFIX) + if encoded == _unmarked(view_id): + return None + try: + model_id: Final = bytes.fromhex(encoded).decode() + except (ValueError, UnicodeDecodeError): + return None + return model_id if _compatibility_id(model_id) == _unmarked(view_id) else None + + +def claude_code_model_id( + model_id: str, + max_input_tokens: float | None, + routing_names: Container[str], +) -> str: + """The collision-free id Claude Code's picker lists a model under.""" + if "*" in model_id: + return model_id + shaped: Final = model_id if CLAUDE_CODE_PICKER_PATTERN.search(model_id) else _compatibility_id(model_id) + one_million: Final = max_input_tokens is not None and max_input_tokens >= _ONE_MILLION_TOKENS + marked: Final = ( + f"{shaped}{_ONE_MILLION_SUFFIX}" if one_million and not shaped.lower().endswith(_ONE_MILLION_SUFFIX) else shaped + ) + return next( + ( + name + for name in (marked, shaped) + if name == model_id or claude_code_group_name(name, routing_names) == model_id + ), + model_id, + ) + + +def claude_code_group_name(view_id: str, routing_names: Container[str]) -> str | None: + """Decode a canonical compatibility id only when no configured route claims it.""" + if view_id in routing_names: + return None + unmarked: Final = _unmarked(view_id) + if unmarked != view_id and unmarked in routing_names: + return unmarked + model_id: Final = _decoded_compatibility_id(view_id) + return model_id if model_id and model_id in routing_names else None + + +def is_claude_code_client(headers: Mapping[str, str]) -> bool: + """Claude Code itself, or a client asking for its view of the listing the way Ramp Router's does""" + from litellm.llms.anthropic.common_utils import is_claude_code_user_agent + + return ( + is_claude_code_user_agent(headers.get("user-agent", "")) + or headers.get(GATEWAY_CLIENT_HEADER, "").lower() == CLAUDE_CODE_CLIENT + ) + + +def claude_code_view_ids( + rows: Sequence[ModelInfoResponse], + headers: Mapping[str, str], + routing_names: Container[str], +) -> Mapping[str, str]: + """served id -> Claude Code id for the requested listing view""" + if not is_claude_code_client(headers): + return MappingProxyType({}) + return MappingProxyType( + {row["id"]: claude_code_model_id(row["id"], row.get("max_input_tokens"), routing_names) for row in rows} + ) + + +@dataclass(frozen=True, slots=True) +class ClaudeCodeRoutingNames: + """Existing routes always own their names, including aliases and wildcard routes.""" + + llm_router: Router | None + team_id: str | None = None + alias_maps: tuple[object, ...] = () + + def __contains__(self, name: object) -> bool: + if not isinstance(name, str): + return False + if name in litellm.model_alias_map or any( + isinstance(aliases, Mapping) and name in aliases for aliases in self.alias_maps + ): + return True + if self.llm_router is None: + return False + return ( + name in self.llm_router.model_group_alias + or self.llm_router.has_model_id(name) + or bool(self.llm_router.get_candidate_model_ids_for_route(name, self.team_id)) + ) + + +def claude_code_requested_group( + requested: str, + llm_router: Router, + team_id: str | None, + alias_maps: tuple[object, ...] = (), +) -> str | None: + return claude_code_group_name(requested, ClaudeCodeRoutingNames(llm_router, team_id, alias_maps)) + + class TeamModelNameTranslator: """Translates internal team routing keys to their public names for the model listing/retrieve responses. Stateless; the live router and general_settings diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3bbe932f5d0..9269fd48e6c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -366,8 +366,11 @@ from litellm.proxy.common_utils.load_config_utils import ( ) from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations from litellm.proxy.common_utils.model_listing_utils import ( + ClaudeCodeRoutingNames, TeamModelNameTranslator, + claude_code_view_ids, configured_display_names, + is_claude_code_client, ) from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, @@ -6639,6 +6642,14 @@ class ProxyConfig: return parsed return None + async def get_hierarchical_router_settings( + self, + user_api_key_dict: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + proxy_logging_obj: ProxyLogging | None = None, + ) -> dict | None: + return await self._get_hierarchical_router_settings(user_api_key_dict, prisma_client, proxy_logging_obj) + async def _get_hierarchical_router_settings( self, user_api_key_dict: Optional["UserAPIKeyAuth"], @@ -8632,6 +8643,7 @@ _STREAM_KEEPALIVE: Final = object() _KEEPALIVE_MIN_SECONDS: Final = 1.0 _KEEPALIVE_MAX_SECONDS: Final = 300.0 _EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) +_EMPTY_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) async def _iter_with_keepalive( @@ -10469,6 +10481,24 @@ async def model_list( wants_anthropic_format: Final = ( http_request is not None and http_request.headers.get("anthropic-version") is not None ) + client_headers: Final[Mapping[str, str]] = http_request.headers if http_request is not None else _EMPTY_HEADERS + view_router_settings: Final = ( + await proxy_config.get_hierarchical_router_settings(user_api_key_dict, prisma_client, proxy_logging_obj) + if wants_anthropic_format and is_claude_code_client(client_headers) + else None + ) + view_aliases: Final = ( + view_router_settings.get("model_group_alias") if isinstance(view_router_settings, Mapping) else None + ) + routing_names: Final = ClaudeCodeRoutingNames( + llm_router, + team_id or user_api_key_dict.team_id, + ( + user_api_key_dict.aliases, + user_api_key_dict.team_model_aliases, + view_aliases, + ), + ) # Validate scope parameter if provided if scope is not None and scope != "expand": @@ -10555,6 +10585,11 @@ async def model_list( return create_anthropic_model_list_response( admin_listing, display_names=configured_display_names(admin_entries, llm_router), + listed_ids=claude_code_view_ids( + admin_listing, + client_headers, + routing_names, + ), ) return dict( @@ -10603,6 +10638,11 @@ async def model_list( return create_anthropic_model_list_response( listing, display_names=configured_display_names(entries, llm_router), + listed_ids=claude_code_view_ids( + listing, + client_headers, + routing_names, + ), ) return dict( 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 91652c89092..e1b39c4ba13 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -2267,3 +2267,25 @@ 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 + + +def test_create_anthropic_model_list_response_lists_ids_as_told(): + """listed_ids renames an entry for the caller while display_name and every other field stay keyed to the served + id, and the envelope's first/last ids follow the renamed entries.""" + from litellm.llms.anthropic.common_utils import ( + create_anthropic_model_list_response, + ) + + response = create_anthropic_model_list_response( + [ + {"id": "gpt-4o", "object": "model", "created": 0, "owned_by": "openai", "max_input_tokens": 1000000}, + {"id": "claude-haiku-4-5", "object": "model", "created": 0, "owned_by": "openai"}, + ], + display_names={"gpt-4o": "GPT 4o"}, + listed_ids={"gpt-4o": "claude-router-gpt-4o[1m]"}, + ) + + gpt, haiku = response["data"] + assert (gpt["id"], gpt["display_name"], gpt["max_input_tokens"]) == ("claude-router-gpt-4o[1m]", "GPT 4o", 1000000) + assert (haiku["id"], haiku["display_name"]) == ("claude-haiku-4-5", "claude-haiku-4-5") + assert (response["first_id"], response["last_id"]) == ("claude-router-gpt-4o[1m]", "claude-haiku-4-5") diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 8e6c41761cc..6cce6d0316b 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -7249,3 +7249,46 @@ async def test_jwt_builder_returns_every_team_grant_the_key_path_gets(is_proxy_a assert token.team_member == Member(user_id="jwt-user", role="admin") assert token.team_member_spend == 1.5 assert token.jwt_claims == {"sub": "jwt-user"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/v1/messages", "/messages", "/v1/chat/completions", "/chat/completions", "/v1/responses", "/responses"]) +async def test_claude_view_normalizes_before_model_access(monkeypatch, route): + from starlette.requests import Request + from litellm.proxy.auth.user_api_key_auth import _enforce_key_and_fallback_model_access + + source = "foo[1m]" + encoded = "claude-router-" + source.encode().hex() + "[1m]" + router = litellm.Router(model_list=[{"model_name": source, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}}]) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + data = {"model": encoded, "messages": [{"role": "user", "content": "hi"}]} + request = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""}) + token = UserAPIKeyAuth(models=[source]) + await _enforce_key_and_fallback_model_access(valid_token=token, request_data=data, route=route, request=request, llm_model_list=router.model_list, llm_router=router) + assert data["model"] == source + assert (await request.json())["model"] == source + assert json.loads(await request.body())["model"] == source + assert request.scope["parsed_body"][1]["model"] == source + with pytest.raises(ProxyException): + await _enforce_key_and_fallback_model_access(valid_token=UserAPIKeyAuth(models=["other"]), request_data=data, route=route, request=request, llm_model_list=router.model_list, llm_router=router) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("layer", ["literal", "global", "router", "key", "hierarchical", "unclaimed"]) +async def test_claude_view_never_reinterprets_explicit_names(monkeypatch, layer): + from starlette.requests import Request + from litellm.proxy.auth.user_api_key_auth import _normalize_claude_model + + encoded = "claude-router-666f6f" + names = ("foo", "other", encoded) if layer == "literal" else ("foo", "other") + alias = {encoded: "other"} + router = litellm.Router(model_list=[{"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} for name in names], model_group_alias=alias if layer == "router" else None) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "model_alias_map", alias if layer == "global" else {}) + token = UserAPIKeyAuth(aliases=alias if layer == "key" else {}, router_settings={"model_group_alias": alias} if layer == "hierarchical" else None) + data = {"model": encoded} + request = Request({"type": "http", "method": "POST", "path": "/v1/messages", "headers": [], "query_string": b""}) + await _normalize_claude_model(data, token, request, "/v1/messages") + assert data["model"] == ("foo" if layer == "unclaimed" else encoded) + await _normalize_claude_model(data, token, request, "/v1/messages") + assert data["model"] == ("foo" if layer == "unclaimed" else encoded) diff --git a/tests/test_litellm/proxy/client/cli/test_configure_commands.py b/tests/test_litellm/proxy/client/cli/test_configure_commands.py index a1808383982..ac7408339fe 100644 --- a/tests/test_litellm/proxy/client/cli/test_configure_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_configure_commands.py @@ -89,7 +89,7 @@ class TestConfigureClaudeWithAVirtualKey: assert "Starting model: claude-auto" in result.output assert "1 of the proxy's 2 models" in result.output assert "lite unconfigure claude" in result.output - assert len(responses.calls) == 1 + assert [call.request.headers.get("x-gateway-client") for call in responses.calls] == ["claude-code"] @responses.activate def test_takes_the_key_from_the_global_option_and_keeps_claude_codes_default(self, runner, paths): @@ -358,3 +358,74 @@ class TestUnconfigureClaude: result = runner.invoke(cli, ["unconfigure", "claude"]) assert result.exit_code != 0 assert "nothing to undo" in result.output + + +class TestClaudeCodeView: + VIEW = {"anthropic-version": "2023-06-01", "x-gateway-client": "claude-code"} + + def _mock(self, rows): + responses.get( + f"{PROXY}/v1/models", + json={"data": rows}, + match=[responses.matchers.header_matcher({"Authorization": f"Bearer {VALID_KEY}", **self.VIEW})], + ) + + @responses.activate + @pytest.mark.parametrize( + "model, pinned", + [ + ("literal-claude-router-source", "emitted-literal"), + ("marked-sibling", "emitted-marked[1m]"), + ("emitted-collision", "emitted-source-priority"), + ("emitted-only", "emitted-only"), + ], + ) + def test_pins_source_identity_before_emitted_id(self, runner, paths, model, pinned): + self._mock( + [ + {"id": "emitted-collision", "source_model": "other-source"}, + {"id": "emitted-source-priority", "source_model": "emitted-collision"}, + {"id": "emitted-marked[1m]", "source_model": "marked-sibling"}, + {"id": "emitted-literal", "source_model": "literal-claude-router-source"}, + {"id": "emitted-only"}, + ] + ) + settings_path, _ = paths + result = _configure(runner, "--api-key", VALID_KEY, "--model", model) + assert result.exit_code == 0, result.output + assert json.loads(settings_path.read_text())["model"] == pinned + assert f"Starting model: {pinned}" in result.output + assert len(responses.calls) == 1 + + @responses.activate + def test_refuses_unknown_short_suffix(self, runner, paths): + self._mock([{"id": "emitted-router-source", "source_model": "literal-router-source"}]) + settings_path, _ = paths + result = _configure(runner, "--api-key", VALID_KEY, "--model", "source") + assert result.exit_code != 0 + assert "'source' is not served" in result.output + assert not settings_path.exists() + + @responses.activate + def test_interactive_picker_uses_source_names(self, paths): + self._mock([{"id": "emitted", "source_model": "source"}]) + settings_path, _ = paths + asked = {} + ctx = click.Context( + configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False} + ) + + def pick_model(listed): + asked["listed"] = tuple(listed) + return "source" + + interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=pick_model) + assert asked["listed"] == ("source",) + assert json.loads(settings_path.read_text())["model"] == "emitted" + + @responses.activate + def test_counts_what_an_older_proxy_lets_the_picker_show(self, runner, paths): + _mock_models() + result = _configure(runner, "--api-key", VALID_KEY) + assert result.exit_code == 0, result.output + assert "/model will list 1 of the proxy's 2 models: Claude Code shows only ids containing" in result.output diff --git a/tests/test_litellm/proxy/client/cli/test_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py index 1c2da514ee2..03e5d7dd197 100644 --- a/tests/test_litellm/proxy/client/cli/test_pi.py +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -13,6 +13,7 @@ from litellm.proxy.client.cli.commands.pi import ( PiSyncError, fetch_model_ids, fetch_model_limits, + fetch_model_listing, models_json_path, provider_block, sync_models_json, @@ -50,6 +51,43 @@ class TestFetchModelIds: assert captured["url"] == "http://localhost:4000/v1/models" assert captured["headers"] == {"Authorization": "Bearer sk-key"} + def test_returns_rows_with_optional_source_model_and_dedups_identical_rows(self): + result = fetch_model_listing( + "http://localhost:4000", + "sk-key", + get=lambda *a, **k: _FakeResponse( + 200, + {"data": [{"id": "emitted", "source_model": "source"}, {"id": "emitted", "source_model": "source"}]}, + ), + ) + assert not isinstance(result, PiSyncError) + assert tuple((model.id, model.source_model) for model in result) == (("emitted", "source"),) + + @pytest.mark.parametrize( + "entry", + [ + {"id": ""}, + {"id": "emitted", "source_model": ""}, + {"id": "emitted", "source_model": 1}, + ], + ) + def test_rejects_invalid_model_identity(self, entry): + result = fetch_model_listing( + "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(200, {"data": [entry]}) + ) + assert isinstance(result, PiSyncError) and result.kind is ListingFailure.BAD_BODY + + def test_rejects_conflicting_emitted_id_mappings(self): + result = fetch_model_listing( + "http://localhost:4000", + "sk-key", + get=lambda *a, **k: _FakeResponse( + 200, + {"data": [{"id": "emitted", "source_model": "one"}, {"id": "emitted", "source_model": "two"}]}, + ), + ) + assert isinstance(result, PiSyncError) and result.kind is ListingFailure.BAD_BODY + def test_network_error_is_a_value(self): def boom(*a, **k): raise requests.ConnectionError("refused") diff --git a/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py b/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py new file mode 100644 index 00000000000..7ef03140093 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py @@ -0,0 +1,109 @@ +"""Model identity survives Claude Code presentation, filtering and configured alias precedence.""" + +from itertools import combinations + +import pytest + +from litellm import Router +from litellm.proxy.common_utils.model_listing_utils import ( + ClaudeCodeRoutingNames, + claude_code_group_name, + claude_code_model_id, + claude_code_requested_group, + claude_code_view_ids, +) + + +def _encoded(name): + return "claude-router-" + name.encode().hex() + + +def _marked(name): + return f"{_encoded(name)}[1m]" + + +def _row(name, limit=1000000): + return {"id": name, "object": "model", "created": 0, "owned_by": "openai", "max_input_tokens": limit} + + +def _router(*names, aliases=None): + return Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} + for name in names + ], + model_group_alias=aliases, + ) + + +@pytest.mark.parametrize("limit", [None, 999999, 1000000]) +@pytest.mark.parametrize("name", ["foo", "foo[1m]", "foo[1M]", "a/b: 世界", "claude-router-foo", "claude-opus-5", "claude-opus-5[1m]"]) +def test_listing_round_trips_entire_source_name(name, limit): + names = frozenset({name}) + view = claude_code_model_id(name, limit, names) + assert (claude_code_group_name(view, names) or view) == name + if "claude" not in name: + assert view.startswith(_encoded(name)) + assert ("[1m]" in view.lower()) == (limit == 1000000 or "[1m]" in name.lower() and "claude" in name) + + +def test_collision_matrix_round_trips_without_duplicate_ids(): + universe = ("foo", "foo[1m]", "claude-router-foo", _encoded("foo"), _encoded("foo") + "[1m]", "claude-opus-5", "claude-opus-5[1m]") + for pair in combinations(universe, 2): + for visible in (pair, pair[:1], pair[1:]): + names = frozenset(pair) + view = claude_code_view_ids(tuple(_row(n) for n in visible), {"user-agent": "claude-code/2.1.267"}, names) + assert len(set(view.values())) == len(visible) + assert all((claude_code_group_name(shown, names) or shown) == source for source, shown in view.items()) + + +@pytest.mark.parametrize("spelling", ["claude-router-foo", "claude-router-ff", "claude-router-66 6f6f", "claude-router-666F6F", "claude-router-", _encoded("missing")]) +def test_unknown_or_noncanonical_ids_are_never_guessed(spelling): + assert claude_code_group_name(spelling, frozenset({"foo"})) is None + + +@pytest.mark.parametrize("headers,enabled", [ + ({"user-agent": "claude-code/2.1.267"}, True), + ({"user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, True), + ({"x-gateway-client": "Claude-Code"}, True), + ({"user-agent": "anthropic-sdk-python/0.40"}, False), + ({}, False), +]) +def test_only_claude_code_gets_the_view(headers, enabled): + rows = (_row("foo"), _row("claude-opus-5")) + view = claude_code_view_ids(rows, headers, frozenset(row["id"] for row in rows)) + assert dict(view) == ({"foo": _encoded("foo") + "[1m]", "claude-opus-5": "claude-opus-5[1m]"} if enabled else {}) + + +@pytest.mark.parametrize("layer", ["literal", "global", "router", "key", "team", "wildcard"]) +def test_configured_names_outrank_generated_ids_even_when_hidden_from_listing(monkeypatch, layer): + import litellm + + encoded = _encoded("foo") + alias = {encoded: "other"} + monkeypatch.setattr(litellm, "model_alias_map", alias if layer == "global" else {}) + router = _router("foo", "other", *( (encoded,) if layer == "literal" else ("*",) if layer == "wildcard" else ()), aliases=alias if layer == "router" else None) + maps = (alias,) if layer in ("key", "team") else () + names = ClaudeCodeRoutingNames(router, None, maps) + assert claude_code_requested_group(encoded, router, None, maps) is None + assert claude_code_view_ids((_row("foo", None),), {"x-gateway-client": "claude-code"}, names)["foo"] == "foo" + + +@pytest.mark.parametrize("source", ["foo", "foo[1m]", "世界"]) +def test_mutation_breaking_the_hex_name_cannot_route_to_the_source(source): + router = _router(source) + encoded = _encoded(source) + malformed = encoded[:-1] + ("0" if encoded[-1] != "0" else "1") + assert claude_code_requested_group(malformed, router, None) is None + assert claude_code_requested_group(_marked(source), router, None) == source + + +def test_team_public_name_uses_the_same_scope_at_list_and_request(): + router = Router(model_list=[{ + "model_name": "model_name_team-a_id", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, + "model_info": {"team_id": "team-a", "team_public_model_name": "shared"}, + }]) + shown = claude_code_view_ids((_row("shared"),), {"user-agent": "claude-code/2.1.267"}, ClaudeCodeRoutingNames(router, "team-a"))["shared"] + assert claude_code_requested_group(shown, router, "team-a") == "shared" + assert claude_code_requested_group(shown, router, "team-b") is None diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_models.py b/tests/test_litellm/proxy/proxy_server/test_routes_models.py index bc6106a06f8..e3d6b46e2d2 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_models.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_models.py @@ -343,3 +343,53 @@ def test_anthropic_format_returns_public_team_model_name( assert response.status_code == 200 assert [m["id"] for m in response.json()["data"]] == ["gpt-4-team"] assert internal_name not in response.text + + +@pytest.mark.parametrize("path", ["/v1/models", "/models"]) +@pytest.mark.parametrize( + "caller_headers", + [ + {"anthropic-version": "2023-06-01", "user-agent": "claude-code/2.1.267"}, + {"anthropic-version": "2023-06-01", "user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, + {"anthropic-version": "2023-06-01", "x-gateway-client": "claude-code"}, + ], +) +def test_anthropic_format_lists_claude_code_view_ids_for_claude_code( + client, auth_as, patched_models, monkeypatch, path, caller_headers +): + """Claude Code drops every id without claude/anthropic in it and reads [1m] as its 1M marker, so for Claude + Code (its discovery fetch's own user agent, its SDK's, or the gateway-client header a launcher sends) every + group is listed under a Claude-shaped id with the marker where the window reaches 1M; the display name stays + the served name.""" + + def _create_model_info_response(model_id, provider="openai", **kwargs): + if model_id != "claude-sonnet": + return _stub_model_info_response(model_id=model_id, provider=provider) + return {**_stub_model_info_response(model_id=model_id, provider=provider), "max_input_tokens": 1000000} + + patched_models.model_group_alias = {} + patched_models.has_model_id.return_value = False + patched_models.get_candidate_model_ids_for_route.side_effect = lambda name, team_id=None: frozenset({name}) if name in ("gpt-4", "claude-sonnet") else frozenset() + monkeypatch.setattr(proxy_utils, "create_model_info_response", _create_model_info_response) + + with auth_as(): + response = client.get(path, headers=caller_headers) + + assert response.status_code == 200 + body = response.json() + assert [(m["id"], m["display_name"]) for m in body["data"]] == [ + ("claude-router-6770742d34", "gpt-4"), + ("claude-sonnet[1m]", "claude-sonnet"), + ] + assert (body["first_id"], body["last_id"]) == ("claude-router-6770742d34", "claude-sonnet[1m]") + assert [row["source_model"] for row in body["data"]] == ["gpt-4", "claude-sonnet"] + + +@pytest.mark.parametrize("path", ["/v1/models", "/models"]) +def test_anthropic_format_keeps_served_ids_for_other_anthropic_clients(client, auth_as, patched_models, path): + """An Anthropic SDK asking for the vendor shape gets the served ids: the view is Claude Code's alone.""" + with auth_as(): + response = client.get(path, headers={"anthropic-version": "2023-06-01", "user-agent": "anthropic-sdk-python/0.40"}) + + assert response.status_code == 200 + assert [m["id"] for m in response.json()["data"]] == ["gpt-4", "claude-sonnet"] From 13837d319daf1a38524a0c08ab929aef1ebc0285 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 9 Sep 2026 22:08:30 -0700 Subject: [PATCH 104/107] fix(openai): drop tool schema regex patterns OpenAI's validator cannot compile (#40485) * fix(openai): drop tool schema regex patterns OpenAI's validator cannot compile OpenAI validates function tool parameters with jsonschema's format checker, which compiles every pattern with Python re. Claude Code's Artifact tool ships an ECMA-262 pattern with \p{..} Unicode property escapes, so any OpenAI target behind /v1/messages, /v1/responses or /v1/chat/completions 400s with "Invalid schema for function 'Artifact': '...' is not a 'regex'" for every model family. Drop only the patterns Python re rejects, keep the rest, at the same seams that already flatten top-level combinators. * fix(openai): walk only schema positions, iteratively, and drop regexes for every openai deployment Review round: the regex sanitizer now walks JSON Schema applicator positions only (properties, items, prefixItems, combinators, $defs, additionalProperties and the rest), so a pattern key inside default, examples, const or a vendor extension is data and stays. It also drops patternProperties keys Python re cannot compile, which OpenAI checks the same way. The walk is level-order and rebuilt deepest level first instead of recursive, so the code-quality recursion gate passes and there is no depth cap below what a JSON parser admits. On the chat wire an openai deployment with a custom api_base now drops such regexes too, since that base is usually a proxy in front of the same validator, while the lossier combinator flattening stays limited to api.openai.com hosts. --- .../prompt_templates/common_utils.py | 123 +++++++++++++- litellm/llms/azure/chat/gpt_transformation.py | 16 +- .../azure/chat/o_series_transformation.py | 4 +- .../llms/openai/chat/gpt_transformation.py | 36 ++-- .../llms/openai/responses/transformation.py | 89 +++++----- ...ore_utils_prompt_templates_common_utils.py | 154 +++++++++++++++++- .../test_azure_chat_gpt_transformation.py | 23 +++ .../response/test_azure_transformation.py | 29 ++++ .../chat/test_openai_gpt_transformation.py | 51 ++++++ .../test_openai_responses_transformation.py | 83 ++++++++++ 10 files changed, 532 insertions(+), 76 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index d70530534da..2485896184e 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -6,8 +6,8 @@ import io import json import mimetypes import re -from collections.abc import Iterable, Iterator, Mapping, Sequence -from itertools import groupby +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from itertools import groupby, islice from os import PathLike from pathlib import Path from types import MappingProxyType @@ -1320,17 +1320,128 @@ def flatten_top_level_schema_combinators(schema: Mapping[str, object]) -> Mappin return _flatten_schema_against_root(schema, schema, frozenset(), 0, {}) # mutable-ok: fresh per-call $ref memo -def tool_with_flattened_parameters(tool: Mapping[str, object]) -> Mapping[str, object]: +_SUBSCHEMA_KEYWORDS: Final = frozenset( + { + "additionalItems", + "additionalProperties", + "contains", + "else", + "if", + "items", + "not", + "propertyNames", + "then", + "unevaluatedItems", + "unevaluatedProperties", + } +) +_SUBSCHEMA_LIST_KEYWORDS: Final = frozenset({"allOf", "anyOf", "items", "oneOf", "prefixItems"}) +_SUBSCHEMA_MAP_KEYWORDS: Final = frozenset( + {"$defs", "definitions", "dependentSchemas", "patternProperties", "properties"} +) + +_MAX_SCHEMA_NESTING: Final = 1024 + + +def drop_non_python_regex_patterns(schema: Mapping[str, object]) -> Mapping[str, object]: + """Drop every regex in a schema position that Python's ``re`` cannot compile. + + OpenAI validates tool ``parameters`` against the 2020-12 metaschema with + ``jsonschema``'s format checker, which hands each ``pattern`` value and each + ``patternProperties`` key to ``re.compile``, so a regex written for an + ECMA-262 engine (Unicode property escapes such as ``\\p{Cc}``, as in Claude + Code's ``Artifact`` tool) is refused with "'...' is not a 'regex'" by every + model family on both the chat and Responses wires. Only schema positions are + walked (properties, items, combinators, ``$defs`` and the other applicators), + so a ``pattern`` key inside ``default``, ``examples``, ``const`` or vendor + extensions is data and stays. Outside strict mode the keyword is only a + hint, so dropping it costs the model a constraint and the caller nothing. + Compilable regexes and everything else pass through, the input is never + mutated, and the same object comes back when nothing was dropped. The walk + is level-order rather than recursive, rebuilt deepest level first, and stops + at more schema levels than a JSON parser admits, so a cyclic schema built in + code cannot spin it. + """ + rebuilt: dict[int, Mapping[str, object]] = {} # mutable-ok: per-call memo of rewritten nodes, deepest level first + for level in reversed(tuple(islice(_schema_levels(schema), _MAX_SCHEMA_NESTING))): + rebuilt.update( + (id(node), rewritten) + for node in level + if (rewritten := _node_without_non_python_regex(node, rebuilt)) is not node + ) + return rebuilt.get(id(schema), schema) + + +def _schema_levels(schema: Mapping[str, object]) -> Iterator[tuple[Mapping[str, object], ...]]: + frontier: tuple[Mapping[str, object], ...] = (schema,) # rebind-ok: level-order cursor, one level a round + while frontier: + yield frontier + frontier = tuple(child for node in frontier for child in _subschemas(node)) + + +def _subschemas(node: Mapping[str, object]) -> Iterator[Mapping[str, object]]: + for key, value in node.items(): + if key in _SUBSCHEMA_MAP_KEYWORDS and isinstance(value, dict): + yield from (sub for sub in value.values() if isinstance(sub, dict)) + elif key in _SUBSCHEMA_LIST_KEYWORDS and isinstance(value, list): + yield from (sub for sub in value if isinstance(sub, dict)) + elif key in _SUBSCHEMA_KEYWORDS and isinstance(value, dict): + yield value + + +def _node_without_non_python_regex( + node: Mapping[str, object], rebuilt: Mapping[int, Mapping[str, object]] +) -> Mapping[str, object]: + kept: Final = { # mutable-ok: tool parameters are JSON dicts + key: _keyword_value_rebuilt(key, value, rebuilt) + for key, value in node.items() + if key != "pattern" or not isinstance(value, str) or _is_python_regex(value) + } + return node if len(kept) == len(node) and all(kept[key] is node[key] for key in kept) else kept + + +def _keyword_value_rebuilt(key: str, value: object, rebuilt: Mapping[int, Mapping[str, object]]) -> object: + if key in _SUBSCHEMA_MAP_KEYWORDS and isinstance(value, dict): + kept: Final = { # mutable-ok: tool parameters are JSON dicts + name: rebuilt.get(id(sub), sub) + for name, sub in value.items() + if key != "patternProperties" or not isinstance(name, str) or _is_python_regex(name) + } + return value if len(kept) == len(value) and all(kept[name] is value[name] for name in kept) else kept + if key in _SUBSCHEMA_LIST_KEYWORDS and isinstance(value, list): + items: Final = [rebuilt.get(id(sub), sub) for sub in value] # mutable-ok: tool parameters are JSON lists + return value if all(new is old for new, old in zip(items, value, strict=True)) else items + if key in _SUBSCHEMA_KEYWORDS and isinstance(value, dict): + return rebuilt.get(id(value), value) + return value + + +def _is_python_regex(pattern: str) -> bool: + try: + re.compile(pattern) + except (re.error, RecursionError): + return False + return True + + +def flatten_combinators_and_drop_non_python_regex_patterns(schema: Mapping[str, object]) -> Mapping[str, object]: + return flatten_top_level_schema_combinators(drop_non_python_regex_patterns(schema)) + + +def tool_with_sanitized_parameters( + tool: Mapping[str, object], + sanitize: Callable[[Mapping[str, object]], Mapping[str, object]], +) -> Mapping[str, object]: function: Final = tool.get("function") if not isinstance(function, dict): return tool parameters: Final = function.get("parameters") if not isinstance(parameters, dict): return tool - flattened: Final = flatten_top_level_schema_combinators(parameters) - if flattened is parameters: + sanitized: Final = sanitize(parameters) + if sanitized is parameters: return tool - return {**tool, "function": {**function, "parameters": flattened}} # mutable-ok: request tools are JSON dicts + return {**tool, "function": {**function, "parameters": sanitized}} # mutable-ok: request tools are JSON dicts def _get_image_mime_type_from_url(url: str) -> str | None: diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 880a51eb584..ed16d7f3de0 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -7,8 +7,9 @@ from httpx._models import Headers, Response import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_tool_reference_parts_from_tool_messages, + flatten_combinators_and_drop_non_python_regex_patterns, hoist_images_from_tool_messages, - tool_with_flattened_parameters, + tool_with_sanitized_parameters, ) from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_azure_openai_messages, @@ -39,14 +40,17 @@ else: _NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) -def flattened_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]: +def sanitized_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]: tools: Final = optional_params.get("tools") if not isinstance(tools, list): return _NO_TOOLS_UPDATE - flattened: Final = [ # mutable-ok: request tools are a JSON list - tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools + sanitized: Final = [ # mutable-ok: request tools are a JSON list + tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) + if isinstance(tool, dict) + else tool + for tool in tools ] - return MappingProxyType({"tools": flattened}) + return MappingProxyType({"tools": sanitized}) class AzureOpenAIConfig(BaseConfig): @@ -278,7 +282,7 @@ class AzureOpenAIConfig(BaseConfig): "model": model, "messages": azure_messages, **optional_params, - **flattened_tools_update(optional_params), + **sanitized_tools_update(optional_params), } def transform_response( diff --git a/litellm/llms/azure/chat/o_series_transformation.py b/litellm/llms/azure/chat/o_series_transformation.py index 246bf69cb5f..09d8075e857 100644 --- a/litellm/llms/azure/chat/o_series_transformation.py +++ b/litellm/llms/azure/chat/o_series_transformation.py @@ -20,7 +20,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.utils import get_model_info, supports_reasoning from ...openai.chat.o_series_transformation import OpenAIOSeriesConfig -from .gpt_transformation import flattened_tools_update +from .gpt_transformation import sanitized_tools_update class AzureOpenAIO1Config(OpenAIOSeriesConfig): @@ -111,6 +111,6 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): model = model.replace("o_series/", "") # handle o_series/my-random-deployment-name flattened_params: Final = { # mutable-ok: transform_request's contract takes a plain JSON params dict **optional_params, - **flattened_tools_update(optional_params), + **sanitized_tools_update(optional_params), } return super().transform_request(model, messages, flattened_params, litellm_params, headers) diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 9afc6331d96..9b410cf073e 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -19,10 +19,12 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo _should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, drop_tool_reference_parts_from_tool_messages, + flatten_combinators_and_drop_non_python_regex_patterns, get_tool_call_names, hoist_images_from_tool_messages, - tool_with_flattened_parameters, + tool_with_sanitized_parameters, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, @@ -432,7 +434,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): custom_llm_provider, api_base ) - def _flattened_tools_update_for_openai( + def _sanitized_tools_update_for_openai( self, optional_params: Mapping[str, object], litellm_params: Mapping[str, object], @@ -440,22 +442,26 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): """ OpenAI's chat completions validator rejects tool `parameters` carrying 'oneOf'/'anyOf'/'allOf'/'enum'/'const'/'not' at the top level for every - model family, unlike the Responses API, where GPT-5+ accepts them. + model family, unlike the Responses API, where GPT-5+ accepts them, and + a `pattern` Python's `re` cannot compile for every model family on both. + A custom api_base on the `openai` provider is usually a proxy in front of + the same validator, so regexes are dropped there too, while the lossier + combinator flattening stays limited to api.openai.com hosts. """ tools: Final = optional_params.get("tools") - if not isinstance(tools, list): - return _NO_TOOLS_UPDATE provider: Final = litellm_params.get("custom_llm_provider") - raw_api_base: Final = litellm_params.get("api_base") - if not self._targets_openai_hosted_endpoint( - provider if isinstance(provider, str) else None, - raw_api_base if isinstance(raw_api_base, str) else None, - ): + if not isinstance(tools, list) or provider != "openai": return _NO_TOOLS_UPDATE - flattened: Final = [ # mutable-ok: request tools are a JSON list - tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools + raw_api_base: Final = litellm_params.get("api_base") + sanitize: Final = ( + flatten_combinators_and_drop_non_python_regex_patterns + if self._targets_openai_hosted_endpoint(provider, raw_api_base if isinstance(raw_api_base, str) else None) + else drop_non_python_regex_patterns + ) + sanitized: Final = [ # mutable-ok: request tools are a JSON list + tool_with_sanitized_parameters(tool, sanitize) if isinstance(tool, dict) else tool for tool in tools ] - return MappingProxyType({"tools": flattened}) + return MappingProxyType({"tools": sanitized}) def transform_request( self, @@ -489,7 +495,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): "model": model, "messages": messages, **optional_params, - **self._flattened_tools_update_for_openai(optional_params, litellm_params), + **self._sanitized_tools_update_for_openai(optional_params, litellm_params), } async def async_transform_request( @@ -521,7 +527,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): "model": model, "messages": transformed_messages, **optional_params, - **self._flattened_tools_update_for_openai(optional_params, litellm_params), + **self._sanitized_tools_update_for_openai(optional_params, litellm_params), } else: ## allow for any object specific behaviour to be handled diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 97e7bcc60b7..666e9b32011 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints @@ -15,6 +15,10 @@ from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + flatten_combinators_and_drop_non_python_regex_patterns, +) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig @@ -40,7 +44,7 @@ else: _NO_TOOL_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) _MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3.5", "chatgpt-4o", "o1", "o3", "o4") -_PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) +_PROVIDERS_WITH_OPENAI_SCHEMA_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) @@ -293,7 +297,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): model=model, input=validated_input, tools=tools ) object_schema_tools: Final = self._tools_with_object_parameters(model=model, tools=stripped_tools) - sanitized_tools: Final = self._flatten_tool_schema_combinators_for_openai( + sanitized_tools: Final = self._sanitized_tool_schemas_for_openai( model=model, tools=object_schema_tools, litellm_params=litellm_params ) return self._drop_foreign_tool_call_item_ids(stripped_input), sanitized_tools @@ -378,35 +382,35 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return item return {key: value for key, value in item.items() if key != "id"} # mutable-ok: outgoing JSON request item - def _flatten_tool_schema_combinators_for_openai( + def _sanitized_tool_schemas_for_openai( self, model: str, tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None, litellm_params: GenericLiteLLMParams, ) -> Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None: - """Flatten top-level schema combinators only where OpenAI's validator rejects them. + """Rewrite tool schemas only where OpenAI's validator rejects them. - OpenAI-compatible backends reusing this config (and the ChatGPT backend - Codex talks to natively) accept them, and so do GPT-5 and later models, - which also call tools better with the union intact. Codex wraps MCP tools - inside namespace entries, so nested ``tools`` arrays are walked too. - Azure OpenAI shares the validator but names deployments arbitrarily, so - the router's declared ``model_info.base_model`` wins over the deployment - name and an unrecognized name without one is left untouched. + Every model family refuses a ``pattern`` Python's ``re`` cannot compile, + while top-level schema combinators are flattened only for the families + whose validator rejects them: OpenAI-compatible backends reusing this + config (and the ChatGPT backend Codex talks to natively) accept them, + and so do GPT-5 and later models, which also call tools better with the + union intact. Codex wraps MCP tools inside namespace entries, so nested + ``tools`` arrays are walked too. Azure OpenAI shares the validator but + names deployments arbitrarily, so the router's declared + ``model_info.base_model`` wins over the deployment name and an + unrecognized name without one keeps its combinators. """ - if tools is None or self.custom_llm_provider not in _PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: + if tools is None or self.custom_llm_provider not in _PROVIDERS_WITH_OPENAI_SCHEMA_VALIDATOR: return tools gate_model: Final = self._combinator_gate_model(model=model, litellm_params=litellm_params) - if not self._rejects_top_level_schema_combinators(gate_model): - return tools - flattened: Final = [ # mutable-ok: request tools are a JSON list - self._flattened_tool_or_passthrough(tool) for tool in tools - ] - return cast("Sequence[ALL_RESPONSES_API_TOOL_PARAMS]", flattened) # cast-ok: spread keeps each tool's shape - - @staticmethod - def _flattened_tool_or_passthrough(tool: object) -> object: - return OpenAIResponsesAPIConfig._flattened_tool_entry(tool) if isinstance(tool, dict) else tool + sanitize: Final = ( + flatten_combinators_and_drop_non_python_regex_patterns + if self._rejects_top_level_schema_combinators(gate_model) + else drop_non_python_regex_patterns + ) + sanitized: Final = self._sanitized_tools(tools, sanitize) + return cast("Sequence[ALL_RESPONSES_API_TOOL_PARAMS]", sanitized) # cast-ok: spread keeps each tool's shape @staticmethod def _rejects_top_level_schema_combinators(model: str) -> bool: @@ -421,35 +425,42 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return base_model if isinstance(base_model, str) and base_model else model @staticmethod - def _flattened_tool_entry( + def _sanitized_tool_entry( entry: Mapping[str, object], - ) -> dict[str, object]: # mutable-ok: request tools are JSON dicts - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - flatten_top_level_schema_combinators, - ) - + sanitize: Callable[[Mapping[str, object]], Mapping[str, object]], + ) -> Mapping[str, object]: parameters: Final = entry.get("parameters") nested_tools: Final = entry.get("tools") + sanitized_parameters: Final = sanitize(parameters) if isinstance(parameters, dict) else parameters + sanitized_nested_tools: Final = ( + OpenAIResponsesAPIConfig._sanitized_tools(nested_tools, sanitize) + if isinstance(nested_tools, list) + else nested_tools + ) parameters_update: Final = ( - MappingProxyType({"parameters": flatten_top_level_schema_combinators(parameters)}) - if isinstance(parameters, dict) + MappingProxyType({"parameters": sanitized_parameters}) + if sanitized_parameters is not parameters else _NO_TOOL_UPDATE ) tools_update: Final = ( - MappingProxyType({"tools": OpenAIResponsesAPIConfig._flattened_nested_tools(nested_tools)}) - if isinstance(nested_tools, list) + MappingProxyType({"tools": sanitized_nested_tools}) + if sanitized_nested_tools is not nested_tools else _NO_TOOL_UPDATE ) + if not parameters_update and not tools_update: + return entry return {**entry, **parameters_update, **tools_update} # mutable-ok: request tools are JSON dicts @staticmethod - def _flattened_nested_tools( - nested_tools: Sequence[object], - ) -> list[object]: # mutable-ok: namespace tools are a JSON list - return [ # mutable-ok: namespace tools are a JSON list - OpenAIResponsesAPIConfig._flattened_tool_entry(item) if isinstance(item, dict) else item - for item in nested_tools + def _sanitized_tools( + tools: Sequence[object], + sanitize: Callable[[Mapping[str, object]], Mapping[str, object]], + ) -> Sequence[object]: + sanitized: Final = [ # mutable-ok: request tools are a JSON list + OpenAIResponsesAPIConfig._sanitized_tool_entry(item, sanitize) if isinstance(item, dict) else item + for item in tools ] + return tools if all(new is old for new, old in zip(sanitized, tools, strict=True)) else sanitized def _validate_input_param(self, input: str | ResponseInputParam) -> str | ResponseInputParam: """ diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 7c1445d79c9..b5890d1a5b0 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -2,11 +2,12 @@ import copy import functools import json import os +import sys +from typing import Final from unittest.mock import MagicMock, patch import pytest - from litellm.litellm_core_utils.prompt_templates.common_utils import ( ENCRYPTED_REASONING_SIGNATURE_PREFIX, TOOL_RESULT_IMAGE_BOUNDARY, @@ -25,6 +26,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( update_messages_with_model_file_ids, ) +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + def test_get_format_from_file_id(): unified_file_id = "litellm_proxy:application/pdf;unified_id,cbbe3534-8bf8-4386-af00-f5f6b7e370bf" @@ -1442,7 +1445,7 @@ class TestFlattenTopLevelSchemaCombinators: assert schema == snapshot -class TestToolWithFlattenedParameters: +class TestToolWithSanitizedParameters: def _anyof_tool(self): return { "type": "function", @@ -1469,11 +1472,12 @@ class TestToolWithFlattenedParameters: def test_flattens_anyof_parameters_into_new_tool(self): from litellm.litellm_core_utils.prompt_templates.common_utils import ( - tool_with_flattened_parameters, + flatten_combinators_and_drop_non_python_regex_patterns, + tool_with_sanitized_parameters, ) tool = self._anyof_tool() - result = tool_with_flattened_parameters(tool) + result = tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) assert result is not tool parameters = result["function"]["parameters"] @@ -1486,7 +1490,8 @@ class TestToolWithFlattenedParameters: def test_clean_parameters_return_the_same_tool_object(self): from litellm.litellm_core_utils.prompt_templates.common_utils import ( - tool_with_flattened_parameters, + flatten_combinators_and_drop_non_python_regex_patterns, + tool_with_sanitized_parameters, ) tool = { @@ -1497,7 +1502,23 @@ class TestToolWithFlattenedParameters: }, } - assert tool_with_flattened_parameters(tool) is tool + assert tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) is tool + + def test_pattern_only_sanitizer_drops_the_regex_and_keeps_the_union(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + tool_with_sanitized_parameters, + ) + + tool = self._anyof_tool() + tool["function"]["parameters"]["properties"]["id"]["pattern"] = _ARTIFACT_FIELD_PATTERN + + result = tool_with_sanitized_parameters(tool, drop_non_python_regex_patterns) + + parameters = result["function"]["parameters"] + assert parameters["properties"]["id"] == {"type": "string"} + assert parameters["anyOf"] == self._anyof_tool()["function"]["parameters"]["anyOf"] + assert tool["function"]["parameters"]["properties"]["id"]["pattern"] == _ARTIFACT_FIELD_PATTERN @pytest.mark.parametrize( "tool", @@ -1510,10 +1531,127 @@ class TestToolWithFlattenedParameters: ) def test_non_dict_function_or_parameters_return_the_same_tool_object(self, tool): from litellm.litellm_core_utils.prompt_templates.common_utils import ( - tool_with_flattened_parameters, + flatten_combinators_and_drop_non_python_regex_patterns, + tool_with_sanitized_parameters, ) - assert tool_with_flattened_parameters(tool) is tool + assert tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) is tool + + +class TestDropNonPythonRegexPatterns: + """Claude Code's Artifact tool declares ECMA-262 ``\\p{..}`` escapes that OpenAI's + validator, which compiles ``pattern`` values and ``patternProperties`` keys with + Python ``re``, refuses as "not a 'regex'".""" + + def _schema(self, pattern): + return { + "type": "object", + "properties": { + "field": {"type": "string", "pattern": pattern}, + "writes": { + "type": "array", + "items": {"properties": {"doc_id": {"type": "string", "pattern": pattern}}}, + }, + "query": {"anyOf": [{"type": "string", "pattern": pattern}, {"type": "null"}]}, + "pair": {"type": "array", "prefixItems": [{"type": "string", "pattern": pattern}]}, + "extra": {"type": "object", "additionalProperties": {"type": "string", "pattern": pattern}}, + "tagged": { + "type": "object", + "patternProperties": {pattern: {"type": "string"}, "^x_": {"type": "integer"}}, + }, + }, + "$defs": {"segment": {"type": "string", "pattern": pattern}}, + "required": ["field"], + } + + def test_drops_every_regex_python_re_rejects_from_every_schema_position(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + schema = self._schema(_ARTIFACT_FIELD_PATTERN) + + result = drop_non_python_regex_patterns(schema) + + assert '"pattern"' not in json.dumps(result) + properties = result["properties"] + assert properties["field"] == {"type": "string"} + assert properties["writes"]["items"]["properties"]["doc_id"] == {"type": "string"} + assert properties["query"]["anyOf"] == [{"type": "string"}, {"type": "null"}] + assert properties["pair"]["prefixItems"] == [{"type": "string"}] + assert properties["extra"]["additionalProperties"] == {"type": "string"} + assert properties["tagged"]["patternProperties"] == {"^x_": {"type": "integer"}} + assert result["$defs"]["segment"] == {"type": "string"} + assert result["required"] == ["field"] + assert schema == self._schema(_ARTIFACT_FIELD_PATTERN) + + def test_keeps_regexes_python_re_compiles_and_returns_the_same_object(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + schema = self._schema(r'^(?!__.*__$)[^"\\./[\]]{1,200}$') + + assert drop_non_python_regex_patterns(schema) is schema + + def test_pattern_keys_inside_data_positions_are_not_regexes(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + schema = { + "type": "object", + "properties": { + "pattern": {"type": "string"}, + "template": {"type": "object", "default": {"pattern": _ARTIFACT_FIELD_PATTERN}}, + "samples": {"type": "array", "examples": [{"pattern": _ARTIFACT_FIELD_PATTERN}]}, + "fixed": {"const": {"pattern": _ARTIFACT_FIELD_PATTERN}}, + "vendor": {"type": "string", "x-litellm": {"pattern": _ARTIFACT_FIELD_PATTERN}}, + }, + "required": ["pattern"], + } + + assert drop_non_python_regex_patterns(schema) is schema + + def test_regex_nested_past_what_python_re_can_parse_is_dropped_not_raised(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + schema = { + "type": "object", + "properties": {"deep": {"type": "string", "pattern": "(" * 2000 + "a" + ")" * 2000}}, + } + + assert drop_non_python_regex_patterns(schema)["properties"]["deep"] == {"type": "string"} + + def test_walks_schemas_deeper_than_the_interpreter_recursion_limit(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + depth = sys.getrecursionlimit() + leaf = {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN} + schema = functools.reduce( + lambda inner, _: {"type": "object", "properties": {"child": inner}}, range(depth), leaf + ) + + result = drop_non_python_regex_patterns(schema) + + assert functools.reduce(lambda node, _: node["properties"]["child"], range(depth), result) == {"type": "string"} + assert functools.reduce(lambda node, _: node["properties"]["child"], range(depth), schema) is leaf + + def test_leaves_levels_past_the_json_nesting_limit_alone(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + leaf = {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN} + schema = functools.reduce( + lambda inner, _: {"type": "object", "properties": {"child": inner}}, range(1100), leaf + ) + + assert drop_non_python_regex_patterns(schema) is schema class TestRequestContainsImageContent: diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py index 4e6b9ed0188..bc6cb0c0fed 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -198,6 +198,9 @@ def test_azure_gpt_5_takes_the_reasoning_path() -> None: assert "reasoning_effort" in supported +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + + class TestAzureToolSchemaCombinatorFlattening: """ Regression tests for LIT-6510: Azure's chat completions validator rejects @@ -259,6 +262,26 @@ class TestAzureToolSchemaCombinatorFlattening: self._transform(AzureOpenAIConfig(), "gpt-4o", [tool]) assert tool == self._anyof_tool() + def test_transform_request_drops_non_python_regex_pattern(self): + tool = { + "type": "function", + "function": { + "name": "Artifact", + "parameters": { + "type": "object", + "properties": {"field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}}, + }, + }, + } + + request = self._transform(AzureOpenAIConfig(), "gpt-4o", [tool]) + + assert request["tools"][0]["function"]["parameters"] == { + "type": "object", + "properties": {"field": {"type": "string"}}, + } + assert tool["function"]["parameters"]["properties"]["field"]["pattern"] == _ARTIFACT_FIELD_PATTERN + def test_clean_object_schema_passes_through_as_same_object(self): tool = { "type": "function", diff --git a/tests/test_litellm/llms/azure/response/test_azure_transformation.py b/tests/test_litellm/llms/azure/response/test_azure_transformation.py index 0cac2705ab0..726c9f65681 100644 --- a/tests/test_litellm/llms/azure/response/test_azure_transformation.py +++ b/tests/test_litellm/llms/azure/response/test_azure_transformation.py @@ -1,4 +1,5 @@ from copy import deepcopy +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -243,6 +244,9 @@ def test_provider_config_manager_o_series_selection(): assert not isinstance(default_config, AzureOpenAIOSeriesResponsesAPIConfig) +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + + class TestAzureResponsesAPIConfig: def setup_method(self): self.config = AzureOpenAIResponsesAPIConfig() @@ -599,6 +603,31 @@ class TestAzureResponsesAPIConfig: assert result["tools"][0] is tool assert "anyOf" in result["tools"][0]["parameters"] + def test_azure_drops_non_python_regex_pattern_while_keeping_gpt5_combinators(self): + tool = { + "type": "function", + "name": "Artifact", + "parameters": { + "type": "object", + "anyOf": [{"properties": {"field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}}}], + "properties": {"field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}}, + }, + } + + result = self.config.transform_responses_api_request( + model="my-eastus-deployment", + input="hi", + response_api_optional_request_params={"tools": [tool]}, + litellm_params=GenericLiteLLMParams(model_info={"base_model": "azure/gpt-5.4-mini"}), + headers={}, + ) + + assert result["tools"][0]["parameters"] == { + "type": "object", + "anyOf": [{"properties": {"field": {"type": "string"}}}], + "properties": {"field": {"type": "string"}}, + } + def test_azure_keeps_combinators_for_unrecognized_deployment_without_base_model(self): tool = self._anyof_tool() diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index 9737d63cc26..b110586ae5b 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py @@ -4,6 +4,7 @@ Tests for OpenAI GPT transformation (litellm/llms/openai/chat/gpt_transformation import pytest +from typing import Final import litellm @@ -1168,6 +1169,9 @@ class TestOpenAIPromptCacheBreakpointChatPath: assert "prompt_cache_options" not in request +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + + class TestToolSchemaCombinatorFlatteningForOpenAI: """ Regression tests for LIT-6488: OpenAI's chat completions validator rejects @@ -1281,3 +1285,50 @@ class TestToolSchemaCombinatorFlatteningForOpenAI: parameters = request["tools"][0]["function"]["parameters"] assert "anyOf" not in parameters assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + + @staticmethod + def _artifact_tool(): + return { + "type": "function", + "function": { + "name": "Artifact", + "parameters": { + "type": "object", + "properties": {"field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}}, + "required": ["field"], + }, + }, + } + + def test_drops_non_python_regex_pattern_for_hosted_openai(self): + tool = self._artifact_tool() + + request = self._transform(self.config, "gpt-4o", {"custom_llm_provider": "openai", "api_base": None}, [tool]) + + assert request["tools"][0]["function"]["parameters"] == { + "type": "object", + "properties": {"field": {"type": "string"}}, + "required": ["field"], + } + assert tool == self._artifact_tool() + + def test_custom_api_base_drops_non_python_regex_pattern_but_keeps_union(self): + tool = self._anyof_tool() + tool["function"]["parameters"]["properties"]["id"]["pattern"] = _ARTIFACT_FIELD_PATTERN + + request = self._transform( + self.config, "gpt-4o", {"custom_llm_provider": "openai", "api_base": "http://localhost:8000/v1"}, [tool] + ) + + parameters = request["tools"][0]["function"]["parameters"] + assert parameters["properties"]["id"] == {"type": "string"} + assert parameters["anyOf"] == self._anyof_tool()["function"]["parameters"]["anyOf"] + + def test_non_openai_provider_keeps_non_python_regex_pattern(self): + tool = self._artifact_tool() + + request = self._transform( + self.config, "some-oss-model", {"custom_llm_provider": "groq", "api_base": None}, [tool] + ) + + assert request["tools"][0] is tool diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index c5902b32a06..4cf8767764b 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -1,5 +1,6 @@ import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx @@ -20,6 +21,8 @@ from litellm.types.llms.openai import ( ) from litellm.types.router import GenericLiteLLMParams +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + class TestOpenAIResponsesAPIConfig: def setup_method(self): @@ -2022,6 +2025,86 @@ class TestFlattenToolSchemaCombinatorsWiring: assert "anyOf" not in result["tools"][1]["parameters"] +class TestToolSchemaRegexPatternWiring: + """Claude Code's Artifact tool reaches /v1/responses (the /v1/messages bridge) with an + ECMA-262 ``pattern``; OpenAI compiles patterns with Python ``re`` and 400s + "'...' is not a 'regex'" for every model family, so the keyword is dropped. + """ + + def _artifact_tool(self): + return { + "type": "function", + "name": "Artifact", + "parameters": { + "type": "object", + "properties": { + "field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}, + "doc_id": {"type": "string", "pattern": r"^(?!\.\.?(?:/|$))[A-Za-z0-9_\-.~:@+]{1,200}$"}, + }, + "required": ["field"], + }, + } + + @pytest.mark.parametrize("model", ["gpt-5.6", "gpt-4o", "o3"]) + def test_openai_drops_only_the_pattern_python_re_rejects_for_every_family(self, model): + tool = self._artifact_tool() + + result = OpenAIResponsesAPIConfig().transform_responses_api_request( + model=model, + input="hi", + response_api_optional_request_params={"tools": [tool]}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + properties = result["tools"][0]["parameters"]["properties"] + assert properties["field"] == {"type": "string"} + assert properties["doc_id"] == tool["parameters"]["properties"]["doc_id"] + assert result["tools"][0]["parameters"]["required"] == ["field"] + assert tool["parameters"]["properties"]["field"]["pattern"] == _ARTIFACT_FIELD_PATTERN + assert json.loads(json.dumps(result["tools"])) == result["tools"] + + def test_openai_drops_patterns_inside_codex_namespace_tools(self): + namespace = {"type": "namespace", "name": "mcp__claude", "tools": [self._artifact_tool()]} + + result = OpenAIResponsesAPIConfig().transform_responses_api_request( + model="gpt-5.6", + input="hi", + response_api_optional_request_params={"tools": [namespace]}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["tools"][0]["tools"][0]["parameters"]["properties"]["field"] == {"type": "string"} + + def test_openai_compact_request_drops_patterns(self): + _, data = OpenAIResponsesAPIConfig().transform_compact_response_api_request( + model="gpt-5.6", + input="hi", + response_api_optional_request_params={"tools": [self._artifact_tool()]}, + api_base="https://api.openai.com/v1/responses", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert data["tools"][0]["parameters"]["properties"]["field"] == {"type": "string"} + + def test_non_openai_subclass_keeps_patterns(self): + from litellm.llms.hosted_vllm.responses.transformation import HostedVLLMResponsesAPIConfig + + tool = self._artifact_tool() + + result = HostedVLLMResponsesAPIConfig().transform_responses_api_request( + model="hosted_vllm/qwen", + input="hi", + response_api_optional_request_params={"tools": [tool]}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["tools"][0] is tool + + class TestReasoningFollowsModelSupport: """Responses API clients like Codex send `reasoning` on every request, and OpenAI 400s it on non-reasoning models like gpt-4o. drop_params must strip it there, the same way the From dde19adde171c3262388e589b6a3f8c4b352433c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 05:44:30 +0000 Subject: [PATCH 105/107] fix(proxy): bound concurrent key and spend-counter DB lookups to stop prisma pool thrash (#40387) * fix(proxy): bound concurrent key and spend-counter DB lookups to stop prisma pool thrash A cache-miss burst fanned every get_key_object DB fallback and SpendCounterReseed point lookup into the prisma query-engine httpx pool at once. httpcore's request assignment is O(queued x connections) per event, so the event loop spent most of its time in pool bookkeeping and the logging worker's 20s wait_for tripped. Callers now wait on a small shared semaphore (PROXY_DB_LOOKUP_MAX_CONCURRENCY, default 25) instead of queueing inside httpcore Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): type the in-flight counting prisma fake Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): drop module docstring from db_lookup_gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): type the in-flight counting table fake Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/auth/auth_checks.py | 62 ++++++++++--------- litellm/proxy/db/db_lookup_gate.py | 23 +++++++ litellm/proxy/db/spend_counter_reseed.py | 50 ++++++++------- .../proxy/auth/test_auth_checks.py | 41 +++++++++++- .../proxy/db/test_spend_counter_reseed.py | 31 ++++++++++ 6 files changed, 154 insertions(+), 54 deletions(-) create mode 100644 litellm/proxy/db/db_lookup_gate.py diff --git a/litellm/constants.py b/litellm/constants.py index d545890dd59..0a7ef363e92 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1682,6 +1682,7 @@ SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INT RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", "3"))) RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2")) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)) +PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, int(os.getenv("PROXY_DB_LOOKUP_MAX_CONCURRENCY", "25"))) DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500"))) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a71a1993064..1efc9611fe6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -94,6 +94,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( team_membership_auth_cache_key, team_membership_reservation_cache_key, ) +from litellm.proxy.db.db_lookup_gate import db_lookup_gate from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( TOOL_CAPABLE_CALL_TYPES, @@ -3450,36 +3451,37 @@ async def _fetch_key_object_from_db_with_reconnect( """ Fetch key object from DB and retry once if a DB connection error can be healed. """ - try: - return await prisma_client.get_data( - token=hashed_token, - table_name="combined_view", - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception as e: - if PrismaDBExceptionHandler.is_database_transport_error(e): - did_reconnect = False - if hasattr(prisma_client, "attempt_db_reconnect"): - auth_reconnect_timeout = getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", 2.0) - if not isinstance(auth_reconnect_timeout, (int, float)): - auth_reconnect_timeout = 2.0 - auth_reconnect_lock_timeout = getattr(prisma_client, "_db_auth_reconnect_lock_timeout_seconds", 0.1) - if not isinstance(auth_reconnect_lock_timeout, (int, float)): - auth_reconnect_lock_timeout = 0.1 - did_reconnect = await prisma_client.attempt_db_reconnect( - reason="auth_get_key_object_lookup_failure", - timeout_seconds=auth_reconnect_timeout, - lock_timeout_seconds=auth_reconnect_lock_timeout, - ) - if did_reconnect: - return await prisma_client.get_data( - token=hashed_token, - table_name="combined_view", - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - raise + async with db_lookup_gate.current(): + try: + return await prisma_client.get_data( + token=hashed_token, + table_name="combined_view", + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: + if PrismaDBExceptionHandler.is_database_transport_error(e): + did_reconnect = False + if hasattr(prisma_client, "attempt_db_reconnect"): + auth_reconnect_timeout = getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", 2.0) + if not isinstance(auth_reconnect_timeout, (int, float)): + auth_reconnect_timeout = 2.0 + auth_reconnect_lock_timeout = getattr(prisma_client, "_db_auth_reconnect_lock_timeout_seconds", 0.1) + if not isinstance(auth_reconnect_lock_timeout, (int, float)): + auth_reconnect_lock_timeout = 0.1 + did_reconnect = await prisma_client.attempt_db_reconnect( + reason="auth_get_key_object_lookup_failure", + timeout_seconds=auth_reconnect_timeout, + lock_timeout_seconds=auth_reconnect_lock_timeout, + ) + if did_reconnect: + return await prisma_client.get_data( + token=hashed_token, + table_name="combined_view", + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + raise def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str) -> str: diff --git a/litellm/proxy/db/db_lookup_gate.py b/litellm/proxy/db/db_lookup_gate.py new file mode 100644 index 00000000000..2fd427687bd --- /dev/null +++ b/litellm/proxy/db/db_lookup_gate.py @@ -0,0 +1,23 @@ +import asyncio +from typing import Final + +from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY + + +class LoopBoundSemaphore: + __slots__ = ("_loop", "_semaphore", "_value") + + def __init__(self, value: int) -> None: + self._value: Final = value + self._loop: asyncio.AbstractEventLoop | None = None + self._semaphore: asyncio.Semaphore | None = None + + def current(self) -> asyncio.Semaphore: + loop: Final = asyncio.get_running_loop() + if self._semaphore is None or self._loop is not loop: + self._semaphore = asyncio.Semaphore(self._value) + self._loop = loop + return self._semaphore + + +db_lookup_gate: Final = LoopBoundSemaphore(PROXY_DB_LOOKUP_MAX_CONCURRENCY) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index a38b8a47dbd..e35b1c8c82b 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -23,6 +23,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.db.db_lookup_gate import db_lookup_gate from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( BudgetWindowSpendRepository, @@ -121,30 +122,33 @@ class SpendCounterReseed: if SpendCounterReseed._is_key_or_team_window_counter(counter_key): return None try: - if counter_key.startswith("spend:key:"): - token: Final = counter_key[len("spend:key:") :] - row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token}) - elif counter_key.startswith("spend:team_member:"): - suffix: Final = counter_key[len("spend:team_member:") :] - if ":" not in suffix: + async with db_lookup_gate.current(): + if counter_key.startswith("spend:key:"): + token: Final = counter_key[len("spend:key:") :] + row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token}) + elif counter_key.startswith("spend:team_member:"): + suffix: Final = counter_key[len("spend:team_member:") :] + if ":" not in suffix: + return None + user_id, team_id = suffix.rsplit(":", 1) + row = await TeamMembershipRepository(prisma_client).table.find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} + ) + elif counter_key.startswith("spend:team:"): + team_id = counter_key[len("spend:team:") :] + row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + elif counter_key.startswith("spend:user:"): + user_id = counter_key[len("spend:user:") :] + row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + elif counter_key.startswith(END_USER_COUNTER_PREFIX) or counter_key.startswith("spend:tag:"): + return None + elif counter_key.startswith("spend:org:"): + org_id: Final = counter_key[len("spend:org:") :] + row = await OrganizationRepository(prisma_client).table.find_unique( + where={"organization_id": org_id} + ) + else: return None - user_id, team_id = suffix.rsplit(":", 1) - row = await TeamMembershipRepository(prisma_client).table.find_unique( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} - ) - elif counter_key.startswith("spend:team:"): - team_id = counter_key[len("spend:team:") :] - row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) - elif counter_key.startswith("spend:user:"): - user_id = counter_key[len("spend:user:") :] - row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) - elif counter_key.startswith(END_USER_COUNTER_PREFIX) or counter_key.startswith("spend:tag:"): - return None - elif counter_key.startswith("spend:org:"): - org_id: Final = counter_key[len("spend:org:") :] - row = await OrganizationRepository(prisma_client).table.find_unique(where={"organization_id": org_id}) - else: - return None except Exception: verbose_proxy_logger.exception("SpendCounterReseed.from_db: failed for %s", counter_key) return None diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5bfef2b6445..8777e24e209 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1,7 +1,7 @@ import asyncio import json from types import SimpleNamespace -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch if TYPE_CHECKING: @@ -37,6 +37,7 @@ from litellm.proxy.auth.auth_checks import ( _can_object_call_vector_stores, _check_end_user_budget, _check_team_member_budget, + _fetch_key_object_from_db_with_reconnect, _get_fuzzy_user_object, _get_team_db_check, _log_budget_lookup_failure, @@ -55,6 +56,7 @@ from litellm.caching.redis_cache import RedisCache from litellm.constants import ( DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + PROXY_DB_LOOKUP_MAX_CONCURRENCY, REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, ) @@ -564,6 +566,43 @@ async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_e assert mock_prisma_client.get_data.await_count == 1 +class _InFlightCountingPrisma: + def __init__(self) -> None: + self.in_flight = 0 + self.max_in_flight = 0 + + async def get_data( + self, token: str, table_name: str, parent_otel_span: None, proxy_logging_obj: None + ) -> UserAPIKeyAuth: + self.in_flight += 1 + self.max_in_flight = max(self.max_in_flight, self.in_flight) + await asyncio.sleep(0.001) + self.in_flight -= 1 + return UserAPIKeyAuth(token=token) + + +@pytest.mark.asyncio +async def test_fetch_key_object_from_db_bounds_in_flight_prisma_requests(): + prisma: Final = _InFlightCountingPrisma() + burst: Final = PROXY_DB_LOOKUP_MAX_CONCURRENCY * 5 + + results: Final = await asyncio.gather( + *( + _fetch_key_object_from_db_with_reconnect( + hashed_token=f"hashed-token-{i}", + prisma_client=prisma, # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient + parent_otel_span=None, + proxy_logging_obj=None, + ) + for i in range(burst) + ) + ) + + assert len(results) == burst + assert {r.token for r in results if r is not None} == {f"hashed-token-{i}" for i in range(burst)} + assert prisma.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY + + def _fake_redis_cache(): fake_redis = MagicMock() fake_redis.async_get_cache = AsyncMock(return_value=None) diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 8cb3fc665eb..0361d5cfe8f 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -7,6 +7,7 @@ allowed to run: only when the row is missing or belongs to an older window. from __future__ import annotations +import asyncio from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final @@ -14,6 +15,7 @@ from typing import Final import pytest from litellm.caching.dual_cache import DualCache +from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) @@ -42,6 +44,19 @@ class _FakeSpendLogsTable: return [{by[0]: where.get(by[0]), "_sum": {"spend": self._total}}] +class _InFlightCountingTable: + def __init__(self) -> None: + self.in_flight = 0 + self.max_in_flight = 0 + + async def find_unique(self, where: dict[str, str]) -> SimpleNamespace: + self.in_flight += 1 + self.max_in_flight = max(self.max_in_flight, self.in_flight) + await asyncio.sleep(0.001) + self.in_flight -= 1 + return SimpleNamespace(token=where["token"], spend=1.0) + + class _FakePrismaClient: def __init__( self, @@ -55,6 +70,7 @@ class _FakePrismaClient: litellm_budgetwindowspend=_FakeFindUniqueTable(row=row, error=error), litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total), litellm_endusertable=_FakeFindUniqueTable(row=end_user_row, error=end_user_error), + litellm_verificationtoken=_InFlightCountingTable(), ) @@ -306,6 +322,21 @@ async def test_end_user_from_db_returns_none_without_a_row_a_client_or_on_db_err ) +@pytest.mark.asyncio +async def test_from_db_bounds_in_flight_prisma_requests_across_counter_keys(): + """Per-counter singleflight only collapses duplicates of one key. A cold-cache burst + over many distinct keys must still not flood the prisma engine HTTP pool (LIT-6435).""" + prisma: Final = _FakePrismaClient() + burst: Final = PROXY_DB_LOOKUP_MAX_CONCURRENCY * 5 + + results: Final = await asyncio.gather( + *(SpendCounterReseed.from_db(prisma_client=prisma, counter_key=f"spend:key:hashed-{i}") for i in range(burst)) + ) + + assert results == [1.0] * burst + assert prisma.db.litellm_verificationtoken.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY + + @pytest.mark.asyncio async def test_from_db_still_never_reads_the_end_user_row(): """A cold end-user counter keeps seeding from the cached end-user object the auth From 6bb60f34e342ebac3bc01cab70820cf83acf9dc2 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 23:14:16 -0700 Subject: [PATCH 106/107] fix(s3_v2): freeze refreshable credentials before signing and retry 403 uploads with a fresh signature (#40187) RefreshableCredentials (IMDS roles) can refresh between the access key, secret and token reads SigV4 performs, producing a mixed-generation signature that S3 rejects with 403 and the log is dropped. Snapshot the credentials with get_frozen_credentials before signing, treat 403 like 500/503 in the upload retry loop, and fetch credentials plus sign again on every attempt in both the async and sync upload paths. Tests load a real botocore credential_process profile and fake only the HTTP boundary with httpx.MockTransport Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/s3_v2.py | 104 ++-- tests/test_litellm/integrations/test_s3_v2.py | 443 ++++++++++++------ 2 files changed, 348 insertions(+), 199 deletions(-) diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 996ff9c75af..972ac79e306 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -10,9 +10,11 @@ import asyncio import time from collections.abc import Mapping from datetime import datetime -from typing import Final, cast +from typing import TYPE_CHECKING, Final, cast from urllib.parse import quote +import httpx + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS @@ -35,6 +37,9 @@ from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload from .custom_batch_logger import CustomBatchLogger +if TYPE_CHECKING: + from botocore.credentials import Credentials + class S3Logger(CustomBatchLogger, BaseAWSLLM): def __init__( @@ -232,6 +237,26 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): f"{get_aws_dns_suffix(self.s3_region_name)}/{encoded_key}" ) + def _sign_put( + self, credentials: "Credentials", url: str, json_string: str, headers: Mapping[str, str] + ) -> dict[str, str]: # mutable-ok: [LIT001] AsyncHTTPHandler.put/HTTPHandler.put only accept dict headers + """ + ``RefreshableCredentials`` (IMDS roles) may refresh between the access key, secret and token + reads SigV4 performs, producing a mixed-generation signature that S3 rejects with 403. + Freezing first makes the three values one atomic snapshot. + """ + from botocore.auth import S3SigV4Auth + from botocore.awsrequest import AWSRequest + from botocore.credentials import RefreshableCredentials + + frozen: Final = ( + credentials.get_frozen_credentials() if isinstance(credentials, RefreshableCredentials) else credentials + ) + aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=dict(headers)) + aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) + S3SigV4Auth(frozen, "s3", aws_region_name).add_auth(aws_request) + return dict(aws_request.headers.items()) + def _sse_headers(self) -> Mapping[str, str]: candidates: Final = { "x-amz-server-side-encryption": self.s3_server_side_encryption, @@ -317,26 +342,12 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): try: import base64 import hashlib - - from botocore.auth import S3SigV4Auth - from botocore.awsrequest import AWSRequest except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") try: from litellm.litellm_core_utils.asyncify import asyncify asyncified_get_credentials: Final = asyncify(self.get_credentials) - credentials: Final = await asyncified_get_credentials( - aws_access_key_id=self.s3_aws_access_key_id, - aws_secret_access_key=self.s3_aws_secret_access_key, - aws_session_token=self.s3_aws_session_token, - aws_region_name=self.s3_region_name, - aws_session_name=self.s3_aws_session_name, - aws_profile_name=self.s3_aws_profile_name, - aws_role_name=self.s3_aws_role_name, - aws_web_identity_token=self.s3_aws_web_identity_token, - aws_sts_endpoint=self.s3_aws_sts_endpoint, - ) verbose_logger.debug("s3_v2 logger - uploading data to s3 - %s", batch_logging_element.s3_object_key) verbose_logger.debug("s3_v2 logger - s3_verify setting: %s", self.s3_verify) @@ -363,19 +374,28 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): **self._sse_headers(), } - # Sign the request - aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers) - aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) - await run_aws_signing(S3SigV4Auth(credentials, "s3", aws_region_name).add_auth, aws_request) + async def signed_put() -> httpx.Response: + credentials: Final = await asyncified_get_credentials( + aws_access_key_id=self.s3_aws_access_key_id, + aws_secret_access_key=self.s3_aws_secret_access_key, + aws_session_token=self.s3_aws_session_token, + aws_region_name=self.s3_region_name, + aws_session_name=self.s3_aws_session_name, + aws_profile_name=self.s3_aws_profile_name, + aws_role_name=self.s3_aws_role_name, + aws_web_identity_token=self.s3_aws_web_identity_token, + aws_sts_endpoint=self.s3_aws_sts_endpoint, + ) + signed_headers: Final = await run_aws_signing(self._sign_put, credentials, url, json_string, headers) + try: + return await self.async_httpx_client.put(url, data=json_string, headers=signed_headers) + except httpx.HTTPStatusError as error: + return error.response - # Prepare the signed headers - signed_headers: Final = dict(aws_request.headers.items()) - - # Make the request with retry for transient S3 errors (500/503) max_retries: Final = 3 for attempt in range(max_retries): - response = await self.async_httpx_client.put(url, data=json_string, headers=signed_headers) - if response.status_code in (500, 503) and attempt < max_retries - 1: + response = await signed_put() + if response.status_code in (403, 500, 503) and attempt < max_retries - 1: wait_time = 2**attempt # 1s, 2s verbose_logger.warning( "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s", @@ -479,20 +499,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): try: import base64 import hashlib - - from botocore.auth import S3SigV4Auth - from botocore.awsrequest import AWSRequest - from botocore.credentials import Credentials except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") try: verbose_logger.debug("s3_v2 logger - uploading data to s3 - %s", batch_logging_element.s3_object_key) - credentials: Final[Credentials] = self.get_credentials( - aws_access_key_id=self.s3_aws_access_key_id, - aws_secret_access_key=self.s3_aws_secret_access_key, - aws_session_token=self.s3_aws_session_token, - aws_region_name=self.s3_region_name, - ) url: Final = self._build_object_url(batch_logging_element.s3_object_key) @@ -516,22 +526,24 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): **self._sse_headers(), } - # Sign the request - aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers) - aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) - S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) - - # Prepare the signed headers - signed_headers: Final = dict(aws_request.headers.items()) - httpx_client: Final = _get_httpx_client( params=({"ssl_verify": self.s3_verify} if self.s3_verify is not None else None) ) - # Make the request with retry for transient S3 errors (500/503) + + def signed_put() -> httpx.Response: + credentials: Final = self.get_credentials( + aws_access_key_id=self.s3_aws_access_key_id, + aws_secret_access_key=self.s3_aws_secret_access_key, + aws_session_token=self.s3_aws_session_token, + aws_region_name=self.s3_region_name, + ) + signed_headers: Final = self._sign_put(credentials, url, json_string, headers) + return httpx_client.put(url, data=json_string, headers=signed_headers) + max_retries: Final = 3 for attempt in range(max_retries): - response = httpx_client.put(url, data=json_string, headers=signed_headers) - if response.status_code in (500, 503) and attempt < max_retries - 1: + response = signed_put() + if response.status_code in (403, 500, 503) and attempt < max_retries - 1: wait_time = 2**attempt # 1s, 2s verbose_logger.warning( "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s", diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index a037284d7c1..08d37297ab1 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1,10 +1,19 @@ import asyncio +import re +import sys +import textwrap +import uuid +from contextlib import asynccontextmanager from datetime import datetime -from unittest.mock import MagicMock, patch +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, call, patch +import httpx import pytest from litellm.integrations.s3_v2 import S3Logger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.integrations.s3_v2 import s3BatchLoggingElement from litellm.types.utils import StandardLoggingPayload @@ -21,9 +30,7 @@ class TestS3V2UnitTests: source_code = inspect.getsource(s3_v2) # Verify that json.dumps is not used directly in the code - assert ( - "json.dumps(" not in source_code - ), "S3 v2 should not use json.dumps directly" + assert "json.dumps(" not in source_code, "S3 v2 should not use json.dumps directly" @patch("asyncio.create_task") @patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") @@ -86,12 +93,8 @@ class TestS3V2UnitTests: call_args_minio = s3_logger_minio.async_httpx_client.put.call_args assert call_args_minio is not None url_minio = call_args_minio[0][0] - expected_minio_url = ( - "https://minio.example.com:9000/litellm-logs/2025-09-14/test-key.json" - ) - assert ( - url_minio == expected_minio_url - ), f"Expected MinIO URL {expected_minio_url}, got {url_minio}" + expected_minio_url = "https://minio.example.com:9000/litellm-logs/2025-09-14/test-key.json" + assert url_minio == expected_minio_url, f"Expected MinIO URL {expected_minio_url}, got {url_minio}" # Test 3: Custom endpoint without bucket name (should fall back to default) s3_logger_no_bucket = S3Logger( @@ -136,12 +139,8 @@ class TestS3V2UnitTests: call_args_sync = mock_sync_client.put.call_args assert call_args_sync is not None url_sync = call_args_sync[0][0] - expected_sync_url = ( - "https://custom.s3.endpoint.com/sync-bucket/2025-09-14/test-key.json" - ) - assert ( - url_sync == expected_sync_url - ), f"Expected sync URL {expected_sync_url}, got {url_sync}" + expected_sync_url = "https://custom.s3.endpoint.com/sync-bucket/2025-09-14/test-key.json" + assert url_sync == expected_sync_url, f"Expected sync URL {expected_sync_url}, got {url_sync}" # Test 5: Download method with custom endpoint s3_logger_download = S3Logger( @@ -158,19 +157,15 @@ class TestS3V2UnitTests: s3_logger_download.async_httpx_client = AsyncMock() s3_logger_download.async_httpx_client.get.return_value = mock_download_response - result = asyncio.run( - s3_logger_download._download_object_from_s3( - "2025-09-14/download-test-key.json" - ) - ) + result = asyncio.run(s3_logger_download._download_object_from_s3("2025-09-14/download-test-key.json")) call_args_download = s3_logger_download.async_httpx_client.get.call_args assert call_args_download is not None url_download = call_args_download[0][0] expected_download_url = "https://download.s3.endpoint.com/download-bucket/2025-09-14/download-test-key.json" - assert ( - url_download == expected_download_url - ), f"Expected download URL {expected_download_url}, got {url_download}" + assert url_download == expected_download_url, ( + f"Expected download URL {expected_download_url}, got {url_download}" + ) assert result == {"downloaded": "data"} @@ -216,12 +211,8 @@ class TestS3V2UnitTests: call_args = s3_logger_virtual.async_httpx_client.put.call_args assert call_args is not None url = call_args[0][0] - expected_url = ( - "https://test-bucket.s3.custom-endpoint.com/2025-09-14/test-key.json" - ) - assert ( - url == expected_url - ), f"Expected virtual-hosted-style URL {expected_url}, got {url}" + expected_url = "https://test-bucket.s3.custom-endpoint.com/2025-09-14/test-key.json" + assert url == expected_url, f"Expected virtual-hosted-style URL {expected_url}, got {url}" # Test 2: Path-style (default behavior with s3_use_virtual_hosted_style=False) s3_logger_path = S3Logger( @@ -241,12 +232,8 @@ class TestS3V2UnitTests: call_args_path = s3_logger_path.async_httpx_client.put.call_args assert call_args_path is not None url_path = call_args_path[0][0] - expected_path_url = ( - "https://s3.custom-endpoint.com/test-bucket/2025-09-14/test-key.json" - ) - assert ( - url_path == expected_path_url - ), f"Expected path-style URL {expected_path_url}, got {url_path}" + expected_path_url = "https://s3.custom-endpoint.com/test-bucket/2025-09-14/test-key.json" + assert url_path == expected_path_url, f"Expected path-style URL {expected_path_url}, got {url_path}" # Test 3: Virtual-hosted-style with http protocol s3_logger_http = S3Logger( @@ -266,12 +253,10 @@ class TestS3V2UnitTests: call_args_http = s3_logger_http.async_httpx_client.put.call_args assert call_args_http is not None url_http = call_args_http[0][0] - expected_http_url = ( - "http://http-bucket.minio.local:9000/2025-09-14/test-key.json" + expected_http_url = "http://http-bucket.minio.local:9000/2025-09-14/test-key.json" + assert url_http == expected_http_url, ( + f"Expected virtual-hosted-style URL with http {expected_http_url}, got {url_http}" ) - assert ( - url_http == expected_http_url - ), f"Expected virtual-hosted-style URL with http {expected_http_url}, got {url_http}" # Test 4: Sync upload method with virtual-hosted-style s3_logger_sync_virtual = S3Logger( @@ -295,12 +280,10 @@ class TestS3V2UnitTests: call_args_sync = mock_sync_client.put.call_args assert call_args_sync is not None url_sync = call_args_sync[0][0] - expected_sync_url = ( - "https://sync-bucket.storage.example.com/2025-09-14/test-key.json" + expected_sync_url = "https://sync-bucket.storage.example.com/2025-09-14/test-key.json" + assert url_sync == expected_sync_url, ( + f"Expected virtual-hosted-style sync URL {expected_sync_url}, got {url_sync}" ) - assert ( - url_sync == expected_sync_url - ), f"Expected virtual-hosted-style sync URL {expected_sync_url}, got {url_sync}" # Test 5: Download method with virtual-hosted-style s3_logger_download_virtual = S3Logger( @@ -316,34 +299,27 @@ class TestS3V2UnitTests: mock_download_response.status_code = 200 mock_download_response.json = MagicMock(return_value={"downloaded": "data"}) s3_logger_download_virtual.async_httpx_client = AsyncMock() - s3_logger_download_virtual.async_httpx_client.get.return_value = ( - mock_download_response - ) + s3_logger_download_virtual.async_httpx_client.get.return_value = mock_download_response - result = asyncio.run( - s3_logger_download_virtual._download_object_from_s3( - "2025-09-14/download-test-key.json" - ) - ) + result = asyncio.run(s3_logger_download_virtual._download_object_from_s3("2025-09-14/download-test-key.json")) call_args_download = s3_logger_download_virtual.async_httpx_client.get.call_args assert call_args_download is not None url_download = call_args_download[0][0] expected_download_url = "https://download-bucket.download.endpoint.com/2025-09-14/download-test-key.json" - assert ( - url_download == expected_download_url - ), f"Expected virtual-hosted-style download URL {expected_download_url}, got {url_download}" + assert url_download == expected_download_url, ( + f"Expected virtual-hosted-style download URL {expected_download_url}, got {url_download}" + ) assert result == {"downloaded": "data"} @patch("asyncio.create_task") @patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") - def test_s3_v2_put_url_encodes_spaces_in_object_key( - self, mock_periodic_flush, mock_create_task - ): - import requests + def test_s3_v2_put_url_encodes_spaces_in_object_key(self, mock_periodic_flush, mock_create_task): from unittest.mock import AsyncMock + import requests + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement mock_periodic_flush.return_value = None @@ -487,9 +463,7 @@ async def test_async_upload_exhausts_retries_on_persistent_503(): # All 3 attempts return 503 response_503 = MagicMock() response_503.status_code = 503 - response_503.raise_for_status = MagicMock( - side_effect=Exception("503 Service Unavailable") - ) + response_503.raise_for_status = MagicMock(side_effect=Exception("503 Service Unavailable")) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = AsyncMock(return_value=response_503) @@ -528,12 +502,12 @@ async def test_async_upload_no_retry_on_4xx(): s3_object_download_filename="test-no-retry.json", ) - response_403 = MagicMock() - response_403.status_code = 403 - response_403.raise_for_status = MagicMock(side_effect=Exception("403 Forbidden")) + response_400 = MagicMock() + response_400.status_code = 400 + response_400.raise_for_status = MagicMock(side_effect=Exception("400 Bad Request")) logger.async_httpx_client = AsyncMock() - logger.async_httpx_client.put = AsyncMock(return_value=response_403) + logger.async_httpx_client.put = AsyncMock(return_value=response_400) with patch.object(logger, "handle_callback_failure") as mock_failure: await logger.async_upload_data_to_s3(test_element) @@ -543,6 +517,190 @@ async def test_async_upload_no_retry_on_4xx(): mock_failure.assert_called_once_with(callback_name="S3Logger") +_SIGV4_ACCESS_KEY = re.compile(r"Credential=(AKIA\d+)/") + + +@pytest.fixture +def rotating_profile(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> str: + """ + A real botocore profile whose credential_process hands out a new key generation on every call and + expires inside the advisory refresh window, so RefreshableCredentials re-runs it on every property read. + """ + counter = tmp_path / "generation" + script = tmp_path / "rotate_credentials.py" + script.write_text( + textwrap.dedent( + f""" + import json, sys + from datetime import datetime, timedelta, timezone + from pathlib import Path + + counter = Path({str(counter)!r}) + generation = int(counter.read_text()) if counter.exists() else 0 + counter.write_text(str(generation + 1)) + expiry = (datetime.now(timezone.utc) + timedelta(minutes=12)).strftime("%Y-%m-%dT%H:%M:%SZ") + json.dump( + {{ + "Version": 1, + "AccessKeyId": f"AKIA{{generation}}", + "SecretAccessKey": f"secret-{{generation}}", + "SessionToken": f"token-{{generation}}", + "Expiration": expiry, + }}, + sys.stdout, + ) + """ + ) + ) + profile = f"rotating-{uuid.uuid4().hex}" + (tmp_path / "config").write_text(f"[profile {profile}]\ncredential_process = {sys.executable} {script}\n") + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) + return profile + + +def _generation(request: httpx.Request) -> tuple[str, str]: + """(access key generation, session token generation) SigV4 baked into one request.""" + access_key = _SIGV4_ACCESS_KEY.search(request.headers["Authorization"]) + assert access_key is not None + return access_key.group(1).removeprefix("AKIA"), request.headers["X-Amz-Security-Token"].removeprefix("token-") + + +@asynccontextmanager +async def _s3_logger_on_production_handler(profile: str, statuses: list[int]): + """ + S3Logger wired to the real AsyncHTTPHandler over an httpx MockTransport that answers with the given + statuses in order, so the handler's own raise_for_status behaviour is exercised end to end. + """ + requests: list[httpx.Request] = [] + replies = iter(statuses) + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(next(replies), request=request, text="SignatureDoesNotMatch") + + handler = AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_region_name="us-east-1", + s3_aws_profile_name=profile, + s3_flush_interval=3600, + ) + logger.async_httpx_client = handler + with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: + yield logger, requests, mock_sleep + await handler.client.aclose() + + +@pytest.mark.asyncio +async def test_async_upload_signs_with_one_frozen_credential_snapshot(rotating_profile: str, caplog): + """ + RefreshableCredentials refreshes on every property read once inside the advisory window, so signing + off the live object would mix the access key of one generation with the token of the next. + """ + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-frozen.json", + payload={"test": "frozen"}, + s3_object_download_filename="test-frozen.json", + ) + async with _s3_logger_on_production_handler(rotating_profile, [200]) as (logger, requests, _): + await logger.async_upload_data_to_s3(test_element) + + assert len(requests) == 1 + access_key_generation, token_generation = _generation(requests[0]) + assert access_key_generation == token_generation + assert "Error uploading to s3" not in caplog.text + + +@pytest.mark.asyncio +async def test_async_upload_retries_403_with_fresh_credentials_and_signature(rotating_profile: str, caplog): + """ + A 403 (SignatureDoesNotMatch after an IMDS rotation) must be retried, and the retry must fetch + credentials again and carry a signature computed from that newer generation. + """ + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-403.json", + payload={"test": "403"}, + s3_object_download_filename="test-403.json", + ) + async with _s3_logger_on_production_handler(rotating_profile, [403, 200]) as (logger, requests, mock_sleep): + await logger.async_upload_data_to_s3(test_element) + + assert len(requests) == 2 + first_key, first_token = _generation(requests[0]) + second_key, second_token = _generation(requests[1]) + assert first_key == first_token + assert second_key == second_token + assert int(second_key) > int(first_key) + assert requests[1].headers["Authorization"] != requests[0].headers["Authorization"] + mock_sleep.assert_awaited_once_with(1) + assert "Error uploading to s3" not in caplog.text + + +@pytest.mark.asyncio +async def test_async_upload_exhausts_403_retries_through_production_http_handler(rotating_profile: str, caplog): + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-403-exhausted.json", + payload={"test": "403-exhausted"}, + s3_object_download_filename="test-403-exhausted.json", + ) + async with _s3_logger_on_production_handler(rotating_profile, [403, 403, 403]) as (logger, requests, mock_sleep): + await logger.async_upload_data_to_s3(test_element) + + assert len(requests) == 3 + assert mock_sleep.await_args_list == [call(1), call(2)] + assert "Error uploading to s3" in caplog.text + + +@pytest.mark.asyncio +async def test_async_upload_does_not_retry_404_through_production_http_handler(rotating_profile: str, caplog): + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-404.json", + payload={"test": "404"}, + s3_object_download_filename="test-404.json", + ) + async with _s3_logger_on_production_handler(rotating_profile, [404]) as (logger, requests, mock_sleep): + await logger.async_upload_data_to_s3(test_element) + + assert len(requests) == 1 + mock_sleep.assert_not_awaited() + assert "Error uploading to s3" in caplog.text + + +def test_sync_upload_retries_403_with_fresh_signature(rotating_profile: str, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_PROFILE", rotating_profile) + logger = S3Logger(s3_bucket_name="test-bucket", s3_region_name="us-east-1", s3_flush_interval=3600) + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-sync-403.json", + payload={"test": "sync-403"}, + s3_object_download_filename="test-sync-403.json", + ) + requests: list[httpx.Request] = [] + replies = iter([403, 200]) + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(next(replies), request=request) + + handler = HTTPHandler() + handler.client = httpx.Client(transport=httpx.MockTransport(respond)) + with ( + patch( # test-quality-ok: sync upload builds its HTTPHandler per call, there is no injection seam for it + "litellm.integrations.s3_v2._get_httpx_client", return_value=handler + ), + patch("time.sleep") as mock_sleep, + ): + logger.upload_data_to_s3(test_element) + + assert len(requests) == 2 + first_key, first_token = _generation(requests[0]) + second_key, second_token = _generation(requests[1]) + assert first_key == first_token + assert second_key == second_token + assert int(second_key) > int(first_key) + mock_sleep.assert_called_once_with(1) + + def test_sync_upload_retries_on_s3_503(): """ Test that the sync upload_data_to_s3 retries on transient S3 503. @@ -626,9 +784,7 @@ async def test_async_log_event_skips_when_standard_logging_object_missing(): # Nothing should have been queued (catches the case where code falls # through without returning and appends None to the queue) - assert ( - len(logger.log_queue) == 0 - ), "log_queue should be empty when standard_logging_object is missing" + assert len(logger.log_queue) == 0, "log_queue should be empty when standard_logging_object is missing" @pytest.mark.asyncio @@ -767,20 +923,18 @@ async def test_s3_verify_false_handling(monkeypatch: pytest.MonkeyPatch): litellm, "s3_callback_params", { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, # This should NOT be ignored - "s3_use_ssl": False, # This should also NOT be ignored - }, + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, # This should NOT be ignored + "s3_use_ssl": False, # This should also NOT be ignored + }, ) with patch("asyncio.create_task"): - with patch( - "litellm.integrations.s3_v2.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.integrations.s3_v2.get_async_httpx_client") as mock_get_client: mock_client = AsyncMock() mock_get_client.return_value = mock_client @@ -788,22 +942,16 @@ async def test_s3_verify_false_handling(monkeypatch: pytest.MonkeyPatch): logger = S3Logger() # Verify s3_verify is False, not None - assert ( - logger.s3_verify is False - ), f"Expected s3_verify=False, got {logger.s3_verify}" - assert ( - logger.s3_use_ssl is False - ), f"Expected s3_use_ssl=False, got {logger.s3_use_ssl}" + assert logger.s3_verify is False, f"Expected s3_verify=False, got {logger.s3_verify}" + assert logger.s3_use_ssl is False, f"Expected s3_use_ssl=False, got {logger.s3_use_ssl}" # Verify that get_async_httpx_client was called with ssl_verify=False mock_get_client.assert_called_once() call_kwargs = mock_get_client.call_args.kwargs - assert ( - "params" in call_kwargs - ), "params should be passed to get_async_httpx_client" - assert call_kwargs["params"] == { - "ssl_verify": False - }, f"Expected ssl_verify=False in params, got {call_kwargs.get('params')}" + assert "params" in call_kwargs, "params should be passed to get_async_httpx_client" + assert call_kwargs["params"] == {"ssl_verify": False}, ( + f"Expected ssl_verify=False in params, got {call_kwargs.get('params')}" + ) @pytest.mark.asyncio @@ -820,17 +968,15 @@ async def test_s3_verify_none_handling(monkeypatch: pytest.MonkeyPatch): litellm, "s3_callback_params", { - "s3_bucket_name": "test-bucket", - "s3_aws_access_key_id": "test-key", - "s3_aws_secret_access_key": "test-secret", - "s3_region_name": "us-east-1", - }, + "s3_bucket_name": "test-bucket", + "s3_aws_access_key_id": "test-key", + "s3_aws_secret_access_key": "test-secret", + "s3_region_name": "us-east-1", + }, ) with patch("asyncio.create_task"): - with patch( - "litellm.integrations.s3_v2.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.integrations.s3_v2.get_async_httpx_client") as mock_get_client: mock_client = AsyncMock() mock_get_client.return_value = mock_client @@ -838,9 +984,7 @@ async def test_s3_verify_none_handling(monkeypatch: pytest.MonkeyPatch): logger = S3Logger() # Verify s3_verify is None (default) - assert ( - logger.s3_verify is None - ), f"Expected s3_verify=None, got {logger.s3_verify}" + assert logger.s3_verify is None, f"Expected s3_verify=None, got {logger.s3_verify}" # Verify that get_async_httpx_client was called mock_get_client.assert_called_once() @@ -868,13 +1012,13 @@ async def test_s3_verify_false_creates_httpx_client_with_verify_false(monkeypatc litellm, "s3_callback_params", { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, - }, + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, + }, ) with patch("asyncio.create_task"): @@ -890,9 +1034,7 @@ async def test_s3_verify_false_creates_httpx_client_with_verify_false(monkeypatc httpx_client = logger.async_httpx_client.client # Check the _verify attribute (httpx internal) if hasattr(httpx_client, "_verify"): - assert ( - httpx_client._verify is False - ), f"Expected httpx client _verify=False, got {httpx_client._verify}" + assert httpx_client._verify is False, f"Expected httpx client _verify=False, got {httpx_client._verify}" @pytest.mark.asyncio @@ -910,13 +1052,13 @@ async def test_s3_verify_false_async_client(monkeypatch: pytest.MonkeyPatch): litellm, "s3_callback_params", { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, - }, + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, + }, ) with patch("asyncio.create_task"): @@ -948,9 +1090,9 @@ async def test_s3_verify_false_async_client(monkeypatch: pytest.MonkeyPatch): if hasattr(logger.async_httpx_client, "client"): httpx_client = logger.async_httpx_client.client if hasattr(httpx_client, "_verify"): - assert ( - httpx_client._verify is False - ), f"Expected async httpx client _verify=False, got {httpx_client._verify}" + assert httpx_client._verify is False, ( + f"Expected async httpx client _verify=False, got {httpx_client._verify}" + ) @pytest.mark.asyncio @@ -1017,9 +1159,7 @@ def patch_asyncio_create_task(): (True, True, None, None, ""), ], ) -def test_s3_object_key_prefix_combinations( - use_team_prefix, use_key_prefix, team_alias, key_alias, expected_prefix -): +def test_s3_object_key_prefix_combinations(use_team_prefix, use_key_prefix, team_alias, key_alias, expected_prefix): """ Validate correct S3 prefix composition for team alias + key alias combinations. """ @@ -1490,9 +1630,7 @@ def test_s3_callback_params_override_does_not_mutate_inputs(monkeypatch): logger = S3Logger(s3_callback_params_override=override) assert logger.s3_bucket_name == "resolved-bucket" assert override["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" - assert ( - litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" - ) + assert litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" def test_s3_callback_params_override_none_falls_back_to_global(monkeypatch): @@ -1520,9 +1658,7 @@ def _expected_content_md5(payload: dict) -> str: from litellm.litellm_core_utils.safe_json_dumps import safe_dumps json_string = safe_dumps(payload) - return base64.b64encode( - hashlib.md5(json_string.encode("utf-8"), usedforsecurity=False).digest() - ).decode() + return base64.b64encode(hashlib.md5(json_string.encode("utf-8"), usedforsecurity=False).digest()).decode() def _require_non_security_md5(monkeypatch): @@ -1658,9 +1794,9 @@ def test_s3_server_side_encryption_read_from_callback_params(monkeypatch): litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "aws:kms" @@ -1789,10 +1925,10 @@ def test_s3_sse_kms_key_id_read_from_callback_params(monkeypatch): litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, ) logger = S3Logger() assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") @@ -1863,10 +1999,10 @@ def test_kms_key_id_dropped_when_algorithm_is_not_kms(monkeypatch): litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "AES256", - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "AES256", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "AES256" @@ -1884,10 +2020,10 @@ def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued(monkeypatch litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": True, - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": True, + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "aws:kms" @@ -1902,10 +2038,10 @@ def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(monkeypatch): litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - "s3_sse_kms_key_id": 12345, - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": 12345, + }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "aws:kms" @@ -2045,6 +2181,7 @@ async def test_download_signs_object_key_with_space_the_way_s3_does(): headers=call.kwargs["headers"], ) + _RESERVED_CHAR_KEYS = ( "2026-08-21/time-05-29-36_resp_bGl0ZWxsbTpjdXN0b20=.json", "session=logs/2026-08-21/time-05-29-36_abc.json", From e5da59336dacf9c69a07c1cb21dae2b9f4f828a4 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 10 Sep 2026 00:29:28 -0700 Subject: [PATCH 107/107] feat(prometheus): bucket latency by input sequence length (#40059) * feat(prometheus): bucket latency by input sequence length * style: format startup resolver call * fix(prometheus): handle unknown input lengths * fix(prometheus): preserve disabled custom input length labels * test(prometheus): seed startup snapshot in mocked logger test * fix(proxy): preserve database setting types during startup * fix(prometheus): distinguish missing usage and preserve config persistence Keep quoted database-storage config values intact for legacy persistence readers while using a local boolean for early callback discovery. Distinguish absent provider usage from an explicitly reported zero when labeling latency metrics. * fix(proxy): normalize input length flag from secret managers * fix(prometheus): isolate input buckets and preserve missing usage Keep built-in buckets on latency histograms, preserve unrelated custom labels, and classify raw incomplete usage and upstream total-only headers as unknown. Cover count conservation, failure callbacks, explicit zero, startup snapshots, and direct caller compatibility. Drop earlier branch budget changes. * fix(proxy): defer Prometheus alerting until stored settings load Reuse successful startup storage resolution and preserve callback deduplication across alerting reloads. * test(prometheus): restore input length flag between tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Revert "test(prometheus): restore input length flag between tests" This reverts commit f302e0b7bdad06e6c445bccdab0735aa8873ac39. * refactor(prometheus): make input length flag config/env only Drop the Admin UI General Settings row, the safe DB override entry, the startup reorder that loaded DB litellm_settings before Prometheus callbacks, and the alerting-only Prometheus path. The flag now behaves like prometheus_emit_stream_label: litellm_settings in config.yaml or an os.environ reference, applied on restart. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 1 + litellm/integrations/prometheus.py | 57 ++- litellm/types/integrations/prometheus.py | 40 +- .../test_prometheus_client_ip_user_agent.py | 1 + ..._prometheus_input_sequence_length_label.py | 428 ++++++++++++++++++ 5 files changed, 518 insertions(+), 9 deletions(-) create mode 100644 tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 5a461801b62..2a85a3d4200 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -475,6 +475,7 @@ prometheus_metrics_config: Optional[List] = None prometheus_exclude_metrics: Optional[List[str]] = None prometheus_exclude_labels: Optional[List[str]] = None prometheus_emit_stream_label: bool = False +prometheus_emit_input_sequence_length_label: bool = False prometheus_deployment_and_latency_caller_identity: Literal[ "api_key_alias", "user_email", diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 6766d246894..540ce6738fc 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -246,6 +246,7 @@ class PrometheusLogger(CustomLogger): # logger so toggling these flags only takes effect after a # restart, keeping init-time and runtime label sets in sync. self._cached_metric_labels: dict[str, list[str]] = {} + self._emit_input_sequence_length_label = litellm.prometheus_emit_input_sequence_length_label is True _custom_buckets: Final = litellm.prometheus_latency_buckets self.latency_buckets = tuple(_custom_buckets) if _custom_buckets is not None else LATENCY_BUCKETS @@ -1522,6 +1523,11 @@ class PrometheusLogger(CustomLogger): # 2. Pyright does not allow us to run isinstance(standard_logging_payload, StandardLoggingPayload) <- this would be ideal enum_values=enum_values, label_context=label_context, + input_sequence_length=( + self._get_input_sequence_length(standard_logging_payload, kwargs, response_obj) + if self._emit_input_sequence_length_label + else None + ), ) # set x-ratelimit headers @@ -2192,6 +2198,36 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_api_key_tokens_for_model.labels(**tokens_labels).set(remaining_tokens) + @staticmethod + def _get_input_sequence_length( + standard_logging_payload: StandardLoggingPayload, + kwargs: Mapping[str, object], + response_obj: object, + ) -> str: + prompt_tokens: Final = standard_logging_payload.get("prompt_tokens") + if prompt_tokens: + return get_input_sequence_length_bucket(prompt_tokens) + combined_usage: Final = kwargs.get("combined_usage_object") + if ( + combined_usage is not None + and getattr(kwargs.get("_litellm_upstream_reported_usage"), "total_tokens", None) is not None + ): + return get_input_sequence_length_bucket(None) + reported_usage: Final = ( + response_obj.get("usage") if isinstance(response_obj, dict) else getattr(response_obj, "usage", None) + ) + if reported_usage is None and combined_usage is None: + return get_input_sequence_length_bucket(None) + usage_metadata: Final = standard_logging_payload["metadata"].get("usage_object") + if isinstance(usage_metadata, Mapping): + return get_input_sequence_length_bucket(usage_metadata.get("prompt_tokens")) + if combined_usage is None and isinstance(response_obj, dict): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + normalized_usage: Final[Mapping[str, object]] = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj) + return get_input_sequence_length_bucket(normalized_usage.get("prompt_tokens")) + return get_input_sequence_length_bucket(prompt_tokens) + def _set_latency_metrics( self, kwargs: dict, @@ -2202,7 +2238,16 @@ class PrometheusLogger(CustomLogger): user_api_team_alias: str | None, enum_values: UserAPIKeyLabelValues, label_context: PrometheusLabelFactoryContext | None = None, + input_sequence_length: str | None = None, ): + latency_enum_values: Final = ( + replace(enum_values, input_sequence_length=input_sequence_length) + if input_sequence_length is not None + else enum_values + ) + latency_label_context: Final = ( + PrometheusLabelFactoryContext(latency_enum_values) if input_sequence_length is not None else label_context + ) # latency metrics end_time: Final[datetime] = kwargs.get("end_time") or datetime.now() start_time: Final[datetime | None] = kwargs.get("start_time") @@ -2220,8 +2265,8 @@ class PrometheusLogger(CustomLogger): supported_enum_labels=self.get_labels_for_metric( metric_name="litellm_llm_api_time_to_first_token_metric" ), - enum_values=enum_values, - label_context=label_context, + enum_values=latency_enum_values, + label_context=latency_label_context, ) self.litellm_llm_api_time_to_first_token_metric.labels(**_ttft_labels).observe(time_to_first_token_seconds) self._track_end_user_metric_series( @@ -2241,8 +2286,8 @@ class PrometheusLogger(CustomLogger): if api_call_total_time_seconds is not None: _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_llm_api_latency_metric"), - enum_values=enum_values, - label_context=label_context, + enum_values=latency_enum_values, + label_context=latency_label_context, ) self.litellm_llm_api_latency_metric.labels(**_labels).observe(api_call_total_time_seconds) self._track_end_user_metric_series( @@ -2272,8 +2317,8 @@ class PrometheusLogger(CustomLogger): ) _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_request_total_latency_metric"), - enum_values=enum_values, - label_context=label_context, + enum_values=latency_enum_values, + label_context=latency_label_context, ) self.litellm_request_total_latency_metric.labels(**_labels).observe(_observed_total_time_seconds) self._track_end_user_metric_series( diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 8498b6f6d00..a024581f600 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -154,6 +154,22 @@ LATENCY_BUCKETS: Final = ( float("inf"), ) +UNKNOWN_INPUT_SEQUENCE_LENGTH: Final = "unknown" +INPUT_SEQUENCE_LENGTH_BUCKETS: Final = ( + (1_000, "0-1k"), + (4_000, "1k-4k"), + (16_000, "4k-16k"), + (64_000, "16k-64k"), + (float("inf"), "64k+"), +) + + +def get_input_sequence_length_bucket(prompt_tokens: object) -> str: + if not isinstance(prompt_tokens, int) or isinstance(prompt_tokens, bool) or prompt_tokens < 0: + return UNKNOWN_INPUT_SEQUENCE_LENGTH + return next(label for upper, label in INPUT_SEQUENCE_LENGTH_BUCKETS if prompt_tokens < upper) + + # Batch jobs can run for minutes to hours; buckets span 1 min → 24 h. BATCH_DURATION_BUCKETS: Final = ( 60.0, @@ -205,6 +221,7 @@ class UserAPIKeyLabelNames(Enum): MCP_TOOL_NAME = "mcp_tool_name" MCP_SERVER_NAME = "mcp_server_name" SERVICE_TIER = "service_tier" + INPUT_SEQUENCE_LENGTH = "input_sequence_length" DEFINED_PROMETHEUS_METRICS = Literal[ @@ -857,6 +874,13 @@ class PrometheusMetricLabels: "litellm_images_generated_metric", } ) + _input_sequence_length_metrics: ClassVar[frozenset[str]] = frozenset( + { + "litellm_llm_api_latency_metric", + "litellm_llm_api_time_to_first_token_metric", + "litellm_request_total_latency_metric", + } + ) # Managed batch metrics _batch_user_labels = [ UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, @@ -955,14 +979,23 @@ class PrometheusMetricLabels: custom_labels.append(label) if label_name in PrometheusMetricLabels._org_label_metrics: - for label in [ + for label in ( UserAPIKeyLabelNames.ORG_ID.value, UserAPIKeyLabelNames.ORG_ALIAS.value, - ]: + ): if label not in default_labels and label not in custom_labels: custom_labels.append(label) - return default_labels + custom_labels + input_sequence_length_labels: Final = ( + (UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value,) + if ( + label_name in PrometheusMetricLabels._input_sequence_length_metrics + and litellm.prometheus_emit_input_sequence_length_label is True + and UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value not in custom_labels + ) + else () + ) + return [*default_labels, *custom_labels, *input_sequence_length_labels] _USER_API_KEY_LABEL_VALUE_INIT_ALIASES: Final[Mapping[str, str]] = MappingProxyType( @@ -1015,6 +1048,7 @@ class UserAPIKeyLabelValues: mcp_tool_name: str | None = None mcp_server_name: str | None = None service_tier: str | None = None + input_sequence_length: str | None = None # Added for test compatibility. def __init__(self, **kwargs: Any) -> None: diff --git a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py index ea661d2ea78..004ac4dbffb 100644 --- a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py +++ b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py @@ -84,6 +84,7 @@ async def test_async_post_call_success_hook_includes_client_ip_user_agent(): "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None ): logger = PrometheusLogger() + logger._emit_input_sequence_length_label = False logger.litellm_proxy_total_requests_metric = MagicMock() logger.get_labels_for_metric = MagicMock( return_value=["client_ip", "user_agent"] diff --git a/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py b/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py new file mode 100644 index 00000000000..bc922061544 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py @@ -0,0 +1,428 @@ +import asyncio +import datetime +from collections.abc import Mapping +from copy import deepcopy +from typing import Final, cast + +import pytest +from prometheus_client import REGISTRY +from prometheus_client.samples import Sample + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import ( + PrometheusMetricLabels, + UserAPIKeyLabelNames, + UserAPIKeyLabelValues, + get_input_sequence_length_bucket, +) +from litellm.types.utils import StandardLoggingPayload + +LATENCY_METRICS: Final = ( + "litellm_llm_api_latency_metric", + "litellm_llm_api_time_to_first_token_metric", + "litellm_request_total_latency_metric", +) +FLAG: Final = "prometheus_emit_input_sequence_length_label" + + +def _clear_prometheus_registry() -> None: + for collector in tuple(REGISTRY._collector_to_names): # pyright: ignore[reportPrivateUsage] # test registry reset + REGISTRY.unregister(collector) + + +@pytest.fixture(autouse=True) +def isolated_registry(monkeypatch: pytest.MonkeyPatch): + _clear_prometheus_registry() + monkeypatch.setattr(litellm, FLAG, False) + yield + _clear_prometheus_registry() + + +@pytest.mark.parametrize("metric", LATENCY_METRICS) +def test_input_sequence_length_label_is_opt_in(monkeypatch: pytest.MonkeyPatch, metric: str): + assert UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value not in PrometheusMetricLabels.get_labels(metric) + + monkeypatch.setattr(litellm, FLAG, True) + assert UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value in PrometheusMetricLabels.get_labels(metric) + + +def test_input_sequence_length_label_stays_off_non_latency_metrics(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, FLAG, True) + assert UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value not in PrometheusMetricLabels.get_labels( + "litellm_proxy_total_requests_metric" + ) + + +@pytest.mark.parametrize( + "prompt_tokens, expected", + [ + (None, "unknown"), + (0, "0-1k"), + (999, "0-1k"), + (1_000, "1k-4k"), + (3_999, "1k-4k"), + (4_000, "4k-16k"), + (15_999, "4k-16k"), + (16_000, "16k-64k"), + (63_999, "16k-64k"), + (64_000, "64k+"), + (10_000_000, "64k+"), + (-1, "unknown"), + ], +) +def test_input_sequence_length_bucket_boundaries(prompt_tokens: int | None, expected: str): + assert get_input_sequence_length_bucket(prompt_tokens) == expected + + +def test_user_api_key_label_values_carries_input_sequence_length(): + values: Final = UserAPIKeyLabelValues(input_sequence_length="4k-16k") + + assert values.input_sequence_length == "4k-16k" + assert values.model_dump()["input_sequence_length"] == "4k-16k" + + +def _assert_latency_metrics(expected: str | None, stream: bool = True, queue_time: float = 0) -> None: + samples: Final = tuple(sample for metric in REGISTRY.collect() for sample in metric.samples) + for metric, duration in zip(LATENCY_METRICS, (2, 1, 3 + queue_time)): + counts: Final = tuple(sample for sample in samples if sample.name == f"{metric}_count") + sums: Final = tuple(sample for sample in samples if sample.name == f"{metric}_sum") + buckets: Final = tuple(sample for sample in samples if sample.name == f"{metric}_bucket") + if not stream and metric == "litellm_llm_api_time_to_first_token_metric": + assert not counts and not sums and not buckets + continue + assert len(counts) == len(sums) == 1 + assert counts[0].value == 1 + assert sums[0].value == pytest.approx(duration) + assert buckets and any(sample.labels["le"] == "+Inf" for sample in buckets) + assert all(sample.value == int(float(sample.labels["le"]) >= duration) for sample in buckets) + assert all(sample.labels.get("input_sequence_length") == expected for sample in (*counts, *sums, *buckets)) + + +def _non_target_samples() -> tuple[Sample, ...]: + return tuple( + sample + for metric in REGISTRY.collect() + if metric.name not in LATENCY_METRICS + for sample in metric.samples + if "input_sequence_length" in sample.labels and not sample.name.endswith("_created") + ) + + +def _standard_logging_payload(now: datetime.datetime, prompt_tokens: int) -> StandardLoggingPayload: + return cast( + StandardLoggingPayload, + { + "id": "t", + "call_type": "completion", + "response_cost": 0.001, + "status": "success", + "total_tokens": prompt_tokens + 20, + "prompt_tokens": prompt_tokens, + "completion_tokens": 20, + "startTime": now - datetime.timedelta(seconds=3), + "endTime": now, + "completionStartTime": now - datetime.timedelta(seconds=1), + "model": "gpt-4o-mini", + "model_id": "model-123", + "model_group": "gpt-4o-mini", + "api_base": "https://api.openai.com", + "custom_llm_provider": "openai", + "request_tags": [], + "stream": True, + "metadata": { + "user_api_key_hash": "h", + "user_api_key_alias": "a", + "user_api_key_team_id": "t", + "user_api_key_team_alias": "ta", + "user_api_key_user_id": "u", + "user_api_key_user_email": "e@x.com", + "user_api_key_org_id": None, + "user_api_key_org_alias": None, + "requester_metadata": None, + "user_api_key_end_user_id": None, + "usage_object": None, + }, + "hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None}, + }, + ) + + +def _success_kwargs( + now: datetime.datetime, prompt_tokens: int, requester_metadata: Mapping[str, object] | None = None +) -> Mapping[str, object]: + payload: Final = _standard_logging_payload(now, prompt_tokens) + return { + "model": "gpt-4o-mini", + "litellm_params": {"metadata": {}}, + "standard_logging_object": { + **payload, + "metadata": {**payload["metadata"], "requester_metadata": requester_metadata}, + }, + "stream": True, + "start_time": now - datetime.timedelta(seconds=3), + "api_call_start_time": now - datetime.timedelta(seconds=2), + "completion_start_time": now - datetime.timedelta(seconds=1), + "end_time": now, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("flag_at_request_time", (True, False)) +async def test_logger_emits_bucket_from_its_startup_label_set( + monkeypatch: pytest.MonkeyPatch, flag_at_request_time: bool +): + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, True) + logger: Final = PrometheusLogger() + monkeypatch.setattr(litellm, FLAG, flag_at_request_time) + + await logger.async_log_success_event(dict(_success_kwargs(now, prompt_tokens=4_000)), None, now, now) + + _assert_latency_metrics("4k-16k") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "combined_usage", "expected"), + ( + ({"id": "moderation", "results": []}, None, "unknown"), + ({"usage": None}, None, "unknown"), + ({"usage": {}}, None, "unknown"), + ({"usage": {"completion_tokens": 3}}, None, "unknown"), + ({"usage": {"total_tokens": 5}}, None, "unknown"), + ({"usage": {"prompt_tokens": 0}}, None, "0-1k"), + ({"usage": {"prompt_tokens": 4_000}}, None, "4k-16k"), + ({"usage": {"input_tokens": 0, "output_tokens": 3, "total_tokens": 3}}, None, "0-1k"), + ({"usage": {"input_tokens": 4_000, "output_tokens": 3, "total_tokens": 4_003}}, None, "4k-16k"), + (litellm.ModelResponse(usage=litellm.Usage(prompt_tokens=0)), None, "0-1k"), + (None, litellm.Usage(prompt_tokens=0), "0-1k"), + ), +) +@pytest.mark.parametrize("include_usage_metadata", (True, False)) +async def test_logger_distinguishes_missing_usage_from_reported_zero( + monkeypatch: pytest.MonkeyPatch, + response: object, + combined_usage: object, + expected: str, + include_usage_metadata: bool, +): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, True) + logger: Final = PrometheusLogger() + usage: Final = StandardLoggingPayloadSetup.get_usage_as_dict( + response_obj=response if isinstance(response, dict) else None + ) + + payload: Final = _standard_logging_payload(now, usage.get("prompt_tokens", 0)) + await logger.async_log_success_event( + { + **_success_kwargs(now, prompt_tokens=usage.get("prompt_tokens", 0)), + "combined_usage_object": combined_usage, + "standard_logging_object": { + **payload, + "metadata": {**payload["metadata"], "usage_object": usage if include_usage_metadata else None}, + }, + }, + response, + now, + now, + ) + + _assert_latency_metrics(expected) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("total_tokens", (None, 0, 5_000)) +@pytest.mark.parametrize("prompt_tokens", (0, 4_000)) +async def test_upstream_total_only_usage_has_unknown_input_length( + monkeypatch: pytest.MonkeyPatch, total_tokens: int | None, prompt_tokens: int +): + import httpx + + from litellm.litellm_core_utils.litellm_logging import Logging, StandardLoggingPayloadSetup + from litellm.proxy.pass_through_endpoints.upstream_usage_headers import apply_upstream_reported_usage + + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, True) + logger: Final = PrometheusLogger() + logging_obj: Final = Logging( + model="gpt-4o-mini", + messages=[], + stream=True, + call_type="pass_through_endpoint", + start_time=now, + litellm_call_id="test-call-id", + function_id="1", + ) + headers: Final = httpx.Headers( + { + "x-litellm-response-cost": "0.001", + **({"x-litellm-total-tokens": str(total_tokens)} if total_tokens is not None else {}), + } + ) + reported: Final = apply_upstream_reported_usage(logging_obj=logging_obj, headers=headers) + assert reported is not None + combined_usage: Final = logging_obj.model_call_details.get("combined_usage_object") + response: Final = {"usage": {"prompt_tokens": prompt_tokens}} + usage: Final = StandardLoggingPayloadSetup.get_usage_as_dict(response, combined_usage) + payload: Final = _standard_logging_payload(now, usage.get("prompt_tokens", 0)) + + await logger.async_log_success_event( + { + **logging_obj.model_call_details, + **_success_kwargs(now, usage.get("prompt_tokens", 0)), + "standard_logging_object": {**payload, "metadata": {**payload["metadata"], "usage_object": usage}}, + }, + response, + now, + now, + ) + + _assert_latency_metrics("unknown" if total_tokens is not None else get_input_sequence_length_bucket(prompt_tokens)) + + +@pytest.mark.asyncio +async def test_logger_built_with_flag_off_emits_no_bucket_label(monkeypatch: pytest.MonkeyPatch): + now: Final = datetime.datetime.now() + logger: Final = PrometheusLogger() + monkeypatch.setattr(litellm, FLAG, True) + + await logger.async_log_success_event(dict(_success_kwargs(now, prompt_tokens=4_000)), None, now, now) + + _assert_latency_metrics(None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("flag_at_startup", (True, False)) +@pytest.mark.parametrize("stream", (True, False)) +@pytest.mark.parametrize( + "metadata", + ( + None, + {}, + {"input_sequence_length": None}, + {"input_sequence_length": False}, + {"input_sequence_length": True}, + {"input_sequence_length": 0}, + {"input_sequence_length": []}, + {"input_sequence_length": {}}, + {"input_sequence_length": ""}, + {"input_sequence_length": "from-metadata"}, + ), +) +async def test_custom_input_length_label_is_scoped_to_target_histograms( + monkeypatch: pytest.MonkeyPatch, flag_at_startup: bool, stream: bool, metadata: Mapping[str, object] | None +): + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, "custom_prometheus_metadata_labels", ["input_sequence_length"]) + kwargs: Final = { + **_success_kwargs(now, prompt_tokens=4_000, requester_metadata=metadata), + "stream": stream, + "litellm_params": {"metadata": {"queue_time_seconds": 0.25}}, + } + original_kwargs: Final = deepcopy(kwargs) + baseline_logger: Final = PrometheusLogger() + await baseline_logger.async_log_success_event(kwargs, None, now, now) + baseline_samples: Final = _non_target_samples() + _clear_prometheus_registry() + monkeypatch.setattr(litellm, FLAG, flag_at_startup) + logger: Final = PrometheusLogger() + monkeypatch.setattr(litellm, FLAG, not flag_at_startup) + + await logger.async_log_success_event(kwargs, None, now, now) + + assert kwargs == original_kwargs + custom_value: Final = (metadata or {}).get("input_sequence_length") + expected: Final = custom_value if isinstance(custom_value, str) else ("4k-16k" if flag_at_startup else "None") + _assert_latency_metrics(expected, stream=stream, queue_time=0.25) + assert all(logger.get_labels_for_metric(metric).count("input_sequence_length") == 1 for metric in LATENCY_METRICS) + non_target_samples: Final = _non_target_samples() + assert { + "litellm_requests_metric_total", + "litellm_spend_metric_total", + "litellm_total_tokens_metric_total", + "litellm_request_queue_time_seconds_count", + "litellm_deployment_success_responses_total", + }.issubset({sample.name for sample in non_target_samples}) + assert non_target_samples == baseline_samples + queue_sum: Final = tuple( + sample for sample in non_target_samples if sample.name == "litellm_request_queue_time_seconds_sum" + ) + assert len(queue_sum) == 1 and queue_sum[0].value == 0.25 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", (True, False)) +async def test_concurrent_requests_keep_independent_buckets(monkeypatch: pytest.MonkeyPatch, enabled: bool): + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, enabled) + logger: Final = PrometheusLogger() + cases: Final = ( + (None, "unknown"), + (0, "0-1k"), + (1_000, "1k-4k"), + (4_000, "4k-16k"), + (16_000, "16k-64k"), + (64_000, "64k+"), + ) + calls: Final = tuple( + ( + {**_success_kwargs(now, prompt_tokens=tokens or 0), "stream": stream}, + {"usage": {"prompt_tokens": tokens}} if tokens is not None else None, + ) + for tokens, _ in cases + for stream in (True, False) + for _ in range(2) + ) + original_calls: Final = deepcopy(calls) + + await asyncio.gather(*(logger.async_log_success_event(kwargs, response, now, now) for kwargs, response in calls)) + + assert calls == original_calls + samples: Final = tuple(sample for metric in REGISTRY.collect() for sample in metric.samples) + for metric, duration in zip(LATENCY_METRICS, (2, 1, 3)): + expected_count: Final = 2 if metric == "litellm_llm_api_time_to_first_token_metric" else 4 + counts: Final = tuple(sample for sample in samples if sample.name == f"{metric}_count") + sums: Final = tuple(sample for sample in samples if sample.name == f"{metric}_sum") + buckets: Final = tuple(sample for sample in samples if sample.name == f"{metric}_bucket") + expected: Final = ( + {bucket: expected_count for _, bucket in cases} if enabled else {None: expected_count * len(cases)} + ) + assert len(counts) == len(sums) == len(expected) + assert {sample.labels.get("input_sequence_length"): sample.value for sample in counts} == expected + assert {sample.labels.get("input_sequence_length"): sample.value for sample in sums} == { + bucket: count * duration for bucket, count in expected.items() + } + assert sum(sample.value for sample in buckets if sample.labels["le"] == "+Inf") == expected_count * len(cases) + assert all( + sample.value + == expected[sample.labels.get("input_sequence_length")] * int(float(sample.labels["le"]) >= duration) + for sample in buckets + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", (True, False)) +async def test_failed_request_does_not_observe_latency(monkeypatch: pytest.MonkeyPatch, enabled: bool): + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, enabled) + monkeypatch.setattr(litellm, "custom_prometheus_metadata_labels", ["input_sequence_length"]) + logger: Final = PrometheusLogger() + kwargs: Final = { + **_success_kwargs(now, prompt_tokens=4_000), + "standard_logging_object": {**_standard_logging_payload(now, 4_000), "status": "failure"}, + "exception": RuntimeError("upstream request failed"), + } + + await logger.async_log_failure_event(kwargs, None, now, now) + + samples: Final = tuple(sample for metric in REGISTRY.collect() for sample in metric.samples) + assert not any(sample.name.startswith(LATENCY_METRICS) for sample in samples) + for metric in ("litellm_llm_api_failed_requests_metric_total", "litellm_deployment_failure_responses_total"): + counts: Final = tuple(sample for sample in samples if sample.name == metric) + assert len(counts) == 1 + assert counts[0].value == 1 + assert counts[0].labels["input_sequence_length"] == "None"