From 676f841534e7c83bcf5d9afb65f5c37bf741af44 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 03:47:38 -0700 Subject: [PATCH 01/32] feat(mistral): add text-to-speech support for /v1/audio/speech --- .../mistral/audio_speech/transformation.py | 209 ++++++++++++++++++ litellm/main.py | 28 +++ ...odel_prices_and_context_window_backup.json | 4 +- litellm/router.py | 4 +- litellm/utils.py | 6 + model_prices_and_context_window.json | 4 +- ...est_mistral_audio_speech_transformation.py | 198 +++++++++++++++++ tests/test_litellm/test_cost_calculator.py | 12 + tests/test_litellm/test_main.py | 28 +++ tests/test_litellm/test_router.py | 26 +++ 10 files changed, 513 insertions(+), 6 deletions(-) create mode 100644 litellm/llms/mistral/audio_speech/transformation.py create mode 100644 tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py diff --git a/litellm/llms/mistral/audio_speech/transformation.py b/litellm/llms/mistral/audio_speech/transformation.py new file mode 100644 index 00000000000..6d1a693268a --- /dev/null +++ b/litellm/llms/mistral/audio_speech/transformation.py @@ -0,0 +1,209 @@ +""" +Support for Mistral Voxtral text-to-speech via ``/v1/audio/speech``. + +API reference: https://docs.mistral.ai/api/#tag/audio/operation/audio_speech_v1_audio_speech_post +""" + +import base64 +import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +import httpx + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.text_to_speech.transformation import ( + BaseTextToSpeechConfig, + TextToSpeechRequestData, +) +from litellm.secret_managers.main import get_secret_str + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.llms.openai import HttpxBinaryResponseContent + + +class MistralTextToSpeechException(BaseLLMException): + pass + + +class MistralTextToSpeechConfig(BaseTextToSpeechConfig): + TTS_BASE_URL: Final[str] = "https://api.mistral.ai/v1" + AUDIO_CONTENT_TYPES: Final[MappingProxyType[str, str]] = MappingProxyType( + { + "mp3": "audio/mpeg", + "wav": "audio/wav", + "pcm": "audio/pcm", + "flac": "audio/flac", + "opus": "audio/ogg", + } + ) + DROPPED_RESPONSE_HEADERS: Final[frozenset[str]] = frozenset( + {"content-encoding", "transfer-encoding", "content-length", "content-type"} + ) + OPENAI_VOICE_ALIASES: Final[MappingProxyType[str, str]] = MappingProxyType( + { + "alloy": "en_paul_neutral", + "echo": "gb_oliver_neutral", + "fable": "en_paul_cheerful", + "onyx": "en_paul_confident", + "nova": "gb_jane_sarcasm", + "shimmer": "gb_jane_sarcasm", + } + ) + + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class contract returns a plain list + return ["voice", "response_format"] # mutable-ok: base class contract returns a plain list + + def _map_openai_voice(self, voice_id: str) -> str: + return self.OPENAI_VOICE_ALIASES.get(voice_id.lower(), voice_id) + + def _resolve_voice_id(self, voice: object) -> str | None: + if isinstance(voice, str) and voice.strip(): + return self._map_openai_voice(voice.strip()) + if isinstance(voice, Mapping): + candidates: Final = (voice.get(key) for key in ("voice_id", "id", "name")) + resolved: Final = next( + (candidate.strip() for candidate in candidates if isinstance(candidate, str) and candidate.strip()), + None, + ) + return self._map_openai_voice(resolved) if resolved else None + return None + + def map_openai_params( + self, + model: str, + optional_params: Mapping[str, object], + voice: object = None, + drop_params: bool = False, + kwargs: Mapping[str, object] | None = None, + ) -> tuple[str | None, dict]: # mutable-ok: base class contract returns a plain dict + response_format: Final = optional_params.get("response_format") + ref_audio: Final = kwargs.get("ref_audio") if kwargs else None + voice_id_kwarg: Final = kwargs.get("voice_id") if kwargs else None + mapped_voice: Final = self._resolve_voice_id(voice) or self._resolve_voice_id(voice_id_kwarg) + mapped_params: Final = { # mutable-ok: base class contract returns a plain dict + key: value + for key, value in (("response_format", response_format), ("ref_audio", ref_audio)) + if isinstance(value, str) + } + return mapped_voice, mapped_params + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: # mutable-ok: base class contract returns a plain dict + resolved_key: Final = api_key or get_secret_str("MISTRAL_API_KEY") + if resolved_key is None: + raise MistralTextToSpeechException( + status_code=401, + message="Mistral API key is required. Set MISTRAL_API_KEY or pass api_key.", + ) + return { # mutable-ok: base class contract returns a plain dict + **headers, + "Authorization": f"Bearer {resolved_key}", + "Content-Type": "application/json", + } + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: Mapping[str, object], + ) -> str: + base_url: Final = api_base or get_secret_str("MISTRAL_API_BASE") or self.TTS_BASE_URL + return f"{base_url.rstrip('/')}/audio/speech" + + def transform_text_to_speech_request( + self, + model: str, + input: str, + voice: str | None, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, str], + ) -> TextToSpeechRequestData: + response_format: Final = optional_params.get("response_format") + ref_audio: Final = optional_params.get("ref_audio") + request_data: Final[TextToSpeechRequestData] = { + "dict_body": { + "model": model, + "input": input, + **({"voice_id": voice} if voice else {}), + **({"response_format": response_format} if isinstance(response_format, str) else {}), + **({"ref_audio": ref_audio} if isinstance(ref_audio, str) else {}), + }, + "headers": {"Content-Type": "application/json"}, + } + return request_data + + def _requested_content_type(self, request: httpx.Request) -> str: + request_body: Final = json.loads(request.content or b"{}") + requested_format: Final = request_body.get("response_format") + if not isinstance(requested_format, str): + return "audio/mpeg" + return self.AUDIO_CONTENT_TYPES.get(requested_format, "audio/mpeg") + + def transform_text_to_speech_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> "HttpxBinaryResponseContent": + from litellm.types.llms.openai import HttpxBinaryResponseContent + + try: + response_json: Final = raw_response.json() + except (json.JSONDecodeError, ValueError): + raise MistralTextToSpeechException( + status_code=raw_response.status_code, + message=f"Non-JSON response from Mistral speech API: {raw_response.text[:500]}", + headers=raw_response.headers, + ) + audio_b64: Final = response_json.get("audio_data") + if not isinstance(audio_b64, str) or not audio_b64: + raise MistralTextToSpeechException( + status_code=500, + message=f"No audio_data in Mistral speech response. Response keys: {tuple(response_json.keys())}", + headers=raw_response.headers, + ) + try: + audio_bytes: Final = base64.b64decode(audio_b64) + except ValueError: + raise MistralTextToSpeechException( + status_code=500, + message="Invalid base64 audio_data in Mistral speech response.", + headers=raw_response.headers, + ) + retained_headers: Final = tuple( + (key, value) + for key, value in raw_response.headers.items() + if key.lower() not in self.DROPPED_RESPONSE_HEADERS + ) + response_headers: Final = retained_headers + ( + ("content-length", str(len(audio_bytes))), + ("content-type", self._requested_content_type(raw_response.request)), + ) + binary_response: Final = httpx.Response( + status_code=200, + headers=response_headers, + content=audio_bytes, + request=raw_response.request, + ) + return HttpxBinaryResponseContent(binary_response) + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict | httpx.Headers, # mutable-ok: BaseLLMException takes a plain dict or httpx.Headers + ) -> BaseLLMException: + return MistralTextToSpeechException( + message=error_message, + status_code=status_code, + headers=headers, + ) diff --git a/litellm/main.py b/litellm/main.py index cafa1e4718f..692df23b3f9 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8367,6 +8367,34 @@ def speech( client=client, _is_async=aspeech or False, ) + elif custom_llm_provider == "mistral": + from litellm.llms.mistral.audio_speech.transformation import ( + MistralTextToSpeechConfig, + ) + + mistral_tts_config: Final = text_to_speech_provider_config or MistralTextToSpeechConfig() + + if api_base is not None: + litellm_params_dict["api_base"] = api_base + if api_key is not None: + litellm_params_dict["api_key"] = api_key + + mistral_voice: Final[str | None] = voice if isinstance(voice, str) else None + + response = base_llm_http_handler.text_to_speech_handler( + model=model, + input=input, + voice=mistral_voice, + text_to_speech_provider_config=mistral_tts_config, + text_to_speech_optional_params=optional_params, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params_dict, + logging_obj=logging_obj, + timeout=timeout, + extra_headers=extra_headers, + client=client, + _is_async=aspeech or False, + ) elif custom_llm_provider == "aws_polly": from litellm.llms.aws_polly.text_to_speech.transformation import ( AWSPollyTextToSpeechConfig, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bebbcc32181..620fe2f8030 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -31126,9 +31126,9 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-latest": { + "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", - "output_cost_per_character": 1.6e-05, "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", "supported_endpoints": [ "/v1/audio/speech" @@ -51677,9 +51677,9 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-2603": { + "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", - "output_cost_per_character": 1.6e-05, "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", "supported_endpoints": [ "/v1/audio/speech" diff --git a/litellm/router.py b/litellm/router.py index c93c1753f0e..d6ec5e57467 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4270,7 +4270,7 @@ class Router: self.fail_calls[model_name] += 1 raise e - async def aspeech(self, model: str, input: str, voice: str, **kwargs): + async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): """ Example Usage: @@ -4322,7 +4322,7 @@ class Router: ) raise e - async def _aspeech(self, model: str, input: str, voice: str, **kwargs): + async def _aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): model_name: Final = model try: verbose_router_logger.debug("Inside _aspeech()- model: %s; kwargs: %s", model, kwargs) diff --git a/litellm/utils.py b/litellm/utils.py index 5cd9bfc5f32..74a9b4ce935 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9408,6 +9408,12 @@ class ProviderConfigManager: ) return MinimaxTextToSpeechConfig() + elif litellm.LlmProviders.MISTRAL == provider: + from litellm.llms.mistral.audio_speech.transformation import ( + MistralTextToSpeechConfig, + ) + + return MistralTextToSpeechConfig() elif litellm.LlmProviders.AWS_POLLY == provider: from litellm.llms.aws_polly.text_to_speech.transformation import ( AWSPollyTextToSpeechConfig, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bebbcc32181..620fe2f8030 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -31126,9 +31126,9 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-latest": { + "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", - "output_cost_per_character": 1.6e-05, "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", "supported_endpoints": [ "/v1/audio/speech" @@ -51677,9 +51677,9 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-2603": { + "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", - "output_cost_per_character": 1.6e-05, "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", "supported_endpoints": [ "/v1/audio/speech" diff --git a/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py b/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py new file mode 100644 index 00000000000..20d07699cb8 --- /dev/null +++ b/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py @@ -0,0 +1,198 @@ +import base64 +from typing import Final +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig +from litellm.llms.mistral.audio_speech.transformation import ( + MistralTextToSpeechConfig, + MistralTextToSpeechException, +) +from litellm.utils import ProviderConfigManager + +SPEECH_URL: Final = "https://api.mistral.ai/v1/audio/speech" + + +def test_mistral_text_to_speech_config_installed(): + config: Final = ProviderConfigManager.get_provider_text_to_speech_config( + model="voxtral-mini-tts-2603", + provider=litellm.LlmProviders.MISTRAL, + ) + assert isinstance(config, BaseTextToSpeechConfig) + assert isinstance(config, MistralTextToSpeechConfig) + + +def test_map_openai_params_drops_speed_and_instructions(): + config: Final = MistralTextToSpeechConfig() + voice, params = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={"response_format": "wav", "speed": 1.5, "instructions": "sound cheerful"}, + voice="en_paul_neutral", + ) + assert voice == "en_paul_neutral" + assert params == {"response_format": "wav"} + + +def test_map_openai_params_accepts_voice_dict_and_ref_audio(): + config: Final = MistralTextToSpeechConfig() + voice, params = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={}, + voice={"voice_id": "1f3a8b0c-voice-uuid"}, + kwargs={"ref_audio": "bXktdm9pY2Utc2FtcGxl"}, + ) + assert voice == "1f3a8b0c-voice-uuid" + assert params == {"ref_audio": "bXktdm9pY2Utc2FtcGxl"} + + +def test_transform_request_builds_mistral_body(): + config: Final = MistralTextToSpeechConfig() + data: Final = config.transform_text_to_speech_request( + model="voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + optional_params={"response_format": "wav"}, + litellm_params={}, + headers={}, + ) + assert data["dict_body"] == { + "model": "voxtral-mini-tts-2603", + "input": "hello from litellm", + "voice_id": "en_paul_neutral", + "response_format": "wav", + } + assert data["headers"] == {"Content-Type": "application/json"} + + +def test_transform_request_omits_voice_for_ref_audio_cloning(): + config: Final = MistralTextToSpeechConfig() + data: Final = config.transform_text_to_speech_request( + model="voxtral-mini-tts-2603", + input="clone me", + voice=None, + optional_params={"ref_audio": "bXktdm9pY2Utc2FtcGxl"}, + litellm_params={}, + headers={}, + ) + assert data["dict_body"] == { + "model": "voxtral-mini-tts-2603", + "input": "clone me", + "ref_audio": "bXktdm9pY2Utc2FtcGxl", + } + + +def test_get_complete_url_default_base(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("MISTRAL_API_BASE", raising=False) + config: Final = MistralTextToSpeechConfig() + url: Final = config.get_complete_url(model="voxtral-mini-tts-2603", api_base=None, litellm_params={}) + assert url == SPEECH_URL + + +def test_get_complete_url_custom_base(): + config: Final = MistralTextToSpeechConfig() + url: Final = config.get_complete_url( + model="voxtral-mini-tts-2603", + api_base="https://custom.api.example.com/v1/", + litellm_params={}, + ) + assert url == "https://custom.api.example.com/v1/audio/speech" + + +def test_validate_environment_sets_bearer_header(): + config: Final = MistralTextToSpeechConfig() + headers: Final = config.validate_environment( + headers={"x-custom": "1"}, + model="voxtral-mini-tts-2603", + api_key="sk-mistral-test", + ) + assert headers == { + "x-custom": "1", + "Authorization": "Bearer sk-mistral-test", + "Content-Type": "application/json", + } + + +def test_validate_environment_requires_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("MISTRAL_API_KEY", raising=False) + config: Final = MistralTextToSpeechConfig() + with pytest.raises(MistralTextToSpeechException, match="MISTRAL_API_KEY"): + config.validate_environment(headers={}, model="voxtral-mini-tts-2603") + + +def test_transform_response_decodes_base64_audio(): + config: Final = MistralTextToSpeechConfig() + audio_bytes: Final = b"RIFF-fake-wav-bytes" + raw_response: Final = httpx.Response( + 200, + json={"audio_data": base64.b64encode(audio_bytes).decode()}, + headers={"x-request-id": "req-123"}, + request=httpx.Request( + "POST", + SPEECH_URL, + json={"model": "voxtral-mini-tts-2603", "input": "hi", "response_format": "wav"}, + ), + ) + result: Final = config.transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + assert result.content == audio_bytes + assert result.response.headers["content-type"] == "audio/wav" + assert result.response.headers["content-length"] == str(len(audio_bytes)) + assert result.response.headers["x-request-id"] == "req-123" + + +def test_transform_response_missing_audio_data_raises(): + config: Final = MistralTextToSpeechConfig() + raw_response: Final = httpx.Response( + 200, + json={"detail": "unexpected"}, + request=httpx.Request("POST", SPEECH_URL, json={"model": "voxtral-mini-tts-2603", "input": "hi"}), + ) + with pytest.raises(MistralTextToSpeechException, match="audio_data"): + config.transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + + +def test_map_openai_params_maps_openai_voice_aliases(): + config: Final = MistralTextToSpeechConfig() + alloy_voice, _ = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={}, + voice="alloy", + ) + nova_voice, _ = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={}, + voice="Nova", + ) + passthrough_voice, _ = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={}, + voice="en_paul_happy", + ) + assert alloy_voice == "en_paul_neutral" + assert nova_voice == "gb_jane_sarcasm" + assert passthrough_voice == "en_paul_happy" + + +def test_transform_response_invalid_base64_raises(): + config: Final = MistralTextToSpeechConfig() + raw_response: Final = httpx.Response( + status_code=200, + json={"audio_data": "!!!not-base64!!!"}, + request=httpx.Request("POST", SPEECH_URL), + ) + with pytest.raises(MistralTextToSpeechException, match="base64"): + config.transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 7c2174018e8..db08ee486bf 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4473,3 +4473,15 @@ def test_explicit_pricing_precedes_private_provider_response_model( ) assert selected == expected + + +def test_cost_per_token_mistral_voxtral_tts_bills_per_input_character(_local_model_cost_map): + prompt_usd, completion_usd = cost_per_token( + model="voxtral-mini-tts-2603", + custom_llm_provider="mistral", + call_type="speech", + prompt_characters=1000, + ) + + assert prompt_usd == pytest.approx(1000 * 1.6e-05) + assert completion_usd == 0.0 diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 8cf878d05d9..47293d9c413 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3181,3 +3181,31 @@ def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent( assert response is not None assert response._hidden_params.get("response_cost") is None + + +def test_speech_mistral_dispatches_and_decodes_audio(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + audio_bytes: Final = b"ID3-fake-mp3-bytes" + mock_route: Final = respx_mock.post("https://api.mistral.ai/v1/audio/speech").mock( + return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = litellm.speech( + model="mistral/voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + response_format="wav", + speed=2, + instructions="sound cheerful", + ) + + assert mock_route.called + request_body: Final = json.loads(mock_route.calls.last.request.content) + assert request_body == { + "model": "voxtral-mini-tts-2603", + "input": "hello from litellm", + "voice_id": "en_paul_neutral", + "response_format": "wav", + } + assert mock_route.calls.last.request.headers["authorization"] == "Bearer sk-mistral-test" + assert response.content == audio_bytes diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 97286017ffe..388011b7d5d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -11530,3 +11530,29 @@ class TestTierParamsTheTargetAccepts: accepted = router._tier_params_the_target_accepts("no-such-group", {"reasoning_effort": "max"}, {}) assert accepted == {"reasoning_effort": "max"} + + +@pytest.mark.asyncio +async def test_router_aspeech_without_voice_dispatches_ref_audio_cloning(respx_mock, monkeypatch): + import base64 + + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + audio_bytes = b"RIFFfake-wav-bytes" + respx_mock.post("https://api.mistral.ai/v1/audio/speech").respond( + json={"audio_data": base64.b64encode(audio_bytes).decode()} + ) + router = Router( + model_list=[ + { + "model_name": "voxtral-tts", + "litellm_params": {"model": "mistral/voxtral-mini-tts-2603"}, + } + ] + ) + + response = await router.aspeech(model="voxtral-tts", input="clone me", ref_audio="ZmFrZQ==") + + request_body = json.loads(respx_mock.calls.last.request.content) + assert request_body == {"model": "voxtral-mini-tts-2603", "input": "clone me", "ref_audio": "ZmFrZQ=="} + assert response.content == audio_bytes From 7c4cf2dcffbbff32b087d7756d4c0c0a4c590a91 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 11:41:24 -0700 Subject: [PATCH 02/32] fix(mistral): reject malformed base64 audio_data with strict validation --- litellm/llms/mistral/audio_speech/transformation.py | 2 +- .../audio_speech/test_mistral_audio_speech_transformation.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/mistral/audio_speech/transformation.py b/litellm/llms/mistral/audio_speech/transformation.py index 6d1a693268a..e7f7d510346 100644 --- a/litellm/llms/mistral/audio_speech/transformation.py +++ b/litellm/llms/mistral/audio_speech/transformation.py @@ -172,7 +172,7 @@ class MistralTextToSpeechConfig(BaseTextToSpeechConfig): headers=raw_response.headers, ) try: - audio_bytes: Final = base64.b64decode(audio_b64) + audio_bytes: Final = base64.b64decode(audio_b64, validate=True) except ValueError: raise MistralTextToSpeechException( status_code=500, diff --git a/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py b/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py index 20d07699cb8..6d250901e50 100644 --- a/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py +++ b/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py @@ -187,7 +187,7 @@ def test_transform_response_invalid_base64_raises(): config: Final = MistralTextToSpeechConfig() raw_response: Final = httpx.Response( status_code=200, - json={"audio_data": "!!!not-base64!!!"}, + json={"audio_data": "QUJD!QUJD"}, request=httpx.Request("POST", SPEECH_URL), ) with pytest.raises(MistralTextToSpeechException, match="base64"): From 7c91b0120fadccd2a97c5a8ae5db77d6d8f59ec8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:32:51 -0700 Subject: [PATCH 03/32] fix(mistral): ensure /v1 on the Voxtral TTS base URL A host-only api_base or MISTRAL_API_BASE (the documented form, https://api.mistral.ai) built https://api.mistral.ai/audio/speech and 404ed. Match the chat and OCR configs by appending /v1 when the configured base does not already end with it. --- .../mistral/audio_speech/transformation.py | 5 +++-- ...est_mistral_audio_speech_transformation.py | 19 +++++++++++++------ 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/litellm/llms/mistral/audio_speech/transformation.py b/litellm/llms/mistral/audio_speech/transformation.py index e7f7d510346..7f5a659bd08 100644 --- a/litellm/llms/mistral/audio_speech/transformation.py +++ b/litellm/llms/mistral/audio_speech/transformation.py @@ -115,8 +115,9 @@ class MistralTextToSpeechConfig(BaseTextToSpeechConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> str: - base_url: Final = api_base or get_secret_str("MISTRAL_API_BASE") or self.TTS_BASE_URL - return f"{base_url.rstrip('/')}/audio/speech" + configured_base: Final = (api_base or get_secret_str("MISTRAL_API_BASE") or self.TTS_BASE_URL).rstrip("/") + versioned_base: Final = configured_base if configured_base.endswith("/v1") else f"{configured_base}/v1" + return f"{versioned_base}/audio/speech" def transform_text_to_speech_request( self, diff --git a/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py b/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py index 6d250901e50..108d4107db1 100644 --- a/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py +++ b/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py @@ -91,16 +91,23 @@ def test_get_complete_url_default_base(monkeypatch: pytest.MonkeyPatch): assert url == SPEECH_URL -def test_get_complete_url_custom_base(): +@pytest.mark.parametrize( + "api_base", + ["https://custom.api.example.com/v1/", "https://custom.api.example.com/v1", "https://custom.api.example.com"], +) +def test_get_complete_url_custom_base_always_versioned(api_base: str): config: Final = MistralTextToSpeechConfig() - url: Final = config.get_complete_url( - model="voxtral-mini-tts-2603", - api_base="https://custom.api.example.com/v1/", - litellm_params={}, - ) + url: Final = config.get_complete_url(model="voxtral-mini-tts-2603", api_base=api_base, litellm_params={}) assert url == "https://custom.api.example.com/v1/audio/speech" +def test_get_complete_url_host_only_env_base_gets_v1(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MISTRAL_API_BASE", "https://api.mistral.ai") + config: Final = MistralTextToSpeechConfig() + url: Final = config.get_complete_url(model="voxtral-mini-tts-2603", api_base=None, litellm_params={}) + assert url == SPEECH_URL + + def test_validate_environment_sets_bearer_header(): config: Final = MistralTextToSpeechConfig() headers: Final = config.validate_environment( From 0a2581c14caaa43f76f1db2399aa66a9812d63bf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:30:07 -0700 Subject: [PATCH 04/32] fix(mistral): keep the deployment voice default and drop the unreachable api base fallback Review turned up two real problems in the TTS path. Router.aspeech forwarded voice=None whenever the caller omitted it, which overwrote a voice set in the deployment's litellm_params, so a configured fallback voice was ignored on voice-less requests. It now leaves the key alone when no voice is passed. get_complete_url also fell back to MISTRAL_API_BASE, but speech() always receives a non-null api_base from get_llm_provider, whose mistral branch only reads MISTRAL_AZURE_API_BASE and otherwise hardcodes the public host. That branch could never run, and its unit test asserted a behavior the real path does not have. The working override is api_base on the deployment, now pinned by an end-to-end test --- .../mistral/audio_speech/transformation.py | 2 +- litellm/router.py | 2 +- ...est_mistral_audio_speech_transformation.py | 10 +--- tests/test_litellm/test_main.py | 18 +++++++ tests/test_litellm/test_router.py | 50 +++++++++++++++++++ 5 files changed, 71 insertions(+), 11 deletions(-) diff --git a/litellm/llms/mistral/audio_speech/transformation.py b/litellm/llms/mistral/audio_speech/transformation.py index 7f5a659bd08..2b3264dc756 100644 --- a/litellm/llms/mistral/audio_speech/transformation.py +++ b/litellm/llms/mistral/audio_speech/transformation.py @@ -115,7 +115,7 @@ class MistralTextToSpeechConfig(BaseTextToSpeechConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> str: - configured_base: Final = (api_base or get_secret_str("MISTRAL_API_BASE") or self.TTS_BASE_URL).rstrip("/") + configured_base: Final = (api_base or self.TTS_BASE_URL).rstrip("/") versioned_base: Final = configured_base if configured_base.endswith("/v1") else f"{configured_base}/v1" return f"{versioned_base}/audio/speech" diff --git a/litellm/router.py b/litellm/router.py index ee3a3ced023..b1357374f5c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4422,7 +4422,7 @@ class Router: **{ **data, "input": input, - "voice": voice, + **({"voice": voice} if voice is not None else {}), "client": model_client, **kwargs, } diff --git a/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py b/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py index 108d4107db1..d819a79cef1 100644 --- a/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py +++ b/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py @@ -84,8 +84,7 @@ def test_transform_request_omits_voice_for_ref_audio_cloning(): } -def test_get_complete_url_default_base(monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv("MISTRAL_API_BASE", raising=False) +def test_get_complete_url_default_base(): config: Final = MistralTextToSpeechConfig() url: Final = config.get_complete_url(model="voxtral-mini-tts-2603", api_base=None, litellm_params={}) assert url == SPEECH_URL @@ -101,13 +100,6 @@ def test_get_complete_url_custom_base_always_versioned(api_base: str): assert url == "https://custom.api.example.com/v1/audio/speech" -def test_get_complete_url_host_only_env_base_gets_v1(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("MISTRAL_API_BASE", "https://api.mistral.ai") - config: Final = MistralTextToSpeechConfig() - url: Final = config.get_complete_url(model="voxtral-mini-tts-2603", api_base=None, litellm_params={}) - assert url == SPEECH_URL - - def test_validate_environment_sets_bearer_header(): config: Final = MistralTextToSpeechConfig() headers: Final = config.validate_environment( diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 9b5122c3f75..e01679048e6 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3287,3 +3287,21 @@ def test_speech_mistral_dispatches_and_decodes_audio(respx_mock: respx.MockRoute } assert mock_route.calls.last.request.headers["authorization"] == "Bearer sk-mistral-test" assert response.content == audio_bytes + + +def test_speech_mistral_routes_to_configured_api_base(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + audio_bytes: Final = b"ID3-gateway-bytes" + gateway_route: Final = respx_mock.post("https://mistral.gateway.internal/v1/audio/speech").mock( + return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = litellm.speech( + model="mistral/voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + api_base="https://mistral.gateway.internal", + ) + + assert gateway_route.called + assert response.content == audio_bytes diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 048536d887f..014937b35cd 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -12278,6 +12278,56 @@ async def test_router_aspeech_without_voice_dispatches_ref_audio_cloning(respx_m assert response.content == audio_bytes +@pytest.mark.asyncio +async def test_router_aspeech_without_voice_keeps_deployment_default_voice(respx_mock, monkeypatch): + import base64 + + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + audio_bytes = b"RIFFfake-wav-bytes" + respx_mock.post("https://api.mistral.ai/v1/audio/speech").respond( + json={"audio_data": base64.b64encode(audio_bytes).decode()} + ) + router = Router( + model_list=[ + { + "model_name": "voxtral-tts", + "litellm_params": {"model": "mistral/voxtral-mini-tts-2603", "voice": "en_paul_neutral"}, + } + ] + ) + + await router.aspeech(model="voxtral-tts", input="use my default") + + request_body = json.loads(respx_mock.calls.last.request.content) + assert request_body["voice_id"] == "en_paul_neutral" + + +@pytest.mark.asyncio +async def test_router_aspeech_request_voice_overrides_deployment_default(respx_mock, monkeypatch): + import base64 + + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + audio_bytes = b"RIFFfake-wav-bytes" + respx_mock.post("https://api.mistral.ai/v1/audio/speech").respond( + json={"audio_data": base64.b64encode(audio_bytes).decode()} + ) + router = Router( + model_list=[ + { + "model_name": "voxtral-tts", + "litellm_params": {"model": "mistral/voxtral-mini-tts-2603", "voice": "en_paul_neutral"}, + } + ] + ) + + await router.aspeech(model="voxtral-tts", input="override me", voice="gb_oliver_neutral") + + request_body = json.loads(respx_mock.calls.last.request.content) + assert request_body["voice_id"] == "gb_oliver_neutral" + + class TestPreRoutingTierDrivesFallbacks: """#38832: a complexity/auto router picks a tier behind the router name, but fallback lookup stayed on the router name, so the tier's configured chain never ran and a From 93fa9892389c6206f6b7fe99b3032096abd10d88 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:42:41 -0700 Subject: [PATCH 05/32] fix(router): fall back to the deployment voice without a conditional dict spread --- litellm/router.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 3b7f2d79823..c211e385d95 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4518,7 +4518,7 @@ class Router: **{ **data, "input": input, - **({"voice": voice} if voice is not None else {}), + "voice": data.get("voice") if voice is None else voice, "client": model_client, **kwargs, } From 4c00a6e189d95f34a72036802937b38769f561b5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:35:32 -0700 Subject: [PATCH 06/32] fix(batches): account a batch's cost once, from the first retrieve that sees it final Every retrieve of a batch through the proxy shares one spend row, the batch id plus the batch cost suffix, and spend log inserts skip duplicates. A poll that landed while the batch was still validating or in progress wrote that row at $0 and no later retrieve could overwrite it, and every completed retrieve after the first added the cost to the key, team, and user counters again with no new row to show for it. The cost callback now writes nothing for a batch retrieve until the batch is final, releasing the poll's budget reservation instead, and once it is final it charges only when no spend row for that batch is queued for flush or already stored. Batch cost rows are flushed to the database right away so a second instance sees them, and the logger prices a batch only once it is final, which also covers a failed batch that never produced an output file. --- litellm/batches/batch_utils.py | 20 ++ litellm/litellm_core_utils/litellm_logging.py | 13 +- litellm/proxy/db/db_spend_update_writer.py | 3 +- .../proxy/hooks/proxy_track_cost_callback.py | 56 ++++- .../openai_files_endpoints/common_utils.py | 8 +- litellm/proxy/utils.py | 10 +- .../test_litellm/batches/test_batch_utils.py | 55 +++- .../test_litellm_logging.py | 82 ++++++ .../proxy/db/test_db_spend_update_writer.py | 10 +- .../hooks/test_proxy_track_cost_callback.py | 238 +++++++++++++----- .../prisma_and_spend/test_spend_functions.py | 15 ++ 11 files changed, 428 insertions(+), 82 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 97be5f77d79..eaac3bf0e9f 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -25,6 +25,26 @@ class BatchCostUsageResult: failed_requests: int +_TERMINAL_BATCH_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "expired"}) + + +def batch_cost_is_final(batch: Batch) -> bool: + """Whether this retrieve of the batch is the one to account its cost from. + + A batch still in flight has nothing to price, and a "completed" batch can report + no output_file_id for a moment before the output populates; pricing either records + $0 under the batch's single spend row and pins it there. Final means a completed + batch whose output file has arrived or whose counts prove no line succeeded, or + any other terminal status (failed, cancelled, expired). + """ + if batch.status not in _TERMINAL_BATCH_STATUSES: + return False + if batch.status != "completed" or batch.output_file_id is not None: + return True + request_counts: Final = batch.request_counts + return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0 + + async def calculate_batch_cost_and_usage( file_content_dictionary: list[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c31c4323157..09ddd1b9720 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -36,7 +36,7 @@ from litellm._logging import ( verbose_logger, ) from litellm._uuid import uuid -from litellm.batches.batch_utils import _handle_completed_batch +from litellm.batches.batch_utils import _handle_completed_batch, batch_cost_is_final from litellm.caching.caching import DualCache, InMemoryCache from litellm.caching.caching_handler import LLMCachingHandler from litellm.constants import ( @@ -2899,13 +2899,6 @@ class Logging(LiteLLMLoggingBaseClass): ): # polling job will query these frequently, don't spam db logs return - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ) - - # check if file id is a unified file id - is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(result.id) - batch_cost: Final = kwargs.get("batch_cost", None) batch_usage = kwargs.get("batch_usage", None) batch_models = kwargs.get("batch_models", None) @@ -2913,9 +2906,7 @@ class Logging(LiteLLMLoggingBaseClass): batch_failed_requests: Final = kwargs.get("batch_failed_requests", None) has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models)) - should_compute_batch_data: Final = ( - not is_base64_unified_file_id or not has_explicit_batch_data and result.status == "completed" - ) + should_compute_batch_data: Final = not has_explicit_batch_data and batch_cost_is_final(result) if has_explicit_batch_data: result._hidden_params["response_cost"] = batch_cost result._hidden_params["batch_models"] = batch_models diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e6880d521f1..ff48b00dc70 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -82,6 +82,7 @@ else: RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) +IMMEDIATE_FLUSH_CALL_TYPES: Final = RESPONSES_SESSION_CALL_TYPES | frozenset({CallTypes.aretrieve_batch.value}) class _SpendBatch(Protocol): @@ -939,7 +940,7 @@ class DBSpendUpdateWriter: from litellm.proxy.utils import enqueue_spend_logs, request_spend_log_flush await enqueue_spend_logs(prisma_client, (payload,)) - if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES: + if payload.get("call_type") in IMMEDIATE_FLUSH_CALL_TYPES: request_spend_log_flush() else: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 7254b05db2e..95e61fdd98c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger +from litellm.batches.batch_utils import batch_cost_is_final from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( @@ -33,17 +34,21 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import ( from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, get_request_model_access_groups, + get_spend_logs_id, ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( CallTypes, + LiteLLMBatch, StandardLoggingPayload, StandardLoggingPayloadErrorInformation, ) from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: - from litellm.proxy.utils import ProxyLogging + from prisma.types import LiteLLM_SpendLogsWhereUniqueInput + + from litellm.proxy.utils import PrismaClient, ProxyLogging _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset( { @@ -224,6 +229,7 @@ class _ProxyDBLogger(CustomLogger): ): from litellm.proxy.proxy_server import ( increment_spend_counters, + prisma_client, proxy_logging_obj, update_cache, ) @@ -248,6 +254,18 @@ class _ProxyDBLogger(CustomLogger): ) _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) + if ( + isinstance(completion_response, LiteLLMBatch) + and kwargs.get("call_type") == CallTypes.aretrieve_batch.value + ): + batch_spend_log_id: Final = get_spend_logs_id( + CallTypes.aretrieve_batch.value, completion_response.model_dump(), kwargs + ) + if not await _batch_cost_is_trackable_now( + batch=completion_response, spend_log_id=batch_spend_log_id, prisma_client=prisma_client + ): + await _release_budget_reservation(budget_reservation=budget_reservation) + return user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) @@ -491,6 +509,42 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: bucket[key] = value +async def _batch_cost_is_trackable_now( + batch: LiteLLMBatch, spend_log_id: str | None, prisma_client: "PrismaClient | None" +) -> bool: + """A batch is billed exactly once, from the first retrieve that sees it final. + + Every retrieve of one batch shares a single spend row (its id plus the batch cost + suffix), so a poll that lands before the output exists would write that row at $0 + and pin it there, and every retrieve after the first would add the cost to the + key, team, and user counters again. + """ + if not batch_cost_is_final(batch): + verbose_proxy_logger.debug("Cost tracking deferred for batch %s still in status %s", batch.id, batch.status) + return False + if prisma_client is None or spend_log_id is None: + return True + if not await _spend_log_already_recorded(prisma_client=prisma_client, request_id=spend_log_id): + return True + verbose_proxy_logger.debug( + "Cost tracking skipped for batch %s: spend row %s already recorded", batch.id, spend_log_id + ) + return False + + +async def _spend_log_already_recorded(prisma_client: "PrismaClient", request_id: str) -> bool: + from litellm.proxy.utils import spend_log_is_queued + + if await spend_log_is_queued(prisma_client, request_id): + return True + spend_log_row: Final[LiteLLM_SpendLogsWhereUniqueInput] = {"request_id": request_id} + try: + return await prisma_client.db.litellm_spendlogs.find_unique(where=spend_log_row) is not None + except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreadable DB must not drop the batch's only spend row + verbose_proxy_logger.warning("Could not check for an existing spend row %s, tracking anyway: %s", request_id, e) + return False + + def _is_unbilled_interaction_response(completion_response: object) -> bool: from litellm.interactions.background_cost_polling import missing_usage_is_expected from litellm.types.interactions import InteractionsAPIResponse diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 15eeddbc489..b1f282a0978 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -15,6 +15,7 @@ from typing import ( runtime_checkable, ) +from litellm.batches.batch_utils import batch_cost_is_final from litellm.proxy._types import ProxyException from litellm.repositories.table_repositories import ( ManagedFileRepository, @@ -1357,12 +1358,7 @@ def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: enumerated the batch and none succeeded. A zero or unknown total means counts are unreported, so stay eligible and let the next poller pass revisit it. (#37713) """ - if response.output_file_id is not None: - return True - request_counts = response.request_counts - if request_counts is None: - return False - return request_counts.total > 0 and request_counts.completed == 0 + return batch_cost_is_final(response) async def update_batch_in_database( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index accf7b720fb..f64b51bc6c6 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6251,7 +6251,9 @@ def request_spend_log_flush() -> None: The Responses API hands the client an id it can chain from straight away, and that lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval. - Repeated requests coalesce into the monitor's next pass, so the batching holds. + A batch's cost row is what every other worker checks before charging the same batch + again, so it cannot wait either. Repeated requests coalesce into the monitor's next + pass, so the batching holds. """ PrismaClient.spend_log_flush_requested.set() @@ -6266,6 +6268,12 @@ async def _wait_for_spend_log_flush_request(interval: float) -> bool: return True +async def spend_log_is_queued(prisma_client: PrismaClient, request_id: str) -> bool: + """Whether a spend log with ``request_id`` is still waiting for the next flush.""" + async with prisma_client._spend_log_transactions_lock: + return any(row.get("request_id") == request_id for row in prisma_client.spend_log_transactions) + + async def dequeue_spend_logs(prisma_client: PrismaClient, limit: int) -> list[dict[str, object]]: """Take up to ``limit`` of the oldest queued spend logs off the queue. diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index c86c7c4df03..8d4f68164b4 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -21,11 +21,12 @@ from types import MappingProxyType import httpx import pytest import respx +from openai.types.batch import BatchRequestCounts import litellm import litellm.batches.batch_utils as bu -from litellm.types.utils import Usage +from litellm.types.utils import LiteLLMBatch, Usage # --------------------------------------------------------------------------- # # Builders for batch OUTPUT file rows. @@ -1718,3 +1719,55 @@ def test_unparsable_bedrock_batch_usage_warns(caplog): assert usage.total_tokens == 0 assert "does not understand" in caplog.text assert "inputTextTokenCount" in caplog.text + + +# --------------------------------------------------------------------------- # +# batch_cost_is_final +# --------------------------------------------------------------------------- # + +def _retrieved_batch( + status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None +) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status=status, + output_file_id=output_file_id, + request_counts=counts, + ) + + +class TestBatchCostIsFinal: + """Every retrieve of one batch writes the same spend row, so the first retrieve + that prices it decides the row for good. A poll before the output exists must + therefore not count as final: pricing it recorded $0 and pinned it (LIT-7048).""" + + @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"]) + def test_in_flight_batch_is_not_final(self, status): + assert bu.batch_cost_is_final(_retrieved_batch(status)) is False + + def test_completed_with_output_is_final(self): + assert bu.batch_cost_is_final(_retrieved_batch("completed", output_file_id="file-out")) is True + + def test_completed_without_output_and_unknown_counts_is_not_final(self): + assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False + + def test_completed_without_output_and_zero_counts_is_not_final(self): + counts = BatchRequestCounts(total=0, completed=0, failed=0) + assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False + + def test_completed_without_output_but_successful_lines_is_not_final(self): + counts = BatchRequestCounts(total=2, completed=2, failed=0) + assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False + + def test_completed_without_output_and_every_line_failed_is_final(self): + counts = BatchRequestCounts(total=2, completed=0, failed=2) + assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is True + + @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) + def test_other_terminal_statuses_are_final(self, status): + assert bu.batch_cost_is_final(_retrieved_batch(status)) is True diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 16a99713a06..4583429bd11 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -632,6 +632,88 @@ class TestRetrieveBatchCostPassesModelIdentity: assert captured["model_info"]["input_cost_per_token"] == 0.0 +class TestRetrieveBatchPricesOnlyFinalBatches: + """Regression (LIT-7048): retrieving a provider-id batch priced it on every poll. + + Every retrieve of one batch logs under the same spend row, so pricing a poll + that landed before the output existed wrote that row at $0 and pinned it there. + Only a final batch gets priced; an in-flight poll carries no cost at all. + """ + + @staticmethod + def _logging_obj() -> LitellmLogging: + obj = LitellmLogging( + model="gpt-5.6-luna", + messages=[{"role": "user", "content": "Hey"}], + stream=False, + call_type="aretrieve_batch", + start_time=time.time(), + litellm_call_id="batch-call-2", + function_id="f", + ) + obj.custom_llm_provider = "openai" + return obj + + @staticmethod + def _batch(status: str, output_file_id: str | None): + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id="batch_6a9c99e185588190877d391f8b9d7f8a", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status=status, + output_file_id=output_file_id, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("status", "output_file_id"), + [("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None)], + ) + async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_file_id) -> None: + from litellm.litellm_core_utils import litellm_logging as logging_module + + handle_completed_batch = AsyncMock() + monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) + batch = self._batch(status, output_file_id) + + with contextlib.suppress(Exception): + await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) + + handle_completed_batch.assert_not_awaited() + assert "response_cost" not in batch._hidden_params + + @pytest.mark.asyncio + async def test_completed_batch_with_output_is_priced(self, monkeypatch) -> None: + from litellm.batches.batch_utils import BatchCostUsageResult + from litellm.litellm_core_utils import litellm_logging as logging_module + from litellm.types.utils import Usage + + handle_completed_batch = AsyncMock( + return_value=BatchCostUsageResult( + cost=8e-06, + usage=Usage(prompt_tokens=26, completion_tokens=9, total_tokens=35), + models=["gpt-5.6-luna"], + successful_requests=2, + failed_requests=0, + ) + ) + monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) + batch = self._batch("completed", "file-out") + + with contextlib.suppress(Exception): + await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) + + handle_completed_batch.assert_awaited_once() + assert batch._hidden_params["response_cost"] == 8e-06 + assert batch.usage is not None + assert batch.usage.total_tokens == 35 + + class TestAnthropicPassthroughCustomPricing: """Verify the Anthropic pass-through handler forwards custom pricing.""" diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 11ef911de3e..e1b2d151c6d 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2934,12 +2934,16 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey @pytest.mark.asyncio @pytest.mark.parametrize( "call_type, expects_flush", - [("aresponses", True), ("responses", True), ("acompletion", False)], + [("aresponses", True), ("responses", True), ("aretrieve_batch", True), ("acompletion", False)], ) -async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(call_type: str, expects_flush: bool): +async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_workers_read_back( + call_type: str, expects_flush: bool +): """ A `previous_response_id` chained straight off the previous turn reads the DB, so a - Responses row cannot sit in this worker's queue until the monitor's next poll. + Responses row cannot sit in this worker's queue until the monitor's next poll. A + batch's cost row is what another worker checks before charging the same batch again + (LIT-7048), so it cannot wait either. """ from litellm.proxy.utils import PrismaClient diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 8043a1aca3f..b7037e8d621 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,4 +1,3 @@ - from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -7,6 +6,7 @@ import pytest from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( + _batch_cost_is_trackable_now, _get_budget_reservation_from_metadata, _ProxyDBLogger, _should_track_cost_callback, @@ -70,9 +70,7 @@ async def test_async_post_call_failure_hook(): # Check that metadata was properly updated assert "litellm_params" in call_args["kwargs"] - assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == { - "request_id": "test_request_id" - } + assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == {"request_id": "test_request_id"} metadata = call_args["kwargs"]["litellm_params"]["metadata"] assert metadata["user_api_key"] == "test_api_key" assert metadata["status"] == "failure" @@ -336,9 +334,7 @@ async def test_should_continue_failure_tracking_when_budget_release_fails(): ) assert mock_invalidate_budget_reservation_counters.await_count == 1 assert ( - mock_invalidate_budget_reservation_counters.await_args.kwargs[ - "budget_reservation" - ] + mock_invalidate_budget_reservation_counters.await_args.kwargs["budget_reservation"] is user_api_key_dict.budget_reservation ) assert user_api_key_dict.budget_reservation["finalized"] is True @@ -433,36 +429,21 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): "entries": [{"counter_key": "spend:key:test_api_key"}], } + assert _get_budget_reservation_from_metadata(metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}) is None assert ( _get_budget_reservation_from_metadata( - metadata={"user_api_key_auth": dict(UserAPIKeyAuth())} - ) - is None - ) - assert ( - _get_budget_reservation_from_metadata( - metadata={ - "user_api_key_auth": UserAPIKeyAuth( - budget_reservation=budget_reservation - ) - } + metadata={"user_api_key_auth": UserAPIKeyAuth(budget_reservation=budget_reservation)} ) == budget_reservation ) assert ( _get_budget_reservation_from_metadata( - metadata={ - "user_api_key_auth": dict( - UserAPIKeyAuth(budget_reservation=budget_reservation) - ) - } + metadata={"user_api_key_auth": dict(UserAPIKeyAuth(budget_reservation=budget_reservation))} ) == budget_reservation ) assert ( - _get_budget_reservation_from_metadata( - metadata={"user_api_key_budget_reservation": budget_reservation} - ) + _get_budget_reservation_from_metadata(metadata={"user_api_key_budget_reservation": budget_reservation}) is budget_reservation ) @@ -470,9 +451,7 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - side_effect=Exception("db unavailable") - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=Exception("db unavailable")) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -508,9 +487,7 @@ async def test_update_database_and_spend_counters_releases_reservation_when_db_u async def test_update_database_and_spend_counters_preserves_db_exception_when_release_fails(): proxy_logging_obj = MagicMock() db_exception = RuntimeError("db unavailable") - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - side_effect=db_exception - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=db_exception) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -554,12 +531,8 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re budget_reservation=budget_reservation, ) assert mock_log_exception.call_count == 2 - mock_log_exception.assert_any_call( - "Failed to release budget reservation after database update failed" - ) - mock_log_exception.assert_any_call( - "Failed to invalidate budget reservation counters after release failed" - ) + mock_log_exception.assert_any_call("Failed to release budget reservation after database update failed") + mock_log_exception.assert_any_call("Failed to invalidate budget reservation counters after release failed") increment_spend_counters.assert_not_awaited() @@ -778,6 +751,169 @@ async def test_track_cost_callback_defers_in_progress_background_interaction(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +def _batch_retrieve_kwargs(call_type: str, reservation: dict | None = None) -> dict: + metadata = { + "user_api_key": "hashed_key", + "user_api_key_user_id": "user-1", + "user_api_key_team_id": "team-1", + **({"user_api_key_budget_reservation": reservation} if reservation is not None else {}), + } + return { + "call_type": call_type, + "model": "gpt-5.6-luna", + "litellm_call_id": "test-call-id", + "litellm_params": {"metadata": metadata}, + "standard_logging_object": {"response_cost": 0.0, "request_tags": None}, + "stream": False, + } + + +def _retrieved_batch(status: str, output_file_id: str | None): + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status=status, + output_file_id=output_file_id, + ) + + +def _prisma_client_with(queued_request_ids: tuple[str, ...], stored_row: object) -> MagicMock: + import asyncio + + prisma_client = MagicMock() + prisma_client._spend_log_transactions_lock = asyncio.Lock() + prisma_client.spend_log_transactions = [{"request_id": request_id} for request_id in queued_request_ids] + prisma_client.db.litellm_spendlogs.find_unique = AsyncMock(return_value=stored_row) + return prisma_client + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status", "output_file_id", "spend_log_id", "prisma_client", "trackable"), + [ + ("in_progress", None, "batch_abc_batch_cost", None, False), + ("in_progress", None, "batch_abc_batch_cost", _prisma_client_with((), None), False), + ("completed", None, "batch_abc_batch_cost", _prisma_client_with((), None), False), + ("completed", "file-out", "batch_abc_batch_cost", None, True), + ("completed", "file-out", None, _prisma_client_with((), None), True), + ("completed", "file-out", "batch_abc_batch_cost", _prisma_client_with(("batch_abc_batch_cost",), None), False), + ( + "completed", + "file-out", + "batch_abc_batch_cost", + _prisma_client_with((), {"request_id": "batch_abc_batch_cost"}), + False, + ), + ("completed", "file-out", "batch_abc_batch_cost", _prisma_client_with((), None), True), + ("failed", None, "batch_abc_batch_cost", _prisma_client_with((), None), True), + ], + ids=[ + "in_progress_without_db", + "in_progress_never_consults_db", + "completed_without_output_yet", + "final_without_db", + "final_without_spend_log_id", + "final_row_queued_for_flush", + "final_row_already_stored", + "final_first_sighting", + "failed_first_sighting", + ], +) +async def test_batch_cost_is_trackable_now(status, output_file_id, spend_log_id, prisma_client, trackable): + """ + A batch is billed from the first retrieve that sees it final and never again: + a poll before that wrote the shared spend row at $0 and pinned it there, and + every completed retrieve after the first charged the key again (LIT-7048). + """ + assert ( + await _batch_cost_is_trackable_now( + batch=_retrieved_batch(status, output_file_id), spend_log_id=spend_log_id, prisma_client=prisma_client + ) + is trackable + ) + + +@pytest.mark.asyncio +async def test_batch_cost_is_trackable_now_when_the_spend_row_lookup_fails(): + """An unreadable spend log table must not drop the batch's only spend row.""" + prisma_client = _prisma_client_with((), None) + prisma_client.db.litellm_spendlogs.find_unique = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert ( + await _batch_cost_is_trackable_now( + batch=_retrieved_batch("completed", "file-out"), + spend_log_id="batch_abc_batch_cost", + prisma_client=prisma_client, + ) + is True + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "status", "output_file_id", "stored_row", "charged"), + [ + ("aretrieve_batch", "in_progress", None, None, False), + ("aretrieve_batch", "completed", "file-out", {"request_id": "batch_abc_batch_cost"}, False), + ("aretrieve_batch", "completed", "file-out", None, True), + ("acreate_batch", "validating", None, None, True), + ], + ids=["retrieve_before_final", "retrieve_already_recorded", "retrieve_first_final", "create_before_final"], +) +async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs and whether the poll's reservation is handed back is the whole observable contract of the gate + call_type, status, output_file_id, stored_row, charged +): + """ + Only retrieves are gated, since creating a batch is its own billable request. + A retrieve that writes nothing hands its budget reservation back instead. + """ + logger = _ProxyDBLogger() + budget_reservation = None if charged else {"reserved_cost": 0.5, "entries": []} + kwargs = _batch_retrieve_kwargs(call_type, reservation=budget_reservation) + + with ( + patch( # test-quality-ok: prisma_client is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.prisma_client", _prisma_client_with((), stored_row) + ), + patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ), + patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), + patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: the release is imported inside the callback's helper, no seam + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", new_callable=AsyncMock + ) as mock_release_budget_reservation, + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=_retrieved_batch(status, output_file_id), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_proxy_logging.failed_tracking_alert.assert_not_called() + if charged: + mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() + mock_release_budget_reservation.assert_not_awaited() + else: + mock_proxy_logging.db_spend_update_writer.update_database.assert_not_called() + mock_release_budget_reservation.assert_awaited_once_with(budget_reservation=budget_reservation) + + def _in_progress_interaction_kwargs(reservation: dict) -> dict: return { "call_type": "acreate_interaction", @@ -1101,10 +1237,7 @@ async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj # standard_logging_object should have been propagated from logging obj assert call_kwargs.get("standard_logging_object") is not None - assert ( - call_kwargs["standard_logging_object"]["trace_id"] - == "trace-id-from-logging-obj" - ) + assert call_kwargs["standard_logging_object"]["trace_id"] == "trace-id-from-logging-obj" # litellm_trace_id should also be propagated as a fallback assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj" @@ -1691,9 +1824,7 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend(): "metadata": {}, "proxy_server_request": {"request_id": "rid"}, "response_cost": 3.5e-05, - "combined_usage_object": Usage( - prompt_tokens=30, completion_tokens=1, total_tokens=31 - ), + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), } with patch( @@ -1772,15 +1903,10 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): assert mock_increment.call_args.kwargs["team_id"] == "team-123" assert mock_increment.call_args.kwargs["org_id"] == "org-456" - update_kwargs = ( - mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs - ) + update_kwargs = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs assert update_kwargs["user_id"] == "mcp-user@example.com" assert update_kwargs["team_id"] == "team-123" - assert ( - kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] - == "mcp-user@example.com" - ) + assert kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] == "mcp-user@example.com" @pytest.mark.parametrize( @@ -1828,9 +1954,7 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect ], ) @pytest.mark.asyncio -async def test_track_cost_callback_logs_unauthenticated_pass_through_request( - call_type, expect_spend_log -): +async def test_track_cost_callback_logs_unauthenticated_pass_through_request(call_type, expect_spend_log): """Regression for LIT-3782: a pass-through request with auth=false reaches the cost callback with no key/user/team/end-user. Before the fix the spend-log write was skipped and the request never appeared in request/usage logs. It @@ -1876,9 +2000,7 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( end_time=datetime.now(), ) - assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( - 1 if expect_spend_log else 0 - ) + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if expect_spend_log else 0) class _FakeDeploymentLookup: diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index a1eb88a7834..fc97d760226 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -6,6 +6,7 @@ Symbols pinned here: - ``update_spend_logs_job`` - ``_monitor_spend_logs_queue`` - ``_raise_failed_update_spend_exception`` + - ``spend_log_is_queued`` """ from __future__ import annotations @@ -22,6 +23,7 @@ from litellm.proxy.utils import ( _monitor_spend_logs_queue, _raise_failed_update_spend_exception, drain_spend_logs_queue, + spend_log_is_queued, update_daily_tag_spend, update_spend, update_spend_logs_job, @@ -629,3 +631,16 @@ def test_raise_failed_update_spend_exception_raises_original_error() -> None: with pytest.raises(ValueError, match="specific"): asyncio.run(_runner()) + + +@pytest.mark.asyncio +async def test_spend_log_is_queued_matches_only_rows_awaiting_flush( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="batch_abc_batch_cost")] + + assert await spend_log_is_queued(mock_prisma_client, "batch_abc_batch_cost") is True + assert await spend_log_is_queued(mock_prisma_client, "batch_abc") is False + + mock_prisma_client.spend_log_transactions = [] + assert await spend_log_is_queued(mock_prisma_client, "batch_abc_batch_cost") is False From 635bb3a2096fbb4f4c8899574807c049bb8e4825 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 17:08:42 -0700 Subject: [PATCH 07/32] feat(cost-map): add azure_ai/gpt-6-astra Foundry pricing A gpt-6-astra deployment on a Foundry project reached through the azure_ai route had no cost map entry of its own, so it resolved to the OpenAI gpt-6-astra card: missing from the azure_ai/* wildcard listing, flex and priority prices and /v1/batch it does not sell, and no none reasoning effort. Add azure_ai/gpt-6-astra mirroring the azure/gpt-6-astra Standard Global sheet the way azure_ai/gpt-5.5 mirrors azure/gpt-5.5, and extend the cost, reasoning-effort, and wildcard listing tests to the Foundry route. --- ...odel_prices_and_context_window_backup.json | 43 +++++++++++++++++++ model_prices_and_context_window.json | 43 +++++++++++++++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 15 +++++-- .../proxy/auth/test_model_checks.py | 19 ++++++++ .../test_reasoning_effort_capability.py | 16 +++++-- 5 files changed, 129 insertions(+), 7 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4273ec54472..ac7407c2608 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3485,6 +3485,49 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "azure_ai/gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "source": "https://ai.azure.com/catalog/models/gpt-6-astra", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure_ai/gpt-5.5": { "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4273ec54472..ac7407c2608 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3485,6 +3485,49 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "azure_ai/gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "source": "https://ai.azure.com/catalog/models/gpt-6-astra", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure_ai/gpt-5.5": { "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index df680b7cb0e..73f1a19d85c 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2008,7 +2008,14 @@ def test_generic_cost_per_token_azure_gpt56(_local_model_cost_map, assert round(completion_cost, 10) == round(output_cost * completion_tokens, 10) -@pytest.mark.parametrize("model,zone_multiplier", [("azure/gpt-6-astra", 1.0), ("azure/us/gpt-6-astra", 1.1)]) +@pytest.mark.parametrize( + "model,custom_llm_provider,zone_multiplier", + [ + ("azure/gpt-6-astra", "azure", 1.0), + ("azure/us/gpt-6-astra", "azure", 1.1), + ("azure_ai/gpt-6-astra", "azure_ai", 1.0), + ], +) @pytest.mark.parametrize( "prompt_tokens,input_side_multiplier,output_multiplier", [(100000, 1.0, 1.0), (300000, 2.0, 1.5)], @@ -2016,6 +2023,7 @@ def test_generic_cost_per_token_azure_gpt56(_local_model_cost_map, def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( _local_model_cost_map, model, + custom_llm_provider, zone_multiplier, prompt_tokens, input_side_multiplier, @@ -2023,7 +2031,8 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( ): """Microsoft Foundry sells gpt-6-astra at the OpenAI rates: $10 input, $1 cache read, $12.50 cache write, $50 output per 1M tokens on Standard Global, with the input side doubling and output 1.5x above 272K - prompt tokens. Standard US Data Zone carries the usual 10% uplift on every rate. + prompt tokens. Standard US Data Zone carries the usual 10% uplift on every rate. A Foundry + deployment reached through the azure_ai route bills the same Standard Global sheet. """ cached_tokens = 50000 cache_write_tokens = 40000 @@ -2041,7 +2050,7 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( prompt_cost, completion_cost = generic_cost_per_token( model=model, usage=usage, - custom_llm_provider="azure", + custom_llm_provider=custom_llm_provider, ) input_side = zone_multiplier * input_side_multiplier diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index d58683fd1e5..eb48f70d5da 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -857,6 +857,25 @@ def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): litellm.add_known_models(model_cost_map={}) assert fake_model not in litellm.models_by_provider["vertex_ai"] + +def test_azure_ai_wildcard_lists_the_foundry_gpt_6_astra_entry(monkeypatch): + """A Foundry (azure_ai) deployment of gpt-6-astra only shows up under an azure_ai/* wildcard + when the cost map carries its own azure_ai/ entry; the azure/ entry from the OpenAI-on-Azure + price sheet never reaches the Foundry provider list (LIT-7081).""" + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + foundry_key = "azure_ai/gpt-6-astra" + local_entry = litellm.get_model_cost_map(url="")[foundry_key] + try: + litellm.add_known_models(model_cost_map={foundry_key: local_entry}) + assert foundry_key in get_known_models_from_wildcard("azure_ai/*") + finally: + litellm.azure_ai_models.discard(foundry_key) + litellm.add_known_models(model_cost_map={}) + + def test_get_complete_model_list_drops_no_default_models_sentinel(): from litellm.proxy.auth.model_checks import get_complete_model_list diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index f181370455d..b7499d1c975 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -389,14 +389,22 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: "max", ) - @pytest.mark.parametrize("model", ["azure/gpt-6-astra", "azure/us/gpt-6-astra"]) - def test_a_foundry_deployment_also_advertises_none(self, local_model_cost_map, model): + @pytest.mark.parametrize( + "model,custom_llm_provider", + [ + ("azure/gpt-6-astra", "azure"), + ("azure/us/gpt-6-astra", "azure"), + ("azure_ai/gpt-6-astra", "azure_ai"), + ], + ) + def test_a_foundry_deployment_also_advertises_none(self, local_model_cost_map, model, custom_llm_provider): """Microsoft Foundry serves the same model but its API accepts reasoning_effort none (verified live: 200 with zero reasoning tokens, and it unlocks temperature), which - OpenAI's rejects, so an Azure deployment offers none on top of low through max.""" + OpenAI's rejects, so an Azure deployment offers none on top of low through max, whether + it is reached through the azure route or the azure_ai (Foundry) route.""" from litellm.utils import _get_model_info_helper - model_info = dict(_get_model_info_helper(model=model, custom_llm_provider="azure")) + model_info = dict(_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)) assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( "none", From b067e836f8df15bb3dbefe345bf503b7d3811b9a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 18:58:11 -0700 Subject: [PATCH 08/32] fix(batches): claim the batch cost spend row in the database before charging The cost callback used to look for an existing `_batch_cost` row before charging a completed batch, which left a window where concurrent retrieves on any instance all charged the key, and it would honor a row any request had written under that id. The spend update writer now inserts the batch cost row itself with `create_many(skip_duplicates=True)` and only the retrieve whose insert lands charges the key, team, and user. An existing row only takes the charge when it is a successful `aretrieve_batch` row, so a client-chosen `x-litellm-call-id` on another endpoint cannot suppress billing. Batch cost rows no longer get their own immediate flush path `batch_cost_is_final` now treats the proxy's normalized `complete` status like `completed`, which the enterprise batch cost poller relies on when it decides whether a completed batch is safe to retire. Tests build that status with `model_copy` since the OpenAI `Batch` model rejects it The `test-quality-ok` markers sit on the `patch(` lines the gate keys on, and the logging tests no longer wrap the priced retrieve in `contextlib.suppress` --- litellm/batches/batch_utils.py | 5 +- litellm/proxy/db/db_spend_update_writer.py | 73 ++++++++-- .../proxy/hooks/proxy_track_cost_callback.py | 68 +++------- litellm/proxy/utils.py | 10 +- .../test_litellm/batches/test_batch_utils.py | 14 +- .../test_litellm_logging.py | 12 +- .../proxy/db/test_db_spend_update_writer.py | 128 +++++++++++++++++- .../hooks/test_proxy_track_cost_callback.py | 116 ++++------------ .../prisma_and_spend/test_spend_functions.py | 15 -- 9 files changed, 249 insertions(+), 192 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index eaac3bf0e9f..959c7498479 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -25,7 +25,8 @@ class BatchCostUsageResult: failed_requests: int -_TERMINAL_BATCH_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "expired"}) +_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) +_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"}) def batch_cost_is_final(batch: Batch) -> bool: @@ -39,7 +40,7 @@ def batch_cost_is_final(batch: Batch) -> bool: """ if batch.status not in _TERMINAL_BATCH_STATUSES: return False - if batch.status != "completed" or batch.output_file_id is not None: + if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None: return True request_counts: Final = batch.request_counts return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0 diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ff48b00dc70..3fad351224b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -82,7 +82,10 @@ else: RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) -IMMEDIATE_FLUSH_CALL_TYPES: Final = RESPONSES_SESSION_CALL_TYPES | frozenset({CallTypes.aretrieve_batch.value}) + + +def _is_batch_cost_row(payload: SpendLogsPayload) -> bool: + return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success" class _SpendBatch(Protocol): @@ -216,7 +219,12 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, - ) -> None: + ) -> bool: + """Record the request's spend, answering whether its cost still needs charging. + + False only for a batch retrieve whose cost row another retrieve already wrote, + so the caller leaves the key, team, and user counters alone (LIT-7048). + """ from litellm.proxy.proxy_server import ( disable_spend_logs, litellm_proxy_budget_name, @@ -233,7 +241,7 @@ class DBSpendUpdateWriter: team_id, ) if ProxyUpdateSpend.disable_spend_updates() is True: - return + return True if token is not None and isinstance(token, str) and token.startswith("sk-"): hashed_token = hash_token(token=token) else: @@ -264,10 +272,8 @@ class DBSpendUpdateWriter: payload["team_id"] = team_id if disable_spend_logs is False: - await self._insert_spend_log_to_db( - payload=payload, - prisma_client=prisma_client, - ) + if not await self._record_spend_log(payload=payload, prisma_client=prisma_client): + return False await self._enqueue_tool_usage_transaction( payload=payload, completion_response=completion_response, @@ -307,6 +313,7 @@ class DBSpendUpdateWriter: ) verbose_proxy_logger.debug("Runs spend update on all tables") + return True except Exception: spend_log_error( "Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue " @@ -319,7 +326,55 @@ class DBSpendUpdateWriter: org_id, end_user_id, ) - return + return True + + async def _record_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None") -> bool: + if prisma_client is None or not _is_batch_cost_row(payload): + await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) + return True + return await self._claim_batch_cost_spend_log(payload=payload, prisma_client=prisma_client) + + async def _claim_batch_cost_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient") -> bool: + """Write the batch's cost row now, or learn that another retrieve already did. + + Every retrieve of one batch shares this row, so the insert that lands first owns + the charge and every later one finds the row and charges nothing (LIT-7048). Only + a row a successful retrieve wrote counts: a failed retrieve, or any request whose + client picked the batch id as its call id, cannot take the charge away. + """ + from litellm.repositories.table_repositories import SpendLogsRepository + + request_id: Final = payload["request_id"] + spend_logs: Final = SpendLogsRepository(prisma_client).table + try: + claimed: Final = await spend_logs.create_many( + data=[prisma_client.jsonify_object(payload)], # mutable-ok: prisma create_many takes a list + skip_duplicates=True, + ) + if claimed == 1: + return True + existing: Final = await spend_logs.find_unique( + where={"request_id": request_id} # mutable-ok: prisma where clause + ) + except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log + verbose_proxy_logger.warning( + "Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e + ) + await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) + return True + if ( + existing is not None + and existing.call_type == CallTypes.aretrieve_batch.value + and existing.status == "success" + ): + verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) + return False + verbose_proxy_logger.warning( + "Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own", + request_id, + getattr(existing, "call_type", None), + ) + return True async def _enqueue_tool_usage_transaction( self, @@ -940,7 +995,7 @@ class DBSpendUpdateWriter: from litellm.proxy.utils import enqueue_spend_logs, request_spend_log_flush await enqueue_spend_logs(prisma_client, (payload,)) - if payload.get("call_type") in IMMEDIATE_FLUSH_CALL_TYPES: + if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES: request_spend_log_flush() else: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 95e61fdd98c..f0c889a8cb4 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -34,7 +34,6 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import ( from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, get_request_model_access_groups, - get_spend_logs_id, ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( @@ -46,9 +45,7 @@ from litellm.types.utils import ( from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: - from prisma.types import LiteLLM_SpendLogsWhereUniqueInput - - from litellm.proxy.utils import PrismaClient, ProxyLogging + from litellm.proxy.utils import ProxyLogging _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset( { @@ -229,7 +226,6 @@ class _ProxyDBLogger(CustomLogger): ): from litellm.proxy.proxy_server import ( increment_spend_counters, - prisma_client, proxy_logging_obj, update_cache, ) @@ -257,15 +253,15 @@ class _ProxyDBLogger(CustomLogger): if ( isinstance(completion_response, LiteLLMBatch) and kwargs.get("call_type") == CallTypes.aretrieve_batch.value + and not batch_cost_is_final(completion_response) ): - batch_spend_log_id: Final = get_spend_logs_id( - CallTypes.aretrieve_batch.value, completion_response.model_dump(), kwargs + verbose_proxy_logger.debug( + "Cost tracking deferred for batch %s still in status %s", + completion_response.id, + completion_response.status, ) - if not await _batch_cost_is_trackable_now( - batch=completion_response, spend_log_id=batch_spend_log_id, prisma_client=prisma_client - ): - await _release_budget_reservation(budget_reservation=budget_reservation) - return + await _release_budget_reservation(budget_reservation=budget_reservation) + return user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) @@ -307,7 +303,7 @@ class _ProxyDBLogger(CustomLogger): call_type=call_type, ): ## UPDATE DATABASE - await _update_database_and_spend_counters( + charged: Final = await _update_database_and_spend_counters( proxy_logging_obj=proxy_logging_obj, increment_spend_counters=increment_spend_counters, user_api_key=user_api_key, @@ -324,6 +320,8 @@ class _ProxyDBLogger(CustomLogger): request_tags=tags, model_access_groups=model_access_groups, ) + if not charged: + return # update cache (fire-and-forget for backward compat: # cached object fields, soft budget alerts, etc.) @@ -509,42 +507,6 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: bucket[key] = value -async def _batch_cost_is_trackable_now( - batch: LiteLLMBatch, spend_log_id: str | None, prisma_client: "PrismaClient | None" -) -> bool: - """A batch is billed exactly once, from the first retrieve that sees it final. - - Every retrieve of one batch shares a single spend row (its id plus the batch cost - suffix), so a poll that lands before the output exists would write that row at $0 - and pin it there, and every retrieve after the first would add the cost to the - key, team, and user counters again. - """ - if not batch_cost_is_final(batch): - verbose_proxy_logger.debug("Cost tracking deferred for batch %s still in status %s", batch.id, batch.status) - return False - if prisma_client is None or spend_log_id is None: - return True - if not await _spend_log_already_recorded(prisma_client=prisma_client, request_id=spend_log_id): - return True - verbose_proxy_logger.debug( - "Cost tracking skipped for batch %s: spend row %s already recorded", batch.id, spend_log_id - ) - return False - - -async def _spend_log_already_recorded(prisma_client: "PrismaClient", request_id: str) -> bool: - from litellm.proxy.utils import spend_log_is_queued - - if await spend_log_is_queued(prisma_client, request_id): - return True - spend_log_row: Final[LiteLLM_SpendLogsWhereUniqueInput] = {"request_id": request_id} - try: - return await prisma_client.db.litellm_spendlogs.find_unique(where=spend_log_row) is not None - except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreadable DB must not drop the batch's only spend row - verbose_proxy_logger.warning("Could not check for an existing spend row %s, tracking anyway: %s", request_id, e) - return False - - def _is_unbilled_interaction_response(completion_response: object) -> bool: from litellm.interactions.background_cost_polling import missing_usage_is_expected from litellm.types.interactions import InteractionsAPIResponse @@ -636,9 +598,9 @@ async def _update_database_and_spend_counters( budget_reservation: dict | None, request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, -) -> None: +) -> bool: try: - await proxy_logging_obj.db_spend_update_writer.update_database( + charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, @@ -663,6 +625,9 @@ async def _update_database_and_spend_counters( "Failed to invalidate budget reservation counters after release failed" ) raise + if not charged: + await _release_budget_reservation(budget_reservation=budget_reservation) + return False try: await increment_spend_counters( @@ -688,6 +653,7 @@ async def _update_database_and_spend_counters( finally: budget_reservation["finalized"] = True raise + return True async def _release_budget_reservation(budget_reservation: dict | None) -> None: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f64b51bc6c6..accf7b720fb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6251,9 +6251,7 @@ def request_spend_log_flush() -> None: The Responses API hands the client an id it can chain from straight away, and that lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval. - A batch's cost row is what every other worker checks before charging the same batch - again, so it cannot wait either. Repeated requests coalesce into the monitor's next - pass, so the batching holds. + Repeated requests coalesce into the monitor's next pass, so the batching holds. """ PrismaClient.spend_log_flush_requested.set() @@ -6268,12 +6266,6 @@ async def _wait_for_spend_log_flush_request(interval: float) -> bool: return True -async def spend_log_is_queued(prisma_client: PrismaClient, request_id: str) -> bool: - """Whether a spend log with ``request_id`` is still waiting for the next flush.""" - async with prisma_client._spend_log_transactions_lock: - return any(row.get("request_id") == request_id for row in prisma_client.spend_log_transactions) - - async def dequeue_spend_logs(prisma_client: PrismaClient, limit: int) -> list[dict[str, object]]: """Take up to ``limit`` of the oldest queued spend logs off the queue. diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 8d4f68164b4..976a96f2db1 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1735,10 +1735,10 @@ def _retrieved_batch( endpoint="/v1/chat/completions", input_file_id="file-in", object="batch", - status=status, + status="validating", output_file_id=output_file_id, request_counts=counts, - ) + ).model_copy(update={"status": status}) class TestBatchCostIsFinal: @@ -1750,8 +1750,9 @@ class TestBatchCostIsFinal: def test_in_flight_batch_is_not_final(self, status): assert bu.batch_cost_is_final(_retrieved_batch(status)) is False - def test_completed_with_output_is_final(self): - assert bu.batch_cost_is_final(_retrieved_batch("completed", output_file_id="file-out")) is True + @pytest.mark.parametrize("status", ["completed", "complete"]) + def test_completed_with_output_is_final(self, status): + assert bu.batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True def test_completed_without_output_and_unknown_counts_is_not_final(self): assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False @@ -1764,9 +1765,10 @@ class TestBatchCostIsFinal: counts = BatchRequestCounts(total=2, completed=2, failed=0) assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False - def test_completed_without_output_and_every_line_failed_is_final(self): + @pytest.mark.parametrize("status", ["completed", "complete"]) + def test_completed_without_output_and_every_line_failed_is_final(self, status): counts = BatchRequestCounts(total=2, completed=0, failed=2) - assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is True + assert bu.batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) def test_other_terminal_statuses_are_final(self, status): diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 4583429bd11..174efaa4679 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -665,14 +665,14 @@ class TestRetrieveBatchPricesOnlyFinalBatches: endpoint="/v1/chat/completions", input_file_id="file-in", object="batch", - status=status, + status="validating", output_file_id=output_file_id, - ) + ).model_copy(update={"status": status}) @pytest.mark.asyncio @pytest.mark.parametrize( ("status", "output_file_id"), - [("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None)], + [("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None), ("complete", None)], ) async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_file_id) -> None: from litellm.litellm_core_utils import litellm_logging as logging_module @@ -681,8 +681,7 @@ class TestRetrieveBatchPricesOnlyFinalBatches: monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) batch = self._batch(status, output_file_id) - with contextlib.suppress(Exception): - await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) + await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) handle_completed_batch.assert_not_awaited() assert "response_cost" not in batch._hidden_params @@ -705,8 +704,7 @@ class TestRetrieveBatchPricesOnlyFinalBatches: monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) batch = self._batch("completed", "file-out") - with contextlib.suppress(Exception): - await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) + await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) handle_completed_batch.assert_awaited_once() assert batch._hidden_params["response_cost"] == 8e-06 diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index e1b2d151c6d..500a0e7bb06 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -7,6 +7,7 @@ import re from collections.abc import Callable from contextlib import asynccontextmanager from datetime import datetime, timezone +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -2934,16 +2935,14 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey @pytest.mark.asyncio @pytest.mark.parametrize( "call_type, expects_flush", - [("aresponses", True), ("responses", True), ("aretrieve_batch", True), ("acompletion", False)], + [("aresponses", True), ("responses", True), ("acompletion", False)], ) async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_workers_read_back( call_type: str, expects_flush: bool ): """ A `previous_response_id` chained straight off the previous turn reads the DB, so a - Responses row cannot sit in this worker's queue until the monitor's next poll. A - batch's cost row is what another worker checks before charging the same batch again - (LIT-7048), so it cannot wait either. + Responses row cannot sit in this worker's queue until the monitor's next poll. """ from litellm.proxy.utils import PrismaClient @@ -2961,6 +2960,127 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_worker PrismaClient.spend_log_flush_requested.clear() +def _batch_cost_payload() -> dict: + return { + **_minimal_spend_payload(), + "request_id": "batch_abc_batch_cost", + "call_type": "aretrieve_batch", + "status": "success", + } + + +def _spend_logs_prisma(inserted: int, existing: object) -> MagicMock: + prisma = _tool_usage_prisma() + prisma.jsonify_object = lambda data: dict(data) + prisma.db.litellm_spendlogs.create_many = AsyncMock(return_value=inserted) + prisma.db.litellm_spendlogs.find_unique = AsyncMock(return_value=existing) + return prisma + + +async def _update_database_with(db_writer: DBSpendUpdateWriter, prisma: MagicMock, payload: dict) -> bool: + with ( + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.disable_spend_logs", False + ), + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.prisma_client", prisma + ), + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" + ), + patch( # test-quality-ok: update_database imports the payload builder inside its body, no seam + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=payload, + ), + ): + charged = await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={"model": "gpt-5.6-luna", "call_type": "aretrieve_batch"}, + completion_response=None, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=0.25, + ) + await asyncio.sleep(0) + return charged + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("inserted", "existing", "charged"), + [ + (1, None, True), + (0, SimpleNamespace(call_type="aretrieve_batch", status="success"), False), + (0, SimpleNamespace(call_type="aretrieve_batch", status="failure"), True), + (0, SimpleNamespace(call_type="aembedding", status="success"), True), + (0, None, True), + ], + ids=[ + "first_retrieve_owns_the_row", + "another_retrieve_already_charged", + "failed_retrieve_holds_the_row", + "client_chosen_call_id_holds_the_row", + "row_gone_between_insert_and_lookup", + ], +) +async def test_update_database_charges_a_batch_only_from_the_retrieve_that_wrote_its_row( + inserted: int, existing: object, charged: bool +): + """ + Every retrieve of one batch shares one spend row, so the insert that lands first is + the charge and every later retrieve must leave the counters alone (LIT-7048). A row + written by anything but a successful retrieve, say a request whose client picked the + batch id as its call id, must not be able to take the charge away. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(inserted, existing) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is charged + + claimed_rows = prisma.db.litellm_spendlogs.create_many.await_args.kwargs + assert claimed_rows["skip_duplicates"] is True + assert [(row["request_id"], row["spend"]) for row in claimed_rows["data"]] == [("batch_abc_batch_cost", 0.25)] + assert prisma.spend_log_transactions == [] + assert db_writer._batch_database_updates.await_count == (1 if charged else 0) + + +@pytest.mark.asyncio +async def test_update_database_queues_a_batch_cost_row_it_could_not_claim(): + """An unreachable DB must not drop the batch's only spend row, nor its charge.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(0, None) + prisma.db.litellm_spendlogs.create_many = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True + + assert [row["request_id"] for row in prisma.spend_log_transactions] == ["batch_abc_batch_cost"] + assert db_writer._batch_database_updates.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [{**_batch_cost_payload(), "call_type": "acompletion"}, {**_batch_cost_payload(), "status": "failure"}], + ids=["not_a_batch_retrieve", "failed_batch_retrieve"], +) +async def test_update_database_queues_every_other_spend_row_for_the_next_flush(payload: dict): + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + + assert await _update_database_with(db_writer, prisma, payload) is True + + prisma.db.litellm_spendlogs.create_many.assert_not_called() + assert prisma.spend_log_transactions == [payload] + assert db_writer._batch_database_updates.await_count == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize( "injected_deployment, attributed", diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index b7037e8d621..2965f8b4006 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,3 +1,4 @@ +import asyncio from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -6,7 +7,6 @@ import pytest from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( - _batch_cost_is_trackable_now, _get_budget_reservation_from_metadata, _ProxyDBLogger, _should_track_cost_callback, @@ -783,110 +783,46 @@ def _retrieved_batch(status: str, output_file_id: str | None): ) -def _prisma_client_with(queued_request_ids: tuple[str, ...], stored_row: object) -> MagicMock: - import asyncio - - prisma_client = MagicMock() - prisma_client._spend_log_transactions_lock = asyncio.Lock() - prisma_client.spend_log_transactions = [{"request_id": request_id} for request_id in queued_request_ids] - prisma_client.db.litellm_spendlogs.find_unique = AsyncMock(return_value=stored_row) - return prisma_client - - @pytest.mark.asyncio @pytest.mark.parametrize( - ("status", "output_file_id", "spend_log_id", "prisma_client", "trackable"), + ("call_type", "status", "output_file_id", "row_claimed", "spend_written", "charged"), [ - ("in_progress", None, "batch_abc_batch_cost", None, False), - ("in_progress", None, "batch_abc_batch_cost", _prisma_client_with((), None), False), - ("completed", None, "batch_abc_batch_cost", _prisma_client_with((), None), False), - ("completed", "file-out", "batch_abc_batch_cost", None, True), - ("completed", "file-out", None, _prisma_client_with((), None), True), - ("completed", "file-out", "batch_abc_batch_cost", _prisma_client_with(("batch_abc_batch_cost",), None), False), - ( - "completed", - "file-out", - "batch_abc_batch_cost", - _prisma_client_with((), {"request_id": "batch_abc_batch_cost"}), - False, - ), - ("completed", "file-out", "batch_abc_batch_cost", _prisma_client_with((), None), True), - ("failed", None, "batch_abc_batch_cost", _prisma_client_with((), None), True), + ("aretrieve_batch", "in_progress", None, True, False, False), + ("aretrieve_batch", "completed", None, True, False, False), + ("aretrieve_batch", "completed", "file-out", False, True, False), + ("aretrieve_batch", "completed", "file-out", True, True, True), + ("aretrieve_batch", "failed", None, True, True, True), + ("acreate_batch", "validating", None, True, True, True), ], ids=[ - "in_progress_without_db", - "in_progress_never_consults_db", - "completed_without_output_yet", - "final_without_db", - "final_without_spend_log_id", - "final_row_queued_for_flush", - "final_row_already_stored", - "final_first_sighting", - "failed_first_sighting", + "retrieve_before_final", + "retrieve_completed_without_output_yet", + "retrieve_after_another_retrieve_charged", + "retrieve_first_final", + "retrieve_failed_batch", + "create_before_final", ], ) -async def test_batch_cost_is_trackable_now(status, output_file_id, spend_log_id, prisma_client, trackable): - """ - A batch is billed from the first retrieve that sees it final and never again: - a poll before that wrote the shared spend row at $0 and pinned it there, and - every completed retrieve after the first charged the key again (LIT-7048). - """ - assert ( - await _batch_cost_is_trackable_now( - batch=_retrieved_batch(status, output_file_id), spend_log_id=spend_log_id, prisma_client=prisma_client - ) - is trackable - ) - - -@pytest.mark.asyncio -async def test_batch_cost_is_trackable_now_when_the_spend_row_lookup_fails(): - """An unreadable spend log table must not drop the batch's only spend row.""" - prisma_client = _prisma_client_with((), None) - prisma_client.db.litellm_spendlogs.find_unique = AsyncMock(side_effect=RuntimeError("db unreachable")) - - assert ( - await _batch_cost_is_trackable_now( - batch=_retrieved_batch("completed", "file-out"), - spend_log_id="batch_abc_batch_cost", - prisma_client=prisma_client, - ) - is True - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("call_type", "status", "output_file_id", "stored_row", "charged"), - [ - ("aretrieve_batch", "in_progress", None, None, False), - ("aretrieve_batch", "completed", "file-out", {"request_id": "batch_abc_batch_cost"}, False), - ("aretrieve_batch", "completed", "file-out", None, True), - ("acreate_batch", "validating", None, None, True), - ], - ids=["retrieve_before_final", "retrieve_already_recorded", "retrieve_first_final", "create_before_final"], -) -async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs and whether the poll's reservation is handed back is the whole observable contract of the gate - call_type, status, output_file_id, stored_row, charged +async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs, whether the counters move, and whether the poll's reservation is handed back is the whole observable contract of the gate + call_type, status, output_file_id, row_claimed, spend_written, charged ): """ - Only retrieves are gated, since creating a batch is its own billable request. - A retrieve that writes nothing hands its budget reservation back instead. + A poll before the batch is final used to pin its shared spend row at $0, and every + completed retrieve after the first charged the key again (LIT-7048). Only retrieves + are gated, since creating a batch is its own billable request, and a retrieve that + charges nothing hands its budget reservation back instead. """ logger = _ProxyDBLogger() budget_reservation = None if charged else {"reserved_cost": 0.5, "entries": []} kwargs = _batch_retrieve_kwargs(call_type, reservation=budget_reservation) with ( - patch( # test-quality-ok: prisma_client is a proxy_server global the callback reads lazily, no seam - "litellm.proxy.proxy_server.prisma_client", _prisma_client_with((), stored_row) - ), patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock - ), + ) as mock_increment_spend_counters, patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock - ), + ) as mock_update_cache, patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam "litellm.proxy.proxy_server.proxy_logging_obj" ) as mock_proxy_logging, @@ -895,7 +831,7 @@ async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # ) as mock_release_budget_reservation, ): mock_proxy_logging.failed_tracking_alert = AsyncMock() - mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=row_claimed) mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() await logger._PROXY_track_cost_callback( @@ -904,13 +840,15 @@ async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # start_time=datetime.now(), end_time=datetime.now(), ) + await asyncio.sleep(0) mock_proxy_logging.failed_tracking_alert.assert_not_called() + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if spend_written else 0) + assert mock_increment_spend_counters.await_count == (1 if charged else 0) + assert mock_update_cache.await_count == (1 if charged else 0) if charged: - mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() mock_release_budget_reservation.assert_not_awaited() else: - mock_proxy_logging.db_spend_update_writer.update_database.assert_not_called() mock_release_budget_reservation.assert_awaited_once_with(budget_reservation=budget_reservation) diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index fc97d760226..a1eb88a7834 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -6,7 +6,6 @@ Symbols pinned here: - ``update_spend_logs_job`` - ``_monitor_spend_logs_queue`` - ``_raise_failed_update_spend_exception`` - - ``spend_log_is_queued`` """ from __future__ import annotations @@ -23,7 +22,6 @@ from litellm.proxy.utils import ( _monitor_spend_logs_queue, _raise_failed_update_spend_exception, drain_spend_logs_queue, - spend_log_is_queued, update_daily_tag_spend, update_spend, update_spend_logs_job, @@ -631,16 +629,3 @@ def test_raise_failed_update_spend_exception_raises_original_error() -> None: with pytest.raises(ValueError, match="specific"): asyncio.run(_runner()) - - -@pytest.mark.asyncio -async def test_spend_log_is_queued_matches_only_rows_awaiting_flush( - mock_prisma_client: Any, make_spend_log_row: Any -) -> None: - mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="batch_abc_batch_cost")] - - assert await spend_log_is_queued(mock_prisma_client, "batch_abc_batch_cost") is True - assert await spend_log_is_queued(mock_prisma_client, "batch_abc") is False - - mock_prisma_client.spend_log_transactions = [] - assert await spend_log_is_queued(mock_prisma_client, "batch_abc_batch_cost") is False From 15372967c6cd5085d5d3d9ebb30c5a158f3f8170 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:06:38 -0700 Subject: [PATCH 09/32] fix(azure_ai): read the azure_ai card for gpt-5 series reasoning effort gates Foundry deployments of gpt-6-astra reached through azure_ai used the bare OpenAI card for the reasoning_effort none gates, so temperature and top_p were refused while the azure_ai card says none is supported. AzureAIStudioConfig now dispatches gpt-5 series params through AzureAIGPT5Config, which looks capabilities up under the azure_ai/ prefix the way the azure route does Also carries the search_context_cost_per_query block azure/gpt-6-astra has, adds a flex service tier cost test that fails at the merge base, and keeps the wildcard test from stripping azure_ai/gpt-6-astra out of the provider set --- litellm/llms/azure_ai/chat/transformation.py | 37 ++++++++++++++++++- ...odel_prices_and_context_window_backup.json | 5 +++ model_prices_and_context_window.json | 5 +++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 15 ++++++++ .../chat/test_azure_ai_transformation.py | 22 +++++++++++ .../proxy/auth/test_model_checks.py | 6 ++- 6 files changed, 87 insertions(+), 3 deletions(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index f2d405e9a17..05abd5882c6 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj +from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error from litellm.llms.openai.openai import OpenAIConfig from litellm.llms.xai.chat.transformation import XAIChatConfig @@ -42,12 +43,25 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ( ) +class AzureAIGPT5Config(OpenAIGPT5Config): + @classmethod + def _model_map_lookup_name(cls, model: str) -> str: + return model if model.startswith("azure_ai/") else f"azure_ai/{model}" + + +azureAIGPT5Config: Final = AzureAIGPT5Config() + + class AzureAIStudioConfig(OpenAIConfig): def get_supported_openai_params(self, model: str) -> list: model_supports_tool_choice = True # azure ai supports this by default if not supports_tool_choice(model=f"azure_ai/{model}"): model_supports_tool_choice = False - supported_params = super().get_supported_openai_params(model) + supported_params = ( + azureAIGPT5Config.get_supported_openai_params(model) + if azureAIGPT5Config.is_model_gpt_5_model(model) + else super().get_supported_openai_params(model) + ) if not model_supports_tool_choice: filtered_supported_params: Final = [] for param in supported_params: @@ -61,6 +75,27 @@ class AzureAIStudioConfig(OpenAIConfig): return supported_params + def map_openai_params( + self, + non_default_params: dict, # mutable-ok: OpenAIConfig.map_openai_params signature + optional_params: dict, # mutable-ok: OpenAIConfig.map_openai_params signature + model: str, + drop_params: bool, + ) -> dict: # mutable-ok: OpenAIConfig.map_openai_params signature + if not azureAIGPT5Config.is_model_gpt_5_model(model): + return super().map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=drop_params, + ) + return azureAIGPT5Config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=drop_params, + ) + def _supports_stop_reason(self, model: str) -> bool: """ Check if the model supports stop tokens. diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ac7407c2608..5b4652d38c7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3499,6 +3499,11 @@ "mode": "chat", "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "source": "https://ai.azure.com/catalog/models/gpt-6-astra", "supported_endpoints": [ "/v1/chat/completions", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ac7407c2608..5b4652d38c7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3499,6 +3499,11 @@ "mode": "chat", "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "source": "https://ai.azure.com/catalog/models/gpt-6-astra", "supported_endpoints": [ "/v1/chat/completions", diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 73f1a19d85c..caf97ba791d 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2060,6 +2060,21 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( assert completion_cost == pytest.approx(zone_multiplier * output_multiplier * completion_tokens * 5e-5) +def test_generic_cost_per_token_azure_ai_gpt_6_astra_flex_bills_the_standard_rate(_local_model_cost_map): + """Foundry sells gpt-6-astra on Standard Global only, so a flex service_tier bills the standard rate. + The bare OpenAI card the azure_ai route fell back to before this entry existed carries flex prices + at half rate (LIT-7081).""" + usage = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100) + + standard = generic_cost_per_token(model="azure_ai/gpt-6-astra", usage=usage, custom_llm_provider="azure_ai") + flex = generic_cost_per_token( + model="azure_ai/gpt-6-astra", usage=usage, custom_llm_provider="azure_ai", service_tier="flex" + ) + + assert flex == standard + assert standard == pytest.approx((1000 * 1e-05, 100 * 5e-05)) + + @pytest.mark.parametrize( "model,expected_none,expected_xhigh,expected_minimal", [ diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 33fbb4e8fc7..25eca3b37ad 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -3,6 +3,8 @@ from unittest.mock import MagicMock, patch import pytest +import litellm +from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map from litellm.llms.azure_ai.azure_model_router.transformation import ( AzureModelRouterConfig, ) @@ -138,6 +140,26 @@ def test_azure_ai_validate_environment_with_azure_ad_token(): assert headers["Content-Type"] == "application/json" +@pytest.fixture +def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url)) + + +def test_foundry_gpt_6_astra_keeps_sampling_params_when_reasoning_effort_is_none(_local_model_cost_map): + """A Foundry deployment reached through azure_ai reads the azure_ai/ card, where gpt-6-astra supports + reasoning_effort none, so temperature and top_p ride along; the bare OpenAI card says none is + unsupported and the route used to refuse temperature and drop top_p (LIT-7081).""" + optional_params = AzureAIStudioConfig().map_openai_params( + non_default_params={"reasoning_effort": "none", "temperature": 0.2, "top_p": 0.9}, + optional_params={}, + model="gpt-6-astra", + drop_params=False, + ) + + assert optional_params == {"reasoning_effort": "none", "temperature": 0.2, "top_p": 0.9} + + def test_azure_ai_grok_stop_parameter_handling(): """ Test that Grok models properly handle stop parameter filtering in Azure AI Studio. diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index eb48f70d5da..56dbcca61f3 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -868,12 +868,14 @@ def test_azure_ai_wildcard_lists_the_foundry_gpt_6_astra_entry(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") foundry_key = "azure_ai/gpt-6-astra" local_entry = litellm.get_model_cost_map(url="")[foundry_key] + registered_before = foundry_key in litellm.azure_ai_models try: litellm.add_known_models(model_cost_map={foundry_key: local_entry}) assert foundry_key in get_known_models_from_wildcard("azure_ai/*") finally: - litellm.azure_ai_models.discard(foundry_key) - litellm.add_known_models(model_cost_map={}) + if not registered_before: + litellm.azure_ai_models.discard(foundry_key) + litellm.add_known_models(model_cost_map={}) def test_get_complete_model_list_drops_no_default_models_sentinel(): From a17fcecf7092d0333afed4168d1954a1e9675d18 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:21:34 -0700 Subject: [PATCH 10/32] refactor(azure_ai): type the Foundry param mapping override and drop test docstrings The AzureAIStudioConfig.map_openai_params override now carries dict[str, object] annotations instead of bare dict, and the docstrings added to the new tests go away since the test names already say what they cover. No behavior change --- litellm/llms/azure_ai/chat/transformation.py | 6 +++--- .../llm_cost_calc/test_llm_cost_calc_utils.py | 3 --- .../llms/azure_ai/chat/test_azure_ai_transformation.py | 3 --- tests/test_litellm/proxy/auth/test_model_checks.py | 3 --- .../router_utils/test_reasoning_effort_capability.py | 3 +-- 5 files changed, 4 insertions(+), 14 deletions(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 05abd5882c6..7c9a26c3f07 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -77,11 +77,11 @@ class AzureAIStudioConfig(OpenAIConfig): def map_openai_params( self, - non_default_params: dict, # mutable-ok: OpenAIConfig.map_openai_params signature - optional_params: dict, # mutable-ok: OpenAIConfig.map_openai_params signature + non_default_params: dict[str, object], # mutable-ok: OpenAIConfig.map_openai_params signature + optional_params: dict[str, object], # mutable-ok: OpenAIConfig.map_openai_params signature model: str, drop_params: bool, - ) -> dict: # mutable-ok: OpenAIConfig.map_openai_params signature + ) -> dict[str, object]: # mutable-ok: OpenAIConfig.map_openai_params signature if not azureAIGPT5Config.is_model_gpt_5_model(model): return super().map_openai_params( non_default_params=non_default_params, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index caf97ba791d..40abb5bfca3 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2061,9 +2061,6 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( def test_generic_cost_per_token_azure_ai_gpt_6_astra_flex_bills_the_standard_rate(_local_model_cost_map): - """Foundry sells gpt-6-astra on Standard Global only, so a flex service_tier bills the standard rate. - The bare OpenAI card the azure_ai route fell back to before this entry existed carries flex prices - at half rate (LIT-7081).""" usage = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100) standard = generic_cost_per_token(model="azure_ai/gpt-6-astra", usage=usage, custom_llm_provider="azure_ai") diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 25eca3b37ad..5ff0b729449 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -147,9 +147,6 @@ def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: def test_foundry_gpt_6_astra_keeps_sampling_params_when_reasoning_effort_is_none(_local_model_cost_map): - """A Foundry deployment reached through azure_ai reads the azure_ai/ card, where gpt-6-astra supports - reasoning_effort none, so temperature and top_p ride along; the bare OpenAI card says none is - unsupported and the route used to refuse temperature and drop top_p (LIT-7081).""" optional_params = AzureAIStudioConfig().map_openai_params( non_default_params={"reasoning_effort": "none", "temperature": 0.2, "top_p": 0.9}, optional_params={}, diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 56dbcca61f3..36bfc4c5dd3 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -859,9 +859,6 @@ def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): def test_azure_ai_wildcard_lists_the_foundry_gpt_6_astra_entry(monkeypatch): - """A Foundry (azure_ai) deployment of gpt-6-astra only shows up under an azure_ai/* wildcard - when the cost map carries its own azure_ai/ entry; the azure/ entry from the OpenAI-on-Azure - price sheet never reaches the Foundry provider list (LIT-7081).""" import litellm from litellm.proxy.auth.model_checks import get_known_models_from_wildcard diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index b7499d1c975..3e1f26b6e1c 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -400,8 +400,7 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: def test_a_foundry_deployment_also_advertises_none(self, local_model_cost_map, model, custom_llm_provider): """Microsoft Foundry serves the same model but its API accepts reasoning_effort none (verified live: 200 with zero reasoning tokens, and it unlocks temperature), which - OpenAI's rejects, so an Azure deployment offers none on top of low through max, whether - it is reached through the azure route or the azure_ai (Foundry) route.""" + OpenAI's rejects, so an Azure deployment offers none on top of low through max.""" from litellm.utils import _get_model_info_helper model_info = dict(_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)) From e8f311429ea9afa09195bed21f078bbd50dd791e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:42:15 -0700 Subject: [PATCH 11/32] fix(cost-map): stop advertising reasoning_effort max on azure_ai/gpt-6-astra Foundry rejects reasoning_effort max on the gpt-6-astra deployment with a 400 that names none, low, medium, high, and xhigh as the supported values, so the card no longer lists max. The request path never gated max (only xhigh is opt-in), so this only changes /model_group/info and router capability gating. The azure/ twin stays as is because it was not verified on an Azure OpenAI host --- .../model_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- .../test_reasoning_effort_capability.py | 14 +++++++++++++- 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5b4652d38c7..b659c3b65e5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3518,7 +3518,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5b4652d38c7..b659c3b65e5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3518,7 +3518,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index 3e1f26b6e1c..fa3a6dcd95a 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -394,7 +394,6 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: [ ("azure/gpt-6-astra", "azure"), ("azure/us/gpt-6-astra", "azure"), - ("azure_ai/gpt-6-astra", "azure_ai"), ], ) def test_a_foundry_deployment_also_advertises_none(self, local_model_cost_map, model, custom_llm_provider): @@ -413,3 +412,16 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: "xhigh", "max", ) + + def test_a_foundry_azure_ai_deployment_advertises_none_but_not_max(self, local_model_cost_map): + from litellm.utils import _get_model_info_helper + + model_info = dict(_get_model_info_helper(model="azure_ai/gpt-6-astra", custom_llm_provider="azure_ai")) + + assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( + "none", + "low", + "medium", + "high", + "xhigh", + ) From 061c25b5cac15212ded745c51e3299aa4d43f056 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 21:25:03 -0700 Subject: [PATCH 12/32] fix(spend): let a batch's charge survive an older proxy's $0 poll row A proxy running the old code wrote _batch_cost at $0 every time it polled a batch that was still running, so after an upgrade the claim found that row and read it as proof the batch had already been charged. Only a row that recorded a charge counts now, which leaves those $0 rows, and any row a client planted under the batch id, to be charged over disable_spend_logs skipped the claim entirely, so under that setting every retrieve of a finished batch charged again. The claim now runs either way and writes the one row per batch that makes the charge exactly once, while the per-request logs stay off --- litellm/proxy/db/db_spend_update_writer.py | 42 +++++++------ .../proxy/db/test_db_spend_update_writer.py | 60 ++++++++++++++++--- 2 files changed, 78 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 3fad351224b..48312c025dd 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -271,9 +271,12 @@ class DBSpendUpdateWriter: if team_id is not None and team_id != "": payload["team_id"] = team_id + if not await self._record_spend_log( + payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs + ): + return False + if disable_spend_logs is False: - if not await self._record_spend_log(payload=payload, prisma_client=prisma_client): - return False await self._enqueue_tool_usage_transaction( payload=payload, completion_response=completion_response, @@ -328,19 +331,23 @@ class DBSpendUpdateWriter: ) return True - async def _record_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None") -> bool: - if prisma_client is None or not _is_batch_cost_row(payload): + async def _record_spend_log( + self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None", disable_spend_logs: bool + ) -> bool: + if prisma_client is not None and _is_batch_cost_row(payload): + return await self._claim_batch_cost_spend_log(payload=payload, prisma_client=prisma_client) + if disable_spend_logs is False: await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) - return True - return await self._claim_batch_cost_spend_log(payload=payload, prisma_client=prisma_client) + return True async def _claim_batch_cost_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient") -> bool: """Write the batch's cost row now, or learn that another retrieve already did. Every retrieve of one batch shares this row, so the insert that lands first owns the charge and every later one finds the row and charges nothing (LIT-7048). Only - a row a successful retrieve wrote counts: a failed retrieve, or any request whose - client picked the batch id as its call id, cannot take the charge away. + a row that recorded a charge counts: a failed retrieve, a request whose client + picked the batch id as its call id, and the $0 row an older proxy left behind + while the batch was still running all leave the charge to be made. """ from litellm.repositories.table_repositories import SpendLogsRepository @@ -362,17 +369,18 @@ class DBSpendUpdateWriter: ) await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) return True - if ( - existing is not None - and existing.call_type == CallTypes.aretrieve_batch.value - and existing.status == "success" - ): + if existing is None or existing.call_type != CallTypes.aretrieve_batch.value or existing.status != "success": + verbose_proxy_logger.warning( + "Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own", + request_id, + getattr(existing, "call_type", None), + ) + return True + if existing.spend > 0: verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) return False - verbose_proxy_logger.warning( - "Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own", - request_id, - getattr(existing, "call_type", None), + verbose_proxy_logger.debug( + "Spend row %s charged nothing for this batch, so this retrieve charges it", request_id ) return True diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 500a0e7bb06..41f2be08545 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2977,10 +2977,12 @@ def _spend_logs_prisma(inserted: int, existing: object) -> MagicMock: return prisma -async def _update_database_with(db_writer: DBSpendUpdateWriter, prisma: MagicMock, payload: dict) -> bool: +async def _update_database_with( + db_writer: DBSpendUpdateWriter, prisma: MagicMock, payload: dict, disable_spend_logs: bool = False +) -> bool: with ( patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam - "litellm.proxy.proxy_server.disable_spend_logs", False + "litellm.proxy.proxy_server.disable_spend_logs", disable_spend_logs ), patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam "litellm.proxy.proxy_server.prisma_client", prisma @@ -3014,14 +3016,16 @@ async def _update_database_with(db_writer: DBSpendUpdateWriter, prisma: MagicMoc ("inserted", "existing", "charged"), [ (1, None, True), - (0, SimpleNamespace(call_type="aretrieve_batch", status="success"), False), - (0, SimpleNamespace(call_type="aretrieve_batch", status="failure"), True), - (0, SimpleNamespace(call_type="aembedding", status="success"), True), + (0, SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25), False), + (0, SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0), True), + (0, SimpleNamespace(call_type="aretrieve_batch", status="failure", spend=0.0), True), + (0, SimpleNamespace(call_type="aembedding", status="success", spend=0.25), True), (0, None, True), ], ids=[ "first_retrieve_owns_the_row", "another_retrieve_already_charged", + "an_older_proxy_left_a_zero_row_while_the_batch_ran", "failed_retrieve_holds_the_row", "client_chosen_call_id_holds_the_row", "row_gone_between_insert_and_lookup", @@ -3033,8 +3037,9 @@ async def test_update_database_charges_a_batch_only_from_the_retrieve_that_wrote """ Every retrieve of one batch shares one spend row, so the insert that lands first is the charge and every later retrieve must leave the counters alone (LIT-7048). A row - written by anything but a successful retrieve, say a request whose client picked the - batch id as its call id, must not be able to take the charge away. + that recorded no charge must not be able to take the charge away: neither one a + client planted under the batch id, nor the $0 row a pre-upgrade proxy wrote every + time it polled the batch while it was still running. """ db_writer = DBSpendUpdateWriter() db_writer._batch_database_updates = AsyncMock() @@ -3049,6 +3054,47 @@ async def test_update_database_charges_a_batch_only_from_the_retrieve_that_wrote assert db_writer._batch_database_updates.await_count == (1 if charged else 0) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("inserted", "existing", "charged"), + [ + (1, None, True), + (0, SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25), False), + ], + ids=["first_retrieve_owns_the_row", "another_retrieve_already_charged"], +) +async def test_update_database_charges_a_batch_once_even_with_spend_logs_disabled( + inserted: int, existing: object, charged: bool +): + """ + disable_spend_logs drops the per-request logs, not the batch's charge, so the one row + that makes a batch chargeable exactly once is still written and still read back. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(inserted, existing) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload(), True) is charged + + assert prisma.db.litellm_spendlogs.create_many.await_count == 1 + assert db_writer._batch_database_updates.await_count == (1 if charged else 0) + + +@pytest.mark.asyncio +async def test_update_database_writes_no_ordinary_spend_row_with_spend_logs_disabled(): + """The batch carve-out above stays a carve-out: every other row still goes unwritten.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + payload = {**_batch_cost_payload(), "call_type": "acompletion"} + + assert await _update_database_with(db_writer, prisma, payload, True) is True + + prisma.db.litellm_spendlogs.create_many.assert_not_called() + assert prisma.spend_log_transactions == [] + assert db_writer._batch_database_updates.await_count == 1 + + @pytest.mark.asyncio async def test_update_database_queues_a_batch_cost_row_it_could_not_claim(): """An unreachable DB must not drop the batch's only spend row, nor its charge.""" From 2e2fce5e583f31889dd4a75bd364cf0c2ba3cbe3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:25:13 -0700 Subject: [PATCH 13/32] fix(router): skip the refusing deployment when retrying a non-transient error BadRequestErrorRetries and ContentPolicyViolationErrorRetries did let a retry happen, but the retry re-picked the deployment that had just refused, since a 400 never puts a deployment in cooldown. On a weighted model group the caller got the same 400 back after every configured retry, and the existing 401/403 "retry on another deployment" rule broke the same way A retry after a non-transient status now carries the deployments that already answered this request in the per-request exclusion list weighted failover already honors, so the next attempt lands on a sibling. Single-deployment groups still retry in place, and 408/429/5xx retries are untouched Adds live e2e coverage for reliability.retry.context_window.succeeds_within_retries and renames the two litellm.utils deployment filters that are now called from outside the module --- basedpyright-code-budget.json | 4 +- litellm/router.py | 45 +++++++++- litellm/utils.py | 4 +- tests/e2e/models.py | 1 + tests/e2e/router/reliability_support.py | 19 +++- .../router/test_reliability_retries_e2e.py | 88 +++++++++++++------ tests/test_litellm/test_router.py | 82 +++++++++++++++++ .../test_router_order_fallback.py | 18 ++-- .../test_router_weighted_failover.py | 16 ++-- 9 files changed, 223 insertions(+), 54 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 0b0a61192e6..57ca267e504 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1808 + "limit": 1804 }, "reportRedeclaration": { "limit": 8 @@ -135,7 +135,7 @@ "limit": 21 }, "reportUnusedFunction": { - "limit": 138 + "limit": 136 }, "reportUnusedImport": { "limit": 542 diff --git a/litellm/router.py b/litellm/router.py index 6d6efad9f42..f72f41313a0 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -401,6 +401,7 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_EXCLUDED_DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]: @@ -7458,6 +7459,28 @@ class Router: Context_Policy_Fallbacks={content_policy_fallbacks}", ) + @staticmethod + def _deployment_ids_to_skip_on_retry( + exception: Exception, + already_skipped: object, + healthy_deployments: list[dict], # mutable-ok: matches the routing filters' list contract + ) -> tuple[str, ...]: + failed_deployment_id: Final[str | None] = getattr(exception, "failed_deployment_id", None) + status_code: Final = getattr(exception, "status_code", None) + if not failed_deployment_id or status_code is None: + return () + if litellm._should_retry(status_code): # pyright: ignore[reportPrivateUsage] # as in should_retry_this_error + return () + already_skipped_ids: Final = _EXCLUDED_DEPLOYMENT_IDS_ADAPTER.validate_python(already_skipped or ()) + skipped: Final = frozenset((*already_skipped_ids, failed_deployment_id)) + same_order_candidates: Final = litellm.utils.get_order_filtered_deployments(healthy_deployments) + if not litellm.utils.get_excluded_filtered_deployments(same_order_candidates, excluded_deployment_ids=skipped): + return () + verbose_router_logger.debug( + "Retry skips deployments that already answered %s to this request: %s", status_code, sorted(skipped) + ) + return tuple(sorted(skipped)) + @tracer.wrap() async def async_function_with_retries(self, *args, **kwargs): verbose_router_logger.debug("Inside async function with retries.") @@ -7553,6 +7576,13 @@ class Router: ## LOGGING if num_retries > 0: kwargs = self.log_retry(kwargs=kwargs, e=original_exception) + skipped_deployment_ids: Final = self._deployment_ids_to_skip_on_retry( + exception=original_exception, + already_skipped=kwargs.get("_excluded_deployment_ids"), + healthy_deployments=_healthy_deployments, + ) + if skipped_deployment_ids: + kwargs["_excluded_deployment_ids"] = skipped_deployment_ids else: raise @@ -7622,6 +7652,13 @@ class Router: except Exception: raise e + retry_skipped_deployment_ids = self._deployment_ids_to_skip_on_retry( + exception=e, + already_skipped=kwargs.get("_excluded_deployment_ids"), + healthy_deployments=_healthy_deployments, + ) + if retry_skipped_deployment_ids: + kwargs["_excluded_deployment_ids"] = retry_skipped_deployment_ids _timeout = self._time_to_sleep_before_retry( e=e, remaining_retries=remaining_retries, @@ -12452,7 +12489,7 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - healthy_deployments = litellm.utils._get_order_filtered_deployments( + healthy_deployments = litellm.utils.get_order_filtered_deployments( cast(list[dict], healthy_deployments), target_order=_target_order ) @@ -12460,7 +12497,7 @@ class Router: ## this request via weighted-failover. Always honored, regardless of the ## router-level flag, so a stale exclusion key on kwargs cannot escape. _excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None) - healthy_deployments = litellm.utils._get_excluded_filtered_deployments( + healthy_deployments = litellm.utils.get_excluded_filtered_deployments( cast(list[dict], healthy_deployments), excluded_deployment_ids=_excluded_deployment_ids, ) @@ -13357,7 +13394,7 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - healthy_deployments = litellm.utils._get_order_filtered_deployments( + healthy_deployments = litellm.utils.get_order_filtered_deployments( healthy_deployments, target_order=_target_order ) @@ -13365,7 +13402,7 @@ class Router: ## this request via weighted-failover. See async counterpart in ## async_get_healthy_deployments for details. _excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None) - healthy_deployments = litellm.utils._get_excluded_filtered_deployments( + healthy_deployments = litellm.utils.get_excluded_filtered_deployments( healthy_deployments, excluded_deployment_ids=_excluded_deployment_ids, ) diff --git a/litellm/utils.py b/litellm/utils.py index 52c1859b525..238225eff99 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4889,7 +4889,7 @@ def _get_deployment_order(deployment: dict | Any) -> int | None: return order -def _get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list: +def get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list: if target_order is not None: return [d for d in healthy_deployments if _get_deployment_order(d) == target_order] @@ -4908,7 +4908,7 @@ def _get_order_filtered_deployments(healthy_deployments: list[dict], target_orde return healthy_deployments -def _get_excluded_filtered_deployments( +def get_excluded_filtered_deployments( healthy_deployments: list[dict], excluded_deployment_ids: Iterable[str] | None = None, ) -> list: diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 2c6c0e9bbd4..c01687d0b31 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -291,6 +291,7 @@ class RouterSettingsOverride(BaseModel): context_window_fallbacks: list[dict[str, list[str]]] | None = None content_policy_fallbacks: list[dict[str, list[str]]] | None = None num_retries: int | None = None + model_group_retry_policy: dict[str, dict[str, int]] | None = None enable_tag_filtering: bool | None = None diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index 5822058003c..1efcb1a045b 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -73,10 +73,25 @@ def create_always_timing_out_deployment(proxy: ProxyClient, name: str) -> str: ) +def create_always_picked_small_context_deployment(proxy: ProxyClient, name: str) -> str: + """The always-picked half of a retry pair on the smallest-context model OpenAI + still serves: it holds all of the model group's shuffle weight, so an oversized + prompt opens on it and earns a real context-window refusal, which never benches + a deployment, so only the retry itself can steer the request off it.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody(model=SMALL_CONTEXT_MODEL, api_key=REAL_KEY, weight=1), + model_info=ModelInfoBody(), + ) + ) + + def create_zero_weight_backup_deployment(proxy: ProxyClient, name: str) -> str: """The other half of a retry pair: healthy, but weight 0, so the weighted shuffle - never opens on it. It is reachable only once its sibling is benched and the - weighted pick falls through to a uniform one over what is left.""" + never opens on it. It is reachable only once its sibling is out of the running, + benched by a cooldown or skipped by the retry, and the weighted pick falls through + to a uniform one over what is left.""" return proxy.register_model( ModelNewBody( model_name=name, diff --git a/tests/e2e/router/test_reliability_retries_e2e.py b/tests/e2e/router/test_reliability_retries_e2e.py index 5441412935c..da45cb46a46 100644 --- a/tests/e2e/router/test_reliability_retries_e2e.py +++ b/tests/e2e/router/test_reliability_retries_e2e.py @@ -1,13 +1,17 @@ """Live e2e: a request that fails on its first deployment is retried inside its own model group and still comes back a completion. -The model group is a pair: an always-timing-out deployment that holds all of the -group's shuffle weight, and a healthy backup at weight 0. The weighted pick always -opens on the timing-out one, its first Timeout benches it (an -`allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`), and the retry falls -through to the only deployment left. So the customer sees a completion and the -proxy reports that it took a retry to get there, with no random first pick in the -middle of it. +Each model group is a pair: a deployment that always refuses and holds all of the +group's shuffle weight, plus a healthy backup at weight 0. The weighted pick always +opens on the refusing one, so the customer sees a completion only if the retry +lands on the backup, and the proxy reports that it took a retry to get there, with +no random first pick in the middle of it. + +The timeout pair relies on cooldown: the first Timeout benches the timing-out +deployment (an `allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`) and the +retry falls through to the only deployment left. The context-window pair cannot: +a 400 never benches a deployment, so the retry policy's `BadRequestErrorRetries` +has to steer the retry off the deployment that just refused the prompt. """ from __future__ import annotations @@ -16,20 +20,48 @@ import pytest from complexity_router_client import ComplexityRouterClient from e2e_config import unique_marker +from e2e_http import StreamingResponse from lifecycle import ResourceManager from models import RouterSettingsOverride from reliability_support import ( chat_override, completion_tokens_of, content_of, + create_always_picked_small_context_deployment, create_always_timing_out_deployment, create_zero_weight_backup_deployment, finish_reason_of, + oversized_prompt, ) pytestmark = pytest.mark.e2e +def assert_retry_landed_on_backup(resp: StreamingResponse) -> None: + assert resp.status_code == 200, ( + f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}" + ) + + attempted = resp.headers.get("x-litellm-attempted-retries") + assert attempted is not None, "response is missing the x-litellm-attempted-retries header" + assert int(attempted) >= 1, ( + f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never " + "opened on the refusing deployment, so this proves nothing about retries" + ) + + content = content_of(resp) + finish_reason = finish_reason_of(resp) + completion_tokens = completion_tokens_of(resp) or 0 + assert isinstance(content, str), ( + f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})" + ) + assert content or (finish_reason == "length" and completion_tokens > 0), ( + f"the retry returned empty content with finish_reason={finish_reason!r}, " + f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget " + f"was spent on non-visible reasoning (body={resp.body[:300]})" + ) + + class TestReliabilityRetries: @pytest.mark.covers("reliability.retry.timeout.succeeds_within_retries") def test_timeout_on_first_deployment_succeeds_on_retry( @@ -49,25 +81,27 @@ class TestReliabilityRetries: override=RouterSettingsOverride(num_retries=2), ) - assert resp.status_code == 200, ( - f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}" + assert_retry_landed_on_backup(resp) + + @pytest.mark.covers("reliability.retry.context_window.succeeds_within_retries") + def test_context_window_refusal_on_first_deployment_succeeds_on_retry( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + group = f"reliability-retry-{unique_marker()}" + small_context = create_always_picked_small_context_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(small_context)) + backup = create_zero_weight_backup_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(backup)) + + resp = chat_override( + client.proxy, + scoped_key, + group, + oversized_prompt(unique_marker()), + override=RouterSettingsOverride( + num_retries=2, + model_group_retry_policy={group: {"BadRequestErrorRetries": 2}}, + ), ) - attempted = resp.headers.get("x-litellm-attempted-retries") - assert attempted is not None, "response is missing the x-litellm-attempted-retries header" - assert int(attempted) >= 1, ( - f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never " - "opened on the timing-out deployment, so this proves nothing about retries" - ) - - content = content_of(resp) - finish_reason = finish_reason_of(resp) - completion_tokens = completion_tokens_of(resp) or 0 - assert isinstance(content, str), ( - f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})" - ) - assert content or (finish_reason == "length" and completion_tokens > 0), ( - f"the retry returned empty content with finish_reason={finish_reason!r}, " - f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget " - f"was spent on non-visible reasoning (body={resp.body[:300]})" - ) + assert_retry_landed_on_backup(resp) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 7c044310e14..12c1516837a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -13116,6 +13116,7 @@ async def test_prompt_management_factory_marks_injection_for_every_deployment(mo ({"DefaultRetries": 0}, 502, litellm.BadGatewayError, 1), ({"DefaultRetries": 0, "ServiceUnavailableErrorRetries": 1}, 503, litellm.ServiceUnavailableError, 2), ({"ServiceUnavailableErrorRetries": 0}, 502, litellm.BadGatewayError, 3), + ({"BadRequestErrorRetries": 2}, 400, litellm.BadRequestError, 3), ], ) async def test_router_retry_policy_controls_upstream_attempt_count( @@ -13152,6 +13153,87 @@ async def test_router_retry_policy_controls_upstream_attempt_count( assert upstream.call_count == expected_upstream_calls +@pytest.mark.asyncio +@pytest.mark.parametrize( + "retry_policy,upstream_error", + [ + ( + {"BadRequestErrorRetries": 2}, + { + "message": "This model's maximum context length is 16385 tokens", + "type": "invalid_request_error", + "code": "context_length_exceeded", + }, + ), + ( + {"ContentPolicyViolationErrorRetries": 2}, + { + "message": "Your request was rejected as a result of our safety system", + "type": "invalid_request_error", + "code": "content_policy_violation", + }, + ), + ], +) +async def test_router_retry_policy_400_retries_on_sibling_deployment( + monkeypatch: pytest.MonkeyPatch, retry_policy, upstream_error +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-fake", + "api_base": "https://rejecting.local/v1", + "weight": 1, + }, + "model_info": {"id": "rejecting"}, + }, + { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-fake", + "api_base": "https://accepting.local/v1", + "weight": 0, + }, + "model_info": {"id": "accepting"}, + }, + ], + num_retries=2, + retry_policy=retry_policy, + disable_cooldowns=True, + ) + + with respx.mock(assert_all_called=True) as respx_mock: + rejecting = respx_mock.post("https://rejecting.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": upstream_error}) + ) + accepting = respx_mock.post("https://accepting.local/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-lit-7036", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.6", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "hi back"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + }, + ) + ) + response = await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) + + assert rejecting.call_count == 1 + assert accepting.call_count == 1 + assert response.choices[0].message.content == "hi back" + assert response._hidden_params["additional_headers"]["x-litellm-attempted-retries"] == 1 + + def _make_failure_logging_obj(): return LiteLLMLogging( model="gpt-5.6", diff --git a/tests/test_litellm/test_router_order_fallback.py b/tests/test_litellm/test_router_order_fallback.py index fde870e5abe..93895bbde08 100644 --- a/tests/test_litellm/test_router_order_fallback.py +++ b/tests/test_litellm/test_router_order_fallback.py @@ -18,10 +18,10 @@ from litellm import Router from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.prompt_caching_cache import PromptCachingCache from litellm.types.router import RouterRateLimitError -from litellm.utils import _get_deployment_order, _get_order_filtered_deployments +from litellm.utils import _get_deployment_order, get_order_filtered_deployments # --------------------------------------------------------------------------- -# Unit tests for _get_order_filtered_deployments +# Unit tests for get_order_filtered_deployments # --------------------------------------------------------------------------- @@ -42,7 +42,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(2, "b"), self._make_deployment(1, "c"), ] - result = _get_order_filtered_deployments(deps) + result = get_order_filtered_deployments(deps) assert len(result) == 2 assert all(d["model_info"]["id"] in ("a", "c") for d in result) @@ -52,7 +52,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(2, "b"), self._make_deployment(3, "c"), ] - result = _get_order_filtered_deployments(deps, target_order=2) + result = get_order_filtered_deployments(deps, target_order=2) assert len(result) == 1 assert result[0]["model_info"]["id"] == "b" @@ -61,7 +61,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(1, "a"), self._make_deployment(2, "b"), ] - result = _get_order_filtered_deployments(deps, target_order=99) + result = get_order_filtered_deployments(deps, target_order=99) assert result == [] def test_target_order_no_match_does_not_reselect_lower_order(self): @@ -70,7 +70,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(2, "b"), ] remaining_after_pre_call = [deps[0]] - result = _get_order_filtered_deployments(remaining_after_pre_call, target_order=2) + result = get_order_filtered_deployments(remaining_after_pre_call, target_order=2) assert result == [] def test_no_order_set_returns_all(self): @@ -78,11 +78,11 @@ class TestGetOrderFilteredDeployments: self._make_deployment(None, "a"), self._make_deployment(None, "b"), ] - result = _get_order_filtered_deployments(deps) + result = get_order_filtered_deployments(deps) assert len(result) == 2 def test_empty_list(self): - result = _get_order_filtered_deployments([]) + result = get_order_filtered_deployments([]) assert result == [] def test_single_order_returns_all_with_that_order(self): @@ -90,7 +90,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(1, "a"), self._make_deployment(1, "b"), ] - result = _get_order_filtered_deployments(deps) + result = get_order_filtered_deployments(deps) assert len(result) == 2 diff --git a/tests/test_litellm/test_router_weighted_failover.py b/tests/test_litellm/test_router_weighted_failover.py index 162312a8c67..9f05654f23f 100644 --- a/tests/test_litellm/test_router_weighted_failover.py +++ b/tests/test_litellm/test_router_weighted_failover.py @@ -15,11 +15,11 @@ import pytest import litellm from litellm import Router -from litellm.utils import _get_excluded_filtered_deployments +from litellm.utils import get_excluded_filtered_deployments # --------------------------------------------------------------------------- -# Unit tests for _get_excluded_filtered_deployments +# Unit tests for get_excluded_filtered_deployments # --------------------------------------------------------------------------- @@ -37,17 +37,17 @@ def _make_dep(dep_id: str, weight: Optional[int] = None) -> dict: class TestGetExcludedFilteredDeployments: def test_no_excluded_returns_all(self): deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=None) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=None) assert len(result) == 2 def test_empty_excluded_returns_all(self): deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=[]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=[]) assert len(result) == 2 def test_drops_excluded(self): deps = [_make_dep("a"), _make_dep("b"), _make_dep("c")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) ids = sorted(d["model_info"]["id"] for d in result) assert ids == ["a", "c"] @@ -57,12 +57,12 @@ class TestGetExcludedFilteredDeployments: # error. Returning the original list here would re-include the # just-failed deployment and let weighted failover re-pick it. deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["a", "b"]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=["a", "b"]) assert result == [] def test_excluded_set_with_unknown_ids(self): deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["zzz"]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=["zzz"]) assert len(result) == 2 def test_handles_missing_model_info(self): @@ -70,7 +70,7 @@ class TestGetExcludedFilteredDeployments: {"model_name": "x", "litellm_params": {"model": "gpt-4o"}}, # no model_info _make_dep("b"), ] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) assert len(result) == 1 From e79f3ec5205d01323093534a6577905a6dfdb7ac Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:31:32 -0700 Subject: [PATCH 14/32] fix(cost-map): stop advertising reasoning_effort max on the azure gpt-6-astra rows Both Azure routes refuse it. A live call to the same deployment through openai/deployments/gpt-6-astra/chat/completions on api-version 2025-04-01-preview answers reasoning_effort max with a 400 unsupported_value naming none, low, medium, high and xhigh as the values it takes, and xhigh returns 200, so azure/gpt-6-astra and azure/us/gpt-6-astra now match the azure_ai row. --- ...odel_prices_and_context_window_backup.json | 4 +-- .../reasoning_effort_capability.py | 4 +-- model_prices_and_context_window.json | 4 +-- .../test_reasoning_effort_capability.py | 26 ++++++------------- 4 files changed, 14 insertions(+), 24 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b659c3b65e5..48c51f8bbf5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -7237,7 +7237,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, @@ -7503,7 +7503,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py index 9185d901a28..7b145c15a07 100644 --- a/litellm/router_utils/reasoning_effort_capability.py +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -10,8 +10,8 @@ opt-in. none is opt-out everywhere except the azure gpt-5 family, whose config r UnsupportedParamsError without an explicit true. xhigh is gated on the request path by the openai and azure gpt-5 configs. max is not gated there at -all: every entry carrying supports_max_reasoning_effort is Claude-family, and -anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort +all: outside the gpt-6-astra rows every entry carrying supports_max_reasoning_effort is Claude-family, +and anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort path maps any level to a thinking budget. Making max opt-in is a deliberate trade, then, since an explicit flag is the only signal that the tier is a real one rather than litellm rounding the level to a budget, and a missing flag costs advisory metadata rather than a rejected request. diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b659c3b65e5..48c51f8bbf5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -7237,7 +7237,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, @@ -7503,7 +7503,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index fa3a6dcd95a..ccd6766b13a 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -394,30 +394,20 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: [ ("azure/gpt-6-astra", "azure"), ("azure/us/gpt-6-astra", "azure"), + ("azure_ai/gpt-6-astra", "azure_ai"), ], ) - def test_a_foundry_deployment_also_advertises_none(self, local_model_cost_map, model, custom_llm_provider): - """Microsoft Foundry serves the same model but its API accepts reasoning_effort none - (verified live: 200 with zero reasoning tokens, and it unlocks temperature), which - OpenAI's rejects, so an Azure deployment offers none on top of low through max.""" + def test_an_azure_hosted_deployment_advertises_none_but_not_max( + self, local_model_cost_map, model, custom_llm_provider + ): + """Microsoft hosts the same model with a different level set than OpenAI does. Verified live + on both Azure routes: none returns 200 with zero reasoning tokens and unlocks temperature, + which OpenAI's API rejects, while max returns 400 unsupported_value naming none through + xhigh as the levels it does take.""" from litellm.utils import _get_model_info_helper model_info = dict(_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)) - assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( - "none", - "low", - "medium", - "high", - "xhigh", - "max", - ) - - def test_a_foundry_azure_ai_deployment_advertises_none_but_not_max(self, local_model_cost_map): - from litellm.utils import _get_model_info_helper - - model_info = dict(_get_model_info_helper(model="azure_ai/gpt-6-astra", custom_llm_provider="azure_ai")) - assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( "none", "low", From fa2b64878b6f7be8fed5139ef95961fe26241f99 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:31:33 -0700 Subject: [PATCH 15/32] fix(azure_ai): redirect a gpt-5 capability lookup only when the map has a foundry row gpt-6-astra is the only gpt-5-family name with an azure_ai row. Prefixing the rest cost them every effort flag, since get_llm_provider sends an azure_ai name down the azure provider when a global AZURE_AI_API_BASE points at an openai.azure.com host and azure/ is not a key either, which turned temperature, top_p and logprobs on azure_ai/gpt-5.1-chat-latest from accepted into an UnsupportedParamsError. --- litellm/llms/azure_ai/chat/transformation.py | 14 ++++++++++- .../chat/test_azure_ai_transformation.py | 23 +++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 7c9a26c3f07..039c462b38a 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -46,7 +46,19 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ( class AzureAIGPT5Config(OpenAIGPT5Config): @classmethod def _model_map_lookup_name(cls, model: str) -> str: - return model if model.startswith("azure_ai/") else f"azure_ai/{model}" + """Normalise a Foundry routing name to its cost-map key, when the map has one. + + A Foundry deployment and its OpenAI-hosted namesake are different products with + different capabilities, so ``azure_ai/`` is the entry to read whenever the map + carries it. Most gpt-5-family names have no ``azure_ai/`` row, though, and prefixing + those anyway costs them every flag: ``get_llm_provider`` re-resolves an ``azure_ai/`` + name to the azure provider when a global AZURE_AI_API_BASE points at an + openai.azure.com host, ``azure/`` is not a key either, so the lookup lands + nowhere and every effort answer degrades to False. A missing key defers to the base + resolver instead. + """ + prefixed: Final = model if model.startswith("azure_ai/") else f"azure_ai/{model}" + return prefixed if prefixed in litellm.model_cost else super()._model_map_lookup_name(model) azureAIGPT5Config: Final = AzureAIGPT5Config() diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 5ff0b729449..9924d77eb39 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -157,6 +157,29 @@ def test_foundry_gpt_6_astra_keeps_sampling_params_when_reasoning_effort_is_none assert optional_params == {"reasoning_effort": "none", "temperature": 0.2, "top_p": 0.9} +def test_a_gpt_5_name_without_a_foundry_row_keeps_reading_its_own_entry( + monkeypatch: pytest.MonkeyPatch, _local_model_cost_map +): + """gpt-6-astra is the only gpt-5-family name with an azure_ai/ row. Reading an azure_ai/ key for + the rest finds nothing, and an openai.azure.com base sends that name down the azure provider, + which has no key for it either, so every effort answer would silently fall back to false and + take temperature, top_p and logprobs down with it.""" + monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.openai.azure.com") + monkeypatch.setenv("AZURE_AI_API_KEY", "placeholder") + + optional_params = litellm.utils.get_optional_params( + model="gpt-5.1-chat-latest", + custom_llm_provider="azure_ai", + temperature=0.2, + top_p=0.9, + logprobs=True, + ) + + assert optional_params["temperature"] == 0.2 + assert optional_params["top_p"] == 0.9 + assert optional_params["logprobs"] is True + + def test_azure_ai_grok_stop_parameter_handling(): """ Test that Grok models properly handle stop parameter filtering in Azure AI Studio. From 0fb3951b2cadca5d9091b7586d0a5a15e4600b42 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:32:06 -0700 Subject: [PATCH 16/32] fix(spend): charge a batch once when an older proxy left its cost row at $0 A proxy without this fix wrote the batch's cost row on every poll while the batch was still running, so that row reads $0 and the insert that claims the charge has nowhere to land. The retrieve that charges the batch now writes its own payload over that row under a where clause that still names spend 0.0, so exactly one retrieve takes it over and every later one reads the charge and charges nothing --- litellm/proxy/db/db_spend_update_writer.py | 42 +++++++++- .../proxy/db/test_db_spend_update_writer.py | 76 ++++++++++++++++++- 2 files changed, 112 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 48312c025dd..ae3bc5663eb 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -14,6 +14,7 @@ import time import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload import litellm @@ -379,9 +380,44 @@ class DBSpendUpdateWriter: if existing.spend > 0: verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) return False - verbose_proxy_logger.debug( - "Spend row %s charged nothing for this batch, so this retrieve charges it", request_id - ) + return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client) + + async def _take_over_uncharged_batch_cost_row( + self, payload: SpendLogsPayload, prisma_client: "PrismaClient" + ) -> bool: + """Take the batch's cost row over from the poll that left it charging nothing. + + A pre-upgrade proxy wrote that row every time it polled the batch while it was still + running, so the charge is still to be made and the row still has to end up carrying + it. The row stops matching the moment it carries a charge, so it is one retrieve that + takes it over and charges, and every later one reads the charge and charges nothing. + """ + from litellm.repositories.table_repositories import SpendLogsRepository + + request_id: Final = payload["request_id"] + if payload["spend"] <= 0: + verbose_proxy_logger.debug( + "Cost tracking skipped: this batch costs nothing and spend row %s says so", request_id + ) + return False + try: + taken_over: Final = await SpendLogsRepository(prisma_client).table.update_many( + data=prisma_client.jsonify_object( + MappingProxyType({field: value for field, value in payload.items() if field != "request_id"}) + ), + where={ # mutable-ok: prisma where clause + "request_id": request_id, + "call_type": CallTypes.aretrieve_batch.value, + "status": "success", + "spend": 0.0, + }, + ) + except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; a row it cannot take over charges the batch + verbose_proxy_logger.warning("Could not take over spend row %s for a batch's cost: %s", request_id, e) + return True + if taken_over == 0: + verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) + return False return True async def _enqueue_tool_usage_transaction( diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 41f2be08545..33b7af06e1a 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2969,16 +2969,21 @@ def _batch_cost_payload() -> dict: } -def _spend_logs_prisma(inserted: int, existing: object) -> MagicMock: +def _spend_logs_prisma(inserted: int, existing: object, taken_over: int = 1) -> MagicMock: prisma = _tool_usage_prisma() prisma.jsonify_object = lambda data: dict(data) prisma.db.litellm_spendlogs.create_many = AsyncMock(return_value=inserted) prisma.db.litellm_spendlogs.find_unique = AsyncMock(return_value=existing) + prisma.db.litellm_spendlogs.update_many = AsyncMock(return_value=taken_over) return prisma async def _update_database_with( - db_writer: DBSpendUpdateWriter, prisma: MagicMock, payload: dict, disable_spend_logs: bool = False + db_writer: DBSpendUpdateWriter, + prisma: MagicMock, + payload: dict, + disable_spend_logs: bool = False, + response_cost: float = 0.25, ) -> bool: with ( patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam @@ -3005,7 +3010,7 @@ async def _update_database_with( completion_response=None, start_time=datetime.now(timezone.utc), end_time=datetime.now(timezone.utc), - response_cost=0.25, + response_cost=response_cost, ) await asyncio.sleep(0) return charged @@ -3054,6 +3059,71 @@ async def test_update_database_charges_a_batch_only_from_the_retrieve_that_wrote assert db_writer._batch_database_updates.await_count == (1 if charged else 0) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("taken_over", "charged"), + [(1, True), (0, False)], + ids=["this_retrieve_takes_it_over", "another_one_got_there_first"], +) +async def test_update_database_charges_a_batch_whose_row_a_pre_upgrade_poll_left_at_zero( + taken_over: int, charged: bool +): + """ + A proxy without this fix wrote the batch's row at $0 on every poll of a running batch, + and the row outlives the upgrade, so the charge has to land on the row itself. Charging + without writing it there would charge again on every later retrieve (LIT-7048). + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0) + prisma = _spend_logs_prisma(0, existing, taken_over) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is charged + + taken = prisma.db.litellm_spendlogs.update_many.await_args.kwargs + assert taken["where"] == { + "request_id": "batch_abc_batch_cost", + "call_type": "aretrieve_batch", + "status": "success", + "spend": 0.0, + } + assert taken["data"]["spend"] == 0.25 + assert "request_id" not in taken["data"] + assert db_writer._batch_database_updates.await_count == (1 if charged else 0) + + +@pytest.mark.asyncio +async def test_update_database_charges_a_batch_whose_zero_row_it_could_not_take_over(): + """A DB that refuses the takeover must not swallow the batch's cost.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0) + prisma = _spend_logs_prisma(0, existing) + prisma.db.litellm_spendlogs.update_many = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True + + assert db_writer._batch_database_updates.await_count == 1 + + +@pytest.mark.asyncio +async def test_update_database_leaves_a_batch_that_cost_nothing_to_the_retrieve_that_wrote_its_row(): + """ + A batch every line of which failed costs $0, so its row reads $0 for the honest reason + and the retrieve that wrote it is still the one that accounted it. Taking that row over + on every later retrieve would count one batch as many requests. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0) + prisma = _spend_logs_prisma(0, existing) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload(), response_cost=0.0) is False + + prisma.db.litellm_spendlogs.update_many.assert_not_called() + assert db_writer._batch_database_updates.await_count == 0 + + @pytest.mark.asyncio @pytest.mark.parametrize( ("inserted", "existing", "charged"), From 7c7810df42a6f90f213b9998e9897da81b8cfb25 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:43:15 -0700 Subject: [PATCH 17/32] fix(router): ignore non-integer status codes when picking retry skips CI's router_code_coverage gate wants every function in router.py called by name from a test file with "router" in its name, and the new helper had no direct caller, so the check-quality job failed on the first tip. Covering it directly also turned up a hole. litellm._should_retry compares the status code to 500, so a provider exception carrying a string status code raises TypeError instead of answering. should_retry_this_error has the same call, but the retry policy path skips it, which is exactly the path this change enables, so the helper was the first to touch that value. Narrowing to int leaves those exceptions on the old retry-in-place behavior. --- litellm/router.py | 2 +- tests/test_litellm/test_router.py | 27 +++++++++++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index f72f41313a0..6f1f5696847 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7467,7 +7467,7 @@ class Router: ) -> tuple[str, ...]: failed_deployment_id: Final[str | None] = getattr(exception, "failed_deployment_id", None) status_code: Final = getattr(exception, "status_code", None) - if not failed_deployment_id or status_code is None: + if not failed_deployment_id or not isinstance(status_code, int): return () if litellm._should_retry(status_code): # pyright: ignore[reportPrivateUsage] # as in should_retry_this_error return () diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 12c1516837a..4bf75ad408c 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -13234,6 +13234,33 @@ async def test_router_retry_policy_400_retries_on_sibling_deployment( assert response._hidden_params["additional_headers"]["x-litellm-attempted-retries"] == 1 +@pytest.mark.parametrize( + "status_code,failed_deployment_id,already_skipped,healthy_deployment_ids,expected", + [ + (400, "rejecting", None, ["rejecting", "accepting"], ("rejecting",)), + (403, "rejecting", None, ["rejecting", "accepting"], ("rejecting",)), + (400, "second", ("first",), ["first", "second", "third"], ("first", "second")), + (429, "rejecting", None, ["rejecting", "accepting"], ()), + (503, "rejecting", None, ["rejecting", "accepting"], ()), + (400, "rejecting", None, ["rejecting"], ()), + (400, None, None, ["rejecting", "accepting"], ()), + (None, "rejecting", None, ["rejecting", "accepting"], ()), + ("400", "rejecting", None, ["rejecting", "accepting"], ()), + ], +) +def test_router_deployment_ids_to_skip_on_retry( + status_code, failed_deployment_id, already_skipped, healthy_deployment_ids, expected +): + exception = Exception("upstream refused this request") + exception.status_code = status_code + exception.failed_deployment_id = failed_deployment_id + healthy_deployments = [{"model_info": {"id": deployment_id}} for deployment_id in healthy_deployment_ids] + + assert ( + litellm.Router._deployment_ids_to_skip_on_retry(exception, already_skipped, healthy_deployments) == expected + ) + + def _make_failure_logging_obj(): return LiteLLMLogging( model="gpt-5.6", From 24f0be80219cd3403ec28cddbed9be946fad1cb0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:47:41 -0700 Subject: [PATCH 18/32] fix(spend): leave a batch uncharged when the database refuses the takeover The takeover of a $0 row an older proxy left behind used to charge the batch when the update could not reach the database. That leaves the row still reading $0, so every later retrieve finds the same row and charges the batch again, which is the repeat charging this PR exists to stop. The retrieve that does take the row over is the one that charges, and a batch nobody retrieves again after that failure is never charged, the same as one whose proxy died inside the write window. --- litellm/proxy/db/db_spend_update_writer.py | 8 +++++--- .../proxy/db/test_db_spend_update_writer.py | 12 ++++++++---- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ae3bc5663eb..ee7802a45d3 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -412,9 +412,11 @@ class DBSpendUpdateWriter: "spend": 0.0, }, ) - except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; a row it cannot take over charges the batch - verbose_proxy_logger.warning("Could not take over spend row %s for a batch's cost: %s", request_id, e) - return True + except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; the next retrieve takes the row over + verbose_proxy_logger.warning( + "Could not take over spend row %s, leaving this batch's cost to the next retrieve: %s", request_id, e + ) + return False if taken_over == 0: verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) return False diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 33b7af06e1a..4efb94b60aa 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -3093,17 +3093,21 @@ async def test_update_database_charges_a_batch_whose_row_a_pre_upgrade_poll_left @pytest.mark.asyncio -async def test_update_database_charges_a_batch_whose_zero_row_it_could_not_take_over(): - """A DB that refuses the takeover must not swallow the batch's cost.""" +async def test_update_database_leaves_a_batch_whose_zero_row_it_could_not_take_over_to_the_next_retrieve(): + """ + A DB that refuses the takeover leaves the row reading $0, so charging here would charge + the batch again on every later retrieve. The retrieve that does take the row over is the + one that charges. + """ db_writer = DBSpendUpdateWriter() db_writer._batch_database_updates = AsyncMock() existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0) prisma = _spend_logs_prisma(0, existing) prisma.db.litellm_spendlogs.update_many = AsyncMock(side_effect=RuntimeError("db unreachable")) - assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is False - assert db_writer._batch_database_updates.await_count == 1 + assert db_writer._batch_database_updates.await_count == 0 @pytest.mark.asyncio From cb1ec76e46245196092a425e8d759f85e47df547 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:02:44 -0700 Subject: [PATCH 19/32] fix(router): keep retry skips on the active order-fallback target The retry-skip guard checks that some other deployment could still answer before it excludes the one that just refused, so a single-deployment group keeps the old retry-in-place behavior. It asked that question at the group's minimum order, but the router picks the retry's deployment at the order the request has already escalated to. So a group with a primary at order 1 and a backup at order 2 answered "yes, order 1 still has a candidate" while the retry was pinned to order 2, and the exclusion left order 2 with nothing. The caller got a no-deployments error in place of the provider's own 400. The helper now takes the active target order and filters by it, which is the same value async_get_healthy_deployments reads off the request. --- litellm/router.py | 8 +++++++- tests/test_litellm/test_router.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 6f1f5696847..4901199902c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -402,6 +402,7 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _EXCLUDED_DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) +_TARGET_ORDER_ADAPTER: Final = TypeAdapter(int | None) def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]: @@ -7464,6 +7465,7 @@ class Router: exception: Exception, already_skipped: object, healthy_deployments: list[dict], # mutable-ok: matches the routing filters' list contract + target_order: object = None, ) -> tuple[str, ...]: failed_deployment_id: Final[str | None] = getattr(exception, "failed_deployment_id", None) status_code: Final = getattr(exception, "status_code", None) @@ -7473,7 +7475,9 @@ class Router: return () already_skipped_ids: Final = _EXCLUDED_DEPLOYMENT_IDS_ADAPTER.validate_python(already_skipped or ()) skipped: Final = frozenset((*already_skipped_ids, failed_deployment_id)) - same_order_candidates: Final = litellm.utils.get_order_filtered_deployments(healthy_deployments) + same_order_candidates: Final = litellm.utils.get_order_filtered_deployments( + healthy_deployments, target_order=_TARGET_ORDER_ADAPTER.validate_python(target_order) + ) if not litellm.utils.get_excluded_filtered_deployments(same_order_candidates, excluded_deployment_ids=skipped): return () verbose_router_logger.debug( @@ -7580,6 +7584,7 @@ class Router: exception=original_exception, already_skipped=kwargs.get("_excluded_deployment_ids"), healthy_deployments=_healthy_deployments, + target_order=kwargs.get("_target_order"), ) if skipped_deployment_ids: kwargs["_excluded_deployment_ids"] = skipped_deployment_ids @@ -7656,6 +7661,7 @@ class Router: exception=e, already_skipped=kwargs.get("_excluded_deployment_ids"), healthy_deployments=_healthy_deployments, + target_order=kwargs.get("_target_order"), ) if retry_skipped_deployment_ids: kwargs["_excluded_deployment_ids"] = retry_skipped_deployment_ids diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 4bf75ad408c..0d7f2e7d3f6 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -13261,6 +13261,36 @@ def test_router_deployment_ids_to_skip_on_retry( ) +@pytest.mark.parametrize( + "target_order,deployment_orders,expected", + [ + (2, {"rejecting": 2, "sibling": 1}, ()), + (2, {"rejecting": 2, "sibling": 2}, ("rejecting",)), + (1, {"rejecting": 1, "sibling": 2}, ()), + (None, {"rejecting": 1, "sibling": 2}, ()), + (None, {"rejecting": 1, "sibling": 1}, ("rejecting",)), + (3, {"rejecting": 2, "sibling": 1}, ()), + ], +) +def test_router_deployment_ids_to_skip_on_retry_honors_order_fallback_target( + target_order, deployment_orders, expected +): + exception = Exception("upstream refused this request") + exception.status_code = 400 + exception.failed_deployment_id = "rejecting" + healthy_deployments = [ + {"model_info": {"id": deployment_id}, "litellm_params": {"order": order}} + for deployment_id, order in deployment_orders.items() + ] + + assert ( + litellm.Router._deployment_ids_to_skip_on_retry( + exception, None, healthy_deployments, target_order=target_order + ) + == expected + ) + + def _make_failure_logging_obj(): return LiteLLMLogging( model="gpt-5.6", From 3dea1ebb32c96560f9f10171d0597ca30d4bea40 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:18:25 -0700 Subject: [PATCH 20/32] fix(cost-map): keep the prompt cache breakpoint flag on the foundry gpt-6-astra row The openai gpt-6-astra card carries supports_prompt_cache_breakpoint, so a Foundry deployment reported it as true until the azure_ai row took over the lookup. The cache control hook still honours breakpoints for that deployment through the bare name, so /model/info was the only thing that changed, and it now agrees with the hook again. --- litellm/model_prices_and_context_window_backup.json | 1 + model_prices_and_context_window.json | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 48c51f8bbf5..2f86deffe54 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3524,6 +3524,7 @@ "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 48c51f8bbf5..2f86deffe54 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3524,6 +3524,7 @@ "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, From fffe0bb0dc94f8b071be85050a0bdaa101c0e2ba Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:18:25 -0700 Subject: [PATCH 21/32] test(azure_ai): pin the tier the messages bridge sends when astra refuses max The /v1/messages adapter lowers a tier the entry does not accept, so dropping max from the astra rows moves that path from Foundry's 400 to a request at xhigh. Nothing pinned that, and the guard test's docstring named gpt-6-astra as the only gpt-5 name with an azure_ai row, which 11 rows contradict. --- ..._handler_reasoning_effort_normalization.py | 19 +++++++++++++++++++ .../chat/test_azure_ai_transformation.py | 8 ++++---- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py index 56b754c3476..af7befecc33 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py @@ -82,3 +82,22 @@ class TestTheNormalizedTierIsTheTierSent: self, local_model_cost_map, model, provider, effort, expected ): assert _reasoning_effort_sent(model, provider, effort) == expected + + @pytest.mark.parametrize( + "model, provider", + [ + ("gpt-6-astra", "azure_ai"), + ("azure_ai/gpt-6-astra", "azure_ai"), + ("gpt-6-astra", "azure"), + ("us/gpt-6-astra", "azure"), + ], + ) + def test_an_azure_hosted_astra_deployment_drops_to_the_tier_it_accepts( + self, local_model_cost_map, model, provider + ): + """The deployment answers ``max`` with a 400 naming ``none`` through ``xhigh``, so the rows + say so and the adapter sends the tier below instead of the rejected one.""" + assert _reasoning_effort_sent(model, provider, "max") == "xhigh" + + def test_the_openai_hosted_twin_still_sends_max(self, local_model_cost_map): + assert _reasoning_effort_sent("gpt-6-astra", "openai", "max") == "max" diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 9924d77eb39..f8cc0b5071e 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -160,10 +160,10 @@ def test_foundry_gpt_6_astra_keeps_sampling_params_when_reasoning_effort_is_none def test_a_gpt_5_name_without_a_foundry_row_keeps_reading_its_own_entry( monkeypatch: pytest.MonkeyPatch, _local_model_cost_map ): - """gpt-6-astra is the only gpt-5-family name with an azure_ai/ row. Reading an azure_ai/ key for - the rest finds nothing, and an openai.azure.com base sends that name down the azure provider, - which has no key for it either, so every effort answer would silently fall back to false and - take temperature, top_p and logprobs down with it.""" + """Most gpt-5-family names have no azure_ai/ row. Reading an azure_ai/ key for those finds + nothing, and an openai.azure.com base sends the name down the azure provider, which has no key + for it either, so every effort answer would silently fall back to false and take temperature, + top_p and logprobs down with it.""" monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.openai.azure.com") monkeypatch.setenv("AZURE_AI_API_KEY", "placeholder") From ecf7e4e766cff606c9bbdef62d1d3a67c95113d0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:36:28 -0700 Subject: [PATCH 22/32] fix(router): keep the provider's error when the retry skip empties the group Before excluding the deployment that just refused, the retry-skip guard asked whether another one could still answer. It asked by re-running a single routing filter, the order filter, while deployment selection also applies cooldowns, the context-window pre-call check, tag routing, and routing plugins. Any filter the guard did not replicate made it answer yes while the real pick was left with nothing. A group narrowed to one deployment by tag routing turned the provider's own 400 into a no-deployments 429. The skip now runs where every filter has already been applied, and it keeps the deployments untouched when skipping would leave none. The caller gets the provider's error either way, and a group with one eligible deployment retries in place as it did before. --- litellm/router.py | 59 ++++--- .../complexity_router/complexity_router.py | 5 +- tests/test_litellm/test_router.py | 153 +++++++++++++----- 3 files changed, 150 insertions(+), 67 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 4901199902c..93178dc8e26 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -401,8 +401,7 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_EXCLUDED_DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) -_TARGET_ORDER_ADAPTER: Final = TypeAdapter(int | None) +_SKIPPED_DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]: @@ -7461,29 +7460,19 @@ class Router: ) @staticmethod - def _deployment_ids_to_skip_on_retry( - exception: Exception, - already_skipped: object, - healthy_deployments: list[dict], # mutable-ok: matches the routing filters' list contract - target_order: object = None, - ) -> tuple[str, ...]: + def _deployment_ids_to_skip_on_retry(exception: Exception, already_skipped: object) -> tuple[str, ...]: failed_deployment_id: Final[str | None] = getattr(exception, "failed_deployment_id", None) status_code: Final = getattr(exception, "status_code", None) if not failed_deployment_id or not isinstance(status_code, int): return () if litellm._should_retry(status_code): # pyright: ignore[reportPrivateUsage] # as in should_retry_this_error return () - already_skipped_ids: Final = _EXCLUDED_DEPLOYMENT_IDS_ADAPTER.validate_python(already_skipped or ()) - skipped: Final = frozenset((*already_skipped_ids, failed_deployment_id)) - same_order_candidates: Final = litellm.utils.get_order_filtered_deployments( - healthy_deployments, target_order=_TARGET_ORDER_ADAPTER.validate_python(target_order) - ) - if not litellm.utils.get_excluded_filtered_deployments(same_order_candidates, excluded_deployment_ids=skipped): - return () + already_skipped_ids: Final = _SKIPPED_DEPLOYMENT_IDS_ADAPTER.validate_python(already_skipped or ()) + skipped: Final = tuple(sorted(frozenset((*already_skipped_ids, failed_deployment_id)))) verbose_router_logger.debug( - "Retry skips deployments that already answered %s to this request: %s", status_code, sorted(skipped) + "Retry skips deployments that already answered %s to this request: %s", status_code, skipped ) - return tuple(sorted(skipped)) + return skipped @tracer.wrap() async def async_function_with_retries(self, *args, **kwargs): @@ -7582,12 +7571,10 @@ class Router: kwargs = self.log_retry(kwargs=kwargs, e=original_exception) skipped_deployment_ids: Final = self._deployment_ids_to_skip_on_retry( exception=original_exception, - already_skipped=kwargs.get("_excluded_deployment_ids"), - healthy_deployments=_healthy_deployments, - target_order=kwargs.get("_target_order"), + already_skipped=kwargs.get("_retry_skipped_deployment_ids"), ) if skipped_deployment_ids: - kwargs["_excluded_deployment_ids"] = skipped_deployment_ids + kwargs["_retry_skipped_deployment_ids"] = skipped_deployment_ids else: raise @@ -7659,12 +7646,10 @@ class Router: retry_skipped_deployment_ids = self._deployment_ids_to_skip_on_retry( exception=e, - already_skipped=kwargs.get("_excluded_deployment_ids"), - healthy_deployments=_healthy_deployments, - target_order=kwargs.get("_target_order"), + already_skipped=kwargs.get("_retry_skipped_deployment_ids"), ) if retry_skipped_deployment_ids: - kwargs["_excluded_deployment_ids"] = retry_skipped_deployment_ids + kwargs["_retry_skipped_deployment_ids"] = retry_skipped_deployment_ids _timeout = self._time_to_sleep_before_retry( e=e, remaining_retries=remaining_retries, @@ -12508,6 +12493,19 @@ class Router: excluded_deployment_ids=_excluded_deployment_ids, ) + ## RETRY SKIP ## -> drop deployments that already refused this request with a + ## non-retryable status, unless that leaves nothing, so the caller still gets + ## the provider's own error instead of a no-deployments error. + _retry_skipped_deployment_ids: Final = ( + request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None + ) + healthy_deployments = ( + litellm.utils.get_excluded_filtered_deployments( + healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids + ) + or healthy_deployments + ) + if len(healthy_deployments) == 0: exception: Final = await async_raise_no_deployment_exception( litellm_router_instance=self, @@ -13413,6 +13411,17 @@ class Router: excluded_deployment_ids=_excluded_deployment_ids, ) + ## RETRY SKIP ## -> see async counterpart in async_get_healthy_deployments. + _retry_skipped_deployment_ids: Final = ( + request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None + ) + healthy_deployments = ( + litellm.utils.get_excluded_filtered_deployments( + healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids + ) + or healthy_deployments + ) + if len(healthy_deployments) == 0: model_ids = self.get_model_ids(model_name=model) _cooldown_time = self.cooldown_cache.get_min_cooldown( diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 98a1eb7ac9e..53a56866e92 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2832,8 +2832,9 @@ class ComplexityRouter(CustomLogger): where the prompt never arrives as messages. Probed on a COPY of request_kwargs because the owner pops routing bookkeeping off the - dict it is handed (`_target_order`, `_excluded_deployment_ids`), and this is a - speculative question about a model that may never be picked. + dict it is handed (`_target_order`, `_excluded_deployment_ids`, + `_retry_skipped_deployment_ids`), and this is a speculative question about a model + that may never be picked. Every way the owner says "nothing here can serve this" is a negative verdict: no healthy deployment for the group at all (BadRequestError, which ContextWindowExceededError diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 0d7f2e7d3f6..34c435af706 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -13234,61 +13234,134 @@ async def test_router_retry_policy_400_retries_on_sibling_deployment( assert response._hidden_params["additional_headers"]["x-litellm-attempted-retries"] == 1 +_UPSTREAM_400 = {"message": "upstream refused this request", "type": "invalid_request_error", "code": "bad_request"} + + +def _retry_skip_deployment(deployment_id, host, litellm_params=None, model_info=None): + return { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-fake", + "api_base": f"https://{host}.local/v1", + **(litellm_params or {}), + }, + "model_info": {"id": deployment_id, **(model_info or {})}, + } + + @pytest.mark.parametrize( - "status_code,failed_deployment_id,already_skipped,healthy_deployment_ids,expected", + "status_code,failed_deployment_id,already_skipped,expected", [ - (400, "rejecting", None, ["rejecting", "accepting"], ("rejecting",)), - (403, "rejecting", None, ["rejecting", "accepting"], ("rejecting",)), - (400, "second", ("first",), ["first", "second", "third"], ("first", "second")), - (429, "rejecting", None, ["rejecting", "accepting"], ()), - (503, "rejecting", None, ["rejecting", "accepting"], ()), - (400, "rejecting", None, ["rejecting"], ()), - (400, None, None, ["rejecting", "accepting"], ()), - (None, "rejecting", None, ["rejecting", "accepting"], ()), - ("400", "rejecting", None, ["rejecting", "accepting"], ()), + (400, "rejecting", None, ("rejecting",)), + (403, "rejecting", None, ("rejecting",)), + (400, "second", ("first",), ("first", "second")), + (400, "first", ("first",), ("first",)), + (429, "rejecting", None, ()), + (503, "rejecting", None, ()), + (408, "rejecting", None, ()), + (400, None, None, ()), + (None, "rejecting", None, ()), + ("400", "rejecting", None, ()), ], ) -def test_router_deployment_ids_to_skip_on_retry( - status_code, failed_deployment_id, already_skipped, healthy_deployment_ids, expected -): +def test_router_deployment_ids_to_skip_on_retry(status_code, failed_deployment_id, already_skipped, expected): exception = Exception("upstream refused this request") exception.status_code = status_code exception.failed_deployment_id = failed_deployment_id - healthy_deployments = [{"model_info": {"id": deployment_id}} for deployment_id in healthy_deployment_ids] - assert ( - litellm.Router._deployment_ids_to_skip_on_retry(exception, already_skipped, healthy_deployments) == expected - ) + assert litellm.Router._deployment_ids_to_skip_on_retry(exception, already_skipped) == expected @pytest.mark.parametrize( - "target_order,deployment_orders,expected", + "deployment_ids,skipped,expected", [ - (2, {"rejecting": 2, "sibling": 1}, ()), - (2, {"rejecting": 2, "sibling": 2}, ("rejecting",)), - (1, {"rejecting": 1, "sibling": 2}, ()), - (None, {"rejecting": 1, "sibling": 2}, ()), - (None, {"rejecting": 1, "sibling": 1}, ("rejecting",)), - (3, {"rejecting": 2, "sibling": 1}, ()), + (["rejecting", "sibling"], ("rejecting",), ["sibling"]), + (["rejecting"], ("rejecting",), ["rejecting"]), + (["rejecting", "sibling"], ("rejecting", "sibling"), ["rejecting", "sibling"]), + (["rejecting", "sibling"], (), ["rejecting", "sibling"]), + (["rejecting", "sibling"], None, ["rejecting", "sibling"]), + (["rejecting", "sibling"], ("absent",), ["rejecting", "sibling"]), ], ) -def test_router_deployment_ids_to_skip_on_retry_honors_order_fallback_target( - target_order, deployment_orders, expected -): - exception = Exception("upstream refused this request") - exception.status_code = 400 - exception.failed_deployment_id = "rejecting" - healthy_deployments = [ - {"model_info": {"id": deployment_id}, "litellm_params": {"order": order}} - for deployment_id, order in deployment_orders.items() - ] - - assert ( - litellm.Router._deployment_ids_to_skip_on_retry( - exception, None, healthy_deployments, target_order=target_order - ) - == expected +@pytest.mark.asyncio +async def test_router_healthy_deployments_keep_the_last_candidate_a_retry_skipped(deployment_ids, skipped, expected): + router = litellm.Router( + model_list=[_retry_skip_deployment(deployment_id, deployment_id) for deployment_id in deployment_ids], + disable_cooldowns=True, ) + request_kwargs = {"_retry_skipped_deployment_ids": skipped} + + healthy_deployments = await router.async_get_healthy_deployments(model="gpt-5.6", request_kwargs=request_kwargs) + + assert sorted(deployment["model_info"]["id"] for deployment in healthy_deployments) == sorted(expected) + assert "_retry_skipped_deployment_ids" not in request_kwargs + + +@pytest.mark.asyncio +async def test_router_retry_policy_400_keeps_upstream_error_on_order_fallback_hop(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + _retry_skip_deployment("order1", "order1", litellm_params={"order": 1}), + _retry_skip_deployment("order2", "order2", litellm_params={"order": 2}), + ], + num_retries=2, + retry_policy={"BadRequestErrorRetries": 2}, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + order1 = respx_mock.post("https://order1.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + order2 = respx_mock.post("https://order2.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + with pytest.raises(litellm.BadRequestError) as raised: + await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) + + assert "upstream refused this request" in str(raised.value) + assert "No deployments available" not in str(raised.value) + assert order1.call_count >= 1 + assert order2.call_count >= 1 + + +@pytest.mark.asyncio +async def test_router_retry_policy_400_keeps_upstream_error_when_tags_narrow_the_group( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + _retry_skip_deployment( + "tagged", "tagged", litellm_params={"tags": ["free"]}, model_info={"enable_tag_filtering": True} + ), + _retry_skip_deployment("untagged", "untagged", model_info={"enable_tag_filtering": True}), + ], + num_retries=2, + retry_policy={"BadRequestErrorRetries": 2}, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + tagged = respx_mock.post("https://tagged.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + untagged = respx_mock.post("https://untagged.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + with pytest.raises(litellm.BadRequestError) as raised: + await router.acompletion( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["free"]}, + ) + + assert "upstream refused this request" in str(raised.value) + assert "No deployments available" not in str(raised.value) + assert tagged.call_count == 3 + assert untagged.call_count == 0 def _make_failure_logging_obj(): From fcb6d2267c09389ae1aa9e80e04a35caa5b8b470 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:51:44 -0700 Subject: [PATCH 23/32] fix(spend): keep a batch's claim row out of the logs a proxy was told not to write disable_spend_logs has to keep meaning that no request gets logged, and the row that makes a batch chargeable exactly once is the one row it cannot drop, so with logging off that row now carries only what tells the retrieves apart. SPEND_LOGS_URL deployments get their copy back too: the claim writes straight to this table, so the row is queued as well when an external writer is the one that takes the spend logs. --- litellm/proxy/db/db_spend_update_writer.py | 49 +++++++-- .../proxy/db/test_db_spend_update_writer.py | 99 +++++++++++++++++++ 2 files changed, 141 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ee7802a45d3..bdc014d7f13 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -89,6 +89,21 @@ def _is_batch_cost_row(payload: SpendLogsPayload) -> bool: return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success" +_BATCH_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"}) + + +def _batch_cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]: + """Reduce a batch's cost row to what tells the retrieves apart when logging is off. + + A proxy run with spend logs disabled still needs one row per batch to charge it once, + so the row is written either way, but it carries no request of its own: no metadata, + no requester IP, no key, model, or token counts (LIT-7048). + """ + if disable_spend_logs is False: + return payload + return MappingProxyType({field: value for field, value in payload.items() if field in _BATCH_COST_CLAIM_FIELDS}) + + class _SpendBatch(Protocol): litellm_usertable: BatchTable litellm_verificationtoken: BatchTable @@ -336,12 +351,16 @@ class DBSpendUpdateWriter: self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None", disable_spend_logs: bool ) -> bool: if prisma_client is not None and _is_batch_cost_row(payload): - return await self._claim_batch_cost_spend_log(payload=payload, prisma_client=prisma_client) + return await self._claim_batch_cost_spend_log( + payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs + ) if disable_spend_logs is False: await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) return True - async def _claim_batch_cost_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient") -> bool: + async def _claim_batch_cost_spend_log( + self, payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool + ) -> bool: """Write the batch's cost row now, or learn that another retrieve already did. Every retrieve of one batch shares this row, so the insert that lands first owns @@ -353,13 +372,17 @@ class DBSpendUpdateWriter: from litellm.repositories.table_repositories import SpendLogsRepository request_id: Final = payload["request_id"] + row: Final = _batch_cost_row_to_write(payload, disable_spend_logs) spend_logs: Final = SpendLogsRepository(prisma_client).table try: claimed: Final = await spend_logs.create_many( - data=[prisma_client.jsonify_object(payload)], # mutable-ok: prisma create_many takes a list + data=[prisma_client.jsonify_object(row)], # mutable-ok: prisma create_many takes a list skip_duplicates=True, ) if claimed == 1: + await self._forward_batch_cost_row( + row=row, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs + ) return True existing: Final = await spend_logs.find_unique( where={"request_id": request_id} # mutable-ok: prisma where clause @@ -368,7 +391,7 @@ class DBSpendUpdateWriter: verbose_proxy_logger.warning( "Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e ) - await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) + await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client) return True if existing is None or existing.call_type != CallTypes.aretrieve_batch.value or existing.status != "success": verbose_proxy_logger.warning( @@ -380,10 +403,22 @@ class DBSpendUpdateWriter: if existing.spend > 0: verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) return False - return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client) + return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client, row=row) + + async def _forward_batch_cost_row( + self, row: Mapping[str, object], prisma_client: "PrismaClient", disable_spend_logs: bool + ) -> None: + """Queue the claimed row for an external spend log writer, which the claim went around. + + With ``SPEND_LOGS_URL`` set the queue posts every spend log to that writer instead of + inserting it, so a batch's cost row reaches it only by being queued here as well. + """ + if disable_spend_logs is True or os.getenv("SPEND_LOGS_URL") is None: + return + await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client) async def _take_over_uncharged_batch_cost_row( - self, payload: SpendLogsPayload, prisma_client: "PrismaClient" + self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object] ) -> bool: """Take the batch's cost row over from the poll that left it charging nothing. @@ -403,7 +438,7 @@ class DBSpendUpdateWriter: try: taken_over: Final = await SpendLogsRepository(prisma_client).table.update_many( data=prisma_client.jsonify_object( - MappingProxyType({field: value for field, value in payload.items() if field != "request_id"}) + MappingProxyType({field: value for field, value in row.items() if field != "request_id"}) ), where={ # mutable-ok: prisma where clause "request_id": request_id, diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 4efb94b60aa..a8aeaced55f 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1,6 +1,7 @@ import asyncio import copy import json +import os import re @@ -3169,6 +3170,104 @@ async def test_update_database_writes_no_ordinary_spend_row_with_spend_logs_disa assert db_writer._batch_database_updates.await_count == 1 +_BATCH_CLAIM_FIELDS = {"request_id", "call_type", "status", "spend", "startTime", "endTime"} + + +def _logged_batch_cost_payload() -> dict: + return { + **_batch_cost_payload(), + "api_key": "0e5b0e9e5f", + "model": "gpt-5.6-luna", + "user": "test-user", + "metadata": '{"batch_models": ["gpt-5.6-luna"]}', + "requester_ip_address": "127.0.0.1", + "proxy_server_request": '{"headers": {"user-agent": "litellm-batch-cost-check"}}', + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("disable_spend_logs", "logs_the_request"), + [(False, True), (True, False)], + ids=["spend_logs_on", "spend_logs_off"], +) +async def test_update_database_claims_a_batch_without_logging_the_request_that_polled_it( + disable_spend_logs: bool, logs_the_request: bool +): + """ + disable_spend_logs has to keep meaning that no request gets logged, and the batch's cost + row is the one row it cannot drop, so with logging off that row carries only what tells + the retrieves apart: no metadata, no requester IP, no key, model, or token counts. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + payload = _logged_batch_cost_payload() + + assert await _update_database_with(db_writer, prisma, payload, disable_spend_logs) is True + + claimed = prisma.db.litellm_spendlogs.create_many.await_args.kwargs["data"][0] + assert set(claimed) == (set(payload) if logs_the_request else _BATCH_CLAIM_FIELDS) + assert claimed["spend"] == 0.25 + assert db_writer._batch_database_updates.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("spend_logs_url", "forwarded"), + [("http://spend-logs.internal", True), (None, False)], + ids=["an_external_writer_takes_the_rows", "rows_are_written_to_this_db"], +) +async def test_update_database_sends_a_claimed_batch_cost_row_on_to_an_external_spend_log_writer( + monkeypatch, spend_logs_url: str | None, forwarded: bool +): + """ + SPEND_LOGS_URL makes the flush post spend logs to that writer instead of inserting them, + and the claim writes straight to this table, so the batch's row reaches the writer only + by being queued as well. Queueing it with no writer configured would insert it twice. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + if spend_logs_url is None: + monkeypatch.delenv("SPEND_LOGS_URL", raising=False) + else: + monkeypatch.setenv("SPEND_LOGS_URL", spend_logs_url) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True + + queued = [row["request_id"] for row in prisma.spend_log_transactions] + assert queued == (["batch_abc_batch_cost"] if forwarded else []) + + +@pytest.mark.asyncio +async def test_update_database_forwards_no_batch_cost_row_a_later_retrieve_had_already_claimed(monkeypatch): + """The retrieve that lost the claim charges nothing, so it must not post a row either.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25) + prisma = _spend_logs_prisma(0, existing) + monkeypatch.setenv("SPEND_LOGS_URL", "http://spend-logs.internal") + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is False + + assert prisma.spend_log_transactions == [] + + +@pytest.mark.asyncio +async def test_update_database_queues_only_the_claim_for_a_batch_it_could_not_write_with_logs_disabled(): + """A refused claim is retried through the queue, so what it queues has to stay unlogged too.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(0, None) + prisma.db.litellm_spendlogs.create_many = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert await _update_database_with(db_writer, prisma, _logged_batch_cost_payload(), True) is True + + assert [set(row) for row in prisma.spend_log_transactions] == [_BATCH_CLAIM_FIELDS] + assert db_writer._batch_database_updates.await_count == 1 + + @pytest.mark.asyncio async def test_update_database_queues_a_batch_cost_row_it_could_not_claim(): """An unreachable DB must not drop the batch's only spend row, nor its charge.""" From 6866eac96feed14e47e051896c45263fee53086a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:53:36 -0700 Subject: [PATCH 24/32] fix(router): ignore a retry skip list the caller sent itself The retry skip travels as a request kwarg, and the router forwards keys it does not recognize, so a client can put _retry_skipped_deployment_ids in its own request body. The value went straight into a pydantic TypeAdapter and then into a set(), so an int or an object raised TypeError and a string, a list, or a dict raised a ValidationError, each of them replacing the 400 the provider had actually returned. Every read now goes through one narrowing function that keeps a tuple of strings and skips nothing otherwise, so a forged value costs the caller nothing beyond the retry landing on the same deployment again. --- litellm/router.py | 11 +++++---- tests/test_litellm/test_router.py | 39 +++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 93178dc8e26..e42bc2af398 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -401,7 +401,10 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_SKIPPED_DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) + + +def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: + return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else () def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]: @@ -7467,7 +7470,7 @@ class Router: return () if litellm._should_retry(status_code): # pyright: ignore[reportPrivateUsage] # as in should_retry_this_error return () - already_skipped_ids: Final = _SKIPPED_DEPLOYMENT_IDS_ADAPTER.validate_python(already_skipped or ()) + already_skipped_ids: Final = _as_retry_skipped_deployment_ids(already_skipped) skipped: Final = tuple(sorted(frozenset((*already_skipped_ids, failed_deployment_id)))) verbose_router_logger.debug( "Retry skips deployments that already answered %s to this request: %s", status_code, skipped @@ -12496,7 +12499,7 @@ class Router: ## RETRY SKIP ## -> drop deployments that already refused this request with a ## non-retryable status, unless that leaves nothing, so the caller still gets ## the provider's own error instead of a no-deployments error. - _retry_skipped_deployment_ids: Final = ( + _retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids( request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None ) healthy_deployments = ( @@ -13412,7 +13415,7 @@ class Router: ) ## RETRY SKIP ## -> see async counterpart in async_get_healthy_deployments. - _retry_skipped_deployment_ids: Final = ( + _retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids( request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None ) healthy_deployments = ( diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 34c435af706..696abcd65a5 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -13263,6 +13263,10 @@ def _retry_skip_deployment(deployment_id, host, litellm_params=None, model_info= (400, None, None, ()), (None, "rejecting", None, ()), ("400", "rejecting", None, ()), + (400, "second", 7, ("second",)), + (400, "second", "first", ("second",)), + (400, "second", ["first"], ("second",)), + (400, "second", ("first", 7), ("first", "second")), ], ) def test_router_deployment_ids_to_skip_on_retry(status_code, failed_deployment_id, already_skipped, expected): @@ -13282,6 +13286,11 @@ def test_router_deployment_ids_to_skip_on_retry(status_code, failed_deployment_i (["rejecting", "sibling"], (), ["rejecting", "sibling"]), (["rejecting", "sibling"], None, ["rejecting", "sibling"]), (["rejecting", "sibling"], ("absent",), ["rejecting", "sibling"]), + (["rejecting", "sibling"], 7, ["rejecting", "sibling"]), + (["rejecting", "sibling"], "rejecting", ["rejecting", "sibling"]), + (["rejecting", "sibling"], ["rejecting"], ["rejecting", "sibling"]), + (["rejecting", "sibling"], {"rejecting": True}, ["rejecting", "sibling"]), + (["rejecting", "sibling"], ("rejecting", 7), ["sibling"]), ], ) @pytest.mark.asyncio @@ -13298,6 +13307,36 @@ async def test_router_healthy_deployments_keep_the_last_candidate_a_retry_skippe assert "_retry_skipped_deployment_ids" not in request_kwargs +@pytest.mark.parametrize("client_supplied", [7, "rejecting", ["rejecting"], {"rejecting": True}, object()]) +@pytest.mark.asyncio +async def test_router_retry_policy_400_keeps_upstream_error_when_a_client_forges_the_skip_list( + monkeypatch: pytest.MonkeyPatch, client_supplied +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[_retry_skip_deployment("rejecting", "rejecting"), _retry_skip_deployment("sibling", "sibling")], + num_retries=2, + retry_policy={"BadRequestErrorRetries": 2}, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + respx_mock.post("https://rejecting.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + respx_mock.post("https://sibling.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + with pytest.raises(litellm.BadRequestError) as raised: + await router.acompletion( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + _retry_skipped_deployment_ids=client_supplied, + ) + + assert "upstream refused this request" in str(raised.value) + + @pytest.mark.asyncio async def test_router_retry_policy_400_keeps_upstream_error_on_order_fallback_hop(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) From defd8661f4e359994bf57a7f9e51ed8c479f17ff Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 00:04:24 -0700 Subject: [PATCH 25/32] refactor(spend): stop queueing a batch's claim row for a writer the proxy never builds SPEND_LOGS_URL only diverts spend logs when db_writer_client is set, and nothing in the proxy ever assigns that global, so the queued copy was only ever skipped as a duplicate by the local insert. --- litellm/proxy/db/db_spend_update_writer.py | 15 ------- .../proxy/db/test_db_spend_update_writer.py | 43 ------------------- 2 files changed, 58 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index bdc014d7f13..9230be8055e 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -380,9 +380,6 @@ class DBSpendUpdateWriter: skip_duplicates=True, ) if claimed == 1: - await self._forward_batch_cost_row( - row=row, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs - ) return True existing: Final = await spend_logs.find_unique( where={"request_id": request_id} # mutable-ok: prisma where clause @@ -405,18 +402,6 @@ class DBSpendUpdateWriter: return False return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client, row=row) - async def _forward_batch_cost_row( - self, row: Mapping[str, object], prisma_client: "PrismaClient", disable_spend_logs: bool - ) -> None: - """Queue the claimed row for an external spend log writer, which the claim went around. - - With ``SPEND_LOGS_URL`` set the queue posts every spend log to that writer instead of - inserting it, so a batch's cost row reaches it only by being queued here as well. - """ - if disable_spend_logs is True or os.getenv("SPEND_LOGS_URL") is None: - return - await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client) - async def _take_over_uncharged_batch_cost_row( self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object] ) -> bool: diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index a8aeaced55f..0bca7c9492c 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1,7 +1,6 @@ import asyncio import copy import json -import os import re @@ -3212,48 +3211,6 @@ async def test_update_database_claims_a_batch_without_logging_the_request_that_p assert db_writer._batch_database_updates.await_count == 1 -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("spend_logs_url", "forwarded"), - [("http://spend-logs.internal", True), (None, False)], - ids=["an_external_writer_takes_the_rows", "rows_are_written_to_this_db"], -) -async def test_update_database_sends_a_claimed_batch_cost_row_on_to_an_external_spend_log_writer( - monkeypatch, spend_logs_url: str | None, forwarded: bool -): - """ - SPEND_LOGS_URL makes the flush post spend logs to that writer instead of inserting them, - and the claim writes straight to this table, so the batch's row reaches the writer only - by being queued as well. Queueing it with no writer configured would insert it twice. - """ - db_writer = DBSpendUpdateWriter() - db_writer._batch_database_updates = AsyncMock() - prisma = _spend_logs_prisma(1, None) - if spend_logs_url is None: - monkeypatch.delenv("SPEND_LOGS_URL", raising=False) - else: - monkeypatch.setenv("SPEND_LOGS_URL", spend_logs_url) - - assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True - - queued = [row["request_id"] for row in prisma.spend_log_transactions] - assert queued == (["batch_abc_batch_cost"] if forwarded else []) - - -@pytest.mark.asyncio -async def test_update_database_forwards_no_batch_cost_row_a_later_retrieve_had_already_claimed(monkeypatch): - """The retrieve that lost the claim charges nothing, so it must not post a row either.""" - db_writer = DBSpendUpdateWriter() - db_writer._batch_database_updates = AsyncMock() - existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25) - prisma = _spend_logs_prisma(0, existing) - monkeypatch.setenv("SPEND_LOGS_URL", "http://spend-logs.internal") - - assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is False - - assert prisma.spend_log_transactions == [] - - @pytest.mark.asyncio async def test_update_database_queues_only_the_claim_for_a_batch_it_could_not_write_with_logs_disabled(): """A refused claim is retried through the queue, so what it queues has to stay unlogged too.""" From 0fcf0fe06c80683da3d4a7b7b63a0c0aa922382b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 00:06:35 -0700 Subject: [PATCH 26/32] test(router): cover the retry skip-list narrowing helper The router code coverage gate reads every function defined in router.py and fails when no test file names it. _as_retry_skipped_deployment_ids was only reached indirectly through the retry path, so the gate went red on this PR's tip. Test it directly instead: a tuple of strings survives, non-string items inside the tuple are dropped, and every other shape a caller could send narrows to an empty skip list. --- tests/test_litellm/test_router.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 696abcd65a5..05ac672691e 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -13277,6 +13277,26 @@ def test_router_deployment_ids_to_skip_on_retry(status_code, failed_deployment_i assert litellm.Router._deployment_ids_to_skip_on_retry(exception, already_skipped) == expected +@pytest.mark.parametrize( + "value,expected", + [ + (("first", "second"), ("first", "second")), + ((), ()), + (("first", 7, None, "second"), ("first", "second")), + (None, ()), + (7, ()), + ("first", ()), + (["first"], ()), + ({"first": True}, ()), + (object(), ()), + ], +) +def test_router_as_retry_skipped_deployment_ids_keeps_only_a_tuple_of_strings(value, expected): + from litellm.router import _as_retry_skipped_deployment_ids + + assert _as_retry_skipped_deployment_ids(value) == expected + + @pytest.mark.parametrize( "deployment_ids,skipped,expected", [ From d664ca139ef6cf2a2d79a8e98bdfb4985bc3ed32 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 00:53:24 -0700 Subject: [PATCH 27/32] chore(router): suppress the retry-skip kwargs writes and correct the filter docstring The two writes that hand the skip list to the next attempt now carry a `# rebind-ok` reason, which is the sanctioned escape hatch for an unavoidable parameter mutation and matches how `log_retry` already writes into the same kwargs dict a few lines above `get_excluded_filtered_deployments`'s docstring said returning the unfiltered list would re-include the deployment that just failed. The retry skip does exactly that on purpose, so the docstring now says each caller decides what an empty result means The reliability registry cell the new e2e test claims is marked `fail_before_fix: proven`: the same config returns 400 at the merge base and 200 off a sibling deployment at the tip --- litellm/router.py | 12 ++++++------ litellm/utils.py | 10 ++++++---- tests/e2e/coverage_registry/reliability.yaml | 2 +- 3 files changed, 13 insertions(+), 11 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index e42bc2af398..85586ccab17 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7572,12 +7572,12 @@ class Router: ## LOGGING if num_retries > 0: kwargs = self.log_retry(kwargs=kwargs, e=original_exception) - skipped_deployment_ids: Final = self._deployment_ids_to_skip_on_retry( + first_skipped_ids: Final = self._deployment_ids_to_skip_on_retry( exception=original_exception, already_skipped=kwargs.get("_retry_skipped_deployment_ids"), ) - if skipped_deployment_ids: - kwargs["_retry_skipped_deployment_ids"] = skipped_deployment_ids + if first_skipped_ids: + kwargs["_retry_skipped_deployment_ids"] = first_skipped_ids # rebind-ok: the next attempt reads it else: raise @@ -7647,12 +7647,12 @@ class Router: except Exception: raise e - retry_skipped_deployment_ids = self._deployment_ids_to_skip_on_retry( + skipped_ids = self._deployment_ids_to_skip_on_retry( exception=e, already_skipped=kwargs.get("_retry_skipped_deployment_ids"), ) - if retry_skipped_deployment_ids: - kwargs["_retry_skipped_deployment_ids"] = retry_skipped_deployment_ids + if skipped_ids: + kwargs["_retry_skipped_deployment_ids"] = skipped_ids # rebind-ok: the next attempt reads it _timeout = self._time_to_sleep_before_retry( e=e, remaining_retries=remaining_retries, diff --git a/litellm/utils.py b/litellm/utils.py index 238225eff99..f99d7e6a4b7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4919,10 +4919,12 @@ def get_excluded_filtered_deployments( across the remaining deployments in the same model group after one of them has failed. - If the filter would leave no deployments, an empty list is returned so the - caller raises its usual no-deployments error and the weighted-failover - helper falls through to the cross-group fallback path. Returning the - original unfiltered list here would re-include the just-failed deployment. + If the filter would leave no deployments, an empty list is returned and the + caller decides what that means. Weighted failover lets it raise the usual + no-deployments error and fall through to the cross-group fallback path; the + retry skip in `async_get_healthy_deployments` deliberately falls back to the + unfiltered list, so a request every deployment refused still comes back with + the provider's own error rather than a no-deployments one. """ if not excluded_deployment_ids: return healthy_deployments diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index b50551ec105..6b69677d490 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -7,7 +7,7 @@ - {id: reliability.retry.timeout.succeeds_within_retries, module: reliability, tier: P0, behavior: retry, variant: timeout, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:44", rationale: "Timeout retried per policy"} - {id: reliability.retry.429.succeeds_within_retries, module: reliability, tier: P0, behavior: retry, variant: "429", assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:46", rationale: "429 retried per RateLimitErrorRetries policy"} - {id: reliability.retry.auth.succeeds_within_retries, module: reliability, tier: P1, behavior: retry, variant: auth, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:42", rationale: "Transient auth glitch retry"} -- {id: reliability.retry.context_window.succeeds_within_retries, module: reliability, tier: P1, behavior: retry, variant: context_window, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:51", rationale: "Multi-attempt on context error"} +- {id: reliability.retry.context_window.succeeds_within_retries, module: reliability, tier: P1, behavior: retry, variant: context_window, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:51", fail_before_fix: proven, rationale: "A context-window 400 under BadRequestErrorRetries retries onto a sibling deployment in the same model group, instead of coming straight back as the 400 the deployment that just refused it returned"} - {id: reliability.cooldown.5xx.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "5xx", assertions: [trips_then_recovers], exercised_on: [chat_completions, messages], source: "cooldown_handlers.py:40", rationale: "Deployment cools after repeated 5xx, recovers after cooldown_time"} - {id: reliability.cooldown.429.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "429", assertions: [trips_then_recovers], exercised_on: [chat_completions, messages], source: "cooldown_handlers.py:69", rationale: "Cools on 429, avoids hammering exhausted provider"} - {id: reliability.cooldown.auth.trips_then_recovers, module: reliability, tier: P1, behavior: cooldown, variant: auth, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages], source: "cooldown_handlers.py:74", rationale: "Cools on 401 auth error"} From a72041b7572fad9e357e67dc4dfb913822d428ce Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 01:03:32 -0700 Subject: [PATCH 28/32] test(router): pin the retry skip list across attempts in a model group --- tests/test_litellm/test_router.py | 48 +++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 05ac672691e..a7d4e66bf73 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -39,6 +39,7 @@ from litellm.router import ( _anthropic_stream_should_drop_pre_content_ping, _is_retriable_anthropic_status, ) +from litellm.router_strategy import simple_shuffle from litellm.types.router import DeploymentTypedDict @@ -13423,6 +13424,53 @@ async def test_router_retry_policy_400_keeps_upstream_error_when_tags_narrow_the assert untagged.call_count == 0 +@pytest.mark.asyncio +async def test_router_retry_policy_400_never_returns_to_a_deployment_that_already_refused( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(simple_shuffle.random, "choice", lambda deployments: deployments[0]) + router = litellm.Router( + model_list=[ + _retry_skip_deployment("first-refuser", "first-refuser", litellm_params={"weight": 1}), + _retry_skip_deployment("second-refuser", "second-refuser", litellm_params={"weight": 0}), + _retry_skip_deployment("accepting", "accepting", litellm_params={"weight": 0}), + ], + num_retries=3, + retry_policy={"BadRequestErrorRetries": 3}, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + first = respx_mock.post("https://first-refuser.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + second = respx_mock.post("https://second-refuser.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + accepting = respx_mock.post("https://accepting.local/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-lit-7036", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.6", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "hi back"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + }, + ) + ) + response = await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) + + assert first.call_count == 1 + assert second.call_count == 1 + assert accepting.call_count == 1 + assert response.choices[0].message.content == "hi back" + + def _make_failure_logging_obj(): return LiteLLMLogging( model="gpt-5.6", From 168a0055a244acdcf97c330c52e085ab40b1424c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 6 Sep 2026 12:44:19 -0700 Subject: [PATCH 29/32] chore(lint): stop ratcheting *-budget.json on PR branches (#39937) --- CLAUDE.md | 2 +- scripts/test_quality_gate.py | 59 +++----------------- tests/test_litellm/test_test_quality_gate.py | 33 ++--------- 3 files changed, 13 insertions(+), 81 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index d9e9e8f1586..a7b9b6b9bc0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -52,7 +52,7 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a Python max line length is 120, not 88 -When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing +Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on `litellm_internal_staging` in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. If your branch already carries a budget edit, drop it before opening the PR `make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice diff --git a/scripts/test_quality_gate.py b/scripts/test_quality_gate.py index 7d34b194f1c..9e22eec29ab 100644 --- a/scripts/test_quality_gate.py +++ b/scripts/test_quality_gate.py @@ -10,17 +10,12 @@ base. Every rule is seeded at exactly its count on the day the gate landed, so the suite's existing debt is grandfathered and any net-new violation trips the gate -immediately. ``--update`` ratchets a limit down by the violations this branch -fixed relative to its branch point (the merge-base), so the ceilings only ever -fall. Base counts are measured with the *current* checker, so a rule introduced -on this branch is counted at the base too and ratchets like every other one. - -Only ever falling is not the same as always falling, so the gate enforces the -second half: a branch that clears violations and leaves the ceiling above its -new count fails, naming the rules and telling the author to run -``make lint-budget-update``. Without that, a removed violation could come back -later under a ceiling nobody lowered. Drift already in the base is never -blamed, so this fires only on the branch that did the clearing. +immediately. ``--update`` ratchets a limit down by the violations fixed relative +to ``--base``, so the ceilings only ever fall. Base counts are measured with the +*current* checker, so a rule introduced on this branch is counted at the base too +and ratchets like every other one. The ratchet runs as a scheduled automation +against litellm_internal_staging, not on PR branches, so concurrent PRs never +race to edit the same limit. The deliberate difference from its sibling: this gate has no headroom anywhere. Type discipline seeded LIT010/LIT011 at 1.5x to leave room for an in-flight @@ -144,21 +139,6 @@ def over_ceiling(head: Mapping[str, int], budget: Mapping[str, Mapping[str, int] ) -def unratcheted( - head: Mapping[str, int], - base: Mapping[str, int], - budget: Mapping[str, Mapping[str, int]], -) -> tuple[Breach, ...]: - """Rules this branch cleared without lowering the ceiling behind them. Requires - both `head < base`, so drift already in the base is never blamed on this change, - and `head < limit`, so a ceiling already at the count is left alone.""" - return tuple(sorted( - Breach(rule, head.get(rule, 0), spec["limit"], head.get(rule, 0) - base.get(rule, 0)) - for rule, spec in budget.items() - if head.get(rule, 0) < base.get(rule, 0) and head.get(rule, 0) < spec["limit"] - )) - - def evaluate( head: Mapping[str, int], base: Mapping[str, int], @@ -198,38 +178,15 @@ def introduced( return tuple(v for v in violations if v.line in changed.get(v.file, frozenset())) -def touches_measured_tree(base_point: str) -> bool: - """Whether this branch changed anything that can move a count. A branch that - touches neither the test tree nor the checker cannot have cleared a violation, - so the base scan is skipped and the gate stays cheap on the common change.""" - changed: Final = _run( - ["git", "diff", "--name-only", base_point, "--", TARGET, str(CHECKER.relative_to(REPO_ROOT))] - ) - return bool(changed.strip()) - - def cmd_check(base: str) -> None: budget: Final = json.loads(BUDGET_PATH.read_text()) head: Final = head_violations() head_counts: Final = count_by_rule(head) - base_point: Final = resolve_base_point(base) - if not over_ceiling(head_counts, budget) and not touches_measured_tree(base_point): + if not over_ceiling(head_counts, budget): print(f"OK: every TQ rule is within its test-suite ceiling (base {base})") return + base_point: Final = resolve_base_point(base) base_at_point: Final = base_counts(base_point) - stale: Final = unratcheted(head_counts, base_at_point, budget) - if stale: - print(f"FAIL: TQ-rule limits were left above the count this branch reached (base {base}):") - for breach in stale: - print( - f" {breach.rule}: this branch cleared {-breach.added} down to {breach.total}, " - f"but the limit is still {breach.cap}" - ) - print( - "Run `make lint-budget-update` and commit the lowered limits, so the " - "violations you cleared cannot come back under a ceiling nobody moved." - ) - raise SystemExit(1) breaches: Final = evaluate(head_counts, base_at_point, budget) if not breaches: print(f"OK: every TQ rule is within its test-suite ceiling (base {base})") diff --git a/tests/test_litellm/test_test_quality_gate.py b/tests/test_litellm/test_test_quality_gate.py index 8cce6bc735a..a8b38ecdf49 100644 --- a/tests/test_litellm/test_test_quality_gate.py +++ b/tests/test_litellm/test_test_quality_gate.py @@ -1,11 +1,9 @@ """Tests for scripts/test_quality_gate.py. -The gate's whole value is that it blames a change only for what it adds, that a limit -can never rise, and that a limit cannot stay above a count the branch pushed below it. -All three live in pure functions, so they are tested directly: `evaluate` for the blame -rule, `ratcheted_budget` for the one-way ratchet, `unratcheted` for the ceiling a branch -left behind, and `parse_changed_lines` for the diff scan that turns a breach into -file:line. +The gate's whole value is that it blames a change only for what it adds and that a +limit can never rise. Both live in pure functions, so they are tested directly: +`evaluate` for the blame rule, `ratcheted_budget` for the one-way ratchet, and +`parse_changed_lines` for the diff scan that turns a breach into file:line. """ import importlib.util @@ -73,29 +71,6 @@ def test_ratchet_lowers_a_rule_introduced_on_this_branch_like_any_other(): assert updated["TQ001"]["limit"] == 4 -def test_a_branch_that_cleared_violations_must_lower_the_ceiling(): - stale = gate.unratcheted({"TQ001": 6}, {"TQ001": 10}, _BUDGET) - assert [(b.rule, b.total, b.cap, b.added) for b in stale] == [("TQ001", 6, 10, -4)] - - -def test_headroom_already_in_the_base_is_not_blamed_on_this_branch(): - assert gate.unratcheted({"TQ001": 6}, {"TQ001": 6}, _BUDGET) == () - - -def test_a_branch_that_cleared_down_to_the_ceiling_exactly_is_clean(): - assert gate.unratcheted({"TQ001": 10}, {"TQ001": 12}, _BUDGET) == () - - -def test_a_branch_that_added_violations_is_not_a_ratchet_finding(): - assert gate.unratcheted({"TQ001": 14}, {"TQ001": 10}, _BUDGET) == () - - -def test_the_ratchet_finding_survives_the_update_that_answers_it(): - cleared = {"TQ001": 6} - updated = gate.ratcheted_budget(_BUDGET, cleared, {"TQ001": 10}) - assert gate.unratcheted(cleared, {"TQ001": 10}, updated) == () - - def test_parse_changed_lines_groups_hunks_under_their_own_file(): diff = ( "diff --git a/tests/a.py b/tests/a.py\n" From 728d0953af94c29959232032c0646bf13f87cfaf Mon Sep 17 00:00:00 2001 From: yujonglee Date: Mon, 7 Sep 2026 09:00:58 -0700 Subject: [PATCH 30/32] ci: simplify Rust checks and remove wheel PR comments (#39975) * ci: limit Rust workflows to Rust directory changes * ci: run Rust checks when their workflow changes * ci: report Rust wheels only for successful Rust changes * ci: keep Rust wheel reports in the workflow summary * ci: group Rust lint and validation jobs * ci: keep Rust job names distinct from required lint and test checks * ci: drop the unused Python setup from the Rust lint job --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../workflows/report-rust-release-wheel.yml | 129 ------------------ .github/workflows/test-rust.yml | 111 +++++++-------- 2 files changed, 50 insertions(+), 190 deletions(-) delete mode 100644 .github/workflows/report-rust-release-wheel.yml diff --git a/.github/workflows/report-rust-release-wheel.yml b/.github/workflows/report-rust-release-wheel.yml deleted file mode 100644 index 74e6be69604..00000000000 --- a/.github/workflows/report-rust-release-wheel.yml +++ /dev/null @@ -1,129 +0,0 @@ -name: Report LiteLLM Rust release wheel - -on: # zizmor: ignore[dangerous-triggers] reporter executes no PR code and consumes no PR artifacts or outputs - workflow_run: - workflows: - - LiteLLM Rust - types: - - completed - -permissions: {} - -concurrency: - group: ${{ github.workflow }}-${{ github.event.workflow_run.pull_requests[0].number || github.event.workflow_run.id }} - cancel-in-progress: false - -jobs: - report-release-wheel: - name: report release wheel - if: >- - github.event.workflow_run.event == 'pull_request' && - github.event.workflow_run.path == '.github/workflows/test-rust.yml' && - github.event.workflow_run.head_repository.full_name == github.repository && - github.event.workflow_run.pull_requests[0].number != null - runs-on: ubuntu-latest - timeout-minutes: 5 - permissions: - pull-requests: write - - steps: - - name: Link release wheel report on PR - uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0 - env: - COMMENT_MARKER: "" - with: - script: | - const marker = process.env.COMMENT_MARKER; - const workflowRun = context.payload.workflow_run; - const allowedConclusions = new Set([ - "action_required", - "cancelled", - "failure", - "neutral", - "skipped", - "stale", - "startup_failure", - "success", - "timed_out", - ]); - if ( - !allowedConclusions.has(workflowRun.conclusion) || - workflowRun.event !== "pull_request" || - workflowRun.path !== ".github/workflows/test-rust.yml" || - workflowRun.head_repository?.full_name !== - `${context.repo.owner}/${context.repo.repo}` || - workflowRun.pull_requests?.length !== 1 - ) { - throw new Error("unexpected source workflow"); - } - const pullRequest = workflowRun.pull_requests[0]; - const pullRequestNumber = pullRequest.number; - const headSha = workflowRun.head_sha; - const runId = workflowRun.id; - if ( - !Number.isSafeInteger(pullRequestNumber) || - pullRequestNumber <= 0 || - !Number.isSafeInteger(runId) || - runId <= 0 || - !/^[0-9a-f]{40}$/.test(headSha) || - pullRequest.head?.sha !== headSha - ) { - throw new Error("invalid source workflow metadata"); - } - const runUrl = - `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}` + - `/actions/runs/${runId}`; - const result = - workflowRun.conclusion === "success" - ? "successfully" - : `with \`${workflowRun.conclusion}\``; - const body = [ - marker, - "## LiteLLM Rust workflow", - "", - `Workflow completed ${result} for \`${headSha}\``, - "", - `[View workflow run](${runUrl})`, - ].join("\n"); - const comments = await github.paginate(github.rest.issues.listComments, { - owner: context.repo.owner, - repo: context.repo.repo, - issue_number: pullRequestNumber, - per_page: 100, - }); - const existing = comments.find( - (comment) => - comment.user?.login === "github-actions[bot]" && - comment.body?.startsWith(marker), - ); - const currentPullRequest = ( - await github.rest.pulls.get({ - owner: context.repo.owner, - repo: context.repo.repo, - pull_number: pullRequestNumber, - }) - ).data; - if ( - currentPullRequest.state !== "open" || - currentPullRequest.head.repo?.full_name !== - `${context.repo.owner}/${context.repo.repo}` || - currentPullRequest.head.sha !== headSha - ) { - core.info("source workflow no longer matches the current pull request head"); - return; - } - if (existing) { - await github.rest.issues.updateComment({ - owner: context.repo.owner, - repo: context.repo.repo, - comment_id: existing.id, - body, - }); - } else { - await github.rest.issues.createComment({ - owner: context.repo.owner, - repo: context.repo.repo, - issue_number: pullRequestNumber, - body, - }); - } diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 9b8b132df62..4f56e78ddee 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -7,6 +7,7 @@ on: - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" + - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" @@ -22,6 +23,7 @@ on: - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" + - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" @@ -34,102 +36,89 @@ concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: true +env: + CARGO_TERM_COLOR: always + jobs: - rust-checks: - name: rustfmt, clippy, test + rust-lint: runs-on: ubuntu-latest timeout-minutes: 10 defaults: run: working-directory: litellm-rust - env: - CARGO_TERM_COLOR: always - steps: - - name: Checkout repository - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: persist-credentials: false - - name: Set up Rust - run: rustup toolchain install + - run: rustup toolchain install --no-self-update - - name: Cache Cargo registry and target - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + - run: cargo fmt --check + + - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 with: path: | ~/.cargo/registry ~/.cargo/git litellm-rust/target - key: ${{ runner.os }}-cargo-${{ hashFiles('rust-toolchain.toml', 'litellm-rust/Cargo.lock') }} + key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }} restore-keys: | - ${{ runner.os }}-cargo- + ${{ runner.os }}-cargo-${{ github.job }}- - - name: Check Rust formatting - run: cargo fmt --check + - run: cargo clippy --workspace --all-targets --locked -- -D warnings - - name: Run Clippy - run: cargo clippy --workspace --all-targets --locked -- -D warnings + - run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings - - name: Run Clippy with Bedrock auth - run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings + - run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings - - name: Run Clippy with all gateway features - run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings - - - name: Run Rust tests - run: cargo test --workspace --locked - - - name: Run core tests with Bedrock auth - run: cargo test -p litellm-core --features bedrock-auth --locked - - # Not --all-features: python-config links libpython, which this job does not install. - - name: Run gateway tests with the server feature - run: cargo test -p litellm-ai-gateway --features server --locked - - release-wheel: - name: release wheel + rust-test: runs-on: ubuntu-latest - timeout-minutes: 20 - permissions: - contents: read - env: - CARGO_TERM_COLOR: always + timeout-minutes: 30 steps: - - name: Checkout repository - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: persist-credentials: false - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 with: python-version: "3.12" - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries + - uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" - - name: Set up Rust - run: rustup toolchain install + - run: rustup toolchain install --no-self-update - - name: Build release wheel - run: uv build --wheel --out-dir dist + - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + litellm-rust/target + key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }} + restore-keys: | + ${{ runner.os }}-cargo-${{ github.job }}- - - name: Build panic contract wheel - run: >- + - run: cargo test --workspace --locked + working-directory: litellm-rust + + - run: cargo test -p litellm-core --features bedrock-auth --locked + working-directory: litellm-rust + + - run: cargo test -p litellm-ai-gateway --features server --locked + working-directory: litellm-rust + + - run: uv build --wheel --out-dir dist + + - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl + env: + RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} + + - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl + + - run: >- uv build --wheel --out-dir panic-dist --config-setting "maturin.build-args=--features panic-test,extension-module" - - name: Smoke-test native panic unwinding - run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl - - - name: Verify stripped native extension - env: - RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl - - - name: Test native route wheel - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl + - run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl From a1e7293fa921f6894aa0bc2fffb2db3f379342f0 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 7 Sep 2026 12:45:16 -0400 Subject: [PATCH 31/32] fix(mcp): apply key and team guardrails to MCP tool calls (#39629) * fix(mcp): apply key and team guardrails to MCP tool calls Guardrails attached to a virtual key or team were only enforced on LLM routes. The synthetic request built for MCP tool call guardrail hooks carried no guardrails in its metadata, so a guardrail with default_on false never ran on tools/call even when the key explicitly listed it. Resolve key, team, and project guardrails onto the synthetic request with the same helper the chat path uses. * fix(mcp): pass project metadata through without a mutable default * fix(mcp): mark the request dict parameter mutable-ok with a reason * test(mcp): explain the premium_user patch and tighten the helper docstring --- litellm/proxy/litellm_pre_call_utils.py | 40 +++++++++------ litellm/proxy/utils.py | 10 +++- .../mcp_server/test_mcp_server_manager.py | 49 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_utils.py | 28 +++++++++++ 4 files changed, 109 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d026c5510e6..56512570448 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2882,6 +2882,28 @@ def _add_guardrails_from_policies_in_metadata( ) +def add_guardrails_from_auth_metadata( + user_api_key_dict: UserAPIKeyAuth, + data: dict, # mutable-ok: writes guardrails into the live request dict, same contract as the helpers it wraps + metadata_variable_name: str, +) -> None: + """Resolve key, team, and project guardrails, direct and via policies, onto the request metadata.""" + _add_guardrails_from_key_or_team_metadata( + key_metadata=user_api_key_dict.metadata, + team_metadata=user_api_key_dict.team_metadata, + project_metadata=user_api_key_dict.project_metadata, + data=data, + metadata_variable_name=metadata_variable_name, + ) + _add_guardrails_from_policies_in_metadata( + key_metadata=user_api_key_dict.metadata, + team_metadata=user_api_key_dict.team_metadata, + project_metadata=user_api_key_dict.project_metadata, + data=data, + metadata_variable_name=metadata_variable_name, + ) + + async def move_guardrails_to_metadata( data: dict, _metadata_variable_name: str, @@ -2914,22 +2936,8 @@ async def move_guardrails_to_metadata( data.pop("policies", None) return - # Check key/team/project-level guardrails - _add_guardrails_from_key_or_team_metadata( - key_metadata=user_api_key_dict.metadata, - team_metadata=user_api_key_dict.team_metadata, - project_metadata=project_metadata, - data=data, - metadata_variable_name=_metadata_variable_name, - ) - - ######################################################################################### - # Add guardrails from policies attached to key/team/project metadata - ######################################################################################### - _add_guardrails_from_policies_in_metadata( - key_metadata=user_api_key_dict.metadata, - team_metadata=user_api_key_dict.team_metadata, - project_metadata=project_metadata, + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_dict, data=data, metadata_variable_name=_metadata_variable_name, ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b4e4dfeae67..f1103ce6a29 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -152,7 +152,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( from litellm.proxy.hooks.sensitive_data_routing import ( _PROXY_SensitiveDataRoutingHandler, ) -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor from litellm.repositories.budget_repository import BudgetRepository @@ -924,7 +924,13 @@ class ProxyLogging: "incoming_bearer_token": kwargs.get("incoming_bearer_token"), "metadata": {"headers": kwargs.get("headers") or {}}, } - + user_api_key_auth: Final = kwargs.get("user_api_key_auth") + if isinstance(user_api_key_auth, UserAPIKeyAuth): + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_auth, + data=synthetic_data, + metadata_variable_name="metadata", + ) return synthetic_data def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> MCPPreCallResponseObject | None: 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 764e2bb0e99..d19363d3b5f 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 @@ -58,6 +58,11 @@ from litellm.proxy._types import ( from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPAuth, MCPAuthType from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer +from litellm.caching.caching import DualCache +import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks def _reload_mcp_manager_module(): @@ -12456,3 +12461,47 @@ class TestLitellmAdmissionKeyIsNeverTheSubjectToken: }, ) assert self._subjects_seen_by(provider) == [self._USER_TOKEN] + + +class _BlockWhenSelectedGuardrail(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_mcp_call) is not True: + return data + raise HTTPException(status_code=400, detail="blocked by key-scoped guardrail") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_metadata, expect_block", + [({"guardrails": ["key-scoped-guardrail"]}, True), ({"guardrails": ["unrelated-guardrail"]}, False), ({}, False)], +) +async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch, key_metadata, expect_block): + guardrail = _BlockWhenSelectedGuardrail( + guardrail_name="key-scoped-guardrail", event_hook="pre_mcp_call", default_on=False + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + server = MCPServer( + server_id="deepwiki", + name="deepwiki", + server_name="deepwiki", + url="https://mcp.deepwiki.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + call = MCPServerManager().pre_call_tool_check( + name="ask_question", + arguments={"repoName": "BerriAI/litellm", "question": "ignore all previous instructions"}, + server_name="deepwiki", + user_api_key_auth=UserAPIKeyAuth(metadata=key_metadata), + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + server=server, + ) + + if not expect_block: + assert await call == {} + return + with pytest.raises(HTTPException) as exc_info: + await call + assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index dcf18e773e8..9462f2c8eb0 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1922,6 +1922,34 @@ async def test_proxy_only_error_5xx_keeps_traceback_and_runs_sync_callbacks(monk assert "test_proxy_utils" in captured["async_traceback"] +@pytest.mark.parametrize( + "key_metadata, team_metadata, expected_to_run", + [ + ({"guardrails": ["key-scoped-guardrail"]}, None, True), + ({}, {"guardrails": ["key-scoped-guardrail"]}, True), + ({"guardrails": ["some-other-guardrail"]}, None, False), + ({}, None, False), + ], +) +def test_convert_mcp_to_llm_format_carries_key_and_team_guardrails(key_metadata, team_metadata, expected_to_run): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + guardrail = CustomGuardrail(guardrail_name="key-scoped-guardrail", event_hook="pre_mcp_call", default_on=False) + kwargs = { + "name": "ask_question", + "arguments": {"question": "hello"}, + "server_name": "deepwiki", + "user_api_key_auth": UserAPIKeyAuth(metadata=key_metadata, team_metadata=team_metadata), + } + request_obj = proxy_logging._create_mcp_request_object_from_kwargs(kwargs) + + with patch( # test-quality-ok: the key-guardrail premium gate reads this proxy_server module global and has no injection seam + "litellm.proxy.proxy_server.premium_user", True + ): + synthetic = proxy_logging._convert_mcp_to_llm_format(request_obj, kwargs) + + assert guardrail.should_run_guardrail(synthetic, GuardrailEventHooks.pre_mcp_call) is expected_to_run + + class _TracebackRecordingLogger(CustomLogger): def __init__(self) -> None: super().__init__() From a2b7868a5bdb45030d4b224189eb1b1010ad1c4f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 7 Sep 2026 09:53:18 -0700 Subject: [PATCH 32/32] test(ui): pin wire contracts for key, model and MCP server forms (#40019) * test(ui): pin wire contracts for key, model and MCP server forms Add vitest cases that pin what the key edit, key create, model edit and MCP server edit forms put on the wire: an edited field reaches the request with its new value, a cleared field reaches it as an explicit null, and the dirty-only body is pinned as an expected failure until each form moves to pickDirty. Model edit also pins the cost-map-derived model_info fields as an expected failure. KeyEditView hands a cleared max_budget to KeyInfoView as an empty string and handleKeyUpdate maps it to null, so the null is pinned at the /key/update boundary in key_info_view.test.tsx and the KeyEditView case is an expected failure. buildEditServerPayload passes a cleared description through as an empty string, so that case is an expected failure too. * test(ui): split masked model_info pins and retarget the create tracker The model_info expected-failure case held three assertions, and it.fails stops at the first one, so a later revamp that fixed max_input_tokens while leaving mode leaking would still report an expected failure. Split it into one case per pinned field group so each flips on its own. The key create tracker asserted a body of only key_alias, which a create can never send: key_type, user_id, duration and metadata are always mounted. Retarget it at the real over-send, which is the Optional Settings section adding fifteen undefined-valued keys when the user opens it without filling anything in. --- .../editServerPayload.differential.test.ts | 43 ++++++++++- .../src/components/model_info_view.test.tsx | 72 +++++++++++++++++++ .../create_key_button.integration.test.tsx | 31 ++++++++ .../templates/key_edit_view.test.tsx | 52 ++++++++++++++ .../templates/key_info_view.test.tsx | 10 +++ 5 files changed, 207 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.test.ts index f9cab181e71..66786fac1c3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.test.ts @@ -12,7 +12,7 @@ import { } from "@/components/mcp_tools/types"; import { AUTH_TYPES_REQUIRING_CREDENTIALS } from "./createServerPayload"; import { TOOL_DISPLAY_NAME_PATTERN, normalizeEnvVars } from "./utils"; -import { buildEditServerPayload, type EditServerUiState } from "./editServerPayload"; +import { buildEditServerPayload, type EditServerFormValues, type EditServerUiState } from "./editServerPayload"; import { CASES, baseUi } from "./editServerPayload.differential.cases"; // GENERATED by scratchpad/emit_test.py. The body below is machine-extracted from @@ -317,6 +317,47 @@ describe("buildEditServerPayload matches the pre-extraction handleSave body", () }); }); +const EDIT_FORM_VALUES: EditServerFormValues = { + server_name: "srv", + alias: "srv_alias", + description: "a server", + transport: "http", + url: "https://example.com/mcp", + auth_type: "none", + mcp_access_groups: [], + extra_headers: [], + static_headers: [], + env_vars: [], + allow_all_keys: false, + available_on_public_internet: true, +}; + +describe("buildEditServerPayload wire contract", () => { + it("carries an edited alias and the server identifier onto the wire", () => { + const result = buildEditServerPayload({ ...EDIT_FORM_VALUES, alias: "renamed" }, baseUi); + + expect(result).toMatchObject({ kind: "ok", payload: { server_id: "srv_1", alias: "renamed" } }); + }); + + it.fails( + "sends description as an explicit null when the field is cleared (expected to fail until the forms revamp, tri-state PATCH tracker: today the cleared field reaches the wire as an empty string)", + () => { + const result = buildEditServerPayload({ ...EDIT_FORM_VALUES, description: "" }, baseUi); + + expect(result).toMatchObject({ kind: "ok", payload: { description: null } }); + }, + ); + + it.fails( + "sends only the server identifier and the edited alias (expected to fail until the forms revamp, tri-state PATCH tracker)", + () => { + const result = buildEditServerPayload({ ...EDIT_FORM_VALUES, alias: "renamed" }, baseUi); + + expect(result).toStrictEqual({ kind: "ok", payload: { server_id: "srv_1", alias: "renamed" } }); + }, + ); +}); + void ADMIN_CONFIG_CREDENTIAL_KEYS; void AUTH_TYPE; void AUTH_TYPES_REQUIRING_CREDENTIALS; diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx index 768183907db..3db9418dfb9 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -1729,5 +1729,77 @@ describe("ModelInfoView", () => { expect(payload.litellm_params.cache_control_injection_points).toEqual([{ location: "message", index: "2" }]); }); }); + + const setInputCost = (value: string) => { + fireEvent.change(screen.getByPlaceholderText("Enter input cost"), { target: { value } }); + }; + + it("carries an edited input cost and the model identifier onto the wire", async () => { + const user = userEvent.setup(); + await enterEditMode(user); + setInputCost("5"); + const payload = await save(user); + + expect(mockModelPatchUpdateCall.mock.calls[0][2]).toBe("123"); + expect(payload.litellm_params.input_cost_per_token).toBe(5 / 1_000_000); + }); + + it.fails( + "sends only the edited input cost (expected to fail until the forms revamp, tri-state PATCH tracker)", + async () => { + const user = userEvent.setup(); + await enterEditMode(user); + setInputCost("5"); + const payload = await save(user); + + expect(payload).toStrictEqual({ litellm_params: { input_cost_per_token: 5 / 1_000_000 } }); + }, + ); + + const savePayloadAfterCostEditOnResolvedModel = async () => { + const resolved = { + ...defaultModelData, + model_info: { + ...defaultModelData.model_info, + max_input_tokens: 128_000, + mode: "chat", + supports_vision: true, + supports_function_calling: true, + }, + }; + mockUseModelsInfo.mockReturnValue({ data: { data: [resolved] }, isLoading: false, error: null }); + mockModelInfoV1Call.mockResolvedValue({ data: [resolved] }); + const user = userEvent.setup(); + await enterEditMode(user); + setInputCost("5"); + return save(user); + }; + + it.fails( + "leaves max_input_tokens off the wire when only the input cost is edited (expected to fail until the forms revamp, tri-state PATCH tracker)", + async () => { + const payload = await savePayloadAfterCostEditOnResolvedModel(); + + expect(payload.model_info).not.toHaveProperty("max_input_tokens"); + }, + ); + + it.fails( + "leaves mode off the wire when only the input cost is edited (expected to fail until the forms revamp, tri-state PATCH tracker)", + async () => { + const payload = await savePayloadAfterCostEditOnResolvedModel(); + + expect(payload.model_info).not.toHaveProperty("mode"); + }, + ); + + it.fails( + "leaves every supports_ capability off the wire when only the input cost is edited (expected to fail until the forms revamp, tri-state PATCH tracker)", + async () => { + const payload = await savePayloadAfterCostEditOnResolvedModel(); + + expect(Object.keys(payload.model_info).filter((key) => key.startsWith("supports_"))).toStrictEqual([]); + }, + ); }); }); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index 045dcfa3ceb..3471ef00eb0 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -544,6 +544,37 @@ describe("CreateKey", () => { expect((await createdPayload()).metadata).toBe('{"team":"research"}'); }); + + it("carries the typed key alias and the chosen team onto the wire", async () => { + state.teams = [{ team_id: "team-1", team_alias: "Team One", models: [] }]; + await openModal({ teams: state.teams as unknown as Team[] }); + await nameTheKey("wire-alias"); + await userEvent.click(await screen.findByLabelText("Team")); + await userEvent.click(await screen.findByRole("option", { name: /Team One/ })); + await submit(); + + expect(await createdPayload()).toMatchObject({ key_alias: "wire-alias", team_id: "team-1" }); + }); + + it("sends team_id as an explicit null when no team is chosen", async () => { + await openModal(); + await nameTheKey(); + await submit(); + + expect(await createdPayload()).toHaveProperty("team_id", null); + }); + + it.fails( + "adds no keys for an Optional Settings section the user opened but never filled (expected to fail until the forms revamp, tri-state PATCH tracker)", + async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await submit(); + + expect(await createdPayload()).toStrictEqual(ALL_CLOSED_PAYLOAD); + }, + ); }); describe("key ownership", () => { diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index 8b506329595..4f732083c3e 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -2300,5 +2300,57 @@ describe("KeyEditView", () => { }); expect(onSubmitMock.mock.calls[0][0]).toHaveProperty("tag_rpm_limit", { "test-tag": 7 }); }); + + const setRpmLimit = (value: string) => { + fireEvent.change(screen.getByLabelText("RPM Limit"), { target: { value } }); + }; + + it("carries an edited RPM limit and the key identifier onto the wire", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmitMock); + await screen.findByRole("button", { name: /save changes/i }); + + setRpmLimit("25"); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalledTimes(1); + }); + expect(onSubmitMock.mock.calls[0][0]).toMatchObject({ token: "test-token-123", rpm_limit: "25" }); + }); + + it.fails( + "sends max_budget as an explicit null when the field is cleared (expected to fail until the forms revamp, tri-state PATCH tracker: today the view hands KeyInfoView an empty string and handleKeyUpdate maps it to null)", + async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmitMock); + await screen.findByRole("button", { name: /save changes/i }); + + await userEvent.clear(screen.getByLabelText("Max Budget (USD)")); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalledTimes(1); + }); + expect(onSubmitMock.mock.calls[0][0]).toHaveProperty("max_budget", null); + }, + ); + + it.fails( + "sends only the key identifier and the edited RPM limit (expected to fail until the forms revamp, tri-state PATCH tracker)", + async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmitMock); + await screen.findByRole("button", { name: /save changes/i }); + + setRpmLimit("25"); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalledTimes(1); + }); + expect(onSubmitMock.mock.calls[0][0]).toStrictEqual({ token: "test-token-123", rpm_limit: "25" }); + }, + ); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index 7817d4e7cea..7a27cc41e5e 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -1015,6 +1015,16 @@ describe("KeyInfoView", () => { expect(keyUpdateCall).toHaveBeenCalledWith(expect.anything(), expect.objectContaining({ policies: [] })); }); + + it("puts the key identifier and an explicit null max_budget on the wire when the edit view hands over a cleared budget", async () => { + await enterEditMode({ ...MOCK_KEY_DATA, user_id: "proxy-admin-user" } as KeyResponse); + await editViewMocks.onSubmit!({ token: MOCK_KEY_DATA.token, max_budget: "" }); + + expect(keyUpdateCall).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ key: "test-token-123", max_budget: null }), + ); + }); }); describe("MCP tool permissions on save", () => {