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 001/164] 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 002/164] 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 318b6a4b36d31c4255c66d7b6289bf782fd37d20 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 29 Aug 2026 14:26:48 -0700
Subject: [PATCH 003/164] fix(mcp): forward staged credentials on
 /mcp-rest/test/connection like /test/tools/list

The connection preview built its temporary MCP client without the credentials the
not-yet-saved server config carries: the Authorization bearer an OAuth2 authorization_code
server had just been granted, the auth_value of an api_key, bearer_token, basic, or
authorization server, and the stored credentials of a saved server being edited. The tools
preview forwarded all three, so the same request succeeded there and failed on the connection
test with the generic "Failed to connect to MCP server" message

Both previews now resolve those credentials through one shared staging step, so they cannot
drift apart again, and the Authorization header is only forwarded upstream when the primary
x-litellm-api-key header carried admission, since otherwise it is the caller's LiteLLM key
---
 .../mcp_server/rest_endpoints.py              |  85 ++++++-----
 .../mcp_server/test_rest_endpoints.py         | 134 ++++++++++++++++++
 2 files changed, 186 insertions(+), 33 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
index 3efb6429326..a1583154916 100644
--- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
+++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
@@ -1,11 +1,13 @@
 import asyncio
 import importlib
 from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, Literal
 
 import httpx
 from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
+from starlette.datastructures import Headers
 
 from litellm._logging import verbose_logger
 from litellm.exceptions import (
@@ -1130,6 +1132,45 @@ if MCP_AVAILABLE:
         scopes: Final[list[str] | None] = scopes_raw if isinstance(scopes_raw, list) else None
         return client_id, client_secret, scopes
 
+    _STAGED_AUTH_VALUE_AUTH_TYPES: Final = frozenset(
+        (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization)
+    )
+
+    @dataclass(frozen=True, slots=True)
+    class _StagedServerTest:
+        request: NewMCPServerRequest
+        mcp_auth_header: str | None
+        oauth2_headers: dict[str, str] | None
+
+    def _stage_server_test(new_mcp_server_request: NewMCPServerRequest, headers: Headers) -> _StagedServerTest:
+        """
+        Resolve the credentials a not-yet-saved server config carries for a preview call.
+
+        Both preview endpoints (``/test/connection`` and ``/test/tools/list``) must hand the
+        temporary client the same credentials, or a server that the saved connection reaches
+        fine fails one of them.
+        """
+        from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
+            MCPRequestHandler,
+        )
+
+        request: Final = _inherit_credentials_from_existing_server(new_mcp_server_request)
+        mcp_auth_header: Final = (
+            request.credentials.get("auth_value")
+            if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES and isinstance(request.credentials, dict)
+            else None
+        )
+        # Authorization doubles as the admission fallback (LITELLM_API_KEY_HEADER_NAME_SECONDARY):
+        # when the primary x-litellm-api-key header is absent, the Authorization value is the
+        # caller's LiteLLM key, not an upstream token, and must never be forwarded upstream.
+        oauth2_headers: Final = (
+            MCPRequestHandler._get_oauth2_headers_from_headers(headers)
+            if request.auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES
+            and headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY)
+            else None
+        )
+        return _StagedServerTest(request=request, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers)
+
     async def _execute_with_mcp_client(
         request: NewMCPServerRequest,
         operation: Callable[..., Awaitable[Mapping[str, object]]],
@@ -1339,6 +1380,8 @@ if MCP_AVAILABLE:
                 },
             )
 
+        staged: Final = _stage_server_test(new_mcp_server_request, request.headers)
+
         async def _test_connection_operation(client):
             async def _noop(session):
                 return "ok"
@@ -1347,8 +1390,10 @@ if MCP_AVAILABLE:
             return {"status": "ok"}
 
         return await _execute_with_mcp_client(
-            new_mcp_server_request,
+            staged.request,
             _test_connection_operation,
+            mcp_auth_header=staged.mcp_auth_header,
+            oauth2_headers=staged.oauth2_headers,
             raw_headers=_safe_get_request_headers(request),
         )
 
@@ -1369,37 +1414,11 @@ if MCP_AVAILABLE:
                 },
             )
 
-        new_mcp_server_request = _inherit_credentials_from_existing_server(new_mcp_server_request)
+        staged: Final = _stage_server_test(new_mcp_server_request, request.headers)
 
         # For OpenAPI spec servers, generate tools from the spec directly
-        if new_mcp_server_request.spec_path:
-            return await _preview_openapi_tools(new_mcp_server_request.spec_path)
-
-        from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
-            MCPRequestHandler,
-        )
-
-        headers: Final = request.headers
-
-        mcp_auth_header: str | None = None
-        if new_mcp_server_request.auth_type in {
-            MCPAuth.api_key,
-            MCPAuth.bearer_token,
-            MCPAuth.basic,
-            MCPAuth.authorization,
-        }:
-            credentials: Final = getattr(new_mcp_server_request, "credentials", None)
-            if isinstance(credentials, dict):
-                mcp_auth_header = credentials.get("auth_value")
-
-        # Authorization doubles as the admission fallback (LITELLM_API_KEY_HEADER_NAME_SECONDARY):
-        # when the primary x-litellm-api-key header is absent, the Authorization value is the
-        # caller's LiteLLM key, not an upstream token, and must never be forwarded upstream.
-        oauth2_headers: dict[str, str] | None = None
-        if new_mcp_server_request.auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and headers.get(
-            MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY
-        ):
-            oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers)
+        if staged.request.spec_path:
+            return await _preview_openapi_tools(staged.request.spec_path)
 
         async def _list_tools_operation(client):
             async def _list_tools_session_operation(session):
@@ -1415,9 +1434,9 @@ if MCP_AVAILABLE:
             }
 
         return await _execute_with_mcp_client(
-            new_mcp_server_request,
+            staged.request,
             _list_tools_operation,
-            mcp_auth_header=mcp_auth_header,
-            oauth2_headers=oauth2_headers,
+            mcp_auth_header=staged.mcp_auth_header,
+            oauth2_headers=staged.oauth2_headers,
             raw_headers=_safe_get_request_headers(request),
         )
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
index ef5631218f3..d32ffc90b55 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
@@ -464,6 +464,140 @@ class TestTestConnection:
         route = _get_route("/mcp-rest/test/connection", "POST")
         assert _route_has_dependency(route, user_api_key_auth)
 
+    @staticmethod
+    def _capture_execute(monkeypatch) -> dict:
+        captured: dict = {}
+
+        async def fake_execute(
+            request,
+            operation,
+            mcp_auth_header=None,
+            oauth2_headers=None,
+            raw_headers=None,
+        ):
+            captured["request"] = request
+            captured["mcp_auth_header"] = mcp_auth_header
+            captured["oauth2_headers"] = oauth2_headers
+            return {"status": "ok"}
+
+        monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False)
+        return captured
+
+    @staticmethod
+    def _oauth2_authorization_code_payload(**overrides) -> NewMCPServerRequest:
+        return NewMCPServerRequest(
+            server_name="github_mcp",
+            url="https://api.githubcopilot.com/mcp/",
+            auth_type=MCPAuth.oauth2,
+            oauth2_flow="authorization_code",
+            authorization_url="https://github.com/login/oauth/authorize",
+            token_url="https://github.com/login/oauth/access_token",
+            **overrides,
+        )
+
+    @pytest.mark.asyncio
+    async def test_forwards_staged_oauth2_bearer(self, monkeypatch):
+        """The just-authorized upstream token rides the request's Authorization header, exactly
+        as /test/tools/list receives it; dropping it makes every authorization_code server fail
+        the connection test that its tools preview passes."""
+        from litellm.proxy._types import LitellmUserRoles
+
+        captured = self._capture_execute(monkeypatch)
+        request = _build_request(
+            {"x-litellm-api-key": "sk-admin-session", "authorization": "Bearer upstream-oauth-token"},
+            path="/mcp-rest/test/connection",
+        )
+
+        result = await rest_endpoints.test_connection(
+            request,
+            self._oauth2_authorization_code_payload(),
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
+        )
+
+        assert result == {"status": "ok"}
+        assert captured["oauth2_headers"] == {"Authorization": "Bearer upstream-oauth-token"}
+        assert captured["mcp_auth_header"] is None
+
+    @pytest.mark.asyncio
+    async def test_forwards_staged_auth_value(self, monkeypatch):
+        from litellm.proxy._types import LitellmUserRoles
+
+        captured = self._capture_execute(monkeypatch)
+        request = _build_request({"x-litellm-api-key": "sk-admin-session"}, path="/mcp-rest/test/connection")
+        payload = NewMCPServerRequest(
+            server_name="example",
+            url="https://example.com/mcp",
+            auth_type=MCPAuth.bearer_token,
+            credentials={"auth_value": "upstream-static-token"},
+        )
+
+        result = await rest_endpoints.test_connection(
+            request,
+            payload,
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
+        )
+
+        assert result == {"status": "ok"}
+        assert captured["mcp_auth_header"] == "upstream-static-token"
+        assert captured["oauth2_headers"] is None
+
+    @pytest.mark.asyncio
+    async def test_inherits_stored_credentials_of_saved_server(self, monkeypatch):
+        """The edit form resends a saved server without its masked credential; the stored one
+        must be used, as /test/tools/list already does."""
+        from litellm.proxy._types import LitellmUserRoles
+        from litellm.types.mcp_server.mcp_server_manager import MCPServer
+
+        captured = self._capture_execute(monkeypatch)
+        saved = MCPServer(
+            server_id="saved-server-id",
+            name="example",
+            url="https://example.com/mcp",
+            transport="http",
+            auth_type=MCPAuth.bearer_token,
+            authentication_token="stored-upstream-token",
+        )
+        monkeypatch.setattr(
+            rest_endpoints.global_mcp_server_manager,
+            "get_mcp_server_by_id",
+            lambda server_id: saved if server_id == "saved-server-id" else None,
+        )
+        request = _build_request({"x-litellm-api-key": "sk-admin-session"}, path="/mcp-rest/test/connection")
+        payload = NewMCPServerRequest(
+            server_id="saved-server-id",
+            server_name="example",
+            url="https://example.com/mcp",
+            auth_type=MCPAuth.bearer_token,
+        )
+
+        result = await rest_endpoints.test_connection(
+            request,
+            payload,
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
+        )
+
+        assert result == {"status": "ok"}
+        assert captured["mcp_auth_header"] == "stored-upstream-token"
+        assert captured["request"].credentials == {"auth_value": "stored-upstream-token"}
+
+    @pytest.mark.asyncio
+    async def test_does_not_forward_authorization_that_satisfied_admission(self, monkeypatch):
+        """With no x-litellm-api-key, the Authorization value is the caller's LiteLLM key and
+        must never reach the upstream."""
+        from litellm.proxy._types import LitellmUserRoles
+
+        captured = self._capture_execute(monkeypatch)
+        request = _build_request({"authorization": "Bearer sk-litellm-admission-key"}, path="/mcp-rest/test/connection")
+
+        result = await rest_endpoints.test_connection(
+            request,
+            self._oauth2_authorization_code_payload(),
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
+        )
+
+        assert result == {"status": "ok"}
+        assert captured["oauth2_headers"] is None
+
 
 class TestTestToolsList:
     pytestmark = pytest.mark.asyncio

From 4ef5db7c91ccfb4b690d811baf7cfad4129ab7ae Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 29 Aug 2026 16:36:05 -0700
Subject: [PATCH 004/164] fix(responses): drop unsupported reasoning param for
 openai non-reasoning models

---
 .../llms/openai/responses/transformation.py   | 29 ++++++++++++
 .../test_openai_responses_transformation.py   | 46 +++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py
index 2fa44cfc2e3..99ce158c4e2 100644
--- a/litellm/llms/openai/responses/transformation.py
+++ b/litellm/llms/openai/responses/transformation.py
@@ -75,6 +75,19 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
 
         return OpenAIGPT5Config.effort_resolves_to_none(model, effort)
 
+    @staticmethod
+    def _is_o_series_name(model: str) -> bool:
+        base: Final = model.split("/")[-1]
+        return len(base) > 1 and base[0] == "o" and base[1].isdigit()
+
+    def _supports_reasoning_param(self, model: str) -> bool:
+        if self._is_gpt_5_model(model=model) or self._is_o_series_name(model=model):
+            return True
+        base: Final = model.split("/")[-1]
+        if base not in litellm.open_ai_chat_completion_models:
+            return True
+        return litellm.supports_reasoning(model=base, custom_llm_provider=self.custom_llm_provider.value)
+
     @staticmethod
     def _enforce_min_max_output_tokens(max_output_tokens: "int | None") -> "int | None":
         """Raise sub-minimum max_output_tokens up to the OpenAI Responses API minimum.
@@ -124,6 +137,22 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
         if "max_output_tokens" in params:
             params["max_output_tokens"] = self._enforce_min_max_output_tokens(params.get("max_output_tokens"))
 
+        if (
+            self.custom_llm_provider == LlmProviders.OPENAI
+            and params.get("reasoning") is not None
+            and not self._supports_reasoning_param(model=model)
+        ):
+            if drop_params or litellm.drop_params:
+                params.pop("reasoning", None)
+            else:
+                raise litellm.UnsupportedParamsError(
+                    message=(
+                        f"{model} doesn't support the `reasoning` parameter. "
+                        "To drop unsupported params set `litellm.drop_params = True`"
+                    ),
+                    status_code=400,
+                )
+
         if self._is_gpt_5_model(model=model):
             temperature: Final = params.get("temperature")
             if temperature is not None and temperature != 1:
diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
index e314b94444b..66d22cf8fb0 100644
--- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
+++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
@@ -1626,3 +1626,49 @@ class TestResponsesSurfaceSharesTheEffortRule:
             drop_params=True,
         )
         assert ("temperature" in mapped) is temperature_survives
+
+
+class TestReasoningFollowsModelSupport:
+    """Responses API clients like Codex send `reasoning` on every request, and OpenAI 400s it
+    on non-reasoning models like gpt-4o. drop_params must strip it there, the same way the
+    chat completions surface already strips reasoning_effort for those models.
+    """
+
+    @pytest.mark.parametrize(
+        "model, reasoning_survives",
+        [
+            ("gpt-4o", False),
+            ("gpt-4.1", False),
+            ("gpt-4o-mini", False),
+            ("gpt-5.6", True),
+            ("o3", True),
+            ("o3-deep-research", True),
+            ("codex-mini-latest", True),
+            ("computer-use-preview", True),
+        ],
+    )
+    def test_drop_params_strips_reasoning_by_model(self, local_model_cost_map, model, reasoning_survives):
+        mapped = OpenAIResponsesAPIConfig().map_openai_params(
+            response_api_optional_params={"reasoning": {"effort": "medium", "summary": "auto"}},
+            model=model,
+            drop_params=True,
+        )
+        assert ("reasoning" in mapped) is reasoning_survives
+
+    def test_without_drop_params_the_error_is_litellms_400(self, local_model_cost_map, monkeypatch):
+        monkeypatch.setattr(litellm, "drop_params", False)
+        with pytest.raises(litellm.UnsupportedParamsError) as excinfo:
+            OpenAIResponsesAPIConfig().map_openai_params(
+                response_api_optional_params={"reasoning": {"effort": "medium"}},
+                model="gpt-4o",
+                drop_params=False,
+            )
+        assert excinfo.value.status_code == 400
+
+    def test_azure_deployments_keep_reasoning(self, local_model_cost_map):
+        mapped = AzureOpenAIResponsesAPIConfig().map_openai_params(
+            response_api_optional_params={"reasoning": {"effort": "medium"}},
+            model="my-o3-deployment",
+            drop_params=True,
+        )
+        assert mapped["reasoning"] == {"effort": "medium"}

From a099be02fda770802b41a6f5b4e072460b82bee9 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 29 Aug 2026 21:35:19 -0700
Subject: [PATCH 005/164] fix(guardrails): resolve generateContent routes and
 async-first passthrough call types

API_ROUTE_TO_CALL_TYPES listed the sync llm_passthrough_route first, so every
call_types[0] consumer resolved /llm_passthrough to a call type with no
guardrail translation handler, and the {model}:generateContent patterns never
matched a concrete route because the placeholder segment carries a literal
suffix the matcher treated as an exact segment. Reorder the passthrough
entries async-first, teach the matcher placeholder-with-suffix segments plus
suffixed multi-segment tails (mirroring FastAPI's {model_name:path}), add the
missing /v1beta generateContent entries, and register a Google GenAI
guardrail translation handler so guardrails actually scan generateContent
requests, responses, and streams.
---
 .../api_route_to_call_types.py                |  43 +++-
 .../guardrail_translation/__init__.py         |  20 ++
 .../guardrail_translation/handler.py          | 237 ++++++++++++++++++
 litellm/types/utils.py                        |  12 +-
 .../test_api_route_to_call_types.py           | 112 +++++++++
 .../llms/gemini/google_genai/__init__.py      |   0
 .../guardrail_translation/__init__.py         |   0
 .../test_google_genai_guardrail_handler.py    | 195 ++++++++++++++
 8 files changed, 609 insertions(+), 10 deletions(-)
 create mode 100644 litellm/llms/gemini/google_genai/guardrail_translation/__init__.py
 create mode 100644 litellm/llms/gemini/google_genai/guardrail_translation/handler.py
 create mode 100644 tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py
 create mode 100644 tests/test_litellm/llms/gemini/google_genai/__init__.py
 create mode 100644 tests/test_litellm/llms/gemini/google_genai/guardrail_translation/__init__.py
 create mode 100644 tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py

diff --git a/litellm/litellm_core_utils/api_route_to_call_types.py b/litellm/litellm_core_utils/api_route_to_call_types.py
index e3562095d7f..428d7563d4a 100644
--- a/litellm/litellm_core_utils/api_route_to_call_types.py
+++ b/litellm/litellm_core_utils/api_route_to_call_types.py
@@ -14,21 +14,48 @@ from typing import Final
 from litellm.types.utils import API_ROUTE_TO_CALL_TYPES, CallTypes
 
 
+def _segment_matches(route_segment: str, pattern_segment: str) -> bool:
+    """
+    Match one concrete path segment against one pattern segment.
+    A bare placeholder ({param}) matches any segment; a placeholder with a
+    literal suffix ({model}:generateContent) requires the segment to end with
+    that suffix and have a non-empty value before it.
+    """
+    if not pattern_segment.startswith("{"):
+        return route_segment == pattern_segment
+    placeholder_end: Final = pattern_segment.find("}")
+    if placeholder_end == -1:
+        return route_segment == pattern_segment
+    literal_suffix: Final = pattern_segment[placeholder_end + 1 :]
+    if not literal_suffix:
+        return True
+    return route_segment.endswith(literal_suffix) and len(route_segment) > len(literal_suffix)
+
+
+def _pattern_tail_spans_segments(pattern_tail: str) -> bool:
+    """
+    Whether the pattern's last segment is a suffixed placeholder
+    ({model}:generateContent) that may absorb extra route segments, mirroring
+    FastAPI's {model_name:path} converter for slash-containing model names.
+    """
+    return pattern_tail.startswith("{") and "}" in pattern_tail and not pattern_tail.endswith("}")
+
+
 def _route_matches_pattern(route: str, pattern: str) -> bool:
     """
     Return True if the concrete route matches the pattern.
-    Pattern segments like {param} match any single path segment.
+    Pattern segments like {param} match any single path segment, and a
+    suffixed placeholder in the last segment may span multiple segments.
     """
     route_parts: Final = route.strip("/").split("/")
     pattern_parts: Final = pattern.strip("/").split("/")
-    if len(route_parts) != len(pattern_parts):
+    if len(route_parts) < len(pattern_parts):
         return False
-    for r, p in zip(route_parts, pattern_parts):
-        if p.startswith("{") and p.endswith("}"):
-            continue
-        if r != p:
-            return False
-    return True
+    if len(route_parts) > len(pattern_parts) and not _pattern_tail_spans_segments(pattern_parts[-1]):
+        return False
+    head_count: Final = len(pattern_parts) - 1
+    merged_parts: Final = (*route_parts[:head_count], "/".join(route_parts[head_count:]))
+    return all(_segment_matches(r, p) for r, p in zip(merged_parts, pattern_parts))
 
 
 def get_call_types_for_route(route: str) -> Sequence[CallTypes] | None:
diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/__init__.py b/litellm/llms/gemini/google_genai/guardrail_translation/__init__.py
new file mode 100644
index 00000000000..494a72d6999
--- /dev/null
+++ b/litellm/llms/gemini/google_genai/guardrail_translation/__init__.py
@@ -0,0 +1,20 @@
+"""Google GenAI generateContent guardrail translation handler."""
+
+from typing import Final
+
+from litellm.llms.gemini.google_genai.guardrail_translation.handler import (
+    GoogleGenAIGenerateContentHandler,
+)
+from litellm.types.utils import CallTypes
+
+guardrail_translation_mappings: Final = {  # mutable-ok: discover_guardrail_translation_mappings only accepts isinstance(mappings, dict)
+    CallTypes.generate_content: GoogleGenAIGenerateContentHandler,
+    CallTypes.agenerate_content: GoogleGenAIGenerateContentHandler,
+    CallTypes.generate_content_stream: GoogleGenAIGenerateContentHandler,
+    CallTypes.agenerate_content_stream: GoogleGenAIGenerateContentHandler,
+}
+
+__all__ = (
+    "GoogleGenAIGenerateContentHandler",
+    "guardrail_translation_mappings",
+)
diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py
new file mode 100644
index 00000000000..dd76cd711d8
--- /dev/null
+++ b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py
@@ -0,0 +1,237 @@
+"""
+Google GenAI generateContent handler for Unified Guardrails.
+
+Extracts text from generateContent requests (contents[].parts[].text) and
+responses (candidates[].content.parts[].text), applies the guardrail, and
+writes the guardrailed text back in place. Requests and responses may be
+dicts (wire format) or google-genai SDK objects; streaming chunks may
+additionally be raw SSE frames, which are scanned for detection (a blocking
+guardrail raises) without rewriting the frames.
+"""
+
+import json
+from collections.abc import Mapping, Sequence
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Final, Optional
+
+from litellm._logging import verbose_proxy_logger
+from litellm.llms.base_llm.guardrail_translation.base_translation import (
+    BaseTranslation,
+    StreamTransformSink,
+)
+from litellm.types.utils import GenericGuardrailAPIInputs
+
+if TYPE_CHECKING:
+    from litellm.integrations.custom_guardrail import CustomGuardrail
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+    from litellm.proxy._types import UserAPIKeyAuth
+
+_EMPTY_REQUEST_DATA: Final[Mapping[str, object]] = MappingProxyType({})
+
+
+def _field(container: object, name: str) -> object | None:
+    if isinstance(container, dict):
+        return container.get(name)
+    return getattr(container, name, None)
+
+
+def _part_text(part: object) -> str | None:
+    text: Final = _field(part, "text")
+    if isinstance(text, str) and text:
+        return text
+    return None
+
+
+def _write_part_text(part: object, text: str) -> None:
+    if isinstance(part, dict):
+        part["text"] = text  # rebind-ok: guardrail write-back rewrites the caller's part in place by handler contract
+        return
+    setattr(part, "text", text)  # noqa: B010  # SDK parts are typed as object here; direct assignment cannot type-check
+
+
+def _content_text_parts(content: object) -> tuple[object, ...]:
+    parts: Final = _field(content, "parts")
+    if not isinstance(parts, (list, tuple)):
+        return ()
+    return tuple(part for part in parts if _part_text(part) is not None)
+
+
+def _request_text_parts(data: Mapping[str, object]) -> tuple[object, ...]:
+    contents: Final = data.get("contents")
+    content_list: Final = (
+        (contents,) if isinstance(contents, dict) else tuple(contents) if isinstance(contents, list) else ()
+    )
+    return tuple(part for content in content_list for part in _content_text_parts(content))
+
+
+def _response_text_parts(response: object) -> tuple[object, ...]:
+    candidates: Final = _field(response, "candidates")
+    if not isinstance(candidates, (list, tuple)):
+        return ()
+    return tuple(part for candidate in candidates for part in _content_text_parts(_field(candidate, "content")))
+
+
+def _part_texts(text_parts: Sequence[object]) -> tuple[str, ...]:
+    return tuple(text for part in text_parts for text in (_part_text(part),) if text is not None)
+
+
+def _texts_payload(
+    texts: Sequence[str],
+) -> list[str]:  # mutable-ok: GenericGuardrailAPIInputs.texts is declared list[str]
+    return list(texts)  # mutable-ok: GenericGuardrailAPIInputs.texts is declared list[str]
+
+
+def _write_back_texts(text_parts: Sequence[object], guardrailed_texts: Sequence[str] | None) -> None:
+    if not guardrailed_texts or len(guardrailed_texts) != len(text_parts):
+        return
+    for part, text in zip(text_parts, guardrailed_texts):
+        _write_part_text(part, text)
+
+
+def _parse_json_dict_or_none(payload: str) -> Mapping[str, object] | None:
+    try:
+        parsed: Final = json.loads(payload)
+    except json.JSONDecodeError:
+        return None
+    if isinstance(parsed, dict):
+        return parsed
+    return None
+
+
+def _sse_payload_texts(sse_text: str) -> tuple[str, ...]:
+    return tuple(
+        text
+        for line in sse_text.splitlines()
+        if line.startswith("data:")
+        for payload in (line[len("data:") :].strip(),)
+        if payload and payload != "[DONE]"
+        for parsed in (_parse_json_dict_or_none(payload),)
+        if parsed is not None
+        for text in _part_texts(_response_text_parts(parsed))
+    )
+
+
+def _chunk_sse_text(chunk: object) -> str | None:
+    if isinstance(chunk, bytes):
+        return chunk.decode("utf-8", errors="replace")
+    if isinstance(chunk, str):
+        return chunk
+    return None
+
+
+def _accumulated_stream_text(responses_so_far: Sequence[object]) -> str:
+    object_texts: Final = tuple(
+        text
+        for chunk in responses_so_far
+        if _chunk_sse_text(chunk) is None
+        for text in _part_texts(_response_text_parts(chunk))
+    )
+    sse_text: Final = "".join(sse for chunk in responses_so_far for sse in (_chunk_sse_text(chunk),) if sse is not None)
+    return "".join(object_texts) + "".join(_sse_payload_texts(sse_text))
+
+
+class GoogleGenAIGenerateContentHandler(BaseTranslation):
+    """
+    Guardrail translation for the google genai generateContent surface
+    (/models/{model}:generateContent, :streamGenerateContent, and the
+    litellm SDK generate_content call types).
+    """
+
+    async def process_input_messages(
+        self,
+        data: dict,  # mutable-ok: base handler contract passes the proxy's request dict through to apply_guardrail
+        guardrail_to_apply: "CustomGuardrail",
+        litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
+    ) -> object:
+        text_parts: Final = _request_text_parts(data)
+        if not text_parts:
+            verbose_proxy_logger.debug("Google GenAI guardrail: no request text found, skipping")
+            return data
+        model: Final = data.get("model")
+        inputs: Final = (
+            GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts)), model=model)
+            if isinstance(model, str)
+            else GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts)))
+        )
+        guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
+            inputs=inputs,
+            request_data=data,
+            input_type="request",
+            logging_obj=litellm_logging_obj,
+        )
+        _write_back_texts(text_parts, guardrailed_inputs.get("texts"))
+        return data
+
+    async def process_output_response(
+        self,
+        response: object,
+        guardrail_to_apply: "CustomGuardrail",
+        litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
+        user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
+        request_data: Mapping[str, object] | None = None,
+    ) -> object:
+        text_parts: Final = _response_text_parts(response)
+        if not text_parts:
+            verbose_proxy_logger.debug("Google GenAI guardrail: no response text found, skipping")
+            return response
+        guardrail_request_data: Final = self._merged_request_data(
+            request_data=request_data,
+            user_api_key_dict=user_api_key_dict,
+            context_key="response",
+            context_value=response,
+        )
+        model: Final = guardrail_request_data.get("model")
+        inputs: Final = (
+            GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts)), model=model)
+            if isinstance(model, str)
+            else GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts)))
+        )
+        guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
+            inputs=inputs,
+            request_data=guardrail_request_data,
+            input_type="response",
+            logging_obj=litellm_logging_obj,
+        )
+        _write_back_texts(text_parts, guardrailed_inputs.get("texts"))
+        return response
+
+    async def process_output_streaming_response(
+        self,
+        responses_so_far: Sequence[object],
+        guardrail_to_apply: "CustomGuardrail",
+        litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
+        user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
+        request_data: Mapping[str, object] | None = None,
+        stream_transform_sink: StreamTransformSink | None = None,
+    ) -> object:
+        accumulated_text: Final = _accumulated_stream_text(responses_so_far)
+        if not accumulated_text:
+            return responses_so_far
+        guardrail_request_data: Final = self._merged_request_data(
+            request_data=request_data,
+            user_api_key_dict=user_api_key_dict,
+            context_key="responses_so_far",
+            context_value=responses_so_far,
+        )
+        _guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
+            inputs=GenericGuardrailAPIInputs(texts=_texts_payload((accumulated_text,))),
+            request_data=guardrail_request_data,
+            input_type="response",
+            logging_obj=litellm_logging_obj,
+        )
+        return responses_so_far
+
+    def _merged_request_data(
+        self,
+        request_data: Mapping[str, object] | None,
+        user_api_key_dict: Optional["UserAPIKeyAuth"],
+        context_key: str,
+        context_value: object,
+    ) -> dict:  # mutable-ok: CustomGuardrail.apply_guardrail requires a plain dict request payload
+        base: Final = request_data if request_data is not None else _EMPTY_REQUEST_DATA
+        user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
+        context_pairs: Final = ((context_key, context_value),) if context_key not in base else ()
+        metadata_pairs: Final = (
+            (("litellm_metadata", user_metadata),) if user_metadata and "litellm_metadata" not in base else ()
+        )
+        return dict((*base.items(), *context_pairs, *metadata_pairs))  # mutable-ok: apply_guardrail takes a plain dict
diff --git a/litellm/types/utils.py b/litellm/types/utils.py
index 4bf8289d725..9e48031dd47 100644
--- a/litellm/types/utils.py
+++ b/litellm/types/utils.py
@@ -919,6 +919,14 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = {
         CallTypes.agenerate_content_stream,
         CallTypes.generate_content_stream,
     ],
+    "/v1beta/models/{model}:generateContent": (
+        CallTypes.agenerate_content,
+        CallTypes.generate_content,
+    ),
+    "/v1beta/models/{model}:streamGenerateContent": (
+        CallTypes.agenerate_content_stream,
+        CallTypes.generate_content_stream,
+    ),
     # MCP (Model Context Protocol)
     "/mcp/call_tool": [CallTypes.call_mcp_tool],
     # A2A (Agent-to-Agent)
@@ -926,12 +934,12 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = {
     "/a2a/{agent_id}/message/send": [CallTypes.asend_message, CallTypes.send_message],
     # Passthrough endpoints
     "/llm_passthrough": [
-        CallTypes.llm_passthrough_route,
         CallTypes.allm_passthrough_route,
+        CallTypes.llm_passthrough_route,
     ],
     "/v1/llm_passthrough": [
-        CallTypes.llm_passthrough_route,
         CallTypes.allm_passthrough_route,
+        CallTypes.llm_passthrough_route,
     ],
     "/v1/messages": [CallTypes.anthropic_messages],
     # OCR
diff --git a/tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py b/tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py
new file mode 100644
index 00000000000..42ca91bfd8f
--- /dev/null
+++ b/tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py
@@ -0,0 +1,112 @@
+"""
+Tests for route -> CallTypes resolution (api_route_to_call_types).
+
+Regression coverage for the guardrail route table bugs:
+- placeholder segments with a literal suffix ({model}:generateContent) never matched
+- the /v1beta generateContent routes were missing from the table
+- /llm_passthrough listed the sync call type first, resolving consumers that
+  take call_types[0] to a handler-less type
+"""
+
+from litellm.litellm_core_utils.api_route_to_call_types import (
+    get_call_types_for_route,
+)
+from litellm.types.utils import API_ROUTE_TO_CALL_TYPES, CallTypes
+
+
+class TestGenerateContentRouteResolution:
+    def test_bare_generate_content_route_resolves(self):
+        call_types = get_call_types_for_route("/models/gemini-2.5-flash:generateContent")
+        assert call_types is not None
+        assert list(call_types) == [CallTypes.agenerate_content, CallTypes.generate_content]
+
+    def test_v1beta_generate_content_route_resolves(self):
+        call_types = get_call_types_for_route("/v1beta/models/gemini-2.5-flash:generateContent")
+        assert call_types is not None
+        assert list(call_types) == [CallTypes.agenerate_content, CallTypes.generate_content]
+
+    def test_bare_stream_generate_content_route_resolves(self):
+        call_types = get_call_types_for_route("/models/gemini-2.5-flash:streamGenerateContent")
+        assert call_types is not None
+        assert list(call_types) == [
+            CallTypes.agenerate_content_stream,
+            CallTypes.generate_content_stream,
+        ]
+
+    def test_v1beta_stream_generate_content_route_resolves(self):
+        call_types = get_call_types_for_route("/v1beta/models/gemini-2.5-flash:streamGenerateContent")
+        assert call_types is not None
+        assert list(call_types) == [
+            CallTypes.agenerate_content_stream,
+            CallTypes.generate_content_stream,
+        ]
+
+    def test_slash_containing_model_name_resolves(self):
+        call_types = get_call_types_for_route("/v1beta/models/gemini/gemini-2.5-flash:generateContent")
+        assert call_types is not None
+        assert list(call_types) == [CallTypes.agenerate_content, CallTypes.generate_content]
+
+    def test_empty_model_name_does_not_match(self):
+        assert get_call_types_for_route("/models/:generateContent") is None
+
+    def test_unrelated_model_action_does_not_match(self):
+        assert get_call_types_for_route("/models/gemini-2.5-flash:countTokens") is None
+
+
+class TestPassthroughRouteOrdering:
+    def test_llm_passthrough_lists_async_call_type_first(self):
+        call_types = get_call_types_for_route("/llm_passthrough")
+        assert call_types is not None
+        assert list(call_types) == [
+            CallTypes.allm_passthrough_route,
+            CallTypes.llm_passthrough_route,
+        ]
+
+    def test_v1_llm_passthrough_lists_async_call_type_first(self):
+        call_types = get_call_types_for_route("/v1/llm_passthrough")
+        assert call_types is not None
+        assert list(call_types) == [
+            CallTypes.allm_passthrough_route,
+            CallTypes.llm_passthrough_route,
+        ]
+
+
+class TestFirstCallTypeHasTranslationHandler:
+    def test_first_call_type_is_translatable_whenever_any_is(self):
+        """
+        Consumers (unified guardrail post-call and streaming resolution) take
+        call_types[0]. A route whose first call type lacks a guardrail
+        translation handler while a later one has it silently skips guardrail
+        scanning, so the table must list a handler-backed call type first.
+        """
+        from litellm.llms import load_guardrail_translation_mappings
+
+        mappings = load_guardrail_translation_mappings()
+        misordered = {
+            route: [call_type.value for call_type in call_types]
+            for route, call_types in API_ROUTE_TO_CALL_TYPES.items()
+            if call_types
+            and call_types[0] not in mappings
+            and any(call_type in mappings for call_type in call_types)
+        }
+        assert misordered == {}
+
+
+class TestExistingRouteResolutionUnchanged:
+    def test_exact_route_still_resolves(self):
+        call_types = get_call_types_for_route("/chat/completions")
+        assert call_types is not None
+        assert CallTypes.acompletion in call_types
+
+    def test_single_segment_placeholder_still_resolves(self):
+        call_types = get_call_types_for_route("/a2a/my-agent/message/send")
+        assert call_types is not None
+        assert list(call_types) == [CallTypes.asend_message, CallTypes.send_message]
+
+    def test_longer_route_does_not_collapse_into_bare_placeholder_pattern(self):
+        call_types = get_call_types_for_route("/responses/resp_123/input_items")
+        assert call_types is not None
+        assert list(call_types) == [CallTypes.alist_input_items]
+
+    def test_unknown_route_returns_none(self):
+        assert get_call_types_for_route("/not/a/real/route") is None
diff --git a/tests/test_litellm/llms/gemini/google_genai/__init__.py b/tests/test_litellm/llms/gemini/google_genai/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/__init__.py b/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py b/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py
new file mode 100644
index 00000000000..42ab8a7431a
--- /dev/null
+++ b/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py
@@ -0,0 +1,195 @@
+"""
+Tests for the Google GenAI generateContent guardrail translation handler.
+"""
+
+import json
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from fastapi import HTTPException
+
+from litellm.llms.gemini.google_genai.guardrail_translation.handler import (
+    GoogleGenAIGenerateContentHandler,
+)
+from litellm.types.utils import CallTypes
+
+
+def _mock_guardrail(returned_texts):
+    guardrail = MagicMock()
+    guardrail.apply_guardrail = AsyncMock(return_value={"texts": returned_texts})
+    return guardrail
+
+
+@pytest.mark.asyncio
+async def test_input_contents_text_is_guardrailed_and_written_back():
+    handler = GoogleGenAIGenerateContentHandler()
+    guardrail = _mock_guardrail(["masked question"])
+    data = {
+        "model": "gemini-2.5-flash",
+        "contents": [{"role": "user", "parts": [{"text": "raw question"}]}],
+    }
+
+    result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
+
+    call_kwargs = guardrail.apply_guardrail.call_args.kwargs
+    assert call_kwargs["inputs"]["texts"] == ["raw question"]
+    assert call_kwargs["inputs"]["model"] == "gemini-2.5-flash"
+    assert call_kwargs["input_type"] == "request"
+    assert result["contents"][0]["parts"][0]["text"] == "masked question"
+
+
+@pytest.mark.asyncio
+async def test_input_without_text_skips_guardrail():
+    handler = GoogleGenAIGenerateContentHandler()
+    guardrail = _mock_guardrail([])
+    data = {"model": "gemini-2.5-flash", "contents": [{"role": "user", "parts": [{"inlineData": {}}]}]}
+
+    result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
+
+    guardrail.apply_guardrail.assert_not_called()
+    assert result is data
+
+
+@pytest.mark.asyncio
+async def test_output_dict_response_text_is_guardrailed_and_written_back():
+    handler = GoogleGenAIGenerateContentHandler()
+    guardrail = _mock_guardrail(["masked answer"])
+    response = {
+        "candidates": [
+            {
+                "content": {"role": "model", "parts": [{"text": "harmful answer"}]},
+                "finishReason": "STOP",
+            }
+        ]
+    }
+
+    result = await handler.process_output_response(
+        response=response,
+        guardrail_to_apply=guardrail,
+        request_data={"model": "gemini-2.5-flash"},
+    )
+
+    call_kwargs = guardrail.apply_guardrail.call_args.kwargs
+    assert call_kwargs["inputs"]["texts"] == ["harmful answer"]
+    assert call_kwargs["input_type"] == "response"
+    assert call_kwargs["request_data"]["response"] is response
+    assert result["candidates"][0]["content"]["parts"][0]["text"] == "masked answer"
+
+
+@pytest.mark.asyncio
+async def test_output_object_response_text_is_guardrailed_and_written_back():
+    handler = GoogleGenAIGenerateContentHandler()
+    guardrail = _mock_guardrail(["masked answer"])
+    part = SimpleNamespace(text="harmful answer")
+    response = SimpleNamespace(
+        candidates=[SimpleNamespace(content=SimpleNamespace(parts=[part]), finish_reason="STOP")]
+    )
+
+    await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
+
+    assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["harmful answer"]
+    assert part.text == "masked answer"
+
+
+@pytest.mark.asyncio
+async def test_output_without_text_skips_guardrail():
+    handler = GoogleGenAIGenerateContentHandler()
+    guardrail = _mock_guardrail([])
+
+    result = await handler.process_output_response(response={"candidates": []}, guardrail_to_apply=guardrail)
+
+    guardrail.apply_guardrail.assert_not_called()
+    assert result == {"candidates": []}
+
+
+@pytest.mark.asyncio
+async def test_output_blocking_guardrail_exception_propagates():
+    handler = GoogleGenAIGenerateContentHandler()
+    guardrail = MagicMock()
+    guardrail.apply_guardrail = AsyncMock(side_effect=HTTPException(status_code=400, detail="blocked"))
+    response = {"candidates": [{"content": {"parts": [{"text": "harmful answer"}]}}]}
+
+    with pytest.raises(HTTPException):
+        await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
+
+
+@pytest.mark.asyncio
+async def test_streaming_dict_chunks_accumulate_text():
+    handler = GoogleGenAIGenerateContentHandler()
+    guardrail = _mock_guardrail(["clean"])
+    chunks = [
+        {"candidates": [{"content": {"parts": [{"text": "harmful "}]}}]},
+        {"candidates": [{"content": {"parts": [{"text": "answer"}]}}]},
+    ]
+
+    result = await handler.process_output_streaming_response(
+        responses_so_far=chunks,
+        guardrail_to_apply=guardrail,
+    )
+
+    call_kwargs = guardrail.apply_guardrail.call_args.kwargs
+    assert call_kwargs["inputs"]["texts"] == ["harmful answer"]
+    assert call_kwargs["input_type"] == "response"
+    assert result is chunks
+
+
+@pytest.mark.asyncio
+async def test_streaming_raw_sse_chunks_accumulate_text_across_split_frames():
+    handler = GoogleGenAIGenerateContentHandler()
+    guardrail = _mock_guardrail(["clean"])
+    frame_one = "data: " + json.dumps({"candidates": [{"content": {"parts": [{"text": "harmful "}]}}]}) + "\r\n\r\n"
+    frame_two = "data: " + json.dumps({"candidates": [{"content": {"parts": [{"text": "answer"}]}}]}) + "\r\n\r\n"
+    split_at = len(frame_one) // 2
+    chunks = [frame_one[:split_at], frame_one[split_at:] + frame_two[:5], frame_two[5:].encode("utf-8")]
+
+    await handler.process_output_streaming_response(
+        responses_so_far=chunks,
+        guardrail_to_apply=guardrail,
+    )
+
+    assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["harmful answer"]
+
+
+@pytest.mark.asyncio
+async def test_streaming_blocking_guardrail_exception_propagates():
+    handler = GoogleGenAIGenerateContentHandler()
+    guardrail = MagicMock()
+    guardrail.apply_guardrail = AsyncMock(side_effect=HTTPException(status_code=400, detail="blocked"))
+    chunks = [{"candidates": [{"content": {"parts": [{"text": "harmful"}]}}]}]
+
+    with pytest.raises(HTTPException):
+        await handler.process_output_streaming_response(responses_so_far=chunks, guardrail_to_apply=guardrail)
+
+
+@pytest.mark.asyncio
+async def test_streaming_without_text_skips_guardrail():
+    handler = GoogleGenAIGenerateContentHandler()
+    guardrail = _mock_guardrail([])
+
+    result = await handler.process_output_streaming_response(responses_so_far=[], guardrail_to_apply=guardrail)
+
+    guardrail.apply_guardrail.assert_not_called()
+    assert result == []
+
+
+def test_generate_content_call_types_are_registered():
+    from litellm.llms.gemini.google_genai.guardrail_translation import (
+        guardrail_translation_mappings,
+    )
+
+    for call_type in (
+        CallTypes.generate_content,
+        CallTypes.agenerate_content,
+        CallTypes.generate_content_stream,
+        CallTypes.agenerate_content_stream,
+    ):
+        assert guardrail_translation_mappings[call_type] is GoogleGenAIGenerateContentHandler
+
+
+def test_discovery_finds_generate_content_handler():
+    from litellm.llms import load_guardrail_translation_mappings
+
+    mappings = load_guardrail_translation_mappings()
+    assert mappings[CallTypes.agenerate_content] is GoogleGenAIGenerateContentHandler
+    assert mappings[CallTypes.agenerate_content_stream] is GoogleGenAIGenerateContentHandler

From 05e4d2f946a2ee8a2beb51d5476b77c4ea4cc027 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 29 Aug 2026 22:11:47 -0700
Subject: [PATCH 006/164] fix(guardrails): scan generateContent
 systemInstruction text and drop fastapi import from handler tests

---
 .../guardrail_translation/handler.py          | 24 +++++++--
 .../test_google_genai_guardrail_handler.py    | 49 +++++++++++++++++--
 2 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py
index dd76cd711d8..e13e1e63cbb 100644
--- a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py
+++ b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py
@@ -1,8 +1,9 @@
 """
 Google GenAI generateContent handler for Unified Guardrails.
 
-Extracts text from generateContent requests (contents[].parts[].text) and
-responses (candidates[].content.parts[].text), applies the guardrail, and
+Extracts text from generateContent requests (systemInstruction.parts[].text
+and contents[].parts[].text) and responses (candidates[].content.parts[].text),
+applies the guardrail, and
 writes the guardrailed text back in place. Requests and responses may be
 dicts (wire format) or google-genai SDK objects; streaming chunks may
 additionally be raw SSE frames, which are scanned for detection (a blocking
@@ -56,12 +57,29 @@ def _content_text_parts(content: object) -> tuple[object, ...]:
     return tuple(part for part in parts if _part_text(part) is not None)
 
 
+def _system_instruction(data: Mapping[str, object]) -> object | None:
+    return next(
+        (
+            value
+            for container in (data, data.get("config"))
+            if container is not None
+            for key in ("systemInstruction", "system_instruction")
+            for value in (_field(container, key),)
+            if value is not None
+        ),
+        None,
+    )
+
+
 def _request_text_parts(data: Mapping[str, object]) -> tuple[object, ...]:
     contents: Final = data.get("contents")
     content_list: Final = (
         (contents,) if isinstance(contents, dict) else tuple(contents) if isinstance(contents, list) else ()
     )
-    return tuple(part for content in content_list for part in _content_text_parts(content))
+    return (
+        *_content_text_parts(_system_instruction(data)),
+        *(part for content in content_list for part in _content_text_parts(content)),
+    )
 
 
 def _response_text_parts(response: object) -> tuple[object, ...]:
diff --git a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py b/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py
index 42ab8a7431a..4119ce99423 100644
--- a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py
+++ b/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py
@@ -7,7 +7,6 @@ from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
-from fastapi import HTTPException
 
 from litellm.llms.gemini.google_genai.guardrail_translation.handler import (
     GoogleGenAIGenerateContentHandler,
@@ -15,6 +14,10 @@ from litellm.llms.gemini.google_genai.guardrail_translation.handler import (
 from litellm.types.utils import CallTypes
 
 
+class GuardrailBlockedError(Exception):
+    pass
+
+
 def _mock_guardrail(returned_texts):
     guardrail = MagicMock()
     guardrail.apply_guardrail = AsyncMock(return_value={"texts": returned_texts})
@@ -39,6 +42,42 @@ async def test_input_contents_text_is_guardrailed_and_written_back():
     assert result["contents"][0]["parts"][0]["text"] == "masked question"
 
 
+@pytest.mark.asyncio
+async def test_input_system_instruction_text_is_scanned_and_written_back():
+    handler = GoogleGenAIGenerateContentHandler()
+    guardrail = _mock_guardrail(["masked instruction", "masked question"])
+    data = {
+        "model": "gemini-2.5-flash",
+        "systemInstruction": {"role": "system", "parts": [{"text": "prohibited instruction"}]},
+        "contents": [{"role": "user", "parts": [{"text": "benign question"}]}],
+    }
+
+    result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
+
+    assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == [
+        "prohibited instruction",
+        "benign question",
+    ]
+    assert result["systemInstruction"]["parts"][0]["text"] == "masked instruction"
+    assert result["contents"][0]["parts"][0]["text"] == "masked question"
+
+
+@pytest.mark.asyncio
+async def test_input_config_nested_snake_case_system_instruction_is_scanned():
+    handler = GoogleGenAIGenerateContentHandler()
+    guardrail = _mock_guardrail(["clean"])
+    instruction_part = SimpleNamespace(text="prohibited instruction")
+    data = {
+        "contents": [],
+        "config": SimpleNamespace(system_instruction=SimpleNamespace(parts=[instruction_part])),
+    }
+
+    await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
+
+    assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["prohibited instruction"]
+    assert instruction_part.text == "clean"
+
+
 @pytest.mark.asyncio
 async def test_input_without_text_skips_guardrail():
     handler = GoogleGenAIGenerateContentHandler()
@@ -107,10 +146,10 @@ async def test_output_without_text_skips_guardrail():
 async def test_output_blocking_guardrail_exception_propagates():
     handler = GoogleGenAIGenerateContentHandler()
     guardrail = MagicMock()
-    guardrail.apply_guardrail = AsyncMock(side_effect=HTTPException(status_code=400, detail="blocked"))
+    guardrail.apply_guardrail = AsyncMock(side_effect=GuardrailBlockedError("blocked"))
     response = {"candidates": [{"content": {"parts": [{"text": "harmful answer"}]}}]}
 
-    with pytest.raises(HTTPException):
+    with pytest.raises(GuardrailBlockedError):
         await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
 
 
@@ -155,10 +194,10 @@ async def test_streaming_raw_sse_chunks_accumulate_text_across_split_frames():
 async def test_streaming_blocking_guardrail_exception_propagates():
     handler = GoogleGenAIGenerateContentHandler()
     guardrail = MagicMock()
-    guardrail.apply_guardrail = AsyncMock(side_effect=HTTPException(status_code=400, detail="blocked"))
+    guardrail.apply_guardrail = AsyncMock(side_effect=GuardrailBlockedError("blocked"))
     chunks = [{"candidates": [{"content": {"parts": [{"text": "harmful"}]}}]}]
 
-    with pytest.raises(HTTPException):
+    with pytest.raises(GuardrailBlockedError):
         await handler.process_output_streaming_response(responses_so_far=chunks, guardrail_to_apply=guardrail)
 
 

From f0a2a2312704df73ba020290ef426c9c4360a130 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 30 Aug 2026 12:32:55 -0700
Subject: [PATCH 007/164] fix(proxy): register SkillsInjectionHook at proxy
 startup instead of import time

---
 litellm/proxy/hooks/litellm_skills/__init__.py   |  6 +-----
 litellm/proxy/hooks/litellm_skills/main.py       | 11 +----------
 .../proxy/hooks/litellm_skills/test_main.py      | 16 ++++++++++++++++
 3 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/litellm/proxy/hooks/litellm_skills/__init__.py b/litellm/proxy/hooks/litellm_skills/__init__.py
index 751122ac51c..d24cdd37161 100644
--- a/litellm/proxy/hooks/litellm_skills/__init__.py
+++ b/litellm/proxy/hooks/litellm_skills/__init__.py
@@ -21,10 +21,7 @@ from litellm.llms.litellm_proxy.skills import (
     code_execution_handler,
     get_litellm_code_execution_tool,
 )
-from litellm.proxy.hooks.litellm_skills.main import (
-    SkillsInjectionHook,
-    skills_injection_hook,
-)
+from litellm.proxy.hooks.litellm_skills.main import SkillsInjectionHook
 
 __all__ = [
     "LITELLM_CODE_EXECUTION_TOOL",
@@ -35,5 +32,4 @@ __all__ = [
     "SkillsSandboxExecutor",
     "code_execution_handler",
     "get_litellm_code_execution_tool",
-    "skills_injection_hook",
 ]
diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py
index 569ec32c1a0..c5e8f03f792 100644
--- a/litellm/proxy/hooks/litellm_skills/main.py
+++ b/litellm/proxy/hooks/litellm_skills/main.py
@@ -29,6 +29,7 @@ import json
 from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, Protocol
 
+import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
 from litellm.integrations.custom_logger import CustomLogger
@@ -475,7 +476,6 @@ class SkillsInjectionHook(CustomLogger):
 
         Returns the final response with generated files inline.
         """
-        import litellm
         from litellm.llms.litellm_proxy.skills.code_execution import (
             LiteLLMInternalTools,
         )
@@ -705,7 +705,6 @@ print('No executable skill module found')
 
         Returns the final response with generated files inline.
         """
-        import litellm
         from litellm.llms.litellm_proxy.skills.code_execution import (
             LiteLLMInternalTools,
         )
@@ -894,11 +893,3 @@ print('No executable skill module found')
         verbose_proxy_logger.debug("SkillsInjectionHook: Attached %s files to response", len(generated_files))
 
         return response
-
-
-# Global instance for registration
-skills_injection_hook: Final = SkillsInjectionHook()
-
-import litellm
-
-litellm.logging_callback_manager.add_litellm_callback(skills_injection_hook)
diff --git a/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py b/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py
index f716a8533d8..c037f60ac25 100644
--- a/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py
+++ b/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py
@@ -1,3 +1,6 @@
+import subprocess
+import sys
+from typing import Final
 from unittest.mock import AsyncMock, patch
 
 import pytest
@@ -65,3 +68,16 @@ async def test_execute_code_loop_dispatches_litellm_skill_tool():
     mock_exec.assert_awaited_once()
     assert mock_exec.await_args.args[0] == SKILL_TOOL_NAME
     assert result is final_response
+
+
+def test_importing_proxy_hooks_does_not_mutate_litellm_callbacks():
+    script: Final = (
+        "import litellm; "
+        "litellm.callbacks = []; "
+        "import litellm.proxy.hooks; "
+        "assert len(litellm.callbacks) == 0, f'mutated: {litellm.callbacks}'"
+    )
+    result: Final = subprocess.run(
+        [sys.executable, "-c", script], capture_output=True, text=True
+    )
+    assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}"

From b8e11a75fa7e656d788855f42b182b9ff862a907 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 30 Aug 2026 12:47:15 -0700
Subject: [PATCH 008/164] test: use local model cost map in import-isolation
 subprocess

---
 tests/test_litellm/proxy/hooks/litellm_skills/test_main.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py b/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py
index c037f60ac25..dc7ebf6e2ec 100644
--- a/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py
+++ b/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py
@@ -1,3 +1,4 @@
+import os
 import subprocess
 import sys
 from typing import Final
@@ -78,6 +79,9 @@ def test_importing_proxy_hooks_does_not_mutate_litellm_callbacks():
         "assert len(litellm.callbacks) == 0, f'mutated: {litellm.callbacks}'"
     )
     result: Final = subprocess.run(
-        [sys.executable, "-c", script], capture_output=True, text=True
+        [sys.executable, "-c", script],
+        capture_output=True,
+        text=True,
+        env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
     )
     assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}"

From 2b8b0eb2024a12ba9b8b152c7d17de37277bacc3 Mon Sep 17 00:00:00 2001
From: David Abutbul <David.a@prompt.security>
Date: Tue, 25 Aug 2026 14:28:53 +0300
Subject: [PATCH 009/164] fix(guardrails): block Prompt Security file
 modifications

---
 .../prompt_security/__init__.py               |   1 +
 .../prompt_security/prompt_security.py        |  39 +++---
 .../guardrail_hooks/prompt_security.py        |   4 +
 .../test_prompt_security_guardrails.py        | 118 ++++++++++++++++++
 4 files changed, 140 insertions(+), 22 deletions(-)

diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py
index 0aaba4016cd..88cf92a4a8c 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py
@@ -21,6 +21,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
         event_hook=litellm_params.mode,
         default_on=litellm_params.default_on,
         file_sanitization_fail_open=getattr(litellm_params, "file_sanitization_fail_open", None),
+        block_on_file_modify=getattr(litellm_params, "block_on_file_modify", None),
     )
     litellm.logging_callback_manager.add_litellm_callback(_prompt_security_callback)
 
diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py
index 84c4f118b00..0954fe1698a 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py
@@ -93,6 +93,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
         check_tool_results: bool | None = None,
         file_sanitization_timeout: float = _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS,
         file_sanitization_fail_open: bool | None = None,
+        block_on_file_modify: bool | None = None,
         **kwargs,
     ):
         kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
@@ -124,6 +125,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
         self.poll_interval = 2  # Seconds between polling attempts
         self.file_sanitization_timeout = file_sanitization_timeout
         self.file_sanitization_fail_open = file_sanitization_fail_open is not False
+        self.block_on_file_modify = block_on_file_modify is not False
 
         super().__init__(**kwargs)
 
@@ -372,13 +374,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
                     result = await self.sanitize_file_content(
                         file_data, filename, user_api_key_alias=user_api_key_alias
                     )
-
-                    if result.get("action") == "block":
-                        violations = result.get("violations", [])
-                        raise HTTPException(
-                            status_code=400,
-                            detail=f"Image blocked by Prompt Security. Violations: {', '.join(violations)}",
-                        )
+                    self._raise_if_file_blocked(result, "Image")
                 except HTTPException:
                     raise
                 except Exception as e:
@@ -408,7 +404,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
         file_data: bytes,
         filename: str,
         user_api_key_alias: str | None = None,
-    ) -> dict:
+    ) -> _SanitizeResult:
         """
         Sanitize file content using Prompt Security API.
         Returns: dict with keys 'action', 'content', 'metadata'
@@ -528,6 +524,17 @@ class PromptSecurityGuardrail(CustomGuardrail):
 
         raise HTTPException(status_code=408, detail="File sanitization timeout")
 
+    def _raise_if_file_blocked(self, sanitization_result: _SanitizeResult, resource_name: str) -> None:
+        action: Final = sanitization_result.get("action")
+        if action != "block" and not (action == "modify" and self.block_on_file_modify):
+            return
+
+        violations: Final = sanitization_result.get("violations", ())
+        raise HTTPException(
+            status_code=400,
+            detail=f"{resource_name} blocked by Prompt Security. Violations: {', '.join(violations)}",
+        )
+
     async def _process_image_url_item(self, item: dict, user_api_key_alias: str | None) -> dict:
         """Process and sanitize image_url items."""
         image_url_data: Final = item.get("image_url", {})
@@ -547,13 +554,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
                 file_data, filename, user_api_key_alias=user_api_key_alias
             )
             action: Final = sanitization_result.get("action")
-
-            if action == "block":
-                violations: Final = sanitization_result.get("violations", [])
-                raise HTTPException(
-                    status_code=400,
-                    detail=f"File blocked by Prompt Security. Violations: {', '.join(violations)}",
-                )
+            self._raise_if_file_blocked(sanitization_result, "File")
 
             if action == "modify":
                 sanitized_content: Final = sanitization_result.get("content", "")
@@ -615,13 +616,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
                 file_data, filename, user_api_key_alias=user_api_key_alias
             )
             action: Final = sanitization_result.get("action")
-
-            if action == "block":
-                violations: Final = sanitization_result.get("violations", [])
-                raise HTTPException(
-                    status_code=400,
-                    detail=f"Document blocked by Prompt Security. Violations: {', '.join(violations)}",
-                )
+            self._raise_if_file_blocked(sanitization_result, "Document")
 
             if action == "modify":
                 sanitized_content: Final = sanitization_result.get("content", "")
diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py b/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py
index 94f8161f44e..29f1b4bdcd6 100644
--- a/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py
+++ b/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py
@@ -16,6 +16,10 @@ class PromptSecurityGuardrailConfigModel(GuardrailConfigModel):
         default=True,
         description="Whether file sanitization timeouts allow the original file through instead of blocking the request.",
     )
+    block_on_file_modify: bool = Field(
+        default=True,
+        description="Whether a file sanitization `modify` verdict blocks the request instead of replacing the file content.",
+    )
 
     @staticmethod
     def ui_friendly_name() -> str:
diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py
index ab4e15ff423..e650f796f29 100644
--- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py
+++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py
@@ -31,6 +31,7 @@ def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch):
                     "mode": "during_call",
                     "default_on": True,
                     "file_sanitization_fail_open": False,
+                    "block_on_file_modify": False,
                 },
             }
         ],
@@ -43,9 +44,11 @@ def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch):
     assert registered[0].default_on is True
     assert registered[0].event_hook == "during_call"
     assert registered[0].file_sanitization_fail_open is False
+    assert registered[0].block_on_file_modify is False
     config_model = registered[0].get_config_model()
     assert config_model is not None
     assert config_model().file_sanitization_fail_open is True
+    assert config_model().block_on_file_modify is True
 
 
 def test_prompt_security_guard_config_no_api_key(monkeypatch: pytest.MonkeyPatch):
@@ -379,6 +382,121 @@ async def test_file_sanitization(monkeypatch: pytest.MonkeyPatch):
     assert result is not None
 
 
+@pytest.mark.asyncio
+async def test_file_sanitization_modify_blocks_by_default(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
+    monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
+
+    guardrail = PromptSecurityGuardrail(
+        guardrail_name="test-guard", event_hook="pre_call", default_on=True
+    )
+    csv_data = b"name,email\nAlice,alice@example.com\n"
+    item = {
+        "type": "file",
+        "file": {
+            "data": base64.b64encode(csv_data).decode(),
+            "mime_type": "text/csv",
+        },
+    }
+    upload_response = Response(
+        json={"jobId": "modify-job"},
+        status_code=200,
+        request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"),
+    )
+    poll_response = Response(
+        json={
+            "status": "done",
+            "content": "name,email\nAlice,[REDACTED]\n",
+            "metadata": {"action": "modify", "violations": ["Sensitive Data"]},
+        },
+        status_code=200,
+        request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"),
+    )
+
+    with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)):
+        with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)):
+            with pytest.raises(HTTPException) as exc_info:
+                await guardrail._process_document_item(item, None)
+
+    assert exc_info.value.status_code == 400
+    assert exc_info.value.detail == "Document blocked by Prompt Security. Violations: Sensitive Data"
+
+
+@pytest.mark.asyncio
+async def test_standalone_image_sanitization_modify_blocks_by_default(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
+    monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
+
+    guardrail = PromptSecurityGuardrail(
+        guardrail_name="test-guard", event_hook="pre_call", default_on=True
+    )
+    guardrail.poll_interval = 0
+    image_url = "data:image/png;base64," + base64.b64encode(b"image-content").decode()
+    upload_response = Response(
+        json={"jobId": "modify-image-job"},
+        status_code=200,
+        request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"),
+    )
+    poll_response = Response(
+        json={
+            "status": "done",
+            "content": "Email: [REDACTED]",
+            "metadata": {"action": "modify", "violations": ["Sensitive Data"]},
+        },
+        status_code=200,
+        request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"),
+    )
+
+    with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)):
+        with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)):
+            with pytest.raises(HTTPException) as exc_info:
+                await guardrail._process_standalone_images([image_url], None)
+
+    assert exc_info.value.status_code == 400
+    assert exc_info.value.detail == "Image blocked by Prompt Security. Violations: Sensitive Data"
+
+
+@pytest.mark.asyncio
+async def test_file_sanitization_modify_can_rewrite_when_blocking_disabled(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
+    monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
+
+    guardrail = PromptSecurityGuardrail(
+        guardrail_name="test-guard",
+        event_hook="pre_call",
+        default_on=True,
+        block_on_file_modify=False,
+    )
+    csv_data = b"name,email\nAlice,alice@example.com\n"
+    item = {
+        "type": "file",
+        "file": {
+            "data": base64.b64encode(csv_data).decode(),
+            "mime_type": "text/csv",
+        },
+    }
+    upload_response = Response(
+        json={"jobId": "modify-job"},
+        status_code=200,
+        request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"),
+    )
+    poll_response = Response(
+        json={
+            "status": "done",
+            "content": "name,email\nAlice,[REDACTED]\n",
+            "metadata": {"action": "modify", "violations": ["Sensitive Data"]},
+        },
+        status_code=200,
+        request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"),
+    )
+
+    with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)):
+        with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)):
+            result = await guardrail._process_document_item(item, None)
+
+    assert base64.b64decode(result["file"]["data"]) == b"name,email\nAlice,[REDACTED]\n"
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "timeout",

From c4982ca407b08a2161e76ebf2c7b3fa4fa3f885f Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 1 Sep 2026 17:04:48 -0700
Subject: [PATCH 010/164] fix(mcp): error instead of silent empty tools when
 scoped MCP access is denied; grant agent MCP servers from the UI

---
 .../proxy/_experimental/mcp_server/server.py  |  57 ++++++
 .../mcp_server/test_mcp_server.py             | 185 ++++++++++++++++++
 .../agents/_components/agent_config.ts        |  23 +++
 .../agent_info.integration.test.tsx           |  41 +++-
 .../agents/_components/agent_info.test.tsx    |  24 +++
 .../agents/_components/agent_info.tsx         |  74 ++++++-
 .../src/components/networking.tsx             |   1 +
 7 files changed, 398 insertions(+), 7 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 989b08b929a..0df38d6d309 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -23,6 +23,7 @@ from pydantic import AnyUrl, ConfigDict
 from starlette.requests import Request as StarletteRequest
 from starlette.responses import JSONResponse
 from starlette.types import Message, Receive, Scope, Send
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_logger
 from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
@@ -816,6 +817,15 @@ if MCP_AVAILABLE:
                 }
             }
             return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta})
+        except HTTPException as e:
+            from mcp.shared.exceptions import McpError
+            from mcp.types import INVALID_REQUEST, ErrorData
+
+            detail: Final = e.detail
+            message: Final = (
+                str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail)
+            )
+            raise McpError(ErrorData(code=INVALID_REQUEST, message=message)) from e
         except Exception as e:
             verbose_logger.exception("Error in list_tools endpoint: %s", e)
             # Return empty list instead of failing completely
@@ -1440,6 +1450,45 @@ if MCP_AVAILABLE:
 
         return allowed_mcp_servers
 
+    class _McpDeniedDetail(TypedDict):
+        error: ReadOnly[str]
+
+    async def _raise_denied_scoped_mcp_access(
+        requested_names: Sequence[str],
+        user_api_key_auth: UserAPIKeyAuth | None,
+        client_ip: str | None = None,
+    ) -> None:
+        """A scoped request (``/mcp/<name>`` path or ``x-mcp-servers`` header) resolved to zero
+        allowed servers. When a requested name IS a registered server visible to this client IP,
+        the denial is a permission outcome and must be loud: a silent 200 with no tools reads as
+        a healthy server with no tools. Names matching no registered server stay fail-closed
+        empty so scoping cannot probe for server existence."""
+        known_targets: Final = tuple(
+            (name, server)
+            for name in requested_names
+            if (server := global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip)) is not None
+        )
+        if not known_targets:
+            return
+        denied_name, denied_server = known_targets[0]
+        agent_id: Final = user_api_key_auth.agent_id if user_api_key_auth else None
+        if user_api_key_auth is not None and agent_id:
+            allowed_without_agent: Final = await global_mcp_server_manager.get_allowed_mcp_servers(
+                user_api_key_auth.model_copy(update=types.MappingProxyType({"agent_id": None}))
+            )
+            if denied_server.server_id in allowed_without_agent:
+                agent_denial: Final[_McpDeniedDetail] = {
+                    "error": (
+                        f"MCP server '{denied_name}' is not available to this key: the key is bound to "
+                        f"agent '{agent_id}', whose MCP grants do not include this server. Add the server "
+                        f"to the agent's object_permission.mcp_servers (edit the agent in the Admin UI or "
+                        f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent."
+                    )
+                }
+                raise HTTPException(status_code=403, detail=agent_denial)
+        key_denial: Final[_McpDeniedDetail] = {"error": f"The key is not allowed to access server {denied_name}"}
+        raise HTTPException(status_code=403, detail=key_denial)
+
     def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool:
         """
         Check if a tool name matches any name in the filter list.
@@ -1964,6 +2013,12 @@ if MCP_AVAILABLE:
                 mcp_servers=mcp_servers,
                 client_ip=client_ip,
             )
+            if mcp_servers is not None and not allowed_mcp_servers:
+                await _raise_denied_scoped_mcp_access(
+                    requested_names=mcp_servers,
+                    user_api_key_auth=user_api_key_auth,
+                    client_ip=client_ip,
+                )
 
             # Pre-fetch OAuth credentials only when at least one server uses OAuth2,
             # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers.
@@ -2388,6 +2443,8 @@ if MCP_AVAILABLE:
             )
             verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools))
             return listing
+        except HTTPException:
+            raise
         except Exception as e:
             verbose_logger.exception("Error getting tools from managed MCP servers: %s", e)
             # Continue with an empty listing instead of failing completely
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
index 82f74cda835..0040149d388 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
@@ -1329,6 +1329,191 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing():
             mock_logger.info.assert_any_call("Successfully fetched %s tools total from all MCP servers", 0)
 
 
+def _denied_scope_manager(known_server_names_to_ids: dict[str, str], allowed_without_agent: list[str]) -> MagicMock:
+    """A manager whose get_mcp_server_by_name knows the given names and whose
+    get_allowed_mcp_servers answers the agent-stripped permission rerun."""
+    servers = {name: MagicMock(server_id=server_id) for name, server_id in known_server_names_to_ids.items()}
+    manager = MagicMock()
+    manager.get_mcp_server_by_name = lambda name, client_ip=None: servers.get(name)
+    manager.get_allowed_mcp_servers = AsyncMock(return_value=allowed_without_agent)
+    return manager
+
+
+@pytest.mark.asyncio
+async def test_scoped_list_denied_by_agent_binding_raises_403_naming_agent():
+    """A scoped tools/list that resolves to zero servers because the key's bound agent lacks the
+    grant must raise a 403 naming the agent, never return a silent 200 with no tools."""
+    try:
+        from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
+    mock_manager = _denied_scope_manager({"github": "srv-github"}, allowed_without_agent=["srv-github"])
+
+    with (
+        patch(  # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
+            AsyncMock(return_value=[]),
+        ),
+        patch(  # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
+            mock_manager,
+        ),
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await _get_tools_from_mcp_servers(
+                user_api_key_auth=user_api_key_auth,
+                mcp_auth_header=None,
+                mcp_servers=["github"],
+            )
+
+    assert exc_info.value.status_code == 403
+    message = exc_info.value.detail["error"]
+    assert "github" in message
+    assert "agent 'agent-123'" in message
+    rerun_auth = mock_manager.get_allowed_mcp_servers.await_args.args[0]
+    assert rerun_auth.agent_id is None
+    assert rerun_auth.user_id == "test_user"
+
+
+@pytest.mark.asyncio
+async def test_scoped_list_denied_for_non_agent_key_raises_generic_403():
+    """A scoped tools/list denied for a key with no agent binding raises the generic 403 and
+    never runs the agent-stripped permission rerun."""
+    try:
+        from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
+    mock_manager = _denied_scope_manager({"github": "srv-github"}, allowed_without_agent=["srv-github"])
+
+    with (
+        patch(  # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
+            AsyncMock(return_value=[]),
+        ),
+        patch(  # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
+            mock_manager,
+        ),
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await _get_tools_from_mcp_servers(
+                user_api_key_auth=user_api_key_auth,
+                mcp_auth_header=None,
+                mcp_servers=["github"],
+            )
+
+    assert exc_info.value.status_code == 403
+    message = exc_info.value.detail["error"]
+    assert "github" in message
+    assert "agent" not in message
+    mock_manager.get_allowed_mcp_servers.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_scoped_list_unknown_server_name_stays_silent_empty():
+    """A scoped request naming no registered server stays fail-closed empty (200, no tools), so
+    scoping cannot probe for server existence."""
+    try:
+        from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
+    mock_manager = _denied_scope_manager({}, allowed_without_agent=["srv-github"])
+
+    with (
+        patch(  # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
+            AsyncMock(return_value=[]),
+        ),
+        patch(  # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
+            mock_manager,
+        ),
+    ):
+        result = await _get_tools_from_mcp_servers(
+            user_api_key_auth=user_api_key_auth,
+            mcp_auth_header=None,
+            mcp_servers=["doesnotexist"],
+        )
+
+    assert result.tools == []
+    assert result.outcomes == {}
+    mock_manager.get_allowed_mcp_servers.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_scoped_list_agent_key_denied_by_key_grants_raises_generic_403():
+    """When the agent-stripped rerun still denies the server, the denial is not the agent's doing,
+    so the 403 stays generic instead of blaming the agent binding."""
+    try:
+        from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
+    mock_manager = _denied_scope_manager({"github": "srv-github"}, allowed_without_agent=[])
+
+    with (
+        patch(  # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
+            AsyncMock(return_value=[]),
+        ),
+        patch(  # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
+            mock_manager,
+        ),
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await _get_tools_from_mcp_servers(
+                user_api_key_auth=user_api_key_auth,
+                mcp_auth_header=None,
+                mcp_servers=["github"],
+            )
+
+    assert exc_info.value.status_code == 403
+    message = exc_info.value.detail["error"]
+    assert "github" in message
+    assert "agent" not in message
+    mock_manager.get_allowed_mcp_servers.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error():
+    """The MCP protocol handler surfaces a permission HTTPException as a clean JSON-RPC error
+    (McpError, INVALID_REQUEST) carrying the denial message, instead of a raw 500."""
+    try:
+        from litellm.proxy._experimental.mcp_server.server import handle_list_tools
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    from mcp.shared.exceptions import McpError
+    from mcp.types import INVALID_REQUEST
+
+    denial_message = "MCP server 'github' is not available to this key: the key is bound to agent 'agent-123'"
+    denial = HTTPException(status_code=403, detail={"error": denial_message})
+
+    with (
+        patch(  # test-quality-ok: the protocol handler reads auth from module context; no injection seam
+            "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context",
+            new=AsyncMock(return_value=(None, None, None, None, None, None, None)),
+        ),
+        patch(  # test-quality-ok: the listing helper is the handler's only collaborator; the suite's seam
+            "litellm.proxy._experimental.mcp_server.server._list_mcp_tools",
+            new=AsyncMock(side_effect=denial),
+        ),
+    ):
+        with pytest.raises(McpError) as exc_info:
+            await handle_list_tools()
+
+    assert exc_info.value.error.code == INVALID_REQUEST
+    assert exc_info.value.error.message == denial_message
+
+
 @pytest.mark.asyncio
 async def test_mcp_server_tool_call_body_with_none_arguments():
     """Test that proxy_server_request body handles None arguments correctly"""
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts
index 442dcd48f66..4e93c0c7a51 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts
@@ -313,6 +313,28 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => {
   return agentData;
 };
 
+/**
+ * Parse MCP grants from an agent's object_permission into the shared MCP form fields
+ */
+export const parseMcpPermissionsForForm = (agent: any) => ({
+  allowed_mcp_servers_and_groups: {
+    servers: agent.object_permission?.mcp_servers ?? [],
+    accessGroups: agent.object_permission?.mcp_access_groups ?? [],
+  },
+  mcp_tool_permissions: agent.object_permission?.mcp_tool_permissions ?? {},
+});
+
+/**
+ * Build the object_permission payload from the shared MCP form fields.
+ * Always includes the MCP keys (empty when cleared) so removals persist;
+ * the proxy merges per key, leaving non-MCP grants untouched.
+ */
+export const buildMcpObjectPermission = (values: any) => ({
+  mcp_servers: values.allowed_mcp_servers_and_groups?.servers ?? [],
+  mcp_access_groups: values.allowed_mcp_servers_and_groups?.accessGroups ?? [],
+  mcp_tool_permissions: values.mcp_tool_permissions ?? {},
+});
+
 /**
  * Parse agent data for form fields
  */
@@ -356,5 +378,6 @@ export const parseAgentForForm = (agent: any) => {
       : [],
     // extra_headers: already an array of strings
     extra_headers: agent.extra_headers ?? [],
+    ...parseMcpPermissionsForForm(agent),
   };
 };
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx
index 79bd2f6a21b..c813c513134 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx
@@ -1,4 +1,5 @@
 import React from "react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event";
 import { describe, it, expect, vi, beforeEach } from "vitest";
@@ -10,6 +11,12 @@ vi.mock("@/components/networking", () => ({
   getAgentInfo: vi.fn(),
   patchAgentCall: vi.fn(),
   getAgentCreateMetadata: vi.fn(),
+  getProxyBaseUrl: vi.fn(() => ""),
+  getUiConfig: vi.fn(async () => ({})),
+  fetchMCPServers: vi.fn(async () => []),
+  fetchMCPAccessGroups: vi.fn(async () => []),
+  fetchMCPToolsets: vi.fn(async () => []),
+  listMCPTools: vi.fn(async () => ({ tools: [] })),
 }));
 
 vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({
@@ -77,7 +84,14 @@ const langgraphInfo: AgentCreateInfo = {
 
 const setup = () => userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never });
 
-const renderView = () => render(<AgentInfoView agentId="agent-1" onClose={vi.fn()} accessToken="tok" isAdmin={true} />);
+const renderView = () => {
+  const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });
+  return render(
+    <QueryClientProvider client={queryClient}>
+      <AgentInfoView agentId="agent-1" onClose={vi.fn()} accessToken="tok" isAdmin={true} />
+    </QueryClientProvider>,
+  );
+};
 
 const openEditor = async (user: ReturnType<typeof setup>) => {
   await user.click(await screen.findByRole("tab", { name: "Settings" }));
@@ -127,6 +141,7 @@ describe("AgentInfoView update payload", () => {
       rpm_limit: 222,
       session_tpm_limit: 333,
       session_rpm_limit: 444,
+      object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_tool_permissions: {} },
     });
   });
 
@@ -167,6 +182,7 @@ describe("AgentInfoView update payload", () => {
       rpm_limit: 222,
       session_tpm_limit: 333,
       session_rpm_limit: 444,
+      object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_tool_permissions: {} },
     });
   });
 
@@ -244,6 +260,29 @@ describe("AgentInfoView update payload", () => {
         api_base: "https://other.example.com",
         model: "langgraph/asst_1",
       },
+      object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_tool_permissions: {} },
+    });
+  });
+
+  it("keeps the agent's existing MCP grants in the update payload", async () => {
+    vi.mocked(networking.getAgentInfo).mockResolvedValue({
+      ...A2A_AGENT,
+      object_permission: {
+        mcp_servers: ["srv-1"],
+        mcp_access_groups: ["grp-a"],
+        mcp_tool_permissions: { "srv-1": ["tool_x"] },
+      },
+    } as never);
+    const user = setup();
+    renderView();
+    await openEditor(user);
+
+    await save(user);
+
+    expect(patchedPayload().object_permission).toEqual({
+      mcp_servers: ["srv-1"],
+      mcp_access_groups: ["grp-a"],
+      mcp_tool_permissions: { "srv-1": ["tool_x"] },
     });
   });
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx
index 0936b8e13db..a351ff2090f 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx
@@ -24,6 +24,18 @@ vi.mock("./agent_form_fields", () => ({
   unmountedA2AFieldNames: () => [],
 }));
 
+vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({
+  useMCPServers: () => ({ data: [{ server_id: "srv-1", server_name: "github" }] }),
+}));
+
+vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({
+  default: () => <div data-testid="mcp-server-selector" />,
+}));
+
+vi.mock("@/components/mcp_server_management/MCPToolPermissions", () => ({
+  default: () => <div data-testid="mcp-tool-permissions" />,
+}));
+
 const agent = {
   agent_id: "agent-1",
   agent_name: "support-agent",
@@ -62,5 +74,17 @@ describe("AgentInfoView settings", () => {
     expect(token).toBe("sk-test");
     expect(agentId).toBe("agent-1");
     expect(payload.tpm_limit).toBe(42);
+    expect(payload.object_permission).toEqual({ mcp_servers: [], mcp_access_groups: [], mcp_tool_permissions: {} });
+  });
+
+  it("shows MCP grants with server names on the overview tab", async () => {
+    vi.mocked(networking.getAgentInfo).mockResolvedValue({
+      ...agent,
+      object_permission: { mcp_servers: ["srv-1"] },
+    } as unknown as Agent);
+
+    render(<AgentInfoView agentId="agent-1" onClose={vi.fn()} accessToken="sk-test" isAdmin={true} />);
+
+    expect(await screen.findByText("github (srv-1)")).toBeInTheDocument();
   });
 });
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx
index eddeeec674b..1592b452455 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx
@@ -15,16 +15,27 @@ import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo }
 import { Agent } from "@/components/agents/types";
 import { KeyResponse } from "@/components/key_team_helpers/key_list";
 import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
+import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers";
 import KeyInfoView from "@/components/templates/key_info_view";
+import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector";
+import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions";
 import AgentVirtualKeys from "./agent_virtual_keys";
 import AgentFormFields, { unmountedA2AFieldNames } from "./agent_form_fields";
 import DynamicAgentFormFields, { buildDynamicAgentData, unmountedDynamicFieldNames } from "./dynamic_agent_form_fields";
-import { AGENT_FORM_CONFIG, buildAgentDataFromForm, parseAgentForForm } from "./agent_config";
+import {
+  AGENT_FORM_CONFIG,
+  buildAgentDataFromForm,
+  buildMcpObjectPermission,
+  parseAgentForForm,
+  parseMcpPermissionsForForm,
+} from "./agent_config";
 import {
   AgentFormField,
   AgentFormValues,
   AgentNumberInput,
   AgentRequestPayload,
+  McpServerSelection,
+  labelWithHint,
   omitFieldValues,
   useCollapsiblePanels,
 } from "./AgentFormKit";
@@ -111,7 +122,7 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
       } else {
         const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType);
         if (typeInfo) {
-          form.reset(parseDynamicAgentForForm(data, typeInfo));
+          form.reset({ ...parseDynamicAgentForForm(data, typeInfo), ...parseMcpPermissionsForForm(data) });
         } else {
           form.reset(parseAgentForForm(data));
         }
@@ -131,7 +142,7 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
       if (agentType !== "a2a") {
         const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType);
         if (typeInfo) {
-          form.reset(parseDynamicAgentForForm(agent, typeInfo));
+          form.reset({ ...parseDynamicAgentForForm(agent, typeInfo), ...parseMcpPermissionsForForm(agent) });
         }
       }
     }
@@ -139,6 +150,14 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
 
   const selectedAgentTypeInfo = agentTypeMetadata.find((t) => t.agent_type === detectedAgentType);
   const watchedFormValues = useWatch({ control: form.control });
+  const mcpSelection = useWatch({ control: form.control, name: "allowed_mcp_servers_and_groups" });
+  const mcpToolPermissions = useWatch({ control: form.control, name: "mcp_tool_permissions" });
+  const { data: mcpServers = [] } = useMCPServers();
+
+  const mcpServerLabel = (serverId: string) => {
+    const server = mcpServers.find((s) => s.server_id === serverId);
+    return server?.server_name ? `${server.server_name} (${serverId})` : serverId;
+  };
 
   const discoveryRequest = useMemo(
     () => buildDiscoveryRequest(detectedAgentType, watchedFormValues || {}, selectedAgentTypeInfo),
@@ -199,7 +218,10 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
         ? overlayDiscoveredCardParams(built, appliedDiscoveredSelection.selected_card)
         : built;
 
-      await patchAgentCall(accessToken, agentId, updateData);
+      await patchAgentCall(accessToken, agentId, {
+        ...updateData,
+        object_permission: buildMcpObjectPermission(values),
+      });
       toast.success("Agent updated successfully");
       setIsEditing(false);
       fetchAgentInfo();
@@ -343,7 +365,13 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
                   <h3 className="text-lg font-medium">MCP Tool Permissions</h3>
                   <DetailList className="mt-4">
                     {agent.object_permission.mcp_servers && agent.object_permission.mcp_servers.length > 0 && (
-                      <DetailItem label="MCP Servers">{agent.object_permission.mcp_servers.join(", ")}</DetailItem>
+                      <DetailItem label="MCP Servers">
+                        <div className="space-y-1">
+                          {agent.object_permission.mcp_servers.map((serverId) => (
+                            <div key={serverId}>{mcpServerLabel(serverId)}</div>
+                          ))}
+                        </div>
+                      </DetailItem>
                     )}
                     {agent.object_permission.mcp_access_groups &&
                       agent.object_permission.mcp_access_groups.length > 0 && (
@@ -357,7 +385,7 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
                           <div className="space-y-1">
                             {Object.entries(agent.object_permission.mcp_tool_permissions).map(([serverId, tools]) => (
                               <div key={serverId}>
-                                <span className="font-medium">{serverId}:</span>{" "}
+                                <span className="font-medium">{mcpServerLabel(serverId)}:</span>{" "}
                                 {Array.isArray(tools) ? tools.join(", ") : String(tools)}
                               </div>
                             ))}
@@ -457,6 +485,40 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
                           {rateLimitField("session_rpm_limit", "Session RPM Limit")}
                         </div>
 
+                        <Separator className="my-6" />
+                        <h3 className="text-lg font-medium mb-4">MCP Servers</h3>
+                        <FieldGroup>
+                          <AgentFormField
+                            name="allowed_mcp_servers_and_groups"
+                            label={labelWithHint(
+                              "Allowed MCP Servers",
+                              "Select which MCP servers or access groups this agent can access. Keys bound to this agent can only reach servers granted here.",
+                            )}
+                          >
+                            {({ value, onChange }) => (
+                              <MCPServerSelector
+                                onChange={onChange}
+                                value={{
+                                  servers: (value as McpServerSelection | undefined)?.servers ?? [],
+                                  accessGroups: (value as McpServerSelection | undefined)?.accessGroups ?? [],
+                                }}
+                                accessToken={accessToken ?? ""}
+                                placeholder="Select MCP servers or access groups (optional)"
+                              />
+                            )}
+                          </AgentFormField>
+                        </FieldGroup>
+                        <div className="mt-4">
+                          <MCPToolPermissions
+                            accessToken={accessToken ?? ""}
+                            selectedServers={mcpSelection?.servers ?? []}
+                            toolPermissions={mcpToolPermissions ?? {}}
+                            onChange={(toolPerms: Record<string, string[]>) =>
+                              form.setValue("mcp_tool_permissions", toolPerms)
+                            }
+                          />
+                        </div>
+
                         <div className="mt-6 flex justify-end gap-2">
                           <Button
                             type="button"
diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx
index b53a67f9bcc..9c9458c81f2 100644
--- a/ui/litellm-dashboard/src/components/networking.tsx
+++ b/ui/litellm-dashboard/src/components/networking.tsx
@@ -6239,6 +6239,7 @@ export const patchAgentCall = async (
     agent_name?: string;
     litellm_params?: Record<string, any>;
     agent_card_params?: Record<string, any>;
+    object_permission?: Record<string, any>;
     tpm_limit?: number | null;
     rpm_limit?: number | null;
     session_tpm_limit?: number | null;

From 659eb0462e251207efb2b69faed16eb47735638d Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 1 Sep 2026 18:37:23 -0700
Subject: [PATCH 011/164] fix(mcp): deny every zero-resolution scoped request
 uniformly and persist agent MCP toolsets from the UI

Any scoped MCP request (/mcp/<name> path or x-mcp-servers header) that resolves to zero allowed servers now returns one generic 403, so unknown, unauthorized, and access-group names are indistinguishable and the error cannot be used to enumerate servers. The agent-attributed variant reruns the same scope resolver with the agent binding stripped and fires only when that rerun resolves, naming the vetoed server or access group and both fix paths

AgentObjectPermission now declares mcp_toolsets so PATCH /v1/agents keeps it, and the agent edit form round-trips toolsets alongside servers and access groups instead of dropping them on save. The two scope-resolver params only iterate their names, so they take Sequence[str] and the LIT001 budget ratchets down by the two annotations this widens
---
 .../proxy/_experimental/mcp_server/server.py  |  60 +++--
 litellm/proxy/_lazy_openapi_snapshot.json     |  14 ++
 litellm/types/agents.py                       |   5 +-
 .../mcp_server/test_mcp_server.py             | 216 +++++++++---------
 type-discipline-budget.json                   |   2 +-
 .../agents/_components/agent_config.ts        |  10 +-
 .../agent_info.integration.test.tsx           |  24 +-
 .../agents/_components/agent_info.test.tsx    |   3 +-
 .../agents/_components/agent_info.tsx         |   5 +
 ui/litellm-dashboard/src/lib/http/schema.d.ts |   2 +
 10 files changed, 189 insertions(+), 152 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 0df38d6d309..9f530ddf862 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -1393,7 +1393,7 @@ if MCP_AVAILABLE:
     ########################################################
 
     async def _get_allowed_mcp_servers_from_mcp_server_names(
-        mcp_servers: list[str] | None,
+        mcp_servers: Sequence[str] | None,
         allowed_mcp_servers: list[MCPServer],
     ) -> list[MCPServer]:
         """
@@ -1459,35 +1459,55 @@ if MCP_AVAILABLE:
         client_ip: str | None = None,
     ) -> None:
         """A scoped request (``/mcp/<name>`` path or ``x-mcp-servers`` header) resolved to zero
-        allowed servers. When a requested name IS a registered server visible to this client IP,
-        the denial is a permission outcome and must be loud: a silent 200 with no tools reads as
-        a healthy server with no tools. Names matching no registered server stay fail-closed
-        empty so scoping cannot probe for server existence."""
-        known_targets: Final = tuple(
-            (name, server)
-            for name in requested_names
-            if (server := global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip)) is not None
-        )
-        if not known_targets:
-            return
-        denied_name, denied_server = known_targets[0]
+        allowed servers, so the denial must be loud: a silent 200 with no tools reads as a healthy
+        server with no tools. Unknown, unauthorized, and access-group names all share one generic
+        error so scoping cannot probe which servers exist; the agent variant fires only when the
+        same request resolves once the agent binding is stripped, proving the binding caused the veto."""
         agent_id: Final = user_api_key_auth.agent_id if user_api_key_auth else None
         if user_api_key_auth is not None and agent_id:
-            allowed_without_agent: Final = await global_mcp_server_manager.get_allowed_mcp_servers(
-                user_api_key_auth.model_copy(update=types.MappingProxyType({"agent_id": None}))
+            resolved_without_agent: Final = await _get_allowed_mcp_servers(
+                user_api_key_auth=user_api_key_auth.model_copy(update=types.MappingProxyType({"agent_id": None})),
+                mcp_servers=requested_names,
+                client_ip=client_ip,
             )
-            if denied_server.server_id in allowed_without_agent:
+            resolved_ids: Final = frozenset(server.server_id for server in resolved_without_agent)
+
+            def _registered_server_id(name: str) -> str | None:
+                server: Final = global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip)
+                return server.server_id if server is not None else None
+
+            vetoed_server: Final = next(
+                (name for name in requested_names if _registered_server_id(name) in resolved_ids), None
+            )
+            if vetoed_server is not None:
                 agent_denial: Final[_McpDeniedDetail] = {
                     "error": (
-                        f"MCP server '{denied_name}' is not available to this key: the key is bound to "
+                        f"MCP server '{vetoed_server}' is not available to this key: the key is bound to "
                         f"agent '{agent_id}', whose MCP grants do not include this server. Add the server "
                         f"to the agent's object_permission.mcp_servers (edit the agent in the Admin UI or "
                         f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent."
                     )
                 }
                 raise HTTPException(status_code=403, detail=agent_denial)
-        key_denial: Final[_McpDeniedDetail] = {"error": f"The key is not allowed to access server {denied_name}"}
-        raise HTTPException(status_code=403, detail=key_denial)
+            vetoed_group: Final = (
+                next((name for name in requested_names if _registered_server_id(name) is None), None)
+                if resolved_ids
+                else None
+            )
+            if vetoed_group is not None:
+                group_denial: Final[_McpDeniedDetail] = {
+                    "error": (
+                        f"MCP access group '{vetoed_group}' is not available to this key: the key is bound to "
+                        f"agent '{agent_id}', whose MCP grants do not include it. Add the group to the "
+                        f"agent's object_permission.mcp_access_groups (edit the agent in the Admin UI or "
+                        f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent."
+                    )
+                }
+                raise HTTPException(status_code=403, detail=group_denial)
+        generic_denial: Final[_McpDeniedDetail] = {
+            "error": f"The key is not allowed to access the requested MCP servers: {', '.join(requested_names)}"
+        }
+        raise HTTPException(status_code=403, detail=generic_denial)
 
     def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool:
         """
@@ -1581,7 +1601,7 @@ if MCP_AVAILABLE:
 
     async def _get_allowed_mcp_servers(
         user_api_key_auth: UserAPIKeyAuth | None,
-        mcp_servers: list[str] | None,
+        mcp_servers: Sequence[str] | None,
         client_ip: str | None = None,
     ) -> list[MCPServer]:
         """Return allowed MCP servers for a request after applying filters.
diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json
index 5af45b29226..f2c4a5e6212 100644
--- a/litellm/proxy/_lazy_openapi_snapshot.json
+++ b/litellm/proxy/_lazy_openapi_snapshot.json
@@ -2576,6 +2576,20 @@
               ],
               "title": "Mcp Tool Permissions"
             },
+            "mcp_toolsets": {
+              "anyOf": [
+                {
+                  "items": {
+                    "type": "string"
+                  },
+                  "type": "array"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Mcp Toolsets"
+            },
             "models": {
               "anyOf": [
                 {
diff --git a/litellm/types/agents.py b/litellm/types/agents.py
index 2cb42ce3fac..dbaaab62d86 100644
--- a/litellm/types/agents.py
+++ b/litellm/types/agents.py
@@ -1,9 +1,9 @@
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, Literal
 
 from pydantic import BaseModel, PrivateAttr, StrictInt
-from typing_extensions import Required, TypedDict
+from typing_extensions import ReadOnly, Required, TypedDict
 
 from litellm.types.llms.base import LiteLLMPydanticObjectBase
 
@@ -172,6 +172,7 @@ class AugmentedAgentCard(AgentCard):
 class AgentObjectPermission(TypedDict, total=False):
     mcp_servers: list[str] | None
     mcp_access_groups: list[str] | None
+    mcp_toolsets: ReadOnly[Sequence[str] | None]
     mcp_tool_permissions: dict[str, list[str]] | None
     models: list[str] | None
     agents: list[str] | None
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
index 0040149d388..fb69ad3b1c4 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
@@ -1329,157 +1329,155 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing():
             mock_logger.info.assert_any_call("Successfully fetched %s tools total from all MCP servers", 0)
 
 
-def _denied_scope_manager(known_server_names_to_ids: dict[str, str], allowed_without_agent: list[str]) -> MagicMock:
-    """A manager whose get_mcp_server_by_name knows the given names and whose
-    get_allowed_mcp_servers answers the agent-stripped permission rerun."""
+def _denied_scope_manager(known_server_names_to_ids: dict[str, str]) -> MagicMock:
     servers = {name: MagicMock(server_id=server_id) for name, server_id in known_server_names_to_ids.items()}
     manager = MagicMock()
     manager.get_mcp_server_by_name = lambda name, client_ip=None: servers.get(name)
-    manager.get_allowed_mcp_servers = AsyncMock(return_value=allowed_without_agent)
     return manager
 
 
+def _scope_resolver(resolved_without_agent: list[str]) -> AsyncMock:
+    async def resolve(user_api_key_auth, mcp_servers, client_ip=None):
+        if user_api_key_auth is not None and user_api_key_auth.agent_id:
+            return []
+        return [MagicMock(server_id=server_id) for server_id in resolved_without_agent]
+
+    return AsyncMock(side_effect=resolve)
+
+
+async def _denied_scoped_list(
+    user_api_key_auth: UserAPIKeyAuth,
+    mcp_servers: list[str],
+    mock_manager: MagicMock,
+    resolver: AsyncMock,
+) -> HTTPException:
+    from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
+
+    with (
+        patch(  # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
+            resolver,
+        ),
+        patch(  # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
+            mock_manager,
+        ),
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await _get_tools_from_mcp_servers(
+                user_api_key_auth=user_api_key_auth,
+                mcp_auth_header=None,
+                mcp_servers=mcp_servers,
+            )
+    return exc_info.value
+
+
 @pytest.mark.asyncio
 async def test_scoped_list_denied_by_agent_binding_raises_403_naming_agent():
-    """A scoped tools/list that resolves to zero servers because the key's bound agent lacks the
-    grant must raise a 403 naming the agent, never return a silent 200 with no tools."""
-    try:
-        from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
-    except ImportError:
-        pytest.skip("MCP server not available")
+    """The agent-binding veto must raise a 403 naming the agent, never a silent 200 with no tools."""
+    pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
 
     user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
-    mock_manager = _denied_scope_manager({"github": "srv-github"}, allowed_without_agent=["srv-github"])
+    resolver = _scope_resolver(resolved_without_agent=["srv-github"])
 
-    with (
-        patch(  # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
-            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
-            AsyncMock(return_value=[]),
-        ),
-        patch(  # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
-            "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
-            mock_manager,
-        ),
-    ):
-        with pytest.raises(HTTPException) as exc_info:
-            await _get_tools_from_mcp_servers(
-                user_api_key_auth=user_api_key_auth,
-                mcp_auth_header=None,
-                mcp_servers=["github"],
-            )
+    denial = await _denied_scoped_list(
+        user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver
+    )
 
-    assert exc_info.value.status_code == 403
-    message = exc_info.value.detail["error"]
-    assert "github" in message
+    assert denial.status_code == 403
+    message = denial.detail["error"]
+    assert "MCP server 'github'" in message
     assert "agent 'agent-123'" in message
-    rerun_auth = mock_manager.get_allowed_mcp_servers.await_args.args[0]
+    assert "mcp_servers" in message
+    rerun_auth = resolver.await_args_list[1].kwargs["user_api_key_auth"]
     assert rerun_auth.agent_id is None
     assert rerun_auth.user_id == "test_user"
+    assert resolver.await_args_list[1].kwargs["mcp_servers"] == ["github"]
 
 
 @pytest.mark.asyncio
 async def test_scoped_list_denied_for_non_agent_key_raises_generic_403():
-    """A scoped tools/list denied for a key with no agent binding raises the generic 403 and
-    never runs the agent-stripped permission rerun."""
-    try:
-        from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
-    except ImportError:
-        pytest.skip("MCP server not available")
+    """A denial for a key with no agent binding stays generic and skips the agent-stripped rerun."""
+    pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
 
     user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
-    mock_manager = _denied_scope_manager({"github": "srv-github"}, allowed_without_agent=["srv-github"])
+    resolver = AsyncMock(return_value=[])
 
-    with (
-        patch(  # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
-            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
-            AsyncMock(return_value=[]),
-        ),
-        patch(  # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
-            "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
-            mock_manager,
-        ),
-    ):
-        with pytest.raises(HTTPException) as exc_info:
-            await _get_tools_from_mcp_servers(
-                user_api_key_auth=user_api_key_auth,
-                mcp_auth_header=None,
-                mcp_servers=["github"],
-            )
+    denial = await _denied_scoped_list(
+        user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver
+    )
 
-    assert exc_info.value.status_code == 403
-    message = exc_info.value.detail["error"]
+    assert denial.status_code == 403
+    message = denial.detail["error"]
     assert "github" in message
     assert "agent" not in message
-    mock_manager.get_allowed_mcp_servers.assert_not_awaited()
+    resolver.assert_awaited_once()
 
 
 @pytest.mark.asyncio
-async def test_scoped_list_unknown_server_name_stays_silent_empty():
-    """A scoped request naming no registered server stays fail-closed empty (200, no tools), so
-    scoping cannot probe for server existence."""
-    try:
-        from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
-    except ImportError:
-        pytest.skip("MCP server not available")
+async def test_scoped_list_unknown_name_raises_same_generic_403_as_unauthorized():
+    """Unknown and registered-but-unauthorized names raise byte-identical generic 403s, so a
+    caller cannot probe which server names exist."""
+    pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
 
     user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
-    mock_manager = _denied_scope_manager({}, allowed_without_agent=["srv-github"])
 
-    with (
-        patch(  # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
-            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
-            AsyncMock(return_value=[]),
-        ),
-        patch(  # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
-            "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
-            mock_manager,
-        ),
-    ):
-        result = await _get_tools_from_mcp_servers(
-            user_api_key_auth=user_api_key_auth,
-            mcp_auth_header=None,
-            mcp_servers=["doesnotexist"],
-        )
+    unknown = await _denied_scoped_list(
+        user_api_key_auth, ["github"], _denied_scope_manager({}), _scope_resolver(resolved_without_agent=[])
+    )
+    unauthorized = await _denied_scoped_list(
+        user_api_key_auth,
+        ["github"],
+        _denied_scope_manager({"github": "srv-github"}),
+        _scope_resolver(resolved_without_agent=[]),
+    )
 
-    assert result.tools == []
-    assert result.outcomes == {}
-    mock_manager.get_allowed_mcp_servers.assert_not_awaited()
+    assert unknown.status_code == unauthorized.status_code == 403
+    assert unknown.detail["error"] == unauthorized.detail["error"]
+    assert "github" in unknown.detail["error"]
+    assert "agent" not in unknown.detail["error"]
+
+
+@pytest.mark.asyncio
+async def test_scoped_list_access_group_vetoed_by_agent_names_agent_and_group():
+    """An access-group scope vetoed by the agent binding raises the 403 naming the agent and the
+    group instead of the silent empty list."""
+    pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
+
+    user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
+
+    denial = await _denied_scoped_list(
+        user_api_key_auth,
+        ["prod-group"],
+        _denied_scope_manager({}),
+        _scope_resolver(resolved_without_agent=["srv-github"]),
+    )
+
+    assert denial.status_code == 403
+    message = denial.detail["error"]
+    assert "access group 'prod-group'" in message
+    assert "agent 'agent-123'" in message
+    assert "mcp_access_groups" in message
 
 
 @pytest.mark.asyncio
 async def test_scoped_list_agent_key_denied_by_key_grants_raises_generic_403():
-    """When the agent-stripped rerun still denies the server, the denial is not the agent's doing,
-    so the 403 stays generic instead of blaming the agent binding."""
-    try:
-        from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
-    except ImportError:
-        pytest.skip("MCP server not available")
+    """When the agent-stripped rerun still resolves nothing, the 403 stays generic instead of
+    blaming the agent binding."""
+    pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
 
     user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
-    mock_manager = _denied_scope_manager({"github": "srv-github"}, allowed_without_agent=[])
+    resolver = _scope_resolver(resolved_without_agent=[])
 
-    with (
-        patch(  # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
-            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
-            AsyncMock(return_value=[]),
-        ),
-        patch(  # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
-            "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
-            mock_manager,
-        ),
-    ):
-        with pytest.raises(HTTPException) as exc_info:
-            await _get_tools_from_mcp_servers(
-                user_api_key_auth=user_api_key_auth,
-                mcp_auth_header=None,
-                mcp_servers=["github"],
-            )
+    denial = await _denied_scoped_list(
+        user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver
+    )
 
-    assert exc_info.value.status_code == 403
-    message = exc_info.value.detail["error"]
+    assert denial.status_code == 403
+    message = denial.detail["error"]
     assert "github" in message
     assert "agent" not in message
-    mock_manager.get_allowed_mcp_servers.assert_awaited_once()
+    assert resolver.await_count == 2
 
 
 @pytest.mark.asyncio
diff --git a/type-discipline-budget.json b/type-discipline-budget.json
index 3d2e97d55a5..3589b007ce5 100644
--- a/type-discipline-budget.json
+++ b/type-discipline-budget.json
@@ -1,6 +1,6 @@
 {
   "LIT001": {
-    "limit": 22367
+    "limit": 22365
   },
   "LIT002": {
     "limit": 26777
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts
index 4e93c0c7a51..e0a5271366f 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts
@@ -313,25 +313,23 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => {
   return agentData;
 };
 
-/**
- * Parse MCP grants from an agent's object_permission into the shared MCP form fields
- */
 export const parseMcpPermissionsForForm = (agent: any) => ({
   allowed_mcp_servers_and_groups: {
     servers: agent.object_permission?.mcp_servers ?? [],
     accessGroups: agent.object_permission?.mcp_access_groups ?? [],
+    toolsets: agent.object_permission?.mcp_toolsets ?? [],
   },
   mcp_tool_permissions: agent.object_permission?.mcp_tool_permissions ?? {},
 });
 
 /**
- * Build the object_permission payload from the shared MCP form fields.
- * Always includes the MCP keys (empty when cleared) so removals persist;
- * the proxy merges per key, leaving non-MCP grants untouched.
+ * Always includes every MCP key (empty when cleared) so removals persist;
+ * the proxy merges object_permission per key, leaving non-MCP grants untouched.
  */
 export const buildMcpObjectPermission = (values: any) => ({
   mcp_servers: values.allowed_mcp_servers_and_groups?.servers ?? [],
   mcp_access_groups: values.allowed_mcp_servers_and_groups?.accessGroups ?? [],
+  mcp_toolsets: values.allowed_mcp_servers_and_groups?.toolsets ?? [],
   mcp_tool_permissions: values.mcp_tool_permissions ?? {},
 });
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx
index c813c513134..5bc8c2a5b1c 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx
@@ -141,7 +141,7 @@ describe("AgentInfoView update payload", () => {
       rpm_limit: 222,
       session_tpm_limit: 333,
       session_rpm_limit: 444,
-      object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_tool_permissions: {} },
+      object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} },
     });
   });
 
@@ -182,7 +182,7 @@ describe("AgentInfoView update payload", () => {
       rpm_limit: 222,
       session_tpm_limit: 333,
       session_rpm_limit: 444,
-      object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_tool_permissions: {} },
+      object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} },
     });
   });
 
@@ -260,18 +260,20 @@ describe("AgentInfoView update payload", () => {
         api_base: "https://other.example.com",
         model: "langgraph/asst_1",
       },
-      object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_tool_permissions: {} },
+      object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} },
     });
   });
 
   it("keeps the agent's existing MCP grants in the update payload", async () => {
+    const existingMcpGrants = {
+      mcp_servers: ["srv-1"],
+      mcp_access_groups: ["grp-a"],
+      mcp_toolsets: ["toolset-1"],
+      mcp_tool_permissions: { "srv-1": ["tool_x"] },
+    };
     vi.mocked(networking.getAgentInfo).mockResolvedValue({
       ...A2A_AGENT,
-      object_permission: {
-        mcp_servers: ["srv-1"],
-        mcp_access_groups: ["grp-a"],
-        mcp_tool_permissions: { "srv-1": ["tool_x"] },
-      },
+      object_permission: existingMcpGrants,
     } as never);
     const user = setup();
     renderView();
@@ -279,11 +281,7 @@ describe("AgentInfoView update payload", () => {
 
     await save(user);
 
-    expect(patchedPayload().object_permission).toEqual({
-      mcp_servers: ["srv-1"],
-      mcp_access_groups: ["grp-a"],
-      mcp_tool_permissions: { "srv-1": ["tool_x"] },
-    });
+    expect(patchedPayload().object_permission).toEqual(existingMcpGrants);
   });
 
   it("reloads the agent and leaves edit mode when the edit is cancelled", async () => {
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx
index a351ff2090f..29d97a20afe 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx
@@ -74,7 +74,8 @@ describe("AgentInfoView settings", () => {
     expect(token).toBe("sk-test");
     expect(agentId).toBe("agent-1");
     expect(payload.tpm_limit).toBe(42);
-    expect(payload.object_permission).toEqual({ mcp_servers: [], mcp_access_groups: [], mcp_tool_permissions: {} });
+    const clearedMcpGrants = { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} };
+    expect(payload.object_permission).toEqual(clearedMcpGrants);
   });
 
   it("shows MCP grants with server names on the overview tab", async () => {
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx
index 1592b452455..6cb99e9692f 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx
@@ -359,6 +359,7 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
             {agent.object_permission &&
               (agent.object_permission.mcp_servers?.length ||
                 agent.object_permission.mcp_access_groups?.length ||
+                agent.object_permission.mcp_toolsets?.length ||
                 (agent.object_permission.mcp_tool_permissions &&
                   Object.keys(agent.object_permission.mcp_tool_permissions).length > 0)) && (
                 <div style={{ marginTop: 24 }}>
@@ -379,6 +380,9 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
                           {agent.object_permission.mcp_access_groups.join(", ")}
                         </DetailItem>
                       )}
+                    {agent.object_permission.mcp_toolsets && agent.object_permission.mcp_toolsets.length > 0 && (
+                      <DetailItem label="MCP Toolsets">{agent.object_permission.mcp_toolsets.join(", ")}</DetailItem>
+                    )}
                     {agent.object_permission.mcp_tool_permissions &&
                       Object.keys(agent.object_permission.mcp_tool_permissions).length > 0 && (
                         <DetailItem label="Tool permissions per server">
@@ -501,6 +505,7 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
                                 value={{
                                   servers: (value as McpServerSelection | undefined)?.servers ?? [],
                                   accessGroups: (value as McpServerSelection | undefined)?.accessGroups ?? [],
+                                  toolsets: (value as McpServerSelection | undefined)?.toolsets ?? [],
                                 }}
                                 accessToken={accessToken ?? ""}
                                 placeholder="Select MCP servers or access groups (optional)"
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 7316b4359cc..6b8b703e5ef 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -22873,6 +22873,8 @@ export interface components {
             mcp_tool_permissions?: {
                 [key: string]: string[];
             } | null;
+            /** Mcp Toolsets */
+            mcp_toolsets?: string[] | null;
             /** Models */
             models?: string[] | null;
         };

From 3c2138f0378ea958d0657af3fb6436d2c8061fb4 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 1 Sep 2026 19:12:24 -0700
Subject: [PATCH 012/164] fix(mcp): count agent toolsets in the agent MCP
 ceiling and tool grants

An agent's object_permission.mcp_toolsets could be persisted through the
new edit form and PATCH /v1/agents but never reached the request-time
checks: _get_allowed_mcp_servers_for_agent read only mcp_servers and
mcp_access_groups, so an agent granted a toolset alone resolved to [] and
placed no ceiling on its keys, and _get_agent_tool_permissions_for_server
ignored the tools those toolsets grant. Both helpers now resolve toolsets
through the shared _toolset_tool_permissions / _toolset_tools_for_server
helpers the key, team, and org levels already use, and a declared toolset
that resolves to nothing raises UnloadableEntitlementError so the
resolver denies instead of reading the agent as unrestricted
---
 .../mcp_server/auth/user_api_key_auth_mcp.py  |  79 +++++-----
 .../auth/test_user_api_key_auth_mcp.py        | 142 ++++++++++++++++++
 2 files changed, 186 insertions(+), 35 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
index 425f82794e6..ce539258fd2 100644
--- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
+++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
@@ -3074,15 +3074,17 @@ class MCPRequestHandler:
     @staticmethod
     async def _get_allowed_mcp_servers_for_agent(
         user_api_key_auth: UserAPIKeyAuth | None = None,
-        agent_object_permission=None,
+        agent_object_permission: LiteLLM_ObjectPermissionTable | None = None,
     ) -> list[str]:
         """
         Get allowed MCP servers for an agent (from the agent's object_permission).
 
-        Returns the MCP servers from the agent's object_permission.
-        If agent has no object_permission, returns [] (no extra restriction). An entitlement the
-        agent LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here so the
-        resolver denies.
+        Returns the agent's direct servers, the servers in its access groups, and the servers reached
+        through its toolsets, exactly as the key, team, and org levels count theirs. If agent has no
+        object_permission, returns [] (no extra restriction). An entitlement the agent LINKS but that
+        cannot be read, or a declared toolset that resolves to no grants, raises
+        ``UnloadableEntitlementError`` out of here so the resolver denies instead of reading the
+        agent as unrestricted.
 
         Args:
             user_api_key_auth: User auth with agent_id
@@ -3092,31 +3094,30 @@ class MCPRequestHandler:
         if not user_api_key_auth or not user_api_key_auth.agent_id:
             return []
 
-        obj_perm = agent_object_permission
-        if obj_perm is None:
-            obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
+        obj_perm: Final = (
+            agent_object_permission
+            if agent_object_permission is not None
+            else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
+        )
         if obj_perm is None:
             return []
 
         try:
-            direct_mcp_servers = getattr(obj_perm, "mcp_servers", None) or []
-            if isinstance(direct_mcp_servers, str):
-                direct_mcp_servers = []
-            mcp_access_groups = getattr(obj_perm, "mcp_access_groups", None) or []
-            if isinstance(mcp_access_groups, str):
-                mcp_access_groups = []
-
-            # Permission entries may be server_ids OR names/aliases — expand to ids.
             from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
                 global_mcp_server_manager,
             )
 
-            expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(list(direct_mcp_servers))
-
-            access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(mcp_access_groups)
-            all_servers: Final = expanded_direct_servers + access_group_servers
-            return list(set(all_servers))
+            expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(
+                obj_perm.mcp_servers or []
+            )
+            access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(
+                obj_perm.mcp_access_groups or []
+            )
+            toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(obj_perm)
+            return list({*expanded_direct_servers, *access_group_servers, *toolset_grants})
         except Exception as e:
+            if isinstance(e, UnloadableEntitlementError):
+                raise
             verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e)
             return []
 
@@ -3124,13 +3125,15 @@ class MCPRequestHandler:
     async def _get_agent_tool_permissions_for_server(
         server_id: str,
         user_api_key_auth: UserAPIKeyAuth | None = None,
-        agent_object_permission=None,
+        agent_object_permission: LiteLLM_ObjectPermissionTable | None = None,
     ) -> list[str] | None:
         """
-        Get allowed tool names for a server from the agent's object_permission.
-        Returns None if agent has no tool restrictions for this server. An entitlement the agent
-        LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here, which the
-        tool resolver turns into deny-all for the server rather than an unrestricted tool list.
+        Get allowed tool names for a server from the agent's object_permission: the union of its
+        direct tool permissions and the tools its toolsets grant on that server, mirroring the key and
+        team levels. Returns None if agent has no tool restrictions for this server. An entitlement the
+        agent LINKS but that cannot be read, or a declared toolset that resolves to no grants, raises
+        ``UnloadableEntitlementError`` out of here, which the tool resolver turns into deny-all for the
+        server rather than an unrestricted tool list.
 
         Args:
             server_id: Server ID to check permissions for
@@ -3141,24 +3144,30 @@ class MCPRequestHandler:
         if not user_api_key_auth or not user_api_key_auth.agent_id:
             return None
 
-        obj_perm = agent_object_permission
-        if obj_perm is None:
-            obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
+        obj_perm: Final = (
+            agent_object_permission
+            if agent_object_permission is not None
+            else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
+        )
         if obj_perm is None:
             return None
 
         try:
-            mcp_tool_permissions: Final = getattr(obj_perm, "mcp_tool_permissions", None)
-            if not mcp_tool_permissions or not isinstance(mcp_tool_permissions, dict):
-                return None
-            # Dict keys may be server_ids OR names/aliases; normalize before lookup.
             from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
                 global_mcp_server_manager,
             )
 
-            tools: Final = global_mcp_server_manager.expand_tool_permissions(mcp_tool_permissions).get(server_id)
-            return list(tools) if tools else None
+            direct_tools: Final = (
+                global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions).get(server_id)
+                if obj_perm.mcp_tool_permissions
+                else None
+            )
+            toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(obj_perm, server_id)
+            agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools)
+            return list(agent_tools) if agent_tools else None
         except Exception as e:
+            if isinstance(e, UnloadableEntitlementError):
+                raise
             verbose_logger.warning("Failed to get agent tool permissions for server: %s", e)
             return None
 
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
index 0144fbb17dd..7caf31ef6f1 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
@@ -11,6 +11,7 @@ from starlette.datastructures import Headers
 
 from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
     MCPRequestHandler,
+    UnloadableEntitlementError,
     _is_mcp_admitted_user_subject,
 )
 from litellm.proxy._types import (
@@ -4358,6 +4359,147 @@ class TestAgentMCPPermissions:
                     )
                     assert sorted(result) == ["tool_a", "tool_b"]
 
+    def _agent_object_permission(self, *, toolset_ids, servers=(), tool_permissions=None):
+        agent_object_permission = MagicMock()
+        agent_object_permission.mcp_servers = list(servers)
+        agent_object_permission.mcp_access_groups = []
+        agent_object_permission.mcp_tool_permissions = tool_permissions
+        agent_object_permission.mcp_toolsets = list(toolset_ids)
+        return agent_object_permission
+
+    def _mock_manager_with_toolsets(self, toolset_perms):
+        mock_manager = MagicMock()
+        mock_manager.expand_permission_list = MagicMock(side_effect=lambda servers: list(servers))
+        mock_manager.expand_tool_permissions = MagicMock(side_effect=lambda perms: perms or {})
+        mock_manager.resolve_toolset_tool_permissions = AsyncMock(return_value=toolset_perms)
+        return mock_manager
+
+    def _agent_toolset_patches(self, agent_object_permission, mock_manager):
+        return (
+            patch.object(  # test-quality-ok: stub the agent perm loader; the resolver reads module globals with no injection seam
+                MCPRequestHandler, "_get_agent_object_permission", AsyncMock(return_value=agent_object_permission)
+            ),
+            patch(  # test-quality-ok: isolate the MCP registry, same seam as the sibling toolset tests
+                "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
+                mock_manager,
+            ),
+            patch.object(  # test-quality-ok: access-group lookup hits the DB, not under test here
+                MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])
+            ),
+        )
+
+    async def test_get_allowed_mcp_servers_for_agent_includes_toolset_servers(self):
+        """An agent granted only mcp_toolsets reaches the toolset's servers, exactly as a
+        key, team, or org granted only toolsets does"""
+        user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets")
+        agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"], servers=["server-direct"])
+        mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]})
+
+        with contextlib.ExitStack() as stack:
+            for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager):
+                stack.enter_context(patcher)
+            result = await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth)
+
+        assert sorted(result) == ["server-a", "server-direct"]
+        mock_manager.resolve_toolset_tool_permissions.assert_awaited_once_with(toolset_ids=["toolset-1"])
+
+    async def test_get_allowed_mcp_servers_toolset_only_agent_caps_key_servers(self):
+        """Regression: an agent whose only grant is a toolset used to resolve to [] and place
+        no ceiling at all, so a key bound to it kept every server the key itself granted"""
+        user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets")
+        agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"])
+        mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]})
+
+        with contextlib.ExitStack() as stack:
+            for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager):
+                stack.enter_context(patcher)
+            stack.enter_context(
+                patch.object(  # test-quality-ok: key resolution has its own tests; pin its grants here
+                    MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"])
+                )
+            )
+            stack.enter_context(
+                patch.object(  # test-quality-ok: team resolution has its own tests; pin it empty here
+                    MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[])
+                )
+            )
+            result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)
+
+        assert result == ["server-a"]
+
+    async def test_get_allowed_mcp_servers_agent_dangling_toolset_denies(self):
+        """An agent toolset that resolves to nothing is a known restriction with unknown
+        contents: deny, never fall through to the key's own servers"""
+        user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets")
+        agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-gone"])
+        mock_manager = self._mock_manager_with_toolsets({})
+
+        with contextlib.ExitStack() as stack:
+            for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager):
+                stack.enter_context(patcher)
+            with pytest.raises(UnloadableEntitlementError):
+                await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth)
+            stack.enter_context(
+                patch.object(  # test-quality-ok: key resolution has its own tests; pin its grants here
+                    MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"])
+                )
+            )
+            stack.enter_context(
+                patch.object(  # test-quality-ok: team resolution has its own tests; pin it empty here
+                    MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[])
+                )
+            )
+            result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)
+
+        assert result == []
+
+    async def test_get_agent_tool_permissions_for_server_unions_direct_and_toolset_tools(self):
+        """The agent's tool ceiling on a server is its direct tool grants plus the tools its
+        toolsets grant there, and None only when neither names the server"""
+        user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets")
+        agent_object_permission = self._agent_object_permission(
+            toolset_ids=["toolset-1"], tool_permissions={"server-a": ["tool_direct"]}
+        )
+        mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_via_toolset"], "server-b": ["tool_b"]})
+
+        with contextlib.ExitStack() as stack:
+            for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager):
+                stack.enter_context(patcher)
+            server_a_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-a", user_api_key_auth)
+            server_b_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-b", user_api_key_auth)
+            server_c_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-c", user_api_key_auth)
+
+        assert sorted(server_a_tools) == ["tool_direct", "tool_via_toolset"]
+        assert server_b_tools == ["tool_b"]
+        assert server_c_tools is None
+
+    async def test_get_allowed_tools_for_server_toolset_only_agent_caps_key_tools(self):
+        """Regression: a key allowing [tool_a, tool_b] bound to an agent whose toolset grants
+        only tool_a on the server ends with [tool_a]; the toolset used to be ignored"""
+        user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets")
+        agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"])
+        mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_a"]})
+        key_perm = MagicMock()
+        key_perm.mcp_tool_permissions = {"server-a": ["tool_a", "tool_b"]}
+        key_perm.mcp_toolsets = []
+
+        with contextlib.ExitStack() as stack:
+            for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager):
+                stack.enter_context(patcher)
+            stack.enter_context(
+                patch.object(  # test-quality-ok: stub the key perm loader; the resolver reads module globals with no injection seam
+                    MCPRequestHandler, "_get_key_object_permission", return_value=key_perm
+                )
+            )
+            stack.enter_context(
+                patch.object(  # test-quality-ok: team resolution has its own tests; pin it absent here
+                    MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)
+                )
+            )
+            result = await MCPRequestHandler.get_allowed_tools_for_server("server-a", user_api_key_auth)
+
+        assert result == ["tool_a"]
+
     async def test_get_agent_object_permission_uses_shared_helper(self):
         """``_get_agent_object_permission`` must resolve the agent's
         ``object_permission_id`` and then defer to the shared

From 222cb293ca61a57d6771bbaca43307f2e15f11fd Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 1 Sep 2026 22:14:18 -0700
Subject: [PATCH 013/164] fix(mcp): name the access group that resolved when an
 agent vetoes a mixed scope

---
 .../proxy/_experimental/mcp_server/server.py  | 12 ++++---
 .../mcp_server/test_mcp_server.py             | 31 +++++++++++++++++--
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 9f530ddf862..b3b574673db 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -1489,10 +1489,14 @@ if MCP_AVAILABLE:
                     )
                 }
                 raise HTTPException(status_code=403, detail=agent_denial)
-            vetoed_group: Final = (
-                next((name for name in requested_names if _registered_server_id(name) is None), None)
-                if resolved_ids
-                else None
+            vetoed_group: Final = next(
+                (
+                    name
+                    for name in requested_names
+                    if _registered_server_id(name) is None
+                    and any(name in (server.access_groups or ()) for server in resolved_without_agent)
+                ),
+                None,
             )
             if vetoed_group is not None:
                 group_denial: Final[_McpDeniedDetail] = {
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
index fb69ad3b1c4..ea64fdf2a5d 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
@@ -1336,11 +1336,13 @@ def _denied_scope_manager(known_server_names_to_ids: dict[str, str]) -> MagicMoc
     return manager
 
 
-def _scope_resolver(resolved_without_agent: list[str]) -> AsyncMock:
+def _scope_resolver(resolved_without_agent: list[str], access_groups: tuple[str, ...] = ()) -> AsyncMock:
     async def resolve(user_api_key_auth, mcp_servers, client_ip=None):
         if user_api_key_auth is not None and user_api_key_auth.agent_id:
             return []
-        return [MagicMock(server_id=server_id) for server_id in resolved_without_agent]
+        return [
+            MagicMock(server_id=server_id, access_groups=list(access_groups)) for server_id in resolved_without_agent
+        ]
 
     return AsyncMock(side_effect=resolve)
 
@@ -1450,7 +1452,7 @@ async def test_scoped_list_access_group_vetoed_by_agent_names_agent_and_group():
         user_api_key_auth,
         ["prod-group"],
         _denied_scope_manager({}),
-        _scope_resolver(resolved_without_agent=["srv-github"]),
+        _scope_resolver(resolved_without_agent=["srv-github"], access_groups=("prod-group",)),
     )
 
     assert denial.status_code == 403
@@ -1460,6 +1462,29 @@ async def test_scoped_list_access_group_vetoed_by_agent_names_agent_and_group():
     assert "mcp_access_groups" in message
 
 
+@pytest.mark.asyncio
+async def test_scoped_list_mixed_unknown_and_vetoed_group_names_the_group_that_resolved():
+    """With an unknown name ahead of the agent-vetoed group in the scope, the 403 must name the group
+    whose servers the key can reach, never the unknown name, or the admin is told to grant a group
+    that does not exist."""
+    pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
+
+    user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
+
+    denial = await _denied_scoped_list(
+        user_api_key_auth,
+        ["no-such-group", "prod-group"],
+        _denied_scope_manager({}),
+        _scope_resolver(resolved_without_agent=["srv-github"], access_groups=("prod-group",)),
+    )
+
+    assert denial.status_code == 403
+    message = denial.detail["error"]
+    assert "access group 'prod-group'" in message
+    assert "no-such-group" not in message
+    assert "agent 'agent-123'" in message
+
+
 @pytest.mark.asyncio
 async def test_scoped_list_agent_key_denied_by_key_grants_raises_generic_403():
     """When the agent-stripped rerun still resolves nothing, the 403 stays generic instead of

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 014/164] 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 5d019fe3925d9d961dff98033f6f4f819b8be172 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 2 Sep 2026 14:27:48 -0700
Subject: [PATCH 015/164] fix(mcp): only claim /mcp/<name> tool URLs the
 gateway serves and guard the provider fallbacks

---
 .../messages/mcp_handler.py                   |   2 +-
 .../proxy/hooks/mcp_semantic_filter/hook.py   |   2 +-
 litellm/responses/main.py                     |   8 +-
 .../responses/mcp/chat_completions_handler.py |   3 +-
 .../mcp/litellm_proxy_mcp_handler.py          | 135 ++++++++++-------
 .../mcp/test_chat_completions_handler.py      |  43 ++++++
 .../mcp/test_litellm_proxy_mcp_handler.py     | 142 ++++++++++++++++++
 7 files changed, 276 insertions(+), 59 deletions(-)

diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py
index c1f10c245f8..d9cc65e730f 100644
--- a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py
+++ b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py
@@ -82,7 +82,7 @@ async def anthropic_messages_with_mcp(
         LiteLLM_Proxy_MCP_Handler,
     )
 
-    mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
+    mcp_references, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
 
     if not mcp_references:
         return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py
index 8b82842353c..5c55b35c968 100644
--- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py
+++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py
@@ -104,7 +104,7 @@ class SemanticToolFilterHook(CustomLogger):
         )
 
         # Parse to separate MCP tools from other tools
-        mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
+        mcp_tools, _ = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
 
         if not mcp_tools:
             return []
diff --git a/litellm/responses/main.py b/litellm/responses/main.py
index f012ec8f07b..bae7c491926 100644
--- a/litellm/responses/main.py
+++ b/litellm/responses/main.py
@@ -175,7 +175,7 @@ async def aresponses_api_with_mcp(
     (
         mcp_tools_with_litellm_proxy,
         other_tools,
-    ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
+    ) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
 
     # Process MCP tools through the complete pipeline (fetch + filter + deduplicate + transform)
     # Extract user_api_key_auth from litellm_metadata (where it's added by add_user_api_key_auth_to_request_metadata)
@@ -234,6 +234,7 @@ async def aresponses_api_with_mcp(
         "timeout": timeout,
         "custom_llm_provider": custom_llm_provider,
         **kwargs,
+        "_skip_mcp_handler": True,
     }
 
     # Handle MCP streaming if requested
@@ -839,13 +840,14 @@ def _responses_try_dispatch_mcp_gateway(
     custom_llm_provider: str | None,
     kwargs: dict[str, object],
     _is_async: bool,
+    skip_mcp_handler: bool,
 ) -> Any | None:
     """Return a response when MCP gateway handles the call; otherwise None."""
     from litellm.responses.mcp.litellm_proxy_mcp_handler import (
         LiteLLM_Proxy_MCP_Handler,
     )
 
-    if not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools):
+    if skip_mcp_handler or not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools):
         return None
     mcp_call_kwargs: Final = {
         "input": input,
@@ -1015,6 +1017,7 @@ def responses(
         litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj")
         litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None)
         _is_async: Final = kwargs.pop("aresponses", False) is True
+        skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False)
         use_chat_completions_api = _pop_use_chat_completions_api_kw(kwargs)
 
         client_headers: Final = kwargs.get("headers")
@@ -1109,6 +1112,7 @@ def responses(
             custom_llm_provider=custom_llm_provider,
             kwargs=kwargs,
             _is_async=_is_async,
+            skip_mcp_handler=skip_mcp_handler,
         )
         if _mcp_dispatch is not None:
             return _mcp_dispatch
diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py
index a75b3768636..ae18d5f6f1b 100644
--- a/litellm/responses/mcp/chat_completions_handler.py
+++ b/litellm/responses/mcp/chat_completions_handler.py
@@ -106,7 +106,7 @@ async def acompletion_with_mcp(
     (
         mcp_tools_with_litellm_proxy,
         other_tools,
-    ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
+    ) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
 
     if not mcp_tools_with_litellm_proxy:
         # No MCP tools, proceed with regular completion
@@ -114,6 +114,7 @@ async def acompletion_with_mcp(
             model=model,
             messages=messages,
             tools=tools,
+            _skip_mcp_handler=True,
             **kwargs,
         )
 
diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py
index 367915156d1..11e87a72813 100644
--- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py
+++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py
@@ -1,6 +1,6 @@
 import re
 import traceback
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload
 
@@ -23,6 +23,7 @@ from litellm.types.llms.openai import (
     ResponsesAPIStreamingResponse,
 )
 from litellm.types.llms.openai import ToolParam as ResponsesToolParam
+from litellm.types.mcp_server.mcp_server_manager import MCPServer
 from litellm.types.utils import (
     CallTypes,
     ChatCompletionMessageCustomToolCall,
@@ -45,6 +46,7 @@ else:
 
 # NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling
 ToolParam: TypeAlias = Mapping[str, object]
+SplitTools: TypeAlias = tuple[list[ToolParam], list[Any]]
 
 
 class MCPToolResult(TypedDict):
@@ -56,14 +58,65 @@ class MCPToolResult(TypedDict):
 LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy"
 LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/"
 
-# Matches any URL whose path ends with /mcp/<server_name> — covers both root-path
-# (http://host:port/mcp/name) and sub-path (http://host/base/mcp/name) proxy deployments.
-# A false-positive match (e.g. an external URL that happens to end with /mcp/<name>) results
-# in a "server not found" error from the internal gateway, not a silent failure or data leak,
-# so this broad pattern is intentional and preferred over anchoring to localhost only.
 _PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$")
 
 
+def _mcp_server_url(tool: ToolParam) -> str | None:
+    if not isinstance(tool, dict) or tool.get("type") != "mcp":
+        return None
+    server_url: Final = tool.get("server_url")
+    return server_url if isinstance(server_url, str) else None
+
+
+def _names_gateway_explicitly(tool: ToolParam) -> bool:
+    return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL)
+
+
+def _proxy_path_mcp_name(tool: ToolParam) -> str | None:
+    server_url: Final = _mcp_server_url(tool)
+    match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url)
+    return None if match is None else match.group(1)
+
+
+def _registered_mcp_servers() -> Collection[MCPServer]:
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+        global_mcp_server_manager,
+    )
+
+    return global_mcp_server_manager.get_registry().values()
+
+
+def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool:
+    return any(
+        name in (server.alias, server.server_name, server.name) or name in (server.access_groups or ())
+        for server in servers
+    )
+
+
+async def _toolset_exists(name: str) -> bool:
+    try:
+        from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+            global_mcp_server_manager,
+        )
+        from litellm.proxy.proxy_server import prisma_client
+
+        if prisma_client is None:
+            return False
+        return await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is not None
+    except Exception as e:
+        verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, e)
+        return False
+
+
+async def _gateway_served_names(
+    names: Collection[str],
+    servers: Callable[[], Collection[MCPServer]] = _registered_mcp_servers,
+    toolset_exists: Callable[[str], Awaitable[bool]] = _toolset_exists,
+) -> frozenset[str]:
+    registered: Final = tuple(servers()) if names else ()
+    return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)])
+
+
 class LiteLLM_Proxy_MCP_Handler:
     """
     Helper class with static methods for MCP integration with Responses API.
@@ -87,57 +140,31 @@ class LiteLLM_Proxy_MCP_Handler:
 
     @staticmethod
     def _should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool:
-        """
-        Returns True if any MCP tool should be handled via the litellm proxy MCP gateway.
-        This includes tools with server_url="litellm_proxy" as well as URLs ending in /mcp/<name>.
-        """
-        if tools:
-            for tool in tools:
-                if isinstance(tool, dict) and tool.get("type") == "mcp":
-                    server_url = tool.get("server_url", "")
-                    if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL):
-                        return True
-                    if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match(server_url):
-                        return True
-        return False
+        """True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in
+        /mcp/<name>. `_split_mcp_tools` then settles which of the latter the gateway actually serves."""
+        return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ())
 
     @staticmethod
-    def _parse_mcp_tools(
+    def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools:
+        items: Final = tuple(tools or ())
+        gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)]
+        other_tools: Final[list[Any]] = [tool for tool in items if not _names_gateway_explicitly(tool)]
+        return gateway_tools, other_tools
+
+    @staticmethod
+    async def _split_mcp_tools(
         tools: Iterable[Mapping[str, object]] | None,
-    ) -> tuple[list[ToolParam], list[Any]]:
-        """
-        Parse tools and separate MCP tools with litellm_proxy from other tools.
-
-        Returns:
-            Tuple of (mcp_tools_with_litellm_proxy, other_tools)
-        """
-        mcp_tools_with_litellm_proxy: Final[list[ToolParam]] = []
-        other_tools: Final[list[Any]] = []
-
-        if tools:
-            for tool in tools:
-                if isinstance(tool, dict) and tool.get("type") == "mcp":
-                    server_url = tool.get("server_url", "")
-                    if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL):
-                        mcp_tools_with_litellm_proxy.append(tool)
-                    elif isinstance(server_url, str):
-                        # Also intercept URLs like http://localhost:4000/mcp/atlassian_test
-                        # by rewriting them to the internal litellm_proxy format.
-                        m = _PROXY_MCP_PATH_RE.match(server_url)
-                        if m:
-                            rewritten = {
-                                **tool,
-                                "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{m.group(1)}",
-                            }
-                            mcp_tools_with_litellm_proxy.append(rewritten)
-                        else:
-                            other_tools.append(tool)
-                    else:
-                        other_tools.append(tool)
-                else:
-                    other_tools.append(tool)
-
-        return mcp_tools_with_litellm_proxy, other_tools
+        served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names,
+    ) -> SplitTools:
+        resolved: Final = tuple((tool, _proxy_path_mcp_name(tool)) for tool in tools or ())
+        names: Final = frozenset(name for _, name in resolved if name is not None)
+        served: Final = await served_names(names) if names else frozenset[str]()
+        return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(
+            [
+                {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} if name in served else tool
+                for tool, name in resolved
+            ]
+        )
 
     @staticmethod
     async def _apply_toolset_permissions(
diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py
index ab4c5185057..2c1845f7b92 100644
--- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py
+++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py
@@ -1,6 +1,13 @@
+import json
+import sys
+import types
+
 import pytest
+import respx
+from httpx import Response
 from unittest.mock import AsyncMock, patch
 
+import litellm
 from litellm.types.utils import ModelResponse
 
 from litellm.responses.mcp import chat_completions_handler
@@ -1344,3 +1351,39 @@ async def test_acompletion_with_mcp_streaming_drains_inner_stream_after_exhausti
 
     assert len(all_chunks) == 3
     assert initial_stream.drained_after_exhaustion is True
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_acompletion_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch):
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
+
+    zapier_tool = {"type": "mcp", "server_label": "zapier", "server_url": "https://mcp.zapier.com/api/mcp/mcp"}
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None))
+    monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {})
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    provider = respx.post("https://api.openai.com/v1/chat/completions").mock(
+        return_value=Response(
+            200,
+            json={
+                "id": "chatcmpl-zapier",
+                "object": "chat.completion",
+                "created": 0,
+                "model": "gpt-4.1",
+                "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
+                "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+            },
+        )
+    )
+
+    result = await acompletion_with_mcp(
+        model="openai/gpt-4.1",
+        messages=[{"role": "user", "content": "hello"}],
+        tools=[zapier_tool],
+        api_key="sk-test",
+        acompletion=True,
+    )
+
+    assert isinstance(result, ModelResponse)
+    assert result.id == "chatcmpl-zapier"
+    assert json.loads(provider.calls.last.request.content)["tools"] == [zapier_tool]
diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
index b33ed3bb581..73d245df651 100644
--- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
+++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
@@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 from fastapi import HTTPException
+from openai.types.responses.tool_param import Mcp
 import importlib
 
 from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing
@@ -719,3 +720,144 @@ def test_extract_tool_call_details_still_prefers_openai_arguments():
     assert name == "get_weather"
     assert call_id == "call_123"
     assert arguments == '{"city": "Paris"}'
+
+
+def _registered(
+    server_id: str,
+    name: str,
+    alias: str | None = None,
+    server_name: str | None = None,
+    access_groups: list[str] | None = None,
+):
+    from litellm.types.mcp import MCPTransport
+    from litellm.types.mcp_server.mcp_server_manager import MCPServer
+
+    return MCPServer(
+        server_id=server_id,
+        name=name,
+        alias=alias,
+        server_name=server_name,
+        transport=MCPTransport.http,
+        access_groups=access_groups,
+    )
+
+
+async def _no_toolset(_: str) -> bool:
+    return False
+
+
+ZAPIER_TOOL: Mcp = {
+    "type": "mcp",
+    "server_label": "zapier",
+    "server_url": "https://mcp.zapier.com/api/mcp/mcp",
+    "require_approval": "never",
+}
+EXPLICIT_GATEWAY_TOOL = {"type": "mcp", "server_label": "github", "server_url": "litellm_proxy/mcp/github"}
+FUNCTION_TOOL = {"type": "function", "name": "get_weather", "parameters": {}}
+
+
+@pytest.mark.asyncio
+async def test_gateway_served_names_matches_alias_server_name_name_access_group_and_toolset():
+    from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names
+
+    servers = (
+        _registered("id-1", "github-name", alias="github", server_name="github-server", access_groups=["prod-group"]),
+        _registered("id-2", "deepwiki"),
+    )
+
+    async def toolset_exists(name: str) -> bool:
+        return name == "my-toolset"
+
+    served = await _gateway_served_names(
+        {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset", "mcp", "nope"},
+        servers=lambda: servers,
+        toolset_exists=toolset_exists,
+    )
+
+    assert served == {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset"}
+
+
+@pytest.mark.asyncio
+async def test_split_mcp_tools_leaves_external_mcp_path_urls_for_the_provider():
+
+    async def served_names(names):
+        assert names == {"mcp"}
+        return frozenset()
+
+    gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(
+        [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names
+    )
+
+    assert gateway_tools == [EXPLICIT_GATEWAY_TOOL]
+    assert other_tools == [ZAPIER_TOOL, FUNCTION_TOOL]
+
+
+@pytest.mark.asyncio
+async def test_split_mcp_tools_repoints_served_proxy_urls_at_the_gateway():
+    served_tool = {
+        "type": "mcp",
+        "server_label": "toolset",
+        "server_url": "http://localhost:4000/mcp/my-toolset",
+        "require_approval": "never",
+        "allowed_tools": ["get_me"],
+    }
+    unserved_tool = {"type": "mcp", "server_label": "typo", "server_url": "http://localhost:4000/mcp/githb"}
+
+    async def served_names(names):
+        return frozenset({"my-toolset"})
+
+    gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(
+        [served_tool, unserved_tool], served_names=served_names
+    )
+
+    assert gateway_tools == [{**served_tool, "server_url": "litellm_proxy/mcp/my-toolset"}]
+    assert other_tools == [unserved_tool]
+
+
+@pytest.mark.asyncio
+async def test_split_mcp_tools_skips_resolution_when_nothing_points_at_the_proxy():
+    async def served_names(names):
+        raise AssertionError("no lookup expected")
+
+    gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(
+        [EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names
+    )
+
+    assert gateway_tools == [EXPLICIT_GATEWAY_TOOL]
+    assert other_tools == [FUNCTION_TOOL]
+
+
+def test_should_use_gateway_still_triggers_on_http_mcp_path():
+    assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([ZAPIER_TOOL]) is True
+    assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([EXPLICIT_GATEWAY_TOOL]) is True
+    assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([FUNCTION_TOOL]) is False
+    assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(None) is False
+
+
+@pytest.mark.asyncio
+async def test_aresponses_api_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch):
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
+    from litellm.responses import main as responses_main
+    from litellm.types.llms.openai import ResponsesAPIResponse
+
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None))
+    monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {})
+    provider_tools: list[object] = []
+
+    def fake_provider(**kwargs: object) -> object:
+        request_params = cast(dict[str, object], kwargs["response_api_optional_request_params"])
+        provider_tools.append(request_params.get("tools"))
+
+        async def respond() -> ResponsesAPIResponse:
+            return ResponsesAPIResponse(id="resp_zapier", created_at=0, output=[])
+
+        return respond()
+
+    monkeypatch.setattr(responses_main.base_llm_http_handler, "response_api_handler", fake_provider)
+
+    response = await responses_main.aresponses_api_with_mcp(
+        input="Reply with the single word ok.", model="openai/gpt-4.1", tools=[ZAPIER_TOOL]
+    )
+
+    assert isinstance(response, ResponsesAPIResponse)
+    assert provider_tools == [[ZAPIER_TOOL]]

From 2a1527c22fc2c4deafb324b925d2c5e453e92580 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 2 Sep 2026 16:03:31 -0700
Subject: [PATCH 016/164] fix(mcp): match gateway-served /mcp/<name> URLs the
 way the gateway resolves them and narrow only gateway references

---
 .../proxy/_experimental/mcp_server/server.py  |  2 +-
 .../proxy/hooks/mcp_semantic_filter/hook.py   | 23 +++++++------
 .../mcp/litellm_proxy_mcp_handler.py          | 32 +++++++++++++++----
 .../mcp_server/test_mcp_server.py             | 27 ++++++++++++++++
 .../mcp_server/test_semantic_tool_filter.py   | 23 +++++++++++++
 .../mcp/test_litellm_proxy_mcp_handler.py     | 31 ++++++++++++++++++
 6 files changed, 122 insertions(+), 16 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index b3b574673db..1100b986f15 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -2037,7 +2037,7 @@ if MCP_AVAILABLE:
                 mcp_servers=mcp_servers,
                 client_ip=client_ip,
             )
-            if mcp_servers is not None and not allowed_mcp_servers:
+            if mcp_servers and not allowed_mcp_servers:
                 await _raise_denied_scoped_mcp_access(
                     requested_names=mcp_servers,
                     user_api_key_auth=user_api_key_auth,
diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py
index 5c55b35c968..afc430173ad 100644
--- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py
+++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py
@@ -5,7 +5,7 @@ Pre-call hook that filters MCP tools semantically before LLM inference.
 Reduces context window size and improves tool selection accuracy.
 """
 
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING, Final, Optional
 
 from fastapi import HTTPException
@@ -173,7 +173,11 @@ class SemanticToolFilterHook(CustomLogger):
         return [name for name in names if name]
 
     @staticmethod
-    def _narrow_mcp_references(tools: Sequence[Mapping[str, object]], selected_tool_names: list[str]) -> list[object]:
+    async def _narrow_mcp_references(
+        tools: Sequence[Mapping[str, object]],
+        selected_tool_names: list[str],
+        served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] | None = None,
+    ) -> list[object]:
         """
         Restrict each litellm_proxy MCP reference to the semantically selected tools.
 
@@ -192,13 +196,14 @@ class SemanticToolFilterHook(CustomLogger):
             LiteLLM_Proxy_MCP_Handler,
         )
 
+        via_gateway: Final = await (
+            LiteLLM_Proxy_MCP_Handler._routes_through_gateway(tools, served_names)
+            if served_names is not None
+            else LiteLLM_Proxy_MCP_Handler._routes_through_gateway(tools)
+        )
         return [
-            (
-                {**tool, "allowed_tools": selected_tool_names}
-                if isinstance(tool, dict) and LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([tool])
-                else tool
-            )
-            for tool in tools
+            {**tool, "allowed_tools": selected_tool_names} if isinstance(tool, dict) and routed else tool
+            for tool, routed in zip(tools, via_gateway, strict=True)
         ]
 
     def _is_mcp_tool(self, tool: object) -> bool:
@@ -325,7 +330,7 @@ class SemanticToolFilterHook(CustomLogger):
                 filtered_expanded_tools = await self._filter_expanded_tools(data=data, expanded_tools=expanded_tools)
 
                 selected_tool_names: Final = self._selected_tool_names(filtered_expanded_tools)
-                narrowed_tools: Final = self._narrow_mcp_references(tools, selected_tool_names)
+                narrowed_tools: Final = await self._narrow_mcp_references(tools, selected_tool_names)
                 data["tools"] = narrowed_tools
                 self._emit_filter_metadata_safe(
                     data=data,
diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py
index 11e87a72813..5db6bac10a1 100644
--- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py
+++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py
@@ -11,6 +11,7 @@ from litellm._logging import verbose_logger
 from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.proxy._experimental.mcp_server.utils import (
+    iter_known_server_prefixes,
     logging_safe_mcp_headers,
     split_server_prefix_from_name,
     strip_known_server_prefix,
@@ -87,8 +88,10 @@ def _registered_mcp_servers() -> Collection[MCPServer]:
 
 
 def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool:
+    requested: Final = name.lower()
     return any(
-        name in (server.alias, server.server_name, server.name) or name in (server.access_groups or ())
+        requested in (known.lower() for known in (*iter_known_server_prefixes(server), server.name))
+        or name in (server.access_groups or ())
         for server in servers
     )
 
@@ -117,6 +120,13 @@ async def _gateway_served_names(
     return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)])
 
 
+async def _served_mcp_path_names(
+    tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]]
+) -> frozenset[str]:
+    names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None)
+    return await served_names(names) if names else frozenset[str]()
+
+
 class LiteLLM_Proxy_MCP_Handler:
     """
     Helper class with static methods for MCP integration with Responses API.
@@ -156,16 +166,26 @@ class LiteLLM_Proxy_MCP_Handler:
         tools: Iterable[Mapping[str, object]] | None,
         served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names,
     ) -> SplitTools:
-        resolved: Final = tuple((tool, _proxy_path_mcp_name(tool)) for tool in tools or ())
-        names: Final = frozenset(name for _, name in resolved if name is not None)
-        served: Final = await served_names(names) if names else frozenset[str]()
+        items: Final = tuple(tools or ())
+        served: Final = await _served_mcp_path_names(items, served_names)
         return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(
             [
-                {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} if name in served else tool
-                for tool, name in resolved
+                {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"}
+                if (name := _proxy_path_mcp_name(tool)) in served
+                else tool
+                for tool in items
             ]
         )
 
+    @staticmethod
+    async def _routes_through_gateway(
+        tools: Iterable[Mapping[str, object]] | None,
+        served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names,
+    ) -> tuple[bool, ...]:
+        items: Final = tuple(tools or ())
+        served: Final = await _served_mcp_path_names(items, served_names)
+        return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items)
+
     @staticmethod
     async def _apply_toolset_permissions(
         resolved_toolset_ids: list[str],
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
index ea64fdf2a5d..52b8fa3d849 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
@@ -1397,6 +1397,33 @@ async def test_scoped_list_denied_by_agent_binding_raises_403_naming_agent():
     assert resolver.await_args_list[1].kwargs["mcp_servers"] == ["github"]
 
 
+@pytest.mark.asyncio
+async def test_empty_scope_lists_nothing_instead_of_raising_a_nameless_denial():
+    """An empty ``x-mcp-servers`` header scopes to no servers; that is an empty listing, not a 403."""
+    pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
+    from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
+
+    user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
+    resolver = AsyncMock(return_value=[])
+
+    with (
+        patch(  # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
+            resolver,
+        ),
+        patch(  # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
+            _denied_scope_manager({"github": "srv-github"}),
+        ),
+    ):
+        listing = await _get_tools_from_mcp_servers(
+            user_api_key_auth=user_api_key_auth, mcp_auth_header=None, mcp_servers=[]
+        )
+
+    assert listing.tools == []
+    resolver.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_scoped_list_denied_for_non_agent_key_raises_generic_403():
     """A denial for a key with no agent binding stays generic and skips the agent-stripped rerun."""
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py
index 054146d474d..7ea90dda971 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py
@@ -2190,3 +2190,26 @@ async def test_top_k_above_router_default_is_respected():
 
     assert len(filtered) == 6
     print("✅ Configured top_k above the semantic-router default of 5 is honored")
+
+
+@pytest.mark.asyncio
+async def test_semantic_filter_hook_narrows_only_references_the_gateway_serves():
+    from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook
+
+    gateway_reference = {"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"}
+    external_tool = {
+        "type": "mcp",
+        "server_label": "zapier",
+        "server_url": "https://mcp.zapier.com/api/mcp/mcp",
+        "allowed_tools": ["zapier_send_email"],
+    }
+
+    async def served_names(names):
+        assert names == {"mcp"}
+        return frozenset()
+
+    narrowed = await SemanticToolFilterHook._narrow_mcp_references(
+        [gateway_reference, external_tool], ["srv-tool_1"], served_names=served_names
+    )
+
+    assert narrowed == [{**gateway_reference, "allowed_tools": ["srv-tool_1"]}, external_tool]
diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
index 73d245df651..4efe4d9bfd9 100644
--- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
+++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
@@ -777,6 +777,37 @@ async def test_gateway_served_names_matches_alias_server_name_name_access_group_
     assert served == {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset"}
 
 
+@pytest.mark.asyncio
+async def test_gateway_served_names_matches_server_id_short_prefix_and_alias_case_like_the_gateway():
+    from litellm.proxy._experimental.mcp_server.utils import compute_short_server_prefix
+    from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names
+
+    server_id = "0b9ae4ca-1bd2-4faa-b183-7dd812597e3b"
+    short_prefix = compute_short_server_prefix(server_id)
+    servers = (_registered(server_id, "github-name", alias="github", access_groups=["prod-group"]),)
+
+    served = await _gateway_served_names(
+        {server_id, short_prefix, "GitHub", "PROD-GROUP", "nope"}, servers=lambda: servers, toolset_exists=_no_toolset
+    )
+
+    assert served == {server_id, short_prefix, "GitHub"}
+
+
+@pytest.mark.asyncio
+async def test_routes_through_gateway_flags_explicit_and_served_tools_only():
+    served_tool = {"type": "mcp", "server_label": "github", "server_url": "http://localhost:4000/mcp/github"}
+
+    async def served_names(names):
+        assert names == {"github", "mcp"}
+        return frozenset({"github"})
+
+    flags = await LiteLLM_Proxy_MCP_Handler._routes_through_gateway(
+        [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, served_tool, FUNCTION_TOOL], served_names=served_names
+    )
+
+    assert flags == (False, True, True, False)
+
+
 @pytest.mark.asyncio
 async def test_split_mcp_tools_leaves_external_mcp_path_urls_for_the_provider():
 

From 0d6948241e97fbd9e30743174d5f6f6799f7897c Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 2 Sep 2026 17:35:51 -0700
Subject: [PATCH 017/164] fix(mcp): attribute the agent veto for differently
 cased server names

---
 .../proxy/_experimental/mcp_server/server.py  | 27 +++++------
 .../mcp_server/test_mcp_server.py             | 45 +++++++++++++++----
 2 files changed, 49 insertions(+), 23 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 1100b986f15..2f6a05481e3 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -1414,13 +1414,10 @@ if MCP_AVAILABLE:
                 server_name_matched = False
 
                 for server in allowed_mcp_servers:
-                    if server:
-                        match_list = [s.lower() for s in iter_known_server_prefixes(server) if s]
-
-                        if server_or_group.lower() in match_list:
-                            filtered_server[server.server_id] = server
-                            server_name_matched = True
-                            break
+                    if server and _server_answers_to(server, server_or_group):
+                        filtered_server[server.server_id] = server
+                        server_name_matched = True
+                        break
 
                 if not server_name_matched:
                     try:
@@ -1450,6 +1447,10 @@ if MCP_AVAILABLE:
 
         return allowed_mcp_servers
 
+    def _server_answers_to(server: MCPServer, name: str) -> bool:
+        requested: Final = name.lower()
+        return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known)
+
     class _McpDeniedDetail(TypedDict):
         error: ReadOnly[str]
 
@@ -1470,15 +1471,11 @@ if MCP_AVAILABLE:
                 mcp_servers=requested_names,
                 client_ip=client_ip,
             )
-            resolved_ids: Final = frozenset(server.server_id for server in resolved_without_agent)
 
-            def _registered_server_id(name: str) -> str | None:
-                server: Final = global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip)
-                return server.server_id if server is not None else None
+            def _resolved_to_server(name: str) -> bool:
+                return any(_server_answers_to(server, name) for server in resolved_without_agent)
 
-            vetoed_server: Final = next(
-                (name for name in requested_names if _registered_server_id(name) in resolved_ids), None
-            )
+            vetoed_server: Final = next((name for name in requested_names if _resolved_to_server(name)), None)
             if vetoed_server is not None:
                 agent_denial: Final[_McpDeniedDetail] = {
                     "error": (
@@ -1493,7 +1490,7 @@ if MCP_AVAILABLE:
                 (
                     name
                     for name in requested_names
-                    if _registered_server_id(name) is None
+                    if not _resolved_to_server(name)
                     and any(name in (server.access_groups or ()) for server in resolved_without_agent)
                 ),
                 None,
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
index 52b8fa3d849..db911364ef3 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
@@ -2,6 +2,7 @@ import asyncio
 import contextvars
 import os
 from datetime import datetime, timedelta
+from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -1336,12 +1337,19 @@ def _denied_scope_manager(known_server_names_to_ids: dict[str, str]) -> MagicMoc
     return manager
 
 
-def _scope_resolver(resolved_without_agent: list[str], access_groups: tuple[str, ...] = ()) -> AsyncMock:
+def _scope_resolver(resolved_without_agent: dict[str, str], access_groups: tuple[str, ...] = ()) -> AsyncMock:
     async def resolve(user_api_key_auth, mcp_servers, client_ip=None):
         if user_api_key_auth is not None and user_api_key_auth.agent_id:
             return []
         return [
-            MagicMock(server_id=server_id, access_groups=list(access_groups)) for server_id in resolved_without_agent
+            SimpleNamespace(
+                server_id=server_id,
+                server_name=server_name,
+                alias=None,
+                short_prefix=None,
+                access_groups=list(access_groups),
+            )
+            for server_name, server_id in resolved_without_agent.items()
         ]
 
     return AsyncMock(side_effect=resolve)
@@ -1380,7 +1388,7 @@ async def test_scoped_list_denied_by_agent_binding_raises_403_naming_agent():
     pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
 
     user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
-    resolver = _scope_resolver(resolved_without_agent=["srv-github"])
+    resolver = _scope_resolver(resolved_without_agent={"github": "srv-github"})
 
     denial = await _denied_scoped_list(
         user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver
@@ -1452,13 +1460,13 @@ async def test_scoped_list_unknown_name_raises_same_generic_403_as_unauthorized(
     user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
 
     unknown = await _denied_scoped_list(
-        user_api_key_auth, ["github"], _denied_scope_manager({}), _scope_resolver(resolved_without_agent=[])
+        user_api_key_auth, ["github"], _denied_scope_manager({}), _scope_resolver(resolved_without_agent={})
     )
     unauthorized = await _denied_scoped_list(
         user_api_key_auth,
         ["github"],
         _denied_scope_manager({"github": "srv-github"}),
-        _scope_resolver(resolved_without_agent=[]),
+        _scope_resolver(resolved_without_agent={}),
     )
 
     assert unknown.status_code == unauthorized.status_code == 403
@@ -1479,7 +1487,7 @@ async def test_scoped_list_access_group_vetoed_by_agent_names_agent_and_group():
         user_api_key_auth,
         ["prod-group"],
         _denied_scope_manager({}),
-        _scope_resolver(resolved_without_agent=["srv-github"], access_groups=("prod-group",)),
+        _scope_resolver(resolved_without_agent={"github": "srv-github"}, access_groups=("prod-group",)),
     )
 
     assert denial.status_code == 403
@@ -1502,7 +1510,7 @@ async def test_scoped_list_mixed_unknown_and_vetoed_group_names_the_group_that_r
         user_api_key_auth,
         ["no-such-group", "prod-group"],
         _denied_scope_manager({}),
-        _scope_resolver(resolved_without_agent=["srv-github"], access_groups=("prod-group",)),
+        _scope_resolver(resolved_without_agent={"github": "srv-github"}, access_groups=("prod-group",)),
     )
 
     assert denial.status_code == 403
@@ -1519,7 +1527,7 @@ async def test_scoped_list_agent_key_denied_by_key_grants_raises_generic_403():
     pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
 
     user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
-    resolver = _scope_resolver(resolved_without_agent=[])
+    resolver = _scope_resolver(resolved_without_agent={})
 
     denial = await _denied_scoped_list(
         user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver
@@ -1532,6 +1540,27 @@ async def test_scoped_list_agent_key_denied_by_key_grants_raises_generic_403():
     assert resolver.await_count == 2
 
 
+@pytest.mark.asyncio
+async def test_scoped_list_agent_veto_attributed_for_differently_cased_server_name():
+    """The scope filter matches `/mcp/GitHub` to a server named `github` case-insensitively, so the
+    agent-attributed 403 must match the same way instead of falling back to the generic denial."""
+    pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
+
+    user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
+
+    denial = await _denied_scoped_list(
+        user_api_key_auth,
+        ["GitHub"],
+        _denied_scope_manager({"github": "srv-github"}),
+        _scope_resolver(resolved_without_agent={"github": "srv-github"}),
+    )
+
+    assert denial.status_code == 403
+    message = denial.detail["error"]
+    assert "MCP server 'GitHub'" in message
+    assert "agent 'agent-123'" in message
+
+
 @pytest.mark.asyncio
 async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error():
     """The MCP protocol handler surfaces a permission HTTPException as a clean JSON-RPC error

From d6b8cb93b3bf9efe07d807f400016eba95cd9385 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 2 Sep 2026 18:18:18 -0700
Subject: [PATCH 018/164] fix(mcp): name the agent veto on scoped tool calls
 too

---
 .../proxy/_experimental/mcp_server/server.py  | 19 +++++--
 .../_experimental/mcp_server/tool_search.py   |  7 +++
 .../mcp_server/test_mcp_server.py             | 55 +++++++++++++++++++
 .../mcp_server/test_mcp_tool_search.py        | 35 ++++++++++++
 4 files changed, 110 insertions(+), 6 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 2f6a05481e3..aa36150b477 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -821,11 +821,7 @@ if MCP_AVAILABLE:
             from mcp.shared.exceptions import McpError
             from mcp.types import INVALID_REQUEST, ErrorData
 
-            detail: Final = e.detail
-            message: Final = (
-                str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail)
-            )
-            raise McpError(ErrorData(code=INVALID_REQUEST, message=message)) from e
+            raise McpError(ErrorData(code=INVALID_REQUEST, message=_http_detail_message(e.detail))) from e
         except Exception as e:
             verbose_logger.exception("Error in list_tools endpoint: %s", e)
             # Return empty list instead of failing completely
@@ -1096,6 +1092,7 @@ if MCP_AVAILABLE:
                     mcp_server_auth_headers=mcp_server_auth_headers,
                     oauth2_headers=oauth2_headers,
                     raw_headers=raw_headers,
+                    client_ip=_client_ip,
                     host_progress_callback=host_progress_callback,
                     **data,  # for logging
                 )
@@ -1129,7 +1126,7 @@ if MCP_AVAILABLE:
             except HTTPException as e:
                 verbose_logger.error("HTTPException in MCP tool call: %s", e)
                 return CallToolResult(
-                    content=[TextContent(text=f"Error: {e.detail}", type="text")],
+                    content=[TextContent(text=f"Error: {_http_detail_message(e.detail)}", type="text")],
                     isError=True,
                 )
             except MCPUpstreamAuthError as e:
@@ -1447,6 +1444,9 @@ if MCP_AVAILABLE:
 
         return allowed_mcp_servers
 
+    def _http_detail_message(detail: object) -> str:
+        return str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail)
+
     def _server_answers_to(server: MCPServer, name: str) -> bool:
         requested: Final = name.lower()
         return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known)
@@ -3148,6 +3148,7 @@ if MCP_AVAILABLE:
         mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
         oauth2_headers: dict[str, str] | None = None,
         raw_headers: dict[str, str] | None = None,
+        client_ip: str | None = None,
         **kwargs: Any,
     ) -> CallToolResult:
         """
@@ -3178,6 +3179,12 @@ if MCP_AVAILABLE:
                 mcp_servers=mcp_servers,
                 allowed_mcp_servers=allowed_mcp_servers,
             )
+            if mcp_servers and not allowed_mcp_servers:
+                await _raise_denied_scoped_mcp_access(
+                    requested_names=mcp_servers,
+                    user_api_key_auth=user_api_key_auth,
+                    client_ip=client_ip,
+                )
             if not allowed_mcp_servers:
                 raise HTTPException(
                     status_code=403,
diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py
index 4f6305d88cf..e624367a887 100644
--- a/litellm/proxy/_experimental/mcp_server/tool_search.py
+++ b/litellm/proxy/_experimental/mcp_server/tool_search.py
@@ -307,6 +307,7 @@ async def handle_mcp_tool_call(
 ) -> CallToolResult:
     from litellm.proxy._experimental.mcp_server.server import (
         _get_allowed_mcp_servers,
+        _raise_denied_scoped_mcp_access,
         execute_mcp_tool,
     )
 
@@ -315,6 +316,12 @@ async def handle_mcp_tool_call(
         mcp_servers=mcp_servers,
         client_ip=client_ip,
     )
+    if mcp_servers and not allowed_mcp_servers:
+        await _raise_denied_scoped_mcp_access(
+            requested_names=mcp_servers,
+            user_api_key_auth=user_api_key_dict,
+            client_ip=client_ip,
+        )
 
     # Reject before dispatch when the key has no accessible servers; otherwise an
     # unprefixed local tool name would fall through to the local registry in
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
index db911364ef3..fd36d739dc5 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
@@ -1593,6 +1593,32 @@ async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error(
     assert exc_info.value.error.message == denial_message
 
 
+@pytest.mark.asyncio
+async def test_mcp_server_tool_call_renders_denial_message_not_detail_dict():
+    try:
+        from litellm.proxy._experimental.mcp_server.server import mcp_server_tool_call
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    denial_message = "MCP server 'github' is not available to this key: the key is bound to agent 'agent-123'"
+    denial = HTTPException(status_code=403, detail={"error": denial_message})
+
+    with (
+        patch(  # test-quality-ok: the protocol handler reads auth from module context; no injection seam
+            "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context",
+            new=AsyncMock(return_value=(None, None, None, None, None, None, None)),
+        ),
+        patch(  # test-quality-ok: the tool-call helper is the handler's only collaborator; the suite's seam
+            "litellm.proxy._experimental.mcp_server.server.call_mcp_tool",
+            new=AsyncMock(side_effect=denial),
+        ),
+    ):
+        result = await mcp_server_tool_call("github-search_issues", {})
+
+    assert result.isError is True
+    assert result.content[0].text == f"Error: {denial_message}"
+
+
 @pytest.mark.asyncio
 async def test_mcp_server_tool_call_body_with_none_arguments():
     """Test that proxy_server_request body handles None arguments correctly"""
@@ -3742,6 +3768,35 @@ async def test_call_mcp_tool_user_unauthorized_access():
         assert "User not allowed to call this tool" in exc_info.value.detail
 
 
+@pytest.mark.asyncio
+async def test_call_mcp_tool_scoped_denial_names_the_binding_agent():
+    from litellm.proxy._experimental.mcp_server.server import call_mcp_tool
+
+    agent_bound_key = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="agent-123")
+
+    with (
+        patch(  # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_allowed_mcp_servers",
+            AsyncMock(return_value=[]),
+        ),
+        patch(  # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
+            _scope_resolver({"github": "srv-github"}),
+        ),
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await call_mcp_tool(
+                name="github-search_issues",
+                arguments={},
+                user_api_key_auth=agent_bound_key,
+                mcp_servers=["github"],
+            )
+
+    assert exc_info.value.status_code == 403
+    assert "MCP server 'github'" in exc_info.value.detail["error"]
+    assert "agent 'agent-123'" in exc_info.value.detail["error"]
+
+
 @pytest.mark.asyncio
 async def test_call_mcp_tool_unauthorized_403_does_not_leak_server_credentials():
     """Regression for LIT-4703 / GH #29936.
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py
index 239f89ebd90..b8e985d59fe 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py
@@ -11,6 +11,7 @@ Covers:
 
 import json
 from collections.abc import Sequence
+from types import SimpleNamespace
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -1229,3 +1230,37 @@ class TestMcpServerToolCallErrorHandling:
 
         assert result.isError is True
         assert "User not allowed to call this tool" in result.content[0].text
+
+
+@pytest.mark.asyncio
+async def test_handle_mcp_tool_call_scoped_denial_names_the_binding_agent() -> None:
+    from fastapi import HTTPException
+
+    from litellm.proxy._experimental.mcp_server.tool_search import handle_mcp_tool_call
+
+    agent_bound_key = UserAPIKeyAuth(api_key="test_key", agent_id="agent-123")
+
+    async def resolve(user_api_key_auth, mcp_servers, client_ip=None):
+        if user_api_key_auth.agent_id:
+            return []
+        return [
+            SimpleNamespace(
+                server_id="srv-github", server_name="github", alias=None, short_prefix=None, access_groups=[]
+            )
+        ]
+
+    with patch(  # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
+        "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
+        new=AsyncMock(side_effect=resolve),
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await handle_mcp_tool_call(
+                tool_name="github-create_issue",
+                arguments={},
+                user_api_key_dict=agent_bound_key,
+                mcp_servers=["github"],
+            )
+
+    assert exc_info.value.status_code == 403
+    assert "MCP server 'github'" in exc_info.value.detail["error"]
+    assert "agent 'agent-123'" in exc_info.value.detail["error"]

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 019/164] 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 c79d1d12ae75a8eaec1de15fe4f0de01330cc553 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 02:15:57 -0700
Subject: [PATCH 020/164] fix(ai-gateway): install a rustls crypto provider
 before dialing upstream WebSockets

The gateway's dependency graph turns on two rustls crypto backends at once:
reqwest's rustls-tls pulls in ring, and litellm-core's bedrock-auth pulls in
aws-lc-rs through aws-config. rustls 0.23 refuses to guess between them, so
ClientConfig::builder panics, and that is exactly how tokio-tungstenite builds
its TLS config. Every outbound WebSocket dial killed its tokio worker and the
client saw the socket vanish with no close frame.

reqwest and the AWS SDK both pick a provider explicitly, so only the tungstenite
path was affected. Route all three dial sites through one helper that installs
ring once per process before connecting.
---
 litellm-rust/Cargo.lock                       |  1 +
 litellm-rust/Cargo.toml                       |  1 +
 litellm-rust/crates/ai-gateway/Cargo.toml     |  3 +
 litellm-rust/crates/ai-gateway/src/io/mod.rs  |  1 +
 .../crates/ai-gateway/src/io/realtime.rs      |  6 +-
 .../crates/ai-gateway/src/io/responses_ws.rs  | 12 ++--
 litellm-rust/crates/ai-gateway/src/io/tls.rs  | 67 +++++++++++++++++++
 7 files changed, 84 insertions(+), 7 deletions(-)
 create mode 100644 litellm-rust/crates/ai-gateway/src/io/tls.rs

diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock
index b3dac5ca935..2ed998174e4 100644
--- a/litellm-rust/Cargo.lock
+++ b/litellm-rust/Cargo.lock
@@ -1415,6 +1415,7 @@ dependencies = [
  "litellm-core",
  "pyo3",
  "reqwest",
+ "rustls 0.23.42",
  "serde",
  "serde_json",
  "sha2 0.10.9",
diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml
index a13dd4c04b0..15074082981 100644
--- a/litellm-rust/Cargo.toml
+++ b/litellm-rust/Cargo.toml
@@ -26,6 +26,7 @@ pythonize = "0.29.0"
 rand = "0.8"
 reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
 rstest = "0.26.1"
+rustls = { version = "0.23", default-features = false, features = ["ring", "std"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 sha2 = "0.10"
diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml
index e3dbdf24ce6..414abc2356d 100644
--- a/litellm-rust/crates/ai-gateway/Cargo.toml
+++ b/litellm-rust/crates/ai-gateway/Cargo.toml
@@ -19,6 +19,9 @@ litellm-core = { workspace = true, features = ["bedrock-auth"] }
 # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
 # Python proxy callbacks API.
 reqwest.workspace = true
+# rustls is a direct dependency so `io::tls` can install a process-level
+# crypto provider; see that module for why the graph needs one.
+rustls.workspace = true
 # `sync` powers the bounded mpsc channel the realtime logger drains.
 tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] }
 tokio-tungstenite.workspace = true
diff --git a/litellm-rust/crates/ai-gateway/src/io/mod.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs
index cce56dd2121..7098d67993f 100644
--- a/litellm-rust/crates/ai-gateway/src/io/mod.rs
+++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs
@@ -3,3 +3,4 @@ pub mod ocr;
 pub mod realtime;
 pub mod realtime_pool;
 pub mod responses_ws;
+pub(crate) mod tls;
diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs
index 662f7328982..53d87848342 100644
--- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs
+++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs
@@ -23,10 +23,12 @@ use tokio_tungstenite::tungstenite::Message;
 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
 use tokio_tungstenite::tungstenite::http::HeaderValue;
 use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
-use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
+use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
 
 use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
 
+use crate::io::tls::connect_upstream;
+
 /// Environment variable holding the OpenAI API key (last-resort fallback).
 const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
 
@@ -84,7 +86,7 @@ pub(crate) async fn dial_upstream(
             .map_err(|err| Error::Auth(err.to_string()))?,
     );
 
-    let (upstream, _response) = connect_async(request)
+    let (upstream, _response) = connect_upstream(request)
         .await
         .map_err(|err| Error::Network(err.to_string()))?;
     Ok(upstream)
diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
index 0b01747b1a5..ee181791413 100644
--- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
+++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
@@ -14,7 +14,9 @@ use tokio_tungstenite::tungstenite::Message;
 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
 use tokio_tungstenite::tungstenite::http::HeaderValue;
 use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName};
-use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
+use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
+
+use crate::io::tls::connect_upstream;
 
 use crate::constants::{
     DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS,
@@ -49,14 +51,14 @@ impl ResponsesWebSocketConnection {
                 .map_err(|error| Error::InvalidRequest(error.to_string()))?;
             request.headers_mut().insert(header_name, header_value);
         }
-        let connect = connect_async(request);
+        let connect = connect_upstream(request);
         let result = match timeout {
             Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
                 Error::Network("Responses WebSocket connection timed out".to_string())
             })?,
             None => connect.await,
         };
-        let (socket, _) = result.map_err(|error| match error {
+        let (socket, _) = result.map_err(|error| match *error {
             tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
                 status: response.status().as_u16(),
                 body: String::new(),
@@ -138,13 +140,13 @@ async fn dial_upstream(
     );
     let result = tokio::time::timeout(
         Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS),
-        connect_async(request),
+        connect_upstream(request),
     )
     .await
     .map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?;
     result
         .map(|(socket, _)| socket)
-        .map_err(|error| match error {
+        .map_err(|error| match *error {
             tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
                 status: response.status().as_u16(),
                 body: String::new(),
diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs
new file mode 100644
index 00000000000..e3754e44002
--- /dev/null
+++ b/litellm-rust/crates/ai-gateway/src/io/tls.rs
@@ -0,0 +1,67 @@
+//! Outbound WebSocket dials, with the rustls crypto provider settled first.
+//!
+//! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth`
+//! enables `rustls/aws-lc-rs`, so `ClientConfig::builder()` — which is how
+//! `tokio-tungstenite` builds its TLS config — panics rather than guess between
+//! them. reqwest and the AWS SDK pick a provider explicitly and never panic.
+//!
+//! Installing from the dial rather than from a `main` also covers the `cdylib`
+//! the Python bridge loads, the tests, and the benches, none of which have one.
+//! ring is what reqwest already falls back to, so installing it changes no
+//! working path, and an embedder that installed its own provider first keeps it.
+
+use std::sync::Once;
+
+use tokio::net::TcpStream;
+use tokio_tungstenite::tungstenite::Error;
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+use tokio_tungstenite::tungstenite::handshake::client::Response;
+use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
+
+static INSTALL_CRYPTO_PROVIDER: Once = Once::new();
+
+pub(crate) fn ensure_crypto_provider() {
+    INSTALL_CRYPTO_PROVIDER.call_once(|| {
+        let _ = rustls::crypto::ring::default_provider().install_default();
+    });
+}
+
+pub(crate) async fn connect_upstream<R>(
+    request: R,
+) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response), Box<Error>>
+where
+    R: IntoClientRequest + Unpin,
+{
+    ensure_crypto_provider();
+    connect_async(request).await.map_err(Box::new)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::ensure_crypto_provider;
+
+    #[test]
+    fn client_config_builder_works_with_both_provider_features_enabled() {
+        ensure_crypto_provider();
+
+        assert!(rustls::crypto::CryptoProvider::get_default().is_some());
+
+        let config = rustls::ClientConfig::builder()
+            .with_root_certificates(rustls::RootCertStore::empty())
+            .with_no_client_auth();
+
+        assert!(!config.crypto_provider().cipher_suites.is_empty());
+    }
+
+    #[test]
+    fn ensure_crypto_provider_is_idempotent() {
+        ensure_crypto_provider();
+        let first = rustls::crypto::CryptoProvider::get_default().cloned();
+
+        ensure_crypto_provider();
+        let second = rustls::crypto::CryptoProvider::get_default().cloned();
+
+        assert!(first.is_some());
+        assert!(std::sync::Arc::ptr_eq(&first.unwrap(), &second.unwrap()));
+    }
+}

From 2c08e7abf8bbf1d32aab48a330f4aadd494b5dd0 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 02:27:00 -0700
Subject: [PATCH 021/164] test(ai-gateway): pin the crypto provider to ring and
 prove the dial installs it

The two tests that shipped with the fix both called ensure_crypto_provider
themselves, so deleting the call from connect_upstream left the whole suite
green, and swapping ring for aws-lc-rs did too.

Adds an integration test, which gets its own process, that dials wss:// at a
local plain-TCP listener through the public Responses WebSocket entrypoint and
asserts an Err plus an installed provider. Without the install in the dial it
panics with the original CryptoProvider message. A unit test now compares the
installed provider's cipher suites and key-exchange groups against ring's, so
the choice of backend is pinned rather than assumed.

Also names tls12 in the workspace rustls features: it already arrives through
reqwest and tokio-rustls, so the graph is unchanged, but a direct dependency
should say it needs TLS 1.2 rather than inherit it.
---
 litellm-rust/Cargo.toml                       |  2 +-
 litellm-rust/crates/ai-gateway/src/io/tls.rs  | 43 +++++++++++++++++--
 .../tests/crypto_provider_wiring.rs           | 41 ++++++++++++++++++
 3 files changed, 81 insertions(+), 5 deletions(-)
 create mode 100644 litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs

diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml
index 15074082981..2e2e8809b7f 100644
--- a/litellm-rust/Cargo.toml
+++ b/litellm-rust/Cargo.toml
@@ -26,7 +26,7 @@ pythonize = "0.29.0"
 rand = "0.8"
 reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
 rstest = "0.26.1"
-rustls = { version = "0.23", default-features = false, features = ["ring", "std"] }
+rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 sha2 = "0.10"
diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs
index e3754e44002..96544adffb7 100644
--- a/litellm-rust/crates/ai-gateway/src/io/tls.rs
+++ b/litellm-rust/crates/ai-gateway/src/io/tls.rs
@@ -8,7 +8,7 @@
 //! Installing from the dial rather than from a `main` also covers the `cdylib`
 //! the Python bridge loads, the tests, and the benches, none of which have one.
 //! ring is what reqwest already falls back to, so installing it changes no
-//! working path, and an embedder that installed its own provider first keeps it.
+//! working path, and whoever installs into this rustls build first still wins.
 
 use std::sync::Once;
 
@@ -38,13 +38,32 @@ where
 
 #[cfg(test)]
 mod tests {
+    use rustls::crypto::CryptoProvider;
+
     use super::ensure_crypto_provider;
 
+    fn fingerprint(
+        provider: &CryptoProvider,
+    ) -> (Vec<rustls::CipherSuite>, Vec<rustls::NamedGroup>) {
+        (
+            provider
+                .cipher_suites
+                .iter()
+                .map(|suite| suite.suite())
+                .collect(),
+            provider
+                .kx_groups
+                .iter()
+                .map(|group| group.name())
+                .collect(),
+        )
+    }
+
     #[test]
     fn client_config_builder_works_with_both_provider_features_enabled() {
         ensure_crypto_provider();
 
-        assert!(rustls::crypto::CryptoProvider::get_default().is_some());
+        assert!(CryptoProvider::get_default().is_some());
 
         let config = rustls::ClientConfig::builder()
             .with_root_certificates(rustls::RootCertStore::empty())
@@ -53,13 +72,29 @@ mod tests {
         assert!(!config.crypto_provider().cipher_suites.is_empty());
     }
 
+    #[test]
+    fn installs_ring_rather_than_aws_lc_rs() {
+        ensure_crypto_provider();
+
+        let installed = CryptoProvider::get_default().expect("a provider is installed");
+
+        assert_eq!(
+            fingerprint(installed),
+            fingerprint(&rustls::crypto::ring::default_provider())
+        );
+        assert_ne!(
+            fingerprint(installed),
+            fingerprint(&rustls::crypto::aws_lc_rs::default_provider())
+        );
+    }
+
     #[test]
     fn ensure_crypto_provider_is_idempotent() {
         ensure_crypto_provider();
-        let first = rustls::crypto::CryptoProvider::get_default().cloned();
+        let first = CryptoProvider::get_default().cloned();
 
         ensure_crypto_provider();
-        let second = rustls::crypto::CryptoProvider::get_default().cloned();
+        let second = CryptoProvider::get_default().cloned();
 
         assert!(first.is_some());
         assert!(std::sync::Arc::ptr_eq(&first.unwrap(), &second.unwrap()));
diff --git a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs
new file mode 100644
index 00000000000..db5698f8460
--- /dev/null
+++ b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs
@@ -0,0 +1,41 @@
+//! Guards the wiring, not just the helper: the dial itself has to install the
+//! rustls provider, in a test binary where nothing else has installed one.
+
+use std::collections::HashMap;
+use std::time::Duration;
+
+use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection;
+use tokio::net::TcpListener;
+
+#[tokio::test]
+async fn dialing_wss_returns_an_error_instead_of_panicking() {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("bind a loopback port");
+    let port = listener
+        .local_addr()
+        .expect("read the bound address")
+        .port();
+
+    tokio::spawn(async move {
+        while let Ok((stream, _peer)) = listener.accept().await {
+            drop(stream);
+        }
+    });
+
+    let result = ResponsesWebSocketConnection::connect_url(
+        &format!("wss://127.0.0.1:{port}/"),
+        &HashMap::new(),
+        Some(Duration::from_secs(10)),
+    )
+    .await;
+
+    assert!(
+        result.is_err(),
+        "a plain TCP server cannot finish a TLS handshake"
+    );
+    assert!(
+        rustls::crypto::CryptoProvider::get_default().is_some(),
+        "the dial is what installs the process-wide provider"
+    );
+}

From 534ab1628dc462912179660774af395fdce50839 Mon Sep 17 00:00:00 2001
From: Riddhi04 <riddhimahajan04@gmail.com>
Date: Wed, 22 Jul 2026 15:58:37 +0400
Subject: [PATCH 022/164] fix(proxy): enforce model access checks on Bedrock
 passthrough routes

get_model_from_request could not resolve a model for /bedrock/... routes since it only checked the JSON body's model field and a small set of URL regexes, none matching Bedrock's passthrough path. This let common_checks skip the key/project model allowlist entirely for any Bedrock passthrough action (invoke, converse, and their streaming variants), while the same model was correctly blocked on /v1/chat/completions

Extract the model from the Bedrock endpoint path using the same helper the passthrough handler itself relies on, so the existing allowlist check applies uniformly across auth methods and call paths
---
 litellm/proxy/auth/auth_utils.py              | 11 ++++
 .../proxy/auth/test_auth_utils.py             | 50 +++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py
index d6007a2d56e..43e60e49e3b 100644
--- a/litellm/proxy/auth/auth_utils.py
+++ b/litellm/proxy/auth/auth_utils.py
@@ -1981,6 +1981,17 @@ def get_model_from_request(
         if vertex_match:
             model = vertex_match.group(1)
 
+    if model is None and route.lower().startswith("/bedrock"):
+        from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
+            _extract_model_from_bedrock_endpoint,
+        )
+
+        bedrock_endpoint = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE)
+        try:
+            model = _extract_model_from_bedrock_endpoint(bedrock_endpoint)
+        except ValueError:
+            model = None
+
     return model
 
 
diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py
index f513f397b64..f3426534c28 100644
--- a/tests/test_litellm/proxy/auth/test_auth_utils.py
+++ b/tests/test_litellm/proxy/auth/test_auth_utils.py
@@ -428,6 +428,56 @@ def test_get_model_from_request_openai_deployment_route_still_works():
     )
 
 
+def test_get_model_from_request_bedrock_converse_passthrough():
+    assert (
+        get_model_from_request(
+            request_data={},
+            route="/bedrock/model/us.anthropic.claude-sonnet-4-6/converse",
+        )
+        == "us.anthropic.claude-sonnet-4-6"
+    )
+
+
+def test_get_model_from_request_bedrock_invoke_passthrough():
+    assert (
+        get_model_from_request(
+            request_data={},
+            route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke",
+        )
+        == "us.anthropic.claude-sonnet-4-6"
+    )
+
+
+def test_get_model_from_request_bedrock_v2_converse_stream_passthrough():
+    assert (
+        get_model_from_request(
+            request_data={},
+            route="/bedrock/v2/model/us.anthropic.claude-sonnet-4-6/converse-stream",
+        )
+        == "us.anthropic.claude-sonnet-4-6"
+    )
+
+
+def test_get_model_from_request_bedrock_model_id_with_slashes():
+    assert (
+        get_model_from_request(
+            request_data={},
+            route="/bedrock/model/aws/anthropic/model-name/invoke",
+        )
+        == "aws/anthropic/model-name"
+    )
+
+
+def test_get_model_from_request_bedrock_unparseable_endpoint_returns_none():
+    assert (
+        get_model_from_request(
+            request_data={},
+            route="/bedrock/agents/some-agent-route",
+        )
+        is None
+    )
+
+
 def test_get_model_from_request_includes_file_endpoint_header_model():
     assert (
         get_model_from_request(

From 6981ccf0ca1c3e1937f71c3fce287c898a15feba Mon Sep 17 00:00:00 2001
From: Riddhi04 <riddhimahajan04@gmail.com>
Date: Fri, 24 Jul 2026 16:38:34 +0400
Subject: [PATCH 023/164] fix(proxy): make URL model authoritative for Bedrock
 path-routed passthrough actions

The allowlist check read model from the request body first, while
bedrock_llm_proxy_route dispatches purely on the path model for invoke,
converse, and their streaming variants. A caller could put an allowed
model in the JSON body while targeting a disallowed model in the URL and
slip past the check. count_tokens keeps reading from the body since its
route has no model segment in the path.
---
 litellm/proxy/auth/auth_utils.py              | 22 +++++-----
 .../proxy/auth/test_auth_utils.py             | 40 +++++++++++++++++++
 2 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py
index 43e60e49e3b..fd746d62e76 100644
--- a/litellm/proxy/auth/auth_utils.py
+++ b/litellm/proxy/auth/auth_utils.py
@@ -1981,16 +1981,20 @@ def get_model_from_request(
         if vertex_match:
             model = vertex_match.group(1)
 
-    if model is None and route.lower().startswith("/bedrock"):
-        from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
-            _extract_model_from_bedrock_endpoint,
-        )
-
+    if route.lower().startswith("/bedrock"):
         bedrock_endpoint = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE)
-        try:
-            model = _extract_model_from_bedrock_endpoint(bedrock_endpoint)
-        except ValueError:
-            model = None
+        is_bedrock_count_tokens_route = (
+            "count_tokens" in bedrock_endpoint.lower() or "count-tokens" in bedrock_endpoint.lower()
+        )
+        if not is_bedrock_count_tokens_route:
+            from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
+                _extract_model_from_bedrock_endpoint,
+            )
+
+            try:
+                model = _extract_model_from_bedrock_endpoint(bedrock_endpoint)
+            except ValueError:
+                pass
 
     return model
 
diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py
index f3426534c28..61d1a308b9d 100644
--- a/tests/test_litellm/proxy/auth/test_auth_utils.py
+++ b/tests/test_litellm/proxy/auth/test_auth_utils.py
@@ -478,6 +478,46 @@ def test_get_model_from_request_bedrock_unparseable_endpoint_returns_none():
     )
 
 
+def test_get_model_from_request_bedrock_url_model_overrides_body_model():
+    assert (
+        get_model_from_request(
+            request_data={"model": "us.anthropic.claude-sonnet-4-6"},
+            route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/converse",
+        )
+        == "us.anthropic.claude-opus-4-6-v1"
+    )
+
+
+def test_get_model_from_request_bedrock_invoke_url_model_overrides_body_model():
+    assert (
+        get_model_from_request(
+            request_data={"model": "us.anthropic.claude-sonnet-4-6"},
+            route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/invoke",
+        )
+        == "us.anthropic.claude-opus-4-6-v1"
+    )
+
+
+def test_get_model_from_request_bedrock_count_tokens_uses_body_model():
+    assert (
+        get_model_from_request(
+            request_data={"model": "us.anthropic.claude-sonnet-4-6"},
+            route="/bedrock/v1/messages/count_tokens",
+        )
+        == "us.anthropic.claude-sonnet-4-6"
+    )
+
+
+def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model():
+    assert (
+        get_model_from_request(
+            request_data={"model": "us.anthropic.claude-sonnet-4-6"},
+            route="/bedrock/agents/some-agent-route",
+        )
+        == "us.anthropic.claude-sonnet-4-6"
+    )
+
+
 def test_get_model_from_request_includes_file_endpoint_header_model():
     assert (
         get_model_from_request(

From 32293295f8fe9811a18179b556e12cdb6c4a11bc Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 15:40:35 -0700
Subject: [PATCH 024/164] refactor(proxy): resolve the Bedrock route model
 through an early return

The Bedrock branch of get_model_from_request reassigned the already resolved
model binding. Move the route parsing into a helper that returns the URL model
or None so the caller picks between it and the body model without rebinding.
---
 litellm/proxy/auth/auth_utils.py | 29 ++++++++++++++++-------------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py
index fd746d62e76..12e6f75889e 100644
--- a/litellm/proxy/auth/auth_utils.py
+++ b/litellm/proxy/auth/auth_utils.py
@@ -1982,23 +1982,26 @@ def get_model_from_request(
             model = vertex_match.group(1)
 
     if route.lower().startswith("/bedrock"):
-        bedrock_endpoint = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE)
-        is_bedrock_count_tokens_route = (
-            "count_tokens" in bedrock_endpoint.lower() or "count-tokens" in bedrock_endpoint.lower()
-        )
-        if not is_bedrock_count_tokens_route:
-            from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
-                _extract_model_from_bedrock_endpoint,
-            )
-
-            try:
-                model = _extract_model_from_bedrock_endpoint(bedrock_endpoint)
-            except ValueError:
-                pass
+        bedrock_model: Final = _model_from_bedrock_route(route)
+        return model if bedrock_model is None else bedrock_model
 
     return model
 
 
+def _model_from_bedrock_route(route: str) -> str | None:
+    bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE)
+    if "count_tokens" in bedrock_endpoint.lower() or "count-tokens" in bedrock_endpoint.lower():
+        return None
+    from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
+        _extract_model_from_bedrock_endpoint,
+    )
+
+    try:
+        return _extract_model_from_bedrock_endpoint(bedrock_endpoint)
+    except ValueError:
+        return None
+
+
 def abbreviate_api_key(api_key: str) -> str:
     if len(api_key) < MINIMUM_CUSTOM_KEY_LENGTH:
         return "sk-..."

From e5d51ee8be3d191b295a94366e790c1cc9b43b08 Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Thu, 3 Sep 2026 19:06:01 -0400
Subject: [PATCH 025/164] fix(vertex_ai): support fine-tuned Gemini endpoints
 in managed batches

Managed batches mangled any Vertex model that is not a plain publisher
model: a fine-tuned Gemini endpoint id was filed under
publishers/google/models/gemini/<id> at upload, then the batch create
parse dropped the id and targeted the nonexistent publisher model
'publishers/google/models/gemini', which Vertex rejects.

Fine-tuned endpoints are now stored under endpoints/<id> in the GCS
object path, and batch create resolves the endpoint to its deployed
tuned model resource (projects/../models/<id>) via GET endpoints/<id>,
which is the only form the v1 batch API accepts for tuned models. The
cost poller's bare-model parse round-trips the endpoint id so unmanaged
batch spend still maps to the configured deployment.

custom_endpoint deployments have no Vertex batch surface, so batch file
uploads and batch creation against them now return a clear 400 instead
of creating a doomed job.

Resolves LIT-6899
---
 litellm/batches/main.py                       |  15 +++
 litellm/llms/vertex_ai/batches/handler.py     |  64 +++++++++-
 .../llms/vertex_ai/batches/transformation.py  |  56 ++++++++-
 litellm/llms/vertex_ai/common_utils.py        |  13 ++
 .../llms/vertex_ai/files/transformation.py    |  28 ++++-
 .../proxy_unit_tests/test_check_batch_cost.py |  29 +++++
 tests/test_litellm/batches/test_main.py       |  10 ++
 .../llms/vertex_ai/batches/test_handler.py    | 119 ++++++++++++++++++
 .../vertex_ai/batches/test_transformation.py  |  58 +++++++++
 .../test_vertex_ai_files_transformation.py    |  55 ++++++++
 10 files changed, 435 insertions(+), 12 deletions(-)

diff --git a/litellm/batches/main.py b/litellm/batches/main.py
index c8360a81c7a..6cf9f070d9c 100644
--- a/litellm/batches/main.py
+++ b/litellm/batches/main.py
@@ -301,6 +301,21 @@ def create_batch(
                 litellm_params=litellm_params,
             )
         elif custom_llm_provider == "vertex_ai":
+            if optional_params.get("custom_endpoint"):
+                raise litellm.exceptions.BadRequestError(
+                    message=(
+                        "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. "
+                        "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; "
+                        "use a publisher model or fine-tuned Gemini endpoint deployment instead."
+                    ),
+                    model=model or "n/a",
+                    llm_provider=custom_llm_provider,
+                    response=httpx.Response(
+                        status_code=400,
+                        content="custom_endpoint deployments do not support vertex_ai batches",
+                        request=httpx.Request(method="create_batch", url="https://github.com/BerriAI/litellm"),
+                    ),
+                )
             api_base = optional_params.api_base or ""
             vertex_ai_project: Final = (
                 optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index 377cd9f3437..1bb8ceb747a 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -12,6 +12,7 @@ from litellm.litellm_core_utils.url_utils import (
     safe_get,
 )
 from litellm.llms.custom_httpx.http_handler import (
+    HTTPHandler,
     _get_httpx_client,
     get_async_httpx_client,
 )
@@ -55,6 +56,20 @@ class _FetchedResponseView(TypedDict):
     response: ReadOnly[httpx.Response]
 
 
+class _VertexEndpointDeployedModel(TypedDict, total=False):
+    model: ReadOnly[str]
+
+
+class _VertexEndpointResponse(TypedDict, total=False):
+    deployedModels: ReadOnly[list[_VertexEndpointDeployedModel]]
+
+
+class _VertexEndpointPayloadView(TypedDict):
+    """Holds one decoded GET endpoints/<id> response so the payload reads back typed."""
+
+    payload: ReadOnly[_VertexEndpointResponse]
+
+
 def _vertex_batch_payload(response: _VertexBatchJsonSource) -> VertexBatchPredictionResponse:
     return response.json()
 
@@ -116,11 +131,19 @@ class VertexAIBatchPrediction(VertexLLM):
             "Authorization": f"Bearer {access_token}",
         }
 
-        vertex_batch_request: Final[VertexAIBatchPredictionJob] = (
+        transformed_batch_request: Final[VertexAIBatchPredictionJob] = (
             VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
-                request=create_batch_data
+                request=create_batch_data,
+                vertex_project=vertex_project or project_id,
+                vertex_location=vertex_location or "us-central1",
             )
         )
+        vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model(
+            vertex_batch_request=transformed_batch_request,
+            headers=headers,
+            sync_handler=sync_handler,
+            vertex_location=vertex_location or "us-central1",
+        )
 
         if _is_async is True:
             return self._async_create_batch(
@@ -142,6 +165,43 @@ class VertexAIBatchPrediction(VertexLLM):
         )
         return vertex_batch_response
 
+    def _resolve_fine_tuned_endpoint_model(
+        self,
+        vertex_batch_request: VertexAIBatchPredictionJob,
+        headers: dict[str, str],
+        sync_handler: HTTPHandler,
+        vertex_location: str,
+    ) -> VertexAIBatchPredictionJob:
+        """
+        A fine-tuned Gemini deployment is configured by its endpoint id, but the v1 batch API only
+        accepts Model resources, so swap the endpoint resource for its deployed tuned model
+        (`projects/../locations/../models/<id>`) read from GET endpoints/<id>.
+        """
+        model: Final = vertex_batch_request.get("model", "")
+        if "/endpoints/" not in model:
+            return vertex_batch_request
+
+        endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}"
+        response: Final = sync_handler.get(url=endpoint_url, headers=headers)
+        if response.status_code != 200:
+            raise VertexAIError(
+                status_code=response.status_code,
+                message=f"Failed to resolve fine-tuned Vertex endpoint '{model}': {response.text}",
+            )
+
+        payload_view: Final[_VertexEndpointPayloadView] = {"payload": response.json()}
+        deployed_models: Final = payload_view["payload"].get("deployedModels") or []
+        deployed_model: Final = deployed_models[0].get("model", "") if deployed_models else ""
+        if not deployed_model:
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    f"Vertex endpoint '{model}' has no deployed model, so there is no tuned model "
+                    "resource to run batch predictions against"
+                ),
+            )
+        return {**vertex_batch_request, "model": deployed_model}
+
     async def _async_create_batch(
         self,
         vertex_batch_request: VertexAIBatchPredictionJob,
diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py
index f284b47292b..daedf6d1959 100644
--- a/litellm/llms/vertex_ai/batches/transformation.py
+++ b/litellm/llms/vertex_ai/batches/transformation.py
@@ -22,6 +22,8 @@ class VertexAIBatchTransformation:
     def transform_openai_batch_request_to_vertex_ai_batch_request(
         cls,
         request: CreateBatchRequest,
+        vertex_project: str | None = None,
+        vertex_location: str | None = None,
     ) -> VertexAIBatchPredictionJob:
         """
         Transforms OpenAI Batch requests to Vertex AI Batch requests
@@ -31,7 +33,11 @@ class VertexAIBatchTransformation:
         if input_file_id is None:
             raise ValueError("input_file_id is required, but not provided")
         input_config: InputConfig = InputConfig(gcsSource=GcsSource(uris=[input_file_id]), instancesFormat="jsonl")
-        model: Final[str] = cls._get_model_from_gcs_file(input_file_id)
+        model: Final[str] = cls._get_batch_job_model(
+            input_file_id=input_file_id,
+            vertex_project=vertex_project,
+            vertex_location=vertex_location,
+        )
         output_config: Final[OutputConfig] = OutputConfig(
             predictionsFormat="jsonl",
             gcsDestination=GcsDestination(outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id)),
@@ -188,6 +194,33 @@ class VertexAIBatchTransformation:
         path_parts: Final = input_file_id.rsplit("/", 1)
         return path_parts[0]
 
+    @classmethod
+    def _get_batch_job_model(
+        cls,
+        input_file_id: str,
+        vertex_project: str | None,
+        vertex_location: str | None,
+    ) -> str:
+        """
+        Returns the `model` for the batchPredictionJobs request: the publisher model path as-is, or
+        the full `projects/../locations/../endpoints/<id>` resource name for a fine-tuned endpoint.
+
+        The v1 batch API only accepts Model resources, so the handler resolves an endpoint resource
+        to its deployed tuned model (`projects/../locations/../models/<id>`) before sending the job.
+        """
+        parsed_model: Final = cls._get_model_from_gcs_file(input_file_id)
+        if not parsed_model.startswith("endpoints/"):
+            return parsed_model
+        if not vertex_project:
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    f"Vertex AI batch jobs against a fine-tuned endpoint ('{parsed_model}') require "
+                    "`vertex_project` to build the endpoint resource name"
+                ),
+            )
+        return f"projects/{vertex_project}/locations/{vertex_location or 'us-central1'}/{parsed_model}"
+
     @classmethod
     def _get_model_from_gcs_file(cls, gcs_file_uri: str) -> str:
         """
@@ -202,6 +235,9 @@ class VertexAIBatchTransformation:
         gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8
         returns: "publishers/google/models/gemini-1.5-flash-001"
 
+        Fine-tuned Gemini endpoints are stored as `endpoints/<numeric id>` in the uri and returned
+        in that form.
+
         Raises a 400 `VertexAIError` when the uri carries no parseable model path.
         """
         model: Final = cls._parse_model_from_gcs_file(gcs_file_uri)
@@ -210,11 +246,13 @@ class VertexAIBatchTransformation:
                 status_code=400,
                 message=(
                     "Vertex AI batch creation requires the model to be part of `input_file_id`, but "
-                    f"'{gcs_file_uri}' contains no 'publishers/<publisher>/models/<model>' path segment. "
+                    f"'{gcs_file_uri}' contains no 'publishers/<publisher>/models/<model>' or "
+                    "'endpoints/<numeric endpoint id>' path segment. "
                     "Either upload the input file through LiteLLM (POST /v1/files with "
                     "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or "
                     "pass a uri of the form "
-                    "gs://<bucket>/<prefix>/publishers/<publisher>/models/<model>/<file>"
+                    "gs://<bucket>/<prefix>/publishers/<publisher>/models/<model>/<file> "
+                    "(or gs://<bucket>/<prefix>/endpoints/<numeric endpoint id>/<file> for fine-tuned models)"
                 ),
             )
         return model
@@ -222,10 +260,16 @@ class VertexAIBatchTransformation:
     @classmethod
     def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None:
         """
-        Returns the `publishers/<publisher>/models/<model>` path from a gcs uri, or None if the uri
-        does not contain one.
+        Returns the `publishers/<publisher>/models/<model>` or `endpoints/<numeric id>` path from a
+        gcs uri, or None if the uri does not contain one.
         """
-        _, separator, model_path = unquote(gcs_file_uri).partition("publishers/")
+        unquoted_uri: Final = unquote(gcs_file_uri)
+        _, endpoint_separator, endpoint_path = unquoted_uri.partition("endpoints/")
+        endpoint_id: Final = endpoint_path.split("/")[0] if endpoint_separator else ""
+        if endpoint_id.isdigit():
+            return f"endpoints/{endpoint_id}"
+
+        _, separator, model_path = unquoted_uri.partition("publishers/")
         if not separator:
             return None
 
diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py
index 970759479fe..ba654f0d851 100644
--- a/litellm/llms/vertex_ai/common_utils.py
+++ b/litellm/llms/vertex_ai/common_utils.py
@@ -310,6 +310,19 @@ def get_vertex_base_model_name(model: str) -> str:
     return model
 
 
+def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None:
+    """
+    Fine-tuned Gemini deployments are addressed by a numeric endpoint id,
+    configured as `vertex_ai/<id>` or `vertex_ai/gemini/<id>`.
+
+    Returns the endpoint id, or None when `model` is a regular publisher model.
+    Mirrors the online chat path in `_get_vertex_url`, which sends numeric
+    models to `endpoints/{id}` instead of `publishers/google/models/{model}`.
+    """
+    candidate: Final = model.split("/")[-1] if "gemini/" in model else model
+    return candidate if candidate.isdigit() else None
+
+
 def validate_vertex_location(vertex_location: str | None) -> str:
     """
     Validate a Vertex AI location before interpolating it into a request host or
diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py
index b6ad9fbcc04..795918f1766 100644
--- a/litellm/llms/vertex_ai/files/transformation.py
+++ b/litellm/llms/vertex_ai/files/transformation.py
@@ -39,6 +39,7 @@ from litellm.llms.base_llm.files.transformation import (
 )
 from litellm.llms.vertex_ai.common_utils import (
     _convert_vertex_datetime_to_openai_datetime,
+    get_vertex_ai_fine_tuned_endpoint_id,
 )
 from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
@@ -712,11 +713,20 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
         Gets a unique GCS object name for the VertexAI batch prediction job
 
         named as: litellm-vertex-{model}-{uuid}
+
+        Fine-tuned Gemini deployments (numeric endpoint ids) are stored under
+        `endpoints/<id>` so the batch transformation can round-trip them into a
+        `projects/../locations/../endpoints/<id>` batch job model instead of a
+        nonexistent publisher model.
         """
-        _model = openai_jsonl_content[0].get("body", {}).get("model", "")
-        if "publishers/google/models" not in _model:
-            _model = f"publishers/google/models/{_model}"
-        safe_model_path: Final = sanitize_cloud_object_path(_model, fallback="model")
+        raw_model: Final = openai_jsonl_content[0].get("body", {}).get("model", "")
+        endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model)
+        model_path: Final = (
+            f"endpoints/{endpoint_id}"
+            if endpoint_id is not None
+            else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}")
+        )
+        safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model")
         object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}"
         return object_name
 
@@ -761,6 +771,16 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
         """
         Get the complete url for the request
         """
+        if data.get("purpose") == "batch" and litellm_params.get("custom_endpoint"):
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. "
+                    "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; "
+                    "remove this deployment from the batch request (e.g. `target_model_names`) or "
+                    "use a publisher model / fine-tuned Gemini endpoint instead."
+                ),
+            )
         bucket_name = self._get_configured_bucket_name(litellm_params)
         bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name)
         file_data: Final = data.get("file")
diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py
index ff5e8f89d64..b69805372ba 100644
--- a/tests/proxy_unit_tests/test_check_batch_cost.py
+++ b/tests/proxy_unit_tests/test_check_batch_cost.py
@@ -1760,6 +1760,35 @@ class TestUnmanagedVertexRouting:
         )
         router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash")
 
+    def test_flag_on_routes_fine_tuned_endpoint_to_vertex_deployment(self):
+        """A fine-tuned Gemini batch stores `endpoints/<id>` in the gs:// path; the bare model
+        (the endpoint id) must round-trip to the deployment configured as
+        `vertex_ai/gemini/<id>` (LIT-6899)."""
+        endpoint_id = "7768560373388541952"
+        router = MagicMock()
+        router.resolve_model_name_from_model_id.return_value = None
+        router.get_model_list.return_value = [
+            {
+                "model_name": "gemini-2.5-flash-dts-usc1",
+                "litellm_params": {
+                    "model": f"vertex_ai/gemini/{endpoint_id}",
+                    "custom_llm_provider": "vertex_ai",
+                },
+                "model_info": {"id": "deploy-ft"},
+            },
+        ]
+        instance = self._instance(track_unmanaged=True, router=router)
+        job = self._job(
+            file_object=_unmanaged_vertex_file_object(
+                input_file_id=f"gs://bucket/litellm-vertex-files/endpoints/{endpoint_id}/abc.jsonl"
+            )
+        )
+
+        with patch(_IS_B64, return_value=False):
+            result = instance._resolve_job_routing(job, MagicMock())
+
+        assert result == ("deploy-ft", "8823717160934178816")
+
     def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self):
         """Flag on, but the only deployment for the model group is a non-vertex_ai
         provider: must not be selected, even though the model group name matches."""
diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py
index c3edb40c819..b10214f884d 100644
--- a/tests/test_litellm/batches/test_main.py
+++ b/tests/test_litellm/batches/test_main.py
@@ -158,6 +158,16 @@ def test_create__vertex_ai_dispatch(seams):
     _assert_only(seams.vertex.create_batch, seams, "create_batch")
 
 
+def test_create__vertex_ai_custom_endpoint_raises_badrequest(seams):
+    """custom_endpoint deployments have no Vertex batch surface; creating a job would target a
+    nonexistent publisher model, so the SDK must 400 before dispatching (LIT-6899)."""
+    with pytest.raises(litellm.exceptions.BadRequestError, match="custom_endpoint"):
+        bm.create_batch(**CREATE_KW, custom_llm_provider="vertex_ai", custom_endpoint=True)
+
+    for m in _all_seam_methods(seams, "create_batch"):
+        m.assert_not_called()
+
+
 def test_create__provider_config_routes_to_base_http_handler(seams):
     """model + a provider batches config (bedrock-style) routes to the generic
     base_llm_http_handler, NOT the per-provider instance."""
diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
index 38fde3caa63..bb1c546614e 100644
--- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
+++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
@@ -178,6 +178,125 @@ def test_create_batch_async_returns_coroutine_and_uses_async_client():
     sync_client.post.assert_not_called()
 
 
+def test_create_batch_sync_does_not_resolve_publisher_models():
+    """Publisher-model jobs must not incur the endpoint-resolution GET."""
+    h = _make_handler()
+    client = MagicMock()
+    client.post.return_value = _http_response()
+
+    with patch(f"{HMOD}._get_httpx_client", return_value=client):
+        h.create_batch(
+            _is_async=False,
+            create_batch_data=CREATE_DATA,
+            api_base=None,
+            vertex_credentials=None,
+            vertex_project=PROJECT,
+            vertex_location=LOCATION,
+            timeout=600.0,
+            max_retries=None,
+        )
+
+    client.get.assert_not_called()
+
+
+ENDPOINT_ID = "7768560373388541952"
+ENDPOINT_CREATE_DATA = {
+    "input_file_id": f"gs://bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/file-uuid"
+}
+TUNED_MODEL_RESOURCE = f"projects/{PROJECT}/locations/{LOCATION}/models/1234509876"
+
+
+def _endpoint_get_response(deployed_models: list | None = None) -> MagicMock:
+    resp = MagicMock()
+    resp.status_code = 200
+    resp.json.return_value = {
+        "name": f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}",
+        "deployedModels": (
+            deployed_models if deployed_models is not None else [{"model": TUNED_MODEL_RESOURCE}]
+        ),
+    }
+    return resp
+
+
+def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model():
+    """A fine-tuned Gemini file id must produce a batch job against the endpoint's deployed
+    tuned model resource; the v1 batch API rejects endpoint resources in `model` (LIT-6899)."""
+    h = _make_handler()
+    client = MagicMock()
+    client.get.return_value = _endpoint_get_response()
+    client.post.return_value = _http_response()
+
+    with patch(f"{HMOD}._get_httpx_client", return_value=client):
+        out = h.create_batch(
+            _is_async=False,
+            create_batch_data=ENDPOINT_CREATE_DATA,
+            api_base=None,
+            vertex_credentials=None,
+            vertex_project=PROJECT,
+            vertex_location=LOCATION,
+            timeout=600.0,
+            max_retries=None,
+        )
+
+    assert isinstance(out, LiteLLMBatch)
+    get_kwargs = client.get.call_args.kwargs
+    assert get_kwargs["url"] == (
+        f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}"
+        f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}"
+    )
+    assert get_kwargs["headers"]["Authorization"] == f"Bearer {TOKEN}"
+    sent = json.loads(client.post.call_args.kwargs["data"])
+    assert sent["model"] == TUNED_MODEL_RESOURCE
+
+
+def test_create_batch_sync_endpoint_resolution_error_raises():
+    h = _make_handler()
+    client = MagicMock()
+    resolve_response = MagicMock()
+    resolve_response.status_code = 404
+    resolve_response.text = "endpoint not found"
+    client.get.return_value = resolve_response
+
+    with patch(f"{HMOD}._get_httpx_client", return_value=client):
+        with pytest.raises(VertexAIError) as exc_info:
+            h.create_batch(
+                _is_async=False,
+                create_batch_data=ENDPOINT_CREATE_DATA,
+                api_base=None,
+                vertex_credentials=None,
+                vertex_project=PROJECT,
+                vertex_location=LOCATION,
+                timeout=600.0,
+                max_retries=None,
+            )
+
+    assert exc_info.value.status_code == 404
+    client.post.assert_not_called()
+
+
+def test_create_batch_sync_endpoint_without_deployed_model_raises_400():
+    h = _make_handler()
+    client = MagicMock()
+    client.get.return_value = _endpoint_get_response(deployed_models=[])
+
+    with patch(f"{HMOD}._get_httpx_client", return_value=client):
+        with pytest.raises(VertexAIError) as exc_info:
+            h.create_batch(
+                _is_async=False,
+                create_batch_data=ENDPOINT_CREATE_DATA,
+                api_base=None,
+                vertex_credentials=None,
+                vertex_project=PROJECT,
+                vertex_location=LOCATION,
+                timeout=600.0,
+                max_retries=None,
+            )
+
+    assert exc_info.value.status_code == 400
+    assert "no deployed model" in str(exc_info.value)
+    client.post.assert_not_called()
+
+
 def test_create_batch_sync_httpstatuserror_propagates():
     """``HTTPHandler.post`` raises for non-2xx via ``raise_for_status``; the
     sync create path must surface that error, not swallow it."""
diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py
index ccb2d7e310d..72ab87bee1a 100644
--- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py
@@ -34,6 +34,12 @@ INPUT_FILE = (
     "models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8"
 )
 
+ENDPOINT_ID = "7768560373388541952"
+ENDPOINT_INPUT_FILE = (
+    f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/"
+    "e9412502-2c91-42a6-8e61-f5c294cc0fc8"
+)
+
 
 # =========================================================================== #
 # transform_openai_batch_request_to_vertex_ai_batch_request
@@ -67,6 +73,41 @@ def test_transform_openai_request_missing_input_file_id_raises():
         T.transform_openai_batch_request_to_vertex_ai_batch_request({})
 
 
+def test_transform_openai_request_fine_tuned_endpoint_builds_endpoint_resource():
+    """A fine-tuned Gemini file id (endpoints/<numeric id>) must target the endpoint resource,
+    not a nonexistent publisher model (LIT-6899)."""
+    job = T.transform_openai_batch_request_to_vertex_ai_batch_request(
+        {"input_file_id": ENDPOINT_INPUT_FILE},
+        vertex_project="my-project",
+        vertex_location="us-central1",
+    )
+    assert job["model"] == f"projects/my-project/locations/us-central1/endpoints/{ENDPOINT_ID}"
+
+
+def test_transform_openai_request_fine_tuned_endpoint_defaults_location():
+    job = T.transform_openai_batch_request_to_vertex_ai_batch_request(
+        {"input_file_id": ENDPOINT_INPUT_FILE},
+        vertex_project="my-project",
+    )
+    assert job["model"] == f"projects/my-project/locations/us-central1/endpoints/{ENDPOINT_ID}"
+
+
+def test_transform_openai_request_fine_tuned_endpoint_without_project_raises_400():
+    with pytest.raises(VertexAIError) as exc_info:
+        T.transform_openai_batch_request_to_vertex_ai_batch_request({"input_file_id": ENDPOINT_INPUT_FILE})
+    assert exc_info.value.status_code == 400
+    assert "vertex_project" in str(exc_info.value)
+
+
+def test_transform_openai_request_publisher_model_ignores_project_and_location():
+    job = T.transform_openai_batch_request_to_vertex_ai_batch_request(
+        {"input_file_id": INPUT_FILE},
+        vertex_project="my-project",
+        vertex_location="europe-west4",
+    )
+    assert job["model"] == "publishers/google/models/gemini-1.5-flash-001"
+
+
 @pytest.mark.parametrize(
     "input_file_id",
     [
@@ -321,6 +362,21 @@ def test_get_model_from_gcs_file_no_publishers_raises_400():
     assert exc_info.value.status_code == 400
 
 
+def test_get_model_from_gcs_file_fine_tuned_endpoint():
+    """The whole endpoint id must survive parsing; the old 3-segment publishers/ parse dropped it."""
+    assert T._get_model_from_gcs_file(ENDPOINT_INPUT_FILE) == f"endpoints/{ENDPOINT_ID}"
+
+
+def test_get_model_from_gcs_file_non_numeric_endpoints_segment_raises_400():
+    with pytest.raises(VertexAIError) as exc_info:
+        T._get_model_from_gcs_file("gs://bucket/endpoints/not-a-number/file-uuid")
+    assert exc_info.value.status_code == 400
+
+
+def test_get_bare_model_name_from_gcs_file_fine_tuned_endpoint():
+    assert T.get_bare_model_name_from_gcs_file(ENDPOINT_INPUT_FILE) == ENDPOINT_ID
+
+
 # =========================================================================== #
 # is_unmanaged_gcs_batch_input_file_id
 # =========================================================================== #
@@ -334,6 +390,8 @@ def test_get_model_from_gcs_file_no_publishers_raises_400():
         ("file-abc123", False),
         ("gs://bucket/no-model-here.jsonl", False),
         ("gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", False),
+        (ENDPOINT_INPUT_FILE, True),
+        ("gs://bucket/endpoints/not-a-number/file-uuid", False),
     ],
 )
 def test_is_unmanaged_gcs_batch_input_file_id(input_file_id, expected):
diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py
index 3c2d56997b7..f0ea61f183c 100644
--- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py
@@ -159,6 +159,61 @@ class TestCreateFileUrl:
         assert "?" not in object_name
 
 
+class TestBatchObjectNaming:
+    def test_should_store_publisher_model_under_publishers_path(self, config):
+        object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "gemini-2.5-flash"}}])
+        assert object_name.startswith("litellm-vertex-files/publishers/google/models/gemini-2.5-flash/")
+
+    def test_should_store_fine_tuned_endpoint_under_endpoints_path(self, config):
+        """A numeric endpoint id must not be filed under publishers/google/models/gemini/<id>,
+        which the batch transformation later mangles into a nonexistent publisher model (LIT-6899)."""
+        object_name = config._get_gcs_object_name_from_batch_jsonl(
+            [{"body": {"model": "gemini/7768560373388541952"}}]
+        )
+        assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/")
+        assert "publishers" not in object_name
+
+    def test_should_store_bare_numeric_endpoint_under_endpoints_path(self, config):
+        object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "7768560373388541952"}}])
+        assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/")
+
+
+class TestCustomEndpointBatchUpload:
+    def test_should_reject_batch_upload_for_custom_endpoint_deployment(self, config):
+        """custom_endpoint deployments have no Vertex batch surface; the upload must 400 instead
+        of staging a file that can only produce a doomed batch job (LIT-6899)."""
+        from litellm.llms.vertex_ai.common_utils import VertexAIError
+
+        with pytest.raises(VertexAIError) as exc_info:
+            config.get_complete_file_url(
+                api_base=None,
+                api_key=None,
+                model="",
+                optional_params={},
+                litellm_params={"gcs_bucket_name": "my-bucket", "custom_endpoint": True},
+                data={
+                    "file": ("batch.jsonl", b'{"body": {"model": "openai/gemma-2-2b-it"}}', "application/jsonl"),
+                    "purpose": "batch",
+                },
+            )
+        assert exc_info.value.status_code == 400
+        assert "custom_endpoint" in str(exc_info.value)
+
+    def test_should_allow_non_batch_upload_for_custom_endpoint_deployment(self, config):
+        url = config.get_complete_file_url(
+            api_base=None,
+            api_key=None,
+            model="",
+            optional_params={},
+            litellm_params={"gcs_bucket_name": "my-bucket", "custom_endpoint": True},
+            data={
+                "file": ("notes.txt", b"hello", "text/plain"),
+                "purpose": "assistants",
+            },
+        )
+        assert "/b/my-bucket/" in url
+
+
 class TestTransformRetrieveFile:
     def test_should_build_correct_gcs_metadata_url(self, config):
         file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl"

From bd9593b74d4a05e265225a43afa0fd6c8d837484 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 16:41:31 -0700
Subject: [PATCH 026/164] fix(proxy): share the Bedrock count-tokens predicate
 between auth and the passthrough handler

---
 litellm/proxy/auth/auth_utils.py                       |  7 ++++---
 .../llm_passthrough_endpoints.py                       |  7 +++++--
 tests/test_litellm/proxy/auth/test_auth_utils.py       | 10 ++++++++++
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py
index 12e6f75889e..1e4836654a1 100644
--- a/litellm/proxy/auth/auth_utils.py
+++ b/litellm/proxy/auth/auth_utils.py
@@ -1989,13 +1989,14 @@ def get_model_from_request(
 
 
 def _model_from_bedrock_route(route: str) -> str | None:
-    bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE)
-    if "count_tokens" in bedrock_endpoint.lower() or "count-tokens" in bedrock_endpoint.lower():
-        return None
     from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
         _extract_model_from_bedrock_endpoint,
+        is_bedrock_count_tokens_endpoint,
     )
 
+    bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE)
+    if is_bedrock_count_tokens_endpoint(bedrock_endpoint):
+        return None
     try:
         return _extract_model_from_bedrock_endpoint(bedrock_endpoint)
     except ValueError:
diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
index 29f216fd450..820312ac0fd 100644
--- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
@@ -703,6 +703,10 @@ BEDROCK_ENDPOINT_ACTIONS: Final = {
 BEDROCK_STREAMING_ACTIONS: Final = {"invoke-with-response-stream", "converse-stream"}
 
 
+def is_bedrock_count_tokens_endpoint(endpoint: str) -> bool:
+    return "count_tokens" in endpoint or "count-tokens" in endpoint
+
+
 def _extract_model_from_bedrock_endpoint(endpoint: str) -> str:
     """
     Extract model name from Bedrock endpoint path.
@@ -977,8 +981,7 @@ async def bedrock_llm_proxy_route(
 
     request_body: Final = await _read_request_body(request=request)
 
-    # Special handling for count_tokens endpoints
-    if "count_tokens" in endpoint or "count-tokens" in endpoint:
+    if is_bedrock_count_tokens_endpoint(endpoint):
         return await handle_bedrock_count_tokens(
             endpoint=endpoint,
             request=request,
diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py
index 61d1a308b9d..a996de4d40c 100644
--- a/tests/test_litellm/proxy/auth/test_auth_utils.py
+++ b/tests/test_litellm/proxy/auth/test_auth_utils.py
@@ -508,6 +508,16 @@ def test_get_model_from_request_bedrock_count_tokens_uses_body_model():
     )
 
 
+def test_get_model_from_request_bedrock_uppercase_count_tokens_segment_is_not_count_tokens():
+    assert (
+        get_model_from_request(
+            request_data={"model": "us.anthropic.claude-haiku-4-5-20251001-v1:0"},
+            route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke/COUNT_TOKENS",
+        )
+        == "us.anthropic.claude-sonnet-4-6"
+    )
+
+
 def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model():
     assert (
         get_model_from_request(

From 1d2ed0bdacc895b9c1e06cef2b97a4646f88f096 Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Thu, 3 Sep 2026 19:49:37 -0400
Subject: [PATCH 027/164] fix(vertex_ai): address batch review findings

Derive the GCS batch object path from the deployment's configured model
when present, so a user-crafted JSONL body.model cannot redirect an
authorized deployment's credentials to a different endpoint; the JSONL
value remains the fallback for direct SDK calls with no deployment
config.

Route the fine-tuned endpoint resolution GET through _check_custom_proxy
so custom api_base deployments do not contact Google directly.

Prefer the publisher model path over an endpoints/ segment when parsing
GCS uris, and use the last endpoints/ occurrence, so a bucket prefix
containing endpoints/<digits> cannot shadow the real model path.

Move the custom_endpoint rejection from the batches dispatcher into the
Vertex batch handler so the provider policy lives in the provider
module.
---
 litellm/batches/main.py                       | 16 +----
 litellm/llms/vertex_ai/batches/handler.py     | 64 +++++++++++++------
 .../llms/vertex_ai/batches/transformation.py  | 22 ++++---
 .../llms/vertex_ai/files/transformation.py    | 25 ++++++--
 tests/test_litellm/batches/test_main.py       | 12 ++--
 .../llms/vertex_ai/batches/test_handler.py    | 26 ++++++++
 .../vertex_ai/batches/test_transformation.py  | 14 ++++
 .../test_vertex_ai_files_transformation.py    | 30 +++++++++
 8 files changed, 152 insertions(+), 57 deletions(-)

diff --git a/litellm/batches/main.py b/litellm/batches/main.py
index 6cf9f070d9c..77a4fdebf16 100644
--- a/litellm/batches/main.py
+++ b/litellm/batches/main.py
@@ -301,21 +301,6 @@ def create_batch(
                 litellm_params=litellm_params,
             )
         elif custom_llm_provider == "vertex_ai":
-            if optional_params.get("custom_endpoint"):
-                raise litellm.exceptions.BadRequestError(
-                    message=(
-                        "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. "
-                        "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; "
-                        "use a publisher model or fine-tuned Gemini endpoint deployment instead."
-                    ),
-                    model=model or "n/a",
-                    llm_provider=custom_llm_provider,
-                    response=httpx.Response(
-                        status_code=400,
-                        content="custom_endpoint deployments do not support vertex_ai batches",
-                        request=httpx.Request(method="create_batch", url="https://github.com/BerriAI/litellm"),
-                    ),
-                )
             api_base = optional_params.api_base or ""
             vertex_ai_project: Final = (
                 optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
@@ -334,6 +319,7 @@ def create_batch(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 create_batch_data=_create_batch_request,
+                custom_endpoint=optional_params.get("custom_endpoint"),
             )
         else:
             raise litellm.exceptions.BadRequestError(
diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index 1bb8ceb747a..ae99eea777c 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -93,7 +93,17 @@ class VertexAIBatchPrediction(VertexLLM):
         vertex_location: str | None,
         timeout: float | httpx.Timeout,
         max_retries: int | None,
+        custom_endpoint: bool | None = None,
     ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]:
+        if custom_endpoint:
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. "
+                    "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; "
+                    "use a publisher model or fine-tuned Gemini endpoint deployment instead."
+                ),
+            )
         sync_handler: Final = _get_httpx_client()
 
         access_token, project_id = self._ensure_access_token(
@@ -102,6 +112,26 @@ class VertexAIBatchPrediction(VertexLLM):
             custom_llm_provider="vertex_ai",
         )
 
+        headers: Final = {
+            "Content-Type": "application/json; charset=utf-8",
+            "Authorization": f"Bearer {access_token}",
+        }
+
+        transformed_batch_request: Final[VertexAIBatchPredictionJob] = (
+            VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
+                request=create_batch_data,
+                vertex_project=vertex_project or project_id,
+                vertex_location=vertex_location or "us-central1",
+            )
+        )
+        vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model(
+            vertex_batch_request=transformed_batch_request,
+            headers=headers,
+            sync_handler=sync_handler,
+            api_base=api_base,
+            vertex_location=vertex_location or "us-central1",
+        )
+
         default_api_base: Final = self.create_vertex_batch_url(
             vertex_location=vertex_location or "us-central1",
             vertex_project=vertex_project or project_id,
@@ -126,25 +156,6 @@ class VertexAIBatchPrediction(VertexLLM):
             vertex_api_version="v1",
         )
 
-        headers: Final = {
-            "Content-Type": "application/json; charset=utf-8",
-            "Authorization": f"Bearer {access_token}",
-        }
-
-        transformed_batch_request: Final[VertexAIBatchPredictionJob] = (
-            VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
-                request=create_batch_data,
-                vertex_project=vertex_project or project_id,
-                vertex_location=vertex_location or "us-central1",
-            )
-        )
-        vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model(
-            vertex_batch_request=transformed_batch_request,
-            headers=headers,
-            sync_handler=sync_handler,
-            vertex_location=vertex_location or "us-central1",
-        )
-
         if _is_async is True:
             return self._async_create_batch(
                 vertex_batch_request=vertex_batch_request,
@@ -170,6 +181,7 @@ class VertexAIBatchPrediction(VertexLLM):
         vertex_batch_request: VertexAIBatchPredictionJob,
         headers: dict[str, str],
         sync_handler: HTTPHandler,
+        api_base: str | None,
         vertex_location: str,
     ) -> VertexAIBatchPredictionJob:
         """
@@ -181,7 +193,19 @@ class VertexAIBatchPrediction(VertexLLM):
         if "/endpoints/" not in model:
             return vertex_batch_request
 
-        endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}"
+        default_endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}"
+        _, endpoint_url = self._check_custom_proxy(
+            api_base=api_base,
+            custom_llm_provider="vertex_ai",
+            gemini_api_key=None,
+            endpoint=(default_endpoint_url.split(":")[-1] if len(default_endpoint_url.split(":")) > 1 else ""),
+            stream=None,
+            auth_header=None,
+            url=default_endpoint_url,
+            model=None,
+            vertex_location=vertex_location,
+            vertex_api_version="v1",
+        )
         response: Final = sync_handler.get(url=endpoint_url, headers=headers)
         if response.status_code != 200:
             raise VertexAIError(
diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py
index daedf6d1959..e63c80dd3cf 100644
--- a/litellm/llms/vertex_ai/batches/transformation.py
+++ b/litellm/llms/vertex_ai/batches/transformation.py
@@ -262,22 +262,24 @@ class VertexAIBatchTransformation:
         """
         Returns the `publishers/<publisher>/models/<model>` or `endpoints/<numeric id>` path from a
         gcs uri, or None if the uri does not contain one.
+
+        A publisher path wins over an `endpoints/` segment, and the last `endpoints/` occurrence is
+        used, so a user-configured bucket prefix that happens to contain `endpoints/<digits>` cannot
+        override the model path LiteLLM appended after it.
         """
         unquoted_uri: Final = unquote(gcs_file_uri)
-        _, endpoint_separator, endpoint_path = unquoted_uri.partition("endpoints/")
+        _, separator, model_path = unquoted_uri.partition("publishers/")
+        if separator:
+            parts: Final = model_path.split("/")
+            if len(parts) >= 3 and parts[1] == "models" and parts[2]:
+                return f"publishers/{'/'.join(parts[:3])}"
+
+        _, endpoint_separator, endpoint_path = unquoted_uri.rpartition("endpoints/")
         endpoint_id: Final = endpoint_path.split("/")[0] if endpoint_separator else ""
         if endpoint_id.isdigit():
             return f"endpoints/{endpoint_id}"
 
-        _, separator, model_path = unquoted_uri.partition("publishers/")
-        if not separator:
-            return None
-
-        parts: Final = model_path.split("/")
-        if len(parts) < 3 or parts[1] != "models" or not parts[2]:
-            return None
-
-        return f"publishers/{'/'.join(parts[:3])}"
+        return None
 
     @classmethod
     def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool:
diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py
index 795918f1766..263956efc9f 100644
--- a/litellm/llms/vertex_ai/files/transformation.py
+++ b/litellm/llms/vertex_ai/files/transformation.py
@@ -708,18 +708,28 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
     def _get_gcs_object_name_from_batch_jsonl(
         self,
         openai_jsonl_content: list[dict[str, Any]],
+        deployment_model: str | None = None,
     ) -> str:
         """
         Gets a unique GCS object name for the VertexAI batch prediction job
 
         named as: litellm-vertex-{model}-{uuid}
 
+        The stored model path decides which Vertex model the batch job later executes against, so
+        `deployment_model` (the deployment's own configured model) wins over the user-supplied
+        JSONL `body.model`; the JSONL value is only a fallback for direct SDK calls that carry no
+        deployment config.
+
         Fine-tuned Gemini deployments (numeric endpoint ids) are stored under
         `endpoints/<id>` so the batch transformation can round-trip them into a
         `projects/../locations/../endpoints/<id>` batch job model instead of a
         nonexistent publisher model.
         """
-        raw_model: Final = openai_jsonl_content[0].get("body", {}).get("model", "")
+        raw_model: Final = (
+            deployment_model.removeprefix("vertex_ai/")
+            if deployment_model
+            else openai_jsonl_content[0].get("body", {}).get("model", "")
+        )
         endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model)
         model_path: Final = (
             f"endpoints/{endpoint_id}"
@@ -730,7 +740,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
         object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}"
         return object_name
 
-    def get_object_name(self, file_data: FileTypes, purpose: str) -> str:
+    def get_object_name(self, file_data: FileTypes, purpose: str, deployment_model: str | None = None) -> str:
         """
         Get the object name for the request.
 
@@ -738,10 +748,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
         upload is never materialized just to derive the GCS object name.
         """
         if purpose == "batch":
-            ## 1. If jsonl, derive the object name from the first entry's model
+            ## 1. If jsonl, derive the object name from the deployment model (or the first entry's)
             first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None)
             if first_entry is not None:
-                return self._get_gcs_object_name_from_batch_jsonl([first_entry])
+                return self._get_gcs_object_name_from_batch_jsonl([first_entry], deployment_model=deployment_model)
 
         ## 2. If not jsonl, store under a server-generated managed object name
         filename, _ = extract_file_metadata(file_data)
@@ -789,7 +799,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
             raise ValueError("file is required")
         if purpose is None:
             raise ValueError("purpose is required")
-        object_name = self.get_object_name(file_data, purpose)
+        configured_model: Final = litellm_params.get("model")
+        object_name = self.get_object_name(
+            file_data,
+            purpose,
+            deployment_model=configured_model if isinstance(configured_model, str) else None,
+        )
         if object_prefix:
             object_name = f"{object_prefix}/{object_name}"
         encoded_object_name: Final = encode_gcs_object_name_for_url(object_name)
diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py
index b10214f884d..b87f9489250 100644
--- a/tests/test_litellm/batches/test_main.py
+++ b/tests/test_litellm/batches/test_main.py
@@ -158,14 +158,12 @@ def test_create__vertex_ai_dispatch(seams):
     _assert_only(seams.vertex.create_batch, seams, "create_batch")
 
 
-def test_create__vertex_ai_custom_endpoint_raises_badrequest(seams):
-    """custom_endpoint deployments have no Vertex batch surface; creating a job would target a
-    nonexistent publisher model, so the SDK must 400 before dispatching (LIT-6899)."""
-    with pytest.raises(litellm.exceptions.BadRequestError, match="custom_endpoint"):
-        bm.create_batch(**CREATE_KW, custom_llm_provider="vertex_ai", custom_endpoint=True)
+def test_create__vertex_ai_forwards_custom_endpoint(seams):
+    """The vertex handler owns the custom_endpoint batch rejection (LIT-6899), so the dispatcher
+    must forward the flag for the handler to act on."""
+    bm.create_batch(**CREATE_KW, custom_llm_provider="vertex_ai", custom_endpoint=True)
 
-    for m in _all_seam_methods(seams, "create_batch"):
-        m.assert_not_called()
+    assert seams.vertex.create_batch.call_args.kwargs["custom_endpoint"] is True
 
 
 def test_create__provider_config_routes_to_base_http_handler(seams):
diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
index bb1c546614e..8db4775a2bc 100644
--- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
+++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
@@ -274,6 +274,32 @@ def test_create_batch_sync_endpoint_resolution_error_raises():
     client.post.assert_not_called()
 
 
+def test_create_batch_custom_endpoint_raises_400_without_io():
+    """custom_endpoint deployments have no Vertex batch surface; creating a job would target a
+    nonexistent publisher model, so the handler must 400 before any auth or HTTP work (LIT-6899)."""
+    h = _make_handler()
+    client = MagicMock()
+
+    with patch(f"{HMOD}._get_httpx_client", return_value=client):
+        with pytest.raises(VertexAIError) as exc_info:
+            h.create_batch(
+                _is_async=False,
+                create_batch_data=CREATE_DATA,
+                api_base=None,
+                vertex_credentials=None,
+                vertex_project=PROJECT,
+                vertex_location=LOCATION,
+                timeout=600.0,
+                max_retries=None,
+                custom_endpoint=True,
+            )
+
+    assert exc_info.value.status_code == 400
+    assert "custom_endpoint" in str(exc_info.value)
+    h._ensure_access_token.assert_not_called()
+    client.post.assert_not_called()
+
+
 def test_create_batch_sync_endpoint_without_deployed_model_raises_400():
     h = _make_handler()
     client = MagicMock()
diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py
index 72ab87bee1a..232c6413e78 100644
--- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py
@@ -367,6 +367,20 @@ def test_get_model_from_gcs_file_fine_tuned_endpoint():
     assert T._get_model_from_gcs_file(ENDPOINT_INPUT_FILE) == f"endpoints/{ENDPOINT_ID}"
 
 
+def test_get_model_from_gcs_file_publisher_path_wins_over_endpoints_prefix():
+    """A bucket prefix containing endpoints/<digits> must not override the publisher model path
+    LiteLLM appended after it."""
+    uri = "gs://bucket/team-endpoints/999/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/uuid"
+    assert T._get_model_from_gcs_file(uri) == "publishers/google/models/gemini-1.5-flash-001"
+
+
+def test_get_model_from_gcs_file_last_endpoints_segment_wins():
+    """With no publisher path, the endpoint id closest to the file (last occurrence) is the one
+    LiteLLM stored; an earlier prefix segment must not shadow it."""
+    uri = f"gs://bucket/endpoints/999/litellm-vertex-files/endpoints/{ENDPOINT_ID}/uuid"
+    assert T._get_model_from_gcs_file(uri) == f"endpoints/{ENDPOINT_ID}"
+
+
 def test_get_model_from_gcs_file_non_numeric_endpoints_segment_raises_400():
     with pytest.raises(VertexAIError) as exc_info:
         T._get_model_from_gcs_file("gs://bucket/endpoints/not-a-number/file-uuid")
diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py
index f0ea61f183c..8a249820cbd 100644
--- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py
@@ -177,6 +177,36 @@ class TestBatchObjectNaming:
         object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "7768560373388541952"}}])
         assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/")
 
+    def test_deployment_model_overrides_jsonl_body_model(self, config):
+        """The stored path decides which Vertex model the batch later runs against with the
+        deployment's credentials, so a user-crafted JSONL body.model must not be able to redirect
+        an authorized deployment to a different endpoint."""
+        object_name = config._get_gcs_object_name_from_batch_jsonl(
+            [{"body": {"model": "9999999999999999999"}}],
+            deployment_model="vertex_ai/gemini/7768560373388541952",
+        )
+        assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/")
+        assert "9999999999999999999" not in object_name
+
+    def test_url_derives_object_path_from_configured_model(self, config):
+        url = config.get_complete_file_url(
+            api_base=None,
+            api_key=None,
+            model="",
+            optional_params={},
+            litellm_params={
+                "gcs_bucket_name": "my-bucket",
+                "model": "vertex_ai/gemini/7768560373388541952",
+            },
+            data={
+                "file": ("batch.jsonl", b'{"body": {"model": "9999999999999999999"}}', "application/jsonl"),
+                "purpose": "batch",
+            },
+        )
+        object_name = parse_qs(urlparse(url).query)["name"][0]
+        assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/")
+        assert "9999999999999999999" not in object_name
+
 
 class TestCustomEndpointBatchUpload:
     def test_should_reject_batch_upload_for_custom_endpoint_deployment(self, config):

From 14d5ff5c2190f2a89a6b1730ce71c15d2332596f Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Thu, 3 Sep 2026 19:58:37 -0400
Subject: [PATCH 028/164] fix(vertex_ai): keep new batch handler code within
 the mutable-collection budget

---
 litellm/llms/vertex_ai/batches/handler.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index ae99eea777c..7f20b197056 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -1,5 +1,5 @@
 import json
-from collections.abc import Coroutine
+from collections.abc import Coroutine, Sequence
 from typing import TYPE_CHECKING, Final, Protocol
 
 import httpx
@@ -61,7 +61,7 @@ class _VertexEndpointDeployedModel(TypedDict, total=False):
 
 
 class _VertexEndpointResponse(TypedDict, total=False):
-    deployedModels: ReadOnly[list[_VertexEndpointDeployedModel]]
+    deployedModels: ReadOnly[Sequence[_VertexEndpointDeployedModel]]
 
 
 class _VertexEndpointPayloadView(TypedDict):
@@ -179,7 +179,7 @@ class VertexAIBatchPrediction(VertexLLM):
     def _resolve_fine_tuned_endpoint_model(
         self,
         vertex_batch_request: VertexAIBatchPredictionJob,
-        headers: dict[str, str],
+        headers: dict[str, str],  # mutable-ok: HTTPHandler.get only accepts dict headers
         sync_handler: HTTPHandler,
         api_base: str | None,
         vertex_location: str,
@@ -214,7 +214,7 @@ class VertexAIBatchPrediction(VertexLLM):
             )
 
         payload_view: Final[_VertexEndpointPayloadView] = {"payload": response.json()}
-        deployed_models: Final = payload_view["payload"].get("deployedModels") or []
+        deployed_models: Final = payload_view["payload"].get("deployedModels") or ()
         deployed_model: Final = deployed_models[0].get("model", "") if deployed_models else ""
         if not deployed_model:
             raise VertexAIError(
@@ -224,7 +224,8 @@ class VertexAIBatchPrediction(VertexLLM):
                     "resource to run batch predictions against"
                 ),
             )
-        return {**vertex_batch_request, "model": deployed_model}
+        resolved_request: Final[VertexAIBatchPredictionJob] = {**vertex_batch_request, "model": deployed_model}
+        return resolved_request
 
     async def _async_create_batch(
         self,

From 988ae7ca80b66dd9930c1a32ca399d18086561fa Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Thu, 3 Sep 2026 20:04:05 -0400
Subject: [PATCH 029/164] test(vertex_ai): assert the publisher-model batch
 payload instead of only mock calls

---
 tests/test_litellm/llms/vertex_ai/batches/test_handler.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
index 8db4775a2bc..d65ba92afd7 100644
--- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
+++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
@@ -179,13 +179,14 @@ def test_create_batch_async_returns_coroutine_and_uses_async_client():
 
 
 def test_create_batch_sync_does_not_resolve_publisher_models():
-    """Publisher-model jobs must not incur the endpoint-resolution GET."""
+    """Publisher-model jobs must not incur the endpoint-resolution GET, and the job model must
+    stay the publisher path untouched."""
     h = _make_handler()
     client = MagicMock()
     client.post.return_value = _http_response()
 
     with patch(f"{HMOD}._get_httpx_client", return_value=client):
-        h.create_batch(
+        out = h.create_batch(
             _is_async=False,
             create_batch_data=CREATE_DATA,
             api_base=None,
@@ -196,6 +197,9 @@ def test_create_batch_sync_does_not_resolve_publisher_models():
             max_retries=None,
         )
 
+    assert isinstance(out, LiteLLMBatch)
+    sent = json.loads(client.post.call_args.kwargs["data"])
+    assert sent["model"] == "publishers/google/models/gemini-1.5-flash-001"
     client.get.assert_not_called()
 
 

From 54fd69beb2525878f35b622374a76c63713f5760 Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Thu, 3 Sep 2026 20:13:51 -0400
Subject: [PATCH 030/164] fix(vertex_ai): build a well-formed
 endpoint-resolution url for path-mounted custom api_base

---
 litellm/llms/vertex_ai/batches/handler.py     | 32 +++++++++++------
 .../llms/vertex_ai/batches/test_handler.py    | 34 ++++++++++++++++++-
 2 files changed, 55 insertions(+), 11 deletions(-)

diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index 7f20b197056..2d5afd4818e 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -1,6 +1,7 @@
 import json
 from collections.abc import Coroutine, Sequence
 from typing import TYPE_CHECKING, Final, Protocol
+from urllib.parse import urlparse
 
 import httpx
 from typing_extensions import ReadOnly, TypedDict
@@ -18,6 +19,7 @@ from litellm.llms.custom_httpx.http_handler import (
 )
 from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
+from litellm.llms.vertex_ai.vertex_llm_base import _graft_default_vertex_path
 from litellm.types.llms.openai import CreateBatchRequest
 from litellm.types.llms.vertex_ai import (
     VERTEX_CREDENTIALS_TYPES,
@@ -176,6 +178,24 @@ class VertexAIBatchPrediction(VertexLLM):
         )
         return vertex_batch_response
 
+    @staticmethod
+    def _build_endpoint_resolution_url(api_base: str | None, model: str, vertex_location: str) -> str:
+        """
+        Builds the GET url for resolving an endpoint resource (`projects/../endpoints/<id>`).
+
+        A custom `api_base` replaces the Google host: its `/v1`/`/v1beta1` path swallows the
+        version segment (matching `_check_custom_proxy`'s grafting), any other path is kept as a
+        mount prefix in front of the full default path. The `:operation` suffix convention from
+        `_check_custom_proxy` does not apply to a plain resource GET.
+        """
+        default_endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}"
+        if not api_base:
+            return default_endpoint_url
+        api_base_path: Final = urlparse(api_base).path.rstrip("/")
+        if api_base_path in ("/v1", "/v1beta1"):
+            return _graft_default_vertex_path(api_base=api_base, default_url=default_endpoint_url)
+        return api_base.rstrip("/") + urlparse(default_endpoint_url).path
+
     def _resolve_fine_tuned_endpoint_model(
         self,
         vertex_batch_request: VertexAIBatchPredictionJob,
@@ -193,18 +213,10 @@ class VertexAIBatchPrediction(VertexLLM):
         if "/endpoints/" not in model:
             return vertex_batch_request
 
-        default_endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}"
-        _, endpoint_url = self._check_custom_proxy(
+        endpoint_url: Final = self._build_endpoint_resolution_url(
             api_base=api_base,
-            custom_llm_provider="vertex_ai",
-            gemini_api_key=None,
-            endpoint=(default_endpoint_url.split(":")[-1] if len(default_endpoint_url.split(":")) > 1 else ""),
-            stream=None,
-            auth_header=None,
-            url=default_endpoint_url,
-            model=None,
+            model=model,
             vertex_location=vertex_location,
-            vertex_api_version="v1",
         )
         response: Final = sync_handler.get(url=endpoint_url, headers=headers)
         if response.status_code != 200:
diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
index d65ba92afd7..bba0eb27ed4 100644
--- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
+++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
@@ -35,7 +35,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import httpx
 import pytest
 
-
 from litellm.llms.vertex_ai.batches.handler import (  # noqa: E402
     VertexAIBatchPrediction,
 )
@@ -253,6 +252,39 @@ def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model():
     assert sent["model"] == TUNED_MODEL_RESOURCE
 
 
+@pytest.mark.parametrize(
+    "api_base, expected",
+    [
+        (
+            None,
+            f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}"
+            f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}",
+        ),
+        (
+            "https://proxy.internal",
+            f"https://proxy.internal/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}",
+        ),
+        (
+            "https://proxy.internal/v1",
+            f"https://proxy.internal/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}",
+        ),
+        (
+            "https://proxy.internal/vertex",
+            f"https://proxy.internal/vertex/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}",
+        ),
+    ],
+)
+def test_build_endpoint_resolution_url(api_base, expected):
+    """A custom api_base must replace the Google host for the endpoint-resolution GET without
+    producing a malformed url (no ':' grafting, no doubled /v1)."""
+    url = VertexAIBatchPrediction._build_endpoint_resolution_url(
+        api_base=api_base,
+        model=f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}",
+        vertex_location=LOCATION,
+    )
+    assert url == expected
+
+
 def test_create_batch_sync_endpoint_resolution_error_raises():
     h = _make_handler()
     client = MagicMock()

From 18373f8f51c3e283d82a0f480dee51d84dc0e9d0 Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Thu, 3 Sep 2026 20:38:24 -0400
Subject: [PATCH 031/164] fix(vertex_ai): route endpoint resolution through
 safe_get to guard caller-supplied api_base

---
 litellm/llms/vertex_ai/batches/handler.py     | 12 +++++++-
 .../llms/vertex_ai/batches/test_handler.py    | 29 ++++++++++++-------
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index 2d5afd4818e..a15ea4d845b 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -218,7 +218,17 @@ class VertexAIBatchPrediction(VertexLLM):
             model=model,
             vertex_location=vertex_location,
         )
-        response: Final = sync_handler.get(url=endpoint_url, headers=headers)
+        # ``api_base`` can come from caller-supplied request kwargs, so wrap the
+        # fetch in ``safe_get``: it rejects DNS-rebind / private / cloud-metadata
+        # targets before the bearer token leaves the process (mirrors retrieve_batch).
+        fetched: Final[_FetchedResponseView] = {
+            "response": safe_get(
+                sync_handler,
+                endpoint_url,
+                headers=headers,
+            )
+        }
+        response: Final = fetched["response"]
         if response.status_code != 200:
             raise VertexAIError(
                 status_code=response.status_code,
diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
index bba0eb27ed4..24df3214da3 100644
--- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
+++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py
@@ -184,7 +184,10 @@ def test_create_batch_sync_does_not_resolve_publisher_models():
     client = MagicMock()
     client.post.return_value = _http_response()
 
-    with patch(f"{HMOD}._get_httpx_client", return_value=client):
+    with (
+        patch(f"{HMOD}._get_httpx_client", return_value=client),
+        patch(f"{HMOD}.safe_get") as safe_get,
+    ):
         out = h.create_batch(
             _is_async=False,
             create_batch_data=CREATE_DATA,
@@ -199,7 +202,7 @@ def test_create_batch_sync_does_not_resolve_publisher_models():
     assert isinstance(out, LiteLLMBatch)
     sent = json.loads(client.post.call_args.kwargs["data"])
     assert sent["model"] == "publishers/google/models/gemini-1.5-flash-001"
-    client.get.assert_not_called()
+    safe_get.assert_not_called()
 
 
 ENDPOINT_ID = "7768560373388541952"
@@ -226,10 +229,12 @@ def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model():
     tuned model resource; the v1 batch API rejects endpoint resources in `model` (LIT-6899)."""
     h = _make_handler()
     client = MagicMock()
-    client.get.return_value = _endpoint_get_response()
     client.post.return_value = _http_response()
 
-    with patch(f"{HMOD}._get_httpx_client", return_value=client):
+    with (
+        patch(f"{HMOD}._get_httpx_client", return_value=client),
+        patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response()) as safe_get,
+    ):
         out = h.create_batch(
             _is_async=False,
             create_batch_data=ENDPOINT_CREATE_DATA,
@@ -242,8 +247,8 @@ def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model():
         )
 
     assert isinstance(out, LiteLLMBatch)
-    get_kwargs = client.get.call_args.kwargs
-    assert get_kwargs["url"] == (
+    get_args, get_kwargs = safe_get.call_args
+    assert get_args[1] == (
         f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}"
         f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}"
     )
@@ -291,9 +296,11 @@ def test_create_batch_sync_endpoint_resolution_error_raises():
     resolve_response = MagicMock()
     resolve_response.status_code = 404
     resolve_response.text = "endpoint not found"
-    client.get.return_value = resolve_response
 
-    with patch(f"{HMOD}._get_httpx_client", return_value=client):
+    with (
+        patch(f"{HMOD}._get_httpx_client", return_value=client),
+        patch(f"{HMOD}.safe_get", return_value=resolve_response),
+    ):
         with pytest.raises(VertexAIError) as exc_info:
             h.create_batch(
                 _is_async=False,
@@ -339,9 +346,11 @@ def test_create_batch_custom_endpoint_raises_400_without_io():
 def test_create_batch_sync_endpoint_without_deployed_model_raises_400():
     h = _make_handler()
     client = MagicMock()
-    client.get.return_value = _endpoint_get_response(deployed_models=[])
 
-    with patch(f"{HMOD}._get_httpx_client", return_value=client):
+    with (
+        patch(f"{HMOD}._get_httpx_client", return_value=client),
+        patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response(deployed_models=[])),
+    ):
         with pytest.raises(VertexAIError) as exc_info:
             h.create_batch(
                 _is_async=False,

From 533a1c959c21f659a78fb130b17c700a3d6d1a06 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 17:42:46 -0700
Subject: [PATCH 032/164] fix(proxy): reserve budget for the Bedrock Converse
 prompt, not the context window

Resolving the model from the /bedrock path sends passthrough calls through
optimistic budget reservation, whose tokenizer cannot walk Converse content
blocks and so fell back to the model's max_input_tokens. Count those messages
as text and read inferenceConfig.maxTokens so a budgeted key reserves the
request's cost.
---
 .../spend_tracking/budget_reservation.py      | 37 ++++++++++++-------
 .../spend_tracking/test_budget_reservation.py | 31 +++++++++++++++-
 2 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py
index 91d2ece7a51..5ed52a327d4 100644
--- a/litellm/proxy/spend_tracking/budget_reservation.py
+++ b/litellm/proxy/spend_tracking/budget_reservation.py
@@ -1362,12 +1362,15 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int:
 def _count_input_tokens(request_body: dict, model: str) -> int | None:
     try:
         if "messages" in request_body:
-            return litellm.token_counter(
-                model=model,
-                messages=request_body.get("messages") or [],
-                tools=request_body.get("tools"),
-                tool_choice=request_body.get("tool_choice"),
-            )
+            try:
+                return litellm.token_counter(
+                    model=model,
+                    messages=request_body.get("messages") or (),
+                    tools=request_body.get("tools"),
+                    tool_choice=request_body.get("tool_choice"),
+                )
+            except ValueError:
+                return _count_text_tokens(model=model, text=request_body.get("messages"))
         if "prompt" in request_body:
             return _count_text_tokens(model=model, text=request_body.get("prompt"))
         if "input" in request_body:
@@ -1415,11 +1418,7 @@ def _estimate_output_tokens(
     if _is_input_only_route(route=route):
         return 0
 
-    requested: int | None = None
-    for key in ("max_completion_tokens", "max_tokens", "max_output_tokens"):
-        requested = _to_int(request_body.get(key))
-        if requested is not None:
-            break
+    requested: Final = _requested_output_tokens(request_body)
 
     # Clamp at min(requested-or-default, model_max-or-default). Two purposes:
     # (1) Without an explicit cap we still need a finite reservation so the
@@ -1430,9 +1429,19 @@ def _estimate_output_tokens(
     #     at the cap — the model can only physically emit max_output_tokens
     #     anyway, so reserving more is both wasteful and a DoS surface.
     model_ceiling: Final = _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK
-    if requested is None:
-        requested = DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK
-    return min(requested, model_ceiling)
+    return min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK if requested is None else requested, model_ceiling)
+
+
+_OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens")
+
+
+def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None:
+    inference_config: Final = request_body.get("inferenceConfig")
+    candidates: Final = (
+        *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS),
+        inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None,
+    )
+    return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None)
 
 
 def _count_text_tokens(model: str, text: object) -> int:
diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
index f65f68812a2..c7de8c943ad 100644
--- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
+++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
@@ -5,7 +5,7 @@ import pytest
 from litellm.caching import DualCache
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
-from litellm.proxy.spend_tracking.budget_reservation import reserve_budget_for_request
+from litellm.proxy.spend_tracking.budget_reservation import estimate_request_max_cost, reserve_budget_for_request
 from litellm.proxy.utils import ProxyLogging
 
 TOKEN_COUNTING_ROUTES: Final = (
@@ -46,3 +46,32 @@ async def test_non_exempt_llm_route_still_reserves_budget():
 
     assert reservation is not None
     assert reservation["reserved_cost"] > 0
+
+
+BEDROCK_SONNET: Final = "us.anthropic.claude-sonnet-4-6"
+CONVERSE_BODY: Final = {
+    "messages": [{"role": "user", "content": [{"text": "Reply with one word: pong"}]}],
+    "inferenceConfig": {"maxTokens": 5},
+}
+INVOKE_BODY: Final = {
+    "anthropic_version": "bedrock-2023-05-31",
+    "max_tokens": 5,
+    "messages": [{"role": "user", "content": "Reply with one word: pong"}],
+}
+
+
+def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window():
+    converse_cost: Final = estimate_request_max_cost(
+        request_body=CONVERSE_BODY,
+        route=f"/bedrock/model/{BEDROCK_SONNET}/converse",
+        llm_router=None,
+        input_token_counts={},
+    )
+    invoke_cost: Final = estimate_request_max_cost(
+        request_body=INVOKE_BODY,
+        route=f"/bedrock/model/{BEDROCK_SONNET}/invoke",
+        llm_router=None,
+        input_token_counts={},
+    )
+    assert converse_cost is not None and invoke_cost is not None
+    assert invoke_cost < converse_cost < 2 * invoke_cost

From f9a5f676a866012c24427fb8afdfa240031f70d9 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 13:16:39 -0700
Subject: [PATCH 033/164] docs(claude): have runs embed their own QA
 screenshots on visual changes

---
 CLAUDE.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CLAUDE.md b/CLAUDE.md
index d9e9e8f1586..9c525ffe420 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -37,7 +37,7 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
 
 If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
 
-Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
+Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
 
 If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
 - don't use emojis

From 4b950cd94f86dbe174d515a468cb9b5288bf0ab8 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 17:01:14 -0700
Subject: [PATCH 034/164] feat(fireworks_ai): add native Responses API config

---
 basedpyright-code-budget.json                 |   4 +-
 litellm/__init__.py                           |   3 +
 litellm/_lazy_imports_registry.py             |   5 +
 litellm/llms/fireworks_ai/common_utils.py     |  44 +--
 .../fireworks_ai/responses/transformation.py  |  68 +++++
 litellm/utils.py                              |   2 +
 ...t_fireworks_ai_responses_transformation.py | 263 ++++++++++++++++++
 .../test_responses_websocket_all_providers.py |   9 +
 type-discipline-budget.json                   |   4 +-
 9 files changed, 382 insertions(+), 20 deletions(-)
 create mode 100644 litellm/llms/fireworks_ai/responses/transformation.py
 create mode 100644 tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py

diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json
index 9b59480a0dc..e7f8aacf835 100644
--- a/basedpyright-code-budget.json
+++ b/basedpyright-code-budget.json
@@ -105,13 +105,13 @@
     "limit": 109
   },
   "reportUnknownMemberType": {
-    "limit": 38309
+    "limit": 38307
   },
   "reportUnknownParameterType": {
     "limit": 19622
   },
   "reportUnknownVariableType": {
-    "limit": 29846
+    "limit": 29840
   },
   "reportUnnecessaryCast": {
     "limit": 111
diff --git a/litellm/__init__.py b/litellm/__init__.py
index 42c0ea881fd..fdc4435e5ff 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -2001,6 +2001,9 @@ if TYPE_CHECKING:
     from .llms.hosted_vllm.responses.transformation import (
         HostedVLLMResponsesAPIConfig as HostedVLLMResponsesAPIConfig,
     )
+    from .llms.fireworks_ai.responses.transformation import (
+        FireworksAIResponsesAPIConfig as FireworksAIResponsesAPIConfig,
+    )
     from .llms.github_copilot.chat.transformation import (
         GithubCopilotConfig as GithubCopilotConfig,
     )
diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py
index e9199e1ec80..dc323c8cc15 100644
--- a/litellm/_lazy_imports_registry.py
+++ b/litellm/_lazy_imports_registry.py
@@ -237,6 +237,7 @@ LLM_CONFIG_NAMES: Final = (
     "XAIResponsesAPIConfig",
     "LiteLLMProxyResponsesAPIConfig",
     "HostedVLLMResponsesAPIConfig",
+    "FireworksAIResponsesAPIConfig",
     "VolcEngineResponsesAPIConfig",
     "PerplexityResponsesConfig",
     "DatabricksResponsesAPIConfig",
@@ -957,6 +958,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
         ".llms.hosted_vllm.responses.transformation",
         "HostedVLLMResponsesAPIConfig",
     ),
+    "FireworksAIResponsesAPIConfig": (
+        ".llms.fireworks_ai.responses.transformation",
+        "FireworksAIResponsesAPIConfig",
+    ),
     "VolcEngineResponsesAPIConfig": (
         ".llms.volcengine.responses.transformation",
         "VolcEngineResponsesAPIConfig",
diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py
index ac934ad0cb5..21a630a76d7 100644
--- a/litellm/llms/fireworks_ai/common_utils.py
+++ b/litellm/llms/fireworks_ai/common_utils.py
@@ -1,3 +1,5 @@
+from collections.abc import Mapping
+from types import MappingProxyType
 from typing import Final
 
 from httpx import Headers
@@ -13,7 +15,7 @@ class FireworksAIException(BaseLLMException):
     pass
 
 
-def get_fireworks_session_id(litellm_params: dict) -> str | None:
+def get_fireworks_session_id(litellm_params: Mapping[str, object]) -> str | None:
     """
     Session id to send as `x-session-affinity`, or None when the caller gave none.
 
@@ -23,19 +25,39 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None:
     """
     params: Final = litellm_params
     metadata: Final = params.get("metadata")
-    if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
+    if isinstance(metadata, Mapping) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
         return None
     for key in ("litellm_session_id", "session_id"):
         value = params.get(key)
         if value:
             return str(value)
-    if isinstance(metadata, dict):
+    if isinstance(metadata, Mapping):
         value = metadata.get("session_id")
         if value:
             return str(value)
     return None
 
 
+def with_fireworks_session_affinity(
+    headers: Mapping[str, str], litellm_params: Mapping[str, object]
+) -> Mapping[str, str]:
+    if any(key.lower() == "x-session-affinity" for key in headers):
+        return headers
+    session_id: Final = get_fireworks_session_id(litellm_params)
+    if not session_id:
+        return headers
+    return MappingProxyType({**headers, "x-session-affinity": session_id})
+
+
+def resolve_fireworks_api_key(api_key: str | None) -> str | None:
+    return api_key or (
+        get_secret_str("FIREWORKS_API_KEY")
+        or get_secret_str("FIREWORKS_AI_API_KEY")
+        or get_secret_str("FIREWORKSAI_API_KEY")
+        or get_secret_str("FIREWORKS_AI_TOKEN")
+    )
+
+
 AZURE_FOUNDRY_FIREWORKS_MODEL_ID_PREFIX: Final = "FW-"
 
 
@@ -63,13 +85,7 @@ class FireworksAIMixin:
         )
 
     def _get_api_key(self, api_key: str | None) -> str | None:
-        dynamic_api_key: Final = api_key or (
-            get_secret_str("FIREWORKS_API_KEY")
-            or get_secret_str("FIREWORKS_AI_API_KEY")
-            or get_secret_str("FIREWORKSAI_API_KEY")
-            or get_secret_str("FIREWORKS_AI_TOKEN")
-        )
-        return dynamic_api_key
+        return resolve_fireworks_api_key(api_key)
 
     def validate_environment(
         self,
@@ -92,9 +108,5 @@ class FireworksAIMixin:
         return self._add_session_affinity_header({**auth_headers, **content_type_header}, litellm_params)
 
     def _add_session_affinity_header(self, headers: dict, litellm_params: dict) -> dict:
-        if any(key.lower() == "x-session-affinity" for key in headers):
-            return headers
-        session_id: Final = get_fireworks_session_id(litellm_params)
-        if not session_id:
-            return headers
-        return {**headers, "x-session-affinity": session_id}
+        pinned: Final = with_fireworks_session_affinity(headers, litellm_params)
+        return dict(pinned)  # mutable-ok: the HTTP handler updates the returned headers in place
diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py
new file mode 100644
index 00000000000..f36030bb50a
--- /dev/null
+++ b/litellm/llms/fireworks_ai/responses/transformation.py
@@ -0,0 +1,68 @@
+from collections.abc import Mapping
+from types import MappingProxyType
+from typing import Final
+
+from litellm.llms.fireworks_ai.common_utils import (
+    resolve_fireworks_api_key,
+    resolve_fireworks_resource_name,
+    with_fireworks_session_affinity,
+)
+from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.llms.openai import ResponseInputParam
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+FIREWORKS_AI_DEFAULT_API_BASE: Final = "https://api.fireworks.ai/inference/v1"
+
+
+def _session_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, object]:
+    extras: Final[Mapping[str, object]] = litellm_params.model_extra or MappingProxyType({})
+    return MappingProxyType(
+        {"litellm_session_id": extras.get("litellm_session_id"), "metadata": extras.get("litellm_metadata")}
+    )
+
+
+class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
+    @property
+    def custom_llm_provider(self) -> LlmProviders:
+        return LlmProviders.FIREWORKS_AI
+
+    def validate_environment(
+        self,
+        headers: Mapping[str, str],
+        model: str,
+        litellm_params: GenericLiteLLMParams | None,
+    ) -> dict:  # mutable-ok: overrides the base class signature
+        params: Final = litellm_params or GenericLiteLLMParams()
+        api_key: Final = resolve_fireworks_api_key(params.api_key)
+        if api_key is None:
+            raise ValueError("FIREWORKS_API_KEY is not set")
+        authorized: Final = MappingProxyType(
+            {"Content-Type": "application/json", **headers, "Authorization": f"Bearer {api_key}"}
+        )
+        pinned: Final = with_fireworks_session_affinity(authorized, _session_params(params))
+        return dict(pinned)  # mutable-ok: the HTTP handler updates the returned headers in place
+
+    def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str:
+        base: Final = (api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE).rstrip("/")
+        return f"{base}/responses"
+
+    def transform_responses_api_request(
+        self,
+        model: str,
+        input: str | ResponseInputParam,
+        response_api_optional_request_params: dict,  # mutable-ok: overrides the base class signature
+        litellm_params: GenericLiteLLMParams,
+        headers: dict,  # mutable-ok: overrides the base class signature
+    ) -> dict:  # mutable-ok: overrides the base class signature
+        return super().transform_responses_api_request(
+            model=resolve_fireworks_resource_name(model),
+            input=input,
+            response_api_optional_request_params=response_api_optional_request_params,
+            litellm_params=litellm_params,
+            headers=headers,
+        )
+
+    def supports_native_websocket(self) -> bool:
+        return False
diff --git a/litellm/utils.py b/litellm/utils.py
index 8b1b32ea328..3043a502aaa 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -8683,6 +8683,8 @@ class ProviderConfigManager:
             return litellm.OpenRouterResponsesAPIConfig()
         elif litellm.LlmProviders.HOSTED_VLLM == provider:
             return litellm.HostedVLLMResponsesAPIConfig()
+        elif litellm.LlmProviders.FIREWORKS_AI == provider:
+            return litellm.FireworksAIResponsesAPIConfig()
         elif litellm.LlmProviders.BEDROCK_MANTLE == provider:
             # Both decisions are data-driven from the model's price-map entry, with
             # no model-name logic. Capability (can it serve Responses?) comes from
diff --git a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
new file mode 100644
index 00000000000..207fd518f89
--- /dev/null
+++ b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
@@ -0,0 +1,263 @@
+import json
+from collections.abc import Mapping
+from types import MappingProxyType
+from typing import Final, TypedDict
+from unittest.mock import MagicMock, patch
+
+import httpx
+import pytest
+from openai.types.responses import (
+    ResponseFunctionToolCall,
+    ResponseOutputMessage,
+    ResponseOutputText,
+    ResponseReasoningItem,
+)
+from openai.types.responses.response_input_param import FunctionCallOutput
+from openai.types.responses.tool_param import Mcp
+from typing_extensions import ReadOnly
+
+import litellm
+from litellm.llms.fireworks_ai.responses.transformation import FireworksAIResponsesAPIConfig
+from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponsesAPIResponse
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+from litellm.utils import ProviderConfigManager
+
+FIREWORKS_RESPONSES_URL: Final = "https://api.fireworks.ai/inference/v1/responses"
+HTTPX_CLIENT_FACTORY: Final = "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client"
+NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({})
+NO_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
+
+
+class _GeneratedSessionMetadata(TypedDict):
+    litellm_session_id_generated: ReadOnly[bool]
+
+
+def _fireworks_response(model: str) -> Mapping[str, object]:
+    return ResponsesAPIResponse(
+        id="resp_0e946f2d46bf4b49bf8b29ff78083583",
+        object="response",
+        created_at=1788550000,
+        model=model,
+        status="completed",
+        output=(
+            ResponseReasoningItem(id="rs_1", summary=(), type="reasoning"),
+            ResponseOutputMessage(
+                id="msg_1",
+                status="completed",
+                role="assistant",
+                type="message",
+                content=(ResponseOutputText(type="output_text", text="Paris is clear and 21C.", annotations=()),),
+            ),
+            ResponseFunctionToolCall(
+                id="fc_1",
+                call_id="call_abc123",
+                name="get_weather",
+                arguments='{"city": "Paris"}',
+                status="completed",
+                type="function_call",
+            ),
+        ),
+        usage=ResponseAPIUsage(
+            input_tokens=179,
+            output_tokens=100,
+            total_tokens=279,
+            input_tokens_details=InputTokensDetails(cached_tokens=0),
+        ),
+    ).model_dump(mode="json", exclude_none=True)
+
+
+def _mock_http_client(response_body: Mapping[str, object]) -> MagicMock:
+    client: Final = MagicMock()
+    response: Final = MagicMock()
+    response.status_code = 200
+    response.headers = httpx.Headers((("content-type", "application/json"),))
+    response.json.return_value = response_body
+    response.text = json.dumps(response_body)
+    client.post.return_value = response
+    return client
+
+
+def _sent_request(client: MagicMock) -> tuple[str, Mapping[str, str], Mapping[str, object]]:
+    kwargs: Final = client.post.call_args.kwargs
+    body: Final = kwargs["json"] if "json" in kwargs else json.loads(kwargs["data"])
+    return kwargs["url"], kwargs["headers"], body
+
+
+@pytest.fixture(autouse=True)
+def fireworks_env(monkeypatch: pytest.MonkeyPatch) -> None:
+    for name in (
+        "FIREWORKS_API_KEY",
+        "FIREWORKS_AI_API_KEY",
+        "FIREWORKSAI_API_KEY",
+        "FIREWORKS_AI_TOKEN",
+        "FIREWORKS_API_BASE",
+    ):
+        monkeypatch.delenv(name, raising=False)
+
+
+def test_fireworks_ai_provider_config_registration() -> None:
+    config: Final = ProviderConfigManager.get_provider_responses_api_config(
+        model="accounts/fireworks/models/kimi-k3", provider=LlmProviders.FIREWORKS_AI
+    )
+    assert isinstance(config, FireworksAIResponsesAPIConfig)
+    assert config.custom_llm_provider == LlmProviders.FIREWORKS_AI
+
+
+def test_responses_call_hits_native_endpoint_with_mcp_tool_untouched() -> None:
+    client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
+    mcp_tool: Final[Mcp] = {
+        "type": "mcp",
+        "server_label": "deepwiki",
+        "server_url": "https://mcp.deepwiki.com/mcp",
+        "require_approval": "never",
+    }
+    with patch(HTTPX_CLIENT_FACTORY, return_value=client):
+        response: Final = litellm.responses(
+            model="fireworks_ai/accounts/fireworks/models/kimi-k3",
+            input="What is litellm?",
+            tools=[mcp_tool],  # mutable-ok: the Responses API takes tools as a JSON list
+            api_key="fw-test-key",
+        )
+    url, headers, body = _sent_request(client)
+    assert url == FIREWORKS_RESPONSES_URL
+    assert headers["Authorization"] == "Bearer fw-test-key"
+    assert body["model"] == "accounts/fireworks/models/kimi-k3"
+    assert tuple(body["tools"]) == (mcp_tool,)
+    assert "messages" not in body
+    assert isinstance(response, ResponsesAPIResponse)
+    function_calls: Final = tuple(item for item in response.output if getattr(item, "type", None) == "function_call")
+    assert getattr(function_calls[0], "call_id", None) == "call_abc123"
+
+
+def test_responses_call_expands_bare_model_name_to_fireworks_resource() -> None:
+    client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/glm-5p3"))
+    with patch(HTTPX_CLIENT_FACTORY, return_value=client):
+        litellm.responses(model="fireworks_ai/glm-5p3", input="hi", api_key="fw-test-key")
+    _, _, body = _sent_request(client)
+    assert body["model"] == "accounts/fireworks/models/glm-5p3"
+
+
+def test_responses_call_forwards_previous_response_id_and_store() -> None:
+    client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
+    tool_output: Final[FunctionCallOutput] = {
+        "type": "function_call_output",
+        "call_id": "call_abc123",
+        "output": "{}",
+    }
+    with patch(HTTPX_CLIENT_FACTORY, return_value=client):
+        litellm.responses(
+            model="fireworks_ai/kimi-k3",
+            input=[tool_output],  # mutable-ok: the Responses API takes input items as a JSON list
+            previous_response_id="resp_0e946f2d46bf4b49bf8b29ff78083583",
+            store=True,
+            api_key="fw-test-key",
+        )
+    _, _, body = _sent_request(client)
+    assert body["previous_response_id"] == "resp_0e946f2d46bf4b49bf8b29ff78083583"
+    assert body["store"] is True
+    assert body["input"][0]["call_id"] == "call_abc123"
+
+
+def test_responses_call_sends_session_affinity_for_caller_session_id() -> None:
+    client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
+    with patch(HTTPX_CLIENT_FACTORY, return_value=client):
+        litellm.responses(model="fireworks_ai/kimi-k3", input="hi", api_key="fw-test-key", litellm_session_id="sess-42")
+    _, headers, _ = _sent_request(client)
+    assert headers["x-session-affinity"] == "sess-42"
+
+
+def test_responses_call_keeps_caller_supplied_session_affinity_header() -> None:
+    client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
+    pinned: Final[Mapping[str, str]] = MappingProxyType({"x-session-affinity": "explicit-node"})
+    with patch(HTTPX_CLIENT_FACTORY, return_value=client):
+        litellm.responses(
+            model="fireworks_ai/kimi-k3",
+            input="hi",
+            api_key="fw-test-key",
+            litellm_session_id="sess-42",
+            extra_headers=pinned,
+        )
+    _, headers, _ = _sent_request(client)
+    assert headers["x-session-affinity"] == "explicit-node"
+
+
+def test_responses_call_maps_provider_errors_to_fireworks_ai() -> None:
+    client: Final = MagicMock()
+    request: Final = httpx.Request("POST", FIREWORKS_RESPONSES_URL)
+    client.post.side_effect = httpx.HTTPStatusError(
+        "unauthorized",
+        request=request,
+        response=httpx.Response(401, text='{"error": {"message": "invalid api key"}}', request=request),
+    )
+    with patch(HTTPX_CLIENT_FACTORY, return_value=client), pytest.raises(litellm.AuthenticationError) as raised:
+        litellm.responses(model="fireworks_ai/kimi-k3", input="hi", api_key="fw-bad-key")
+    assert raised.value.llm_provider == "fireworks_ai"
+    assert raised.value.status_code == 401
+    assert "invalid api key" in str(raised.value)
+
+
+def test_responses_call_skips_session_affinity_for_proxy_generated_session_id() -> None:
+    client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
+    generated: Final[_GeneratedSessionMetadata] = {"litellm_session_id_generated": True}
+    with patch(HTTPX_CLIENT_FACTORY, return_value=client):
+        litellm.responses(
+            model="fireworks_ai/kimi-k3",
+            input="hi",
+            api_key="fw-test-key",
+            litellm_session_id="generated-1",
+            litellm_metadata=generated,
+        )
+    _, headers, _ = _sent_request(client)
+    assert "x-session-affinity" not in headers
+
+
+@pytest.mark.parametrize(
+    "api_base, expected",
+    (
+        (None, FIREWORKS_RESPONSES_URL),
+        ("https://api.fireworks.ai/inference/v1", FIREWORKS_RESPONSES_URL),
+        ("https://api.fireworks.ai/inference/v1/", FIREWORKS_RESPONSES_URL),
+        ("https://gateway.example.com/fireworks", "https://gateway.example.com/fireworks/responses"),
+    ),
+)
+def test_get_complete_url(api_base: str | None, expected: str) -> None:
+    assert FireworksAIResponsesAPIConfig().get_complete_url(api_base=api_base, litellm_params=NO_PARAMS) == expected
+
+
+def test_responses_call_reads_fireworks_api_base_env(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("FIREWORKS_API_BASE", "https://self-hosted.example.com/v1")
+    client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
+    with patch(HTTPX_CLIENT_FACTORY, return_value=client):
+        litellm.responses(model="fireworks_ai/kimi-k3", input="hi", api_key="fw-test-key")
+    url, _, _ = _sent_request(client)
+    assert url == "https://self-hosted.example.com/v1/responses"
+
+
+@pytest.mark.parametrize(
+    "env_name", ("FIREWORKS_API_KEY", "FIREWORKS_AI_API_KEY", "FIREWORKSAI_API_KEY", "FIREWORKS_AI_TOKEN")
+)
+def test_validate_environment_reads_every_fireworks_key_name(monkeypatch: pytest.MonkeyPatch, env_name: str) -> None:
+    monkeypatch.setenv(env_name, "env-key")
+    headers: Final = FireworksAIResponsesAPIConfig().validate_environment(
+        headers=NO_HEADERS, model="accounts/fireworks/models/kimi-k3", litellm_params=GenericLiteLLMParams()
+    )
+    assert headers["Authorization"] == "Bearer env-key"
+    assert headers["Content-Type"] == "application/json"
+
+
+def test_validate_environment_prefers_explicit_api_key(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("FIREWORKS_API_KEY", "env-key")
+    headers: Final = FireworksAIResponsesAPIConfig().validate_environment(
+        headers=NO_HEADERS,
+        model="accounts/fireworks/models/kimi-k3",
+        litellm_params=GenericLiteLLMParams(api_key="explicit"),
+    )
+    assert headers["Authorization"] == "Bearer explicit"
+
+
+def test_validate_environment_without_any_key_raises() -> None:
+    with pytest.raises(ValueError, match="FIREWORKS_API_KEY"):
+        FireworksAIResponsesAPIConfig().validate_environment(
+            headers=NO_HEADERS, model="accounts/fireworks/models/kimi-k3", litellm_params=None
+        )
diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py
index e8333214ea8..fe3c4a0640d 100644
--- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py
+++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py
@@ -17,6 +17,9 @@ from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPICon
 from litellm.llms.databricks.responses.transformation import (
     DatabricksResponsesAPIConfig,
 )
+from litellm.llms.fireworks_ai.responses.transformation import (
+    FireworksAIResponsesAPIConfig,
+)
 from litellm.llms.github_copilot.responses.transformation import (
     GithubCopilotResponsesAPIConfig,
 )
@@ -102,6 +105,12 @@ class TestResponsesAPIWebSocketSupport:
     def test_openai_model_in_websocket_url_default(self):
         assert OpenAIResponsesAPIConfig().model_in_websocket_url() is True
 
+    def test_fireworks_ai_uses_managed_websocket(self):
+        """Fireworks AI should use managed websocket handler"""
+        assert (
+            FireworksAIResponsesAPIConfig().supports_native_websocket() is False
+        ), "Fireworks AI should use managed websocket handler"
+
     def test_xai_uses_managed_websocket(self):
         """XAI should use managed websocket handler"""
         config = XAIResponsesAPIConfig()
diff --git a/type-discipline-budget.json b/type-discipline-budget.json
index 8589a9451cf..835b24b2b3c 100644
--- a/type-discipline-budget.json
+++ b/type-discipline-budget.json
@@ -1,9 +1,9 @@
 {
   "LIT001": {
-    "limit": 22328
+    "limit": 22327
   },
   "LIT002": {
-    "limit": 26748
+    "limit": 26747
   },
   "LIT003": {
     "limit": 261

From 61bed7956689959c940fbb75d610c0dc7d5ffb44 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 17:18:17 -0700
Subject: [PATCH 035/164] feat(ci): add the cost map guard check

Replace test-model-map.yml with a pull_request_target guard that validates the
cost map, its backup, and its generated schema on every PR, and additionally
enforces the sync bot contract on litellm_cost_map_sync_* branches: only the
three cost map files may change, no model or field is removed, and the special
root keys stay untouched.
---
 .github/workflows/cost-map-guard.yml      |  45 +++++
 .github/workflows/test-model-map.yml      |  37 ----
 ci_cd/cost_map_guard.py                   | 146 ++++++++++++++++
 tests/test_litellm/test_cost_map_guard.py | 195 ++++++++++++++++++++++
 4 files changed, 386 insertions(+), 37 deletions(-)
 create mode 100644 .github/workflows/cost-map-guard.yml
 delete mode 100644 .github/workflows/test-model-map.yml
 create mode 100644 ci_cd/cost_map_guard.py
 create mode 100644 tests/test_litellm/test_cost_map_guard.py

diff --git a/.github/workflows/cost-map-guard.yml b/.github/workflows/cost-map-guard.yml
new file mode 100644
index 00000000000..61a56f1f47e
--- /dev/null
+++ b/.github/workflows/cost-map-guard.yml
@@ -0,0 +1,45 @@
+name: Cost map guard
+
+on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's cost map files are read as data and never executed
+  pull_request_target:
+    branches:
+      - main
+      - litellm_internal_staging
+      - litellm_oss_staging
+      - "litellm_**"
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
+  cancel-in-progress: true
+
+jobs:
+  cost-map-guard:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
+      - name: Fetch the pull request head and its merge base
+        id: revisions
+        env:
+          GH_TOKEN: ${{ github.token }}
+          BASE_SHA: ${{ github.event.pull_request.base.sha }}
+          HEAD_SHA: ${{ github.event.pull_request.head.sha }}
+        run: |
+          merge_base="$(gh api "repos/${GITHUB_REPOSITORY}/compare/${BASE_SHA}...${HEAD_SHA}" --jq '.merge_base_commit.sha')"
+          git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA"
+          echo "merge_base=$merge_base" >> "$GITHUB_OUTPUT"
+      - name: Set up uv
+        uses: ./.github/actions/setup-uv-with-retries
+        with:
+          version: "0.10.9"
+      - name: Run the guard
+        env:
+          MERGE_BASE: ${{ steps.revisions.outputs.merge_base }}
+          HEAD_SHA: ${{ github.event.pull_request.head.sha }}
+          HEAD_REF: ${{ github.event.pull_request.head.ref }}
+        run: |
+          uv run --frozen python ci_cd/cost_map_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --head-ref "$HEAD_REF"
diff --git a/.github/workflows/test-model-map.yml b/.github/workflows/test-model-map.yml
deleted file mode 100644
index c2770e5da4c..00000000000
--- a/.github/workflows/test-model-map.yml
+++ /dev/null
@@ -1,37 +0,0 @@
-name: Validate model_prices_and_context_window.json
-
-on:
-  pull_request:
-    branches:
-      - main
-      - litellm_internal_staging
-      - litellm_oss_staging
-      - "litellm_**"
-
-permissions:
-  contents: read
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
-  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
-
-jobs:
-  validate-model-prices-json:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
-        with:
-          persist-credentials: false
-
-      - name: Validate model_prices_and_context_window.json
-        run: |
-          jq empty model_prices_and_context_window.json
-
-      - name: Set up uv
-        uses: ./.github/actions/setup-uv-with-retries
-        with:
-          version: "0.10.9"
-
-      - name: Check model_prices_and_context_window.schema.json is in sync
-        run: |
-          uv run --frozen python ci_cd/generate_model_prices_schema.py --check
diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py
new file mode 100644
index 00000000000..50aa40ba220
--- /dev/null
+++ b/ci_cd/cost_map_guard.py
@@ -0,0 +1,146 @@
+"""Guard the cost map on pull requests.
+
+Every pull request gets the file checks: the three cost map files parse, the backup copy matches the root file,
+and the JSON schema is in sync and validates the map. Pull requests from the cost map sync bot (branches named
+litellm_cost_map_sync_*) additionally may only touch those three files and may only add or update models.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+import sys
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Final
+
+from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors
+
+COST_MAP_PATH: Final = "model_prices_and_context_window.json"
+BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
+SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
+GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH)
+BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
+
+CostMap = dict[str, object]
+
+
+@dataclass(frozen=True, slots=True)
+class Snapshot:
+    cost_map: str
+    backup: str
+    schema: str
+
+
+def _parse_object(text: str, path: str) -> CostMap | str:
+    try:
+        parsed: Final = json.loads(text)
+    except json.JSONDecodeError as error:
+        return f"{path} is not valid JSON: {error}"
+    return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root"
+
+
+def _rendered_schema(cost_map: CostMap) -> str:
+    try:
+        return render(build_schema(cost_map))
+    except SystemExit as error:
+        return str(error)
+
+
+def _file_failures(head: Snapshot, head_map: CostMap) -> tuple[str, ...]:
+    schema_text: Final = _rendered_schema(head_map)
+    if not schema_text.startswith("{"):
+        return (schema_text,)
+    backup_failure: Final = (
+        ()
+        if head.backup == head.cost_map
+        else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",)
+    )
+    schema_failure: Final = (
+        ()
+        if head.schema == schema_text
+        else (
+            f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; "
+            "run `python ci_cd/generate_model_prices_schema.py` and commit the result",
+        )
+    )
+    return (
+        *backup_failure,
+        *schema_failure,
+        *(
+            f"{COST_MAP_PATH} does not validate against its schema: {error}"
+            for error in validation_errors(head_map, json.loads(schema_text))[:20]
+        ),
+    )
+
+
+def _entries(cost_map: CostMap) -> dict[str, dict[str, object]]:
+    return {key: entry for key, entry in cost_map.items() if isinstance(entry, dict)}
+
+
+def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str]) -> tuple[str, ...]:
+    base_map: Final = _parse_object(base.cost_map, COST_MAP_PATH)
+    if isinstance(base_map, str):
+        return (f"merge base: {base_map}",)
+    base_entries: Final = _entries(base_map)
+    head_entries: Final = _entries(head_map)
+    removed_fields: Final = tuple(
+        f"{key}.{field}"
+        for key, entry in base_entries.items()
+        if key in head_entries
+        for field in entry
+        if field not in head_entries[key]
+    )
+    return (
+        *(
+            f"bot PRs may only change the cost map files, not {path}"
+            for path in changed_files
+            if path not in GUARDED_PATHS
+        ),
+        *(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map),
+        *(f"bot PRs may not remove fields: {ref}" for ref in removed_fields),
+        *(
+            f"bot PRs may not change {key}"
+            for key in sorted(SPECIAL_ROOT_KEYS)
+            if base_map.get(key) != head_map.get(key)
+        ),
+    )
+
+
+def guard_failures(base: Snapshot, head: Snapshot, changed_files: Sequence[str], bot: bool) -> tuple[str, ...]:
+    head_map: Final = _parse_object(head.cost_map, COST_MAP_PATH)
+    if isinstance(head_map, str):
+        return (head_map,)
+    return (*_file_failures(head, head_map), *(_bot_failures(base, head_map, changed_files) if bot else ()))
+
+
+def _git(*args: str) -> str:
+    result: Final = subprocess.run(("git", *args), check=False, capture_output=True, text=True)
+    return result.stdout if result.returncode == 0 else ""
+
+
+def snapshot(revision: str) -> Snapshot:
+    return Snapshot(*(_git("show", f"{revision}:{path}") for path in GUARDED_PATHS))
+
+
+def main(argv: Sequence[str]) -> int:
+    parser: Final = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--base", required=True, help="merge base of the pull request")
+    parser.add_argument("--head", required=True, help="head commit of the pull request")
+    parser.add_argument("--head-ref", required=True, help="head branch name of the pull request")
+    args: Final = parser.parse_args(argv)
+    bot: Final = args.head_ref.startswith(BOT_BRANCH_PREFIX)
+    changed_files: Final = tuple(_git("diff", "--name-only", args.base, args.head).splitlines())
+    failures: Final = guard_failures(snapshot(args.base), snapshot(args.head), changed_files, bot)
+    contract: Final = "bot contract enforced" if bot else "human PR, file checks only"
+    if failures:
+        print(f"cost map guard failed ({contract}):")
+        print("\n".join(f"- {failure}" for failure in failures))
+        return 1
+    print(f"cost map guard passed ({contract})")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))
diff --git a/tests/test_litellm/test_cost_map_guard.py b/tests/test_litellm/test_cost_map_guard.py
new file mode 100644
index 00000000000..1b4330ed62c
--- /dev/null
+++ b/tests/test_litellm/test_cost_map_guard.py
@@ -0,0 +1,195 @@
+import importlib.util
+import json
+import subprocess
+import sys
+from pathlib import Path
+from types import ModuleType
+from typing import Final
+
+import pytest
+
+ROOT: Final = Path(__file__).resolve().parents[2]
+CI_CD: Final = ROOT / "ci_cd"
+
+
+def _load(name: str) -> ModuleType:
+    spec = importlib.util.spec_from_file_location(name, CI_CD / f"{name}.py")
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+schema_module: Final = _load("generate_model_prices_schema")
+guard: Final = _load("cost_map_guard")
+
+MAP_FILES: Final = (guard.COST_MAP_PATH,)
+BOT_REF: Final = "litellm_cost_map_sync_2026-09-04T12-00Z"
+
+
+def _entry(price: float = 1e-06, **extra: object) -> dict[str, object]:
+    return {
+        "input_cost_per_token": price,
+        "output_cost_per_token": price * 2,
+        "litellm_provider": "openrouter",
+        "mode": "chat",
+        "max_tokens": 4096,
+        **extra,
+    }
+
+
+BASE_MAP: Final = {
+    "sample_spec": {"input_cost_per_token": "USD per prompt token"},
+    "fallback_generalizations": {"rules": [{"name": "r", "pattern": "^x"}]},
+    "openrouter/a": _entry(supports_vision=True),
+    "openrouter/b": _entry(2e-06),
+}
+
+
+def _serialize(cost_map: dict[str, object]) -> str:
+    return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n"
+
+
+def _snapshot(cost_map: dict[str, object], backup: str | None = None, schema: str | None = None) -> object:
+    text = _serialize(cost_map)
+    rendered = schema_module.render(schema_module.build_schema(cost_map))
+    return guard.Snapshot(
+        cost_map=text, backup=text if backup is None else backup, schema=rendered if schema is None else schema
+    )
+
+
+BASE: Final = _snapshot(BASE_MAP)
+
+
+def _failures(head: object, changed_files: tuple[str, ...] = MAP_FILES, bot: bool = True) -> tuple[str, ...]:
+    return guard.guard_failures(BASE, head, changed_files, bot)
+
+
+def test_in_sync_files_pass_for_humans_and_bots() -> None:
+    assert _failures(BASE, bot=False) == ()
+    assert _failures(BASE, bot=True) == ()
+
+
+def test_bot_may_add_and_reprice_models() -> None:
+    head = _snapshot({**BASE_MAP, "openrouter/a": _entry(9e-06, supports_vision=True), "openrouter/c": _entry()})
+    assert _failures(head) == ()
+
+
+def test_broken_json_is_reported() -> None:
+    head = guard.Snapshot(cost_map="{not json", backup="{not json", schema="{}")
+    assert _failures(head, bot=False) == (
+        f"{guard.COST_MAP_PATH} is not valid JSON: Expecting property name enclosed in double quotes: line 1 column 2 (char 1)",
+    )
+
+
+def test_non_object_root_is_reported() -> None:
+    head = guard.Snapshot(cost_map="[]", backup="[]", schema="{}")
+    assert _failures(head, bot=False) == (f"{guard.COST_MAP_PATH} must be a JSON object at the root",)
+
+
+def test_backup_drift_is_reported() -> None:
+    head = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)}))
+    assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.BACKUP_PATH)]
+
+
+def test_schema_out_of_sync_is_reported() -> None:
+    head = _snapshot({**BASE_MAP, "openrouter/c": _entry(supports_audio_input=True)}, schema=BASE.schema)
+    assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.SCHEMA_PATH)]
+
+
+def test_schema_validation_errors_are_reported() -> None:
+    head = _snapshot({**BASE_MAP, "openrouter/c": _entry(-1e-06)})
+    prefix = f"{guard.COST_MAP_PATH} does not validate against its schema: openrouter/c."
+    assert [failure.removeprefix(prefix).split(":")[0] for failure in _failures(head, bot=False)] == [
+        "input_cost_per_token",
+        "output_cost_per_token",
+    ]
+
+
+def test_unclassified_entry_key_is_reported() -> None:
+    text = _serialize({**BASE_MAP, "openrouter/c": _entry(weird_thing=1)})
+    head = guard.Snapshot(cost_map=text, backup=text, schema=BASE.schema)
+    (failure,) = _failures(head, bot=False)
+    assert "Unclassified keys" in failure and "weird_thing" in failure
+
+
+def test_bot_may_only_touch_the_cost_map_files() -> None:
+    changed = (*guard.GUARDED_PATHS, "litellm/utils.py", ".github/workflows/cost-map-guard.yml")
+    assert _failures(BASE, changed_files=changed, bot=False) == ()
+    assert _failures(BASE, changed_files=changed) == (
+        "bot PRs may only change the cost map files, not litellm/utils.py",
+        "bot PRs may only change the cost map files, not .github/workflows/cost-map-guard.yml",
+    )
+
+
+def test_bot_may_not_remove_models() -> None:
+    head = _snapshot({key: value for key, value in BASE_MAP.items() if key != "openrouter/b"})
+    assert _failures(head, bot=False) == ()
+    assert _failures(head) == ("bot PRs may not remove models: openrouter/b",)
+
+
+def test_bot_may_not_remove_fields() -> None:
+    head = _snapshot({**BASE_MAP, "openrouter/a": _entry()})
+    assert _failures(head, bot=False) == ()
+    assert _failures(head) == ("bot PRs may not remove fields: openrouter/a.supports_vision",)
+
+
+def test_bot_may_not_change_special_root_keys() -> None:
+    head = _snapshot({**BASE_MAP, "fallback_generalizations": {"rules": []}})
+    assert _failures(head, bot=False) == ()
+    assert _failures(head) == ("bot PRs may not change fallback_generalizations",)
+
+
+def _commit(repo: Path, cost_map: dict[str, object], message: str) -> str:
+    text = _serialize(cost_map)
+    (repo / guard.COST_MAP_PATH).write_text(text)
+    (repo / guard.BACKUP_PATH).parent.mkdir(exist_ok=True)
+    (repo / guard.BACKUP_PATH).write_text(text)
+    (repo / guard.SCHEMA_PATH).write_text(schema_module.render(schema_module.build_schema(cost_map)))
+    subprocess.run(("git", "add", "-A"), cwd=repo, check=True)
+    subprocess.run(
+        ("git", "-c", "user.name=t", "-c", "user.email=t@example.com", "commit", "-q", "-m", message),
+        cwd=repo,
+        check=True,
+    )
+    return subprocess.run(
+        ("git", "rev-parse", "HEAD"), cwd=repo, check=True, capture_output=True, text=True
+    ).stdout.strip()
+
+
+def _run_guard(repo: Path, base: str, head: str, head_ref: str) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        (sys.executable, str(CI_CD / "cost_map_guard.py"), "--base", base, "--head", head, "--head-ref", head_ref),
+        cwd=repo,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+
+
+@pytest.mark.parametrize(
+    ("head_ref", "expected_code", "expected_line"),
+    [
+        (BOT_REF, 1, "- bot PRs may not remove models: openrouter/b"),
+        ("litellm_fix_pricing", 0, "cost map guard passed (human PR, file checks only)"),
+    ],
+)
+def test_main_reads_both_revisions_from_git(
+    tmp_path: Path, head_ref: str, expected_code: int, expected_line: str
+) -> None:
+    subprocess.run(("git", "init", "-q", str(tmp_path)), check=True)
+    base = _commit(tmp_path, BASE_MAP, "base")
+    head = _commit(tmp_path, {key: value for key, value in BASE_MAP.items() if key != "openrouter/b"}, "head")
+    result = _run_guard(tmp_path, base, head, head_ref)
+    assert result.returncode == expected_code, result.stdout + result.stderr
+    assert expected_line in result.stdout.splitlines()
+
+
+def test_main_rejects_a_bot_pr_that_edits_code(tmp_path: Path) -> None:
+    subprocess.run(("git", "init", "-q", str(tmp_path)), check=True)
+    base = _commit(tmp_path, BASE_MAP, "base")
+    (tmp_path / "litellm" / "utils.py").write_text("print('hi')\n")
+    head = _commit(tmp_path, {**BASE_MAP, "openrouter/c": _entry()}, "head")
+    assert _run_guard(tmp_path, base, head, BOT_REF).returncode == 1
+    assert _run_guard(tmp_path, base, head, "litellm_fix_pricing").returncode == 0

From 217b5c404c42163add869108d672926758d257ba Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 17:24:09 -0700
Subject: [PATCH 036/164] fix(fireworks_ai): map the Fireworks delete response
 body to DeleteResponseResult

---
 .../fireworks_ai/responses/transformation.py    | 17 ++++++++++++++++-
 ...est_fireworks_ai_responses_transformation.py | 14 ++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py
index f36030bb50a..7f29e03ff0f 100644
--- a/litellm/llms/fireworks_ai/responses/transformation.py
+++ b/litellm/llms/fireworks_ai/responses/transformation.py
@@ -1,6 +1,9 @@
 from collections.abc import Mapping
 from types import MappingProxyType
-from typing import Final
+from typing import TYPE_CHECKING, Final
+from urllib.parse import unquote
+
+import httpx
 
 from litellm.llms.fireworks_ai.common_utils import (
     resolve_fireworks_api_key,
@@ -10,9 +13,13 @@ from litellm.llms.fireworks_ai.common_utils import (
 from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import ResponseInputParam
+from litellm.types.responses.main import DeleteResponseResult
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import LlmProviders
 
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+
 FIREWORKS_AI_DEFAULT_API_BASE: Final = "https://api.fireworks.ai/inference/v1"
 
 
@@ -64,5 +71,13 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
             headers=headers,
         )
 
+    def transform_delete_response_api_response(
+        self,
+        raw_response: httpx.Response,
+        logging_obj: "LiteLLMLoggingObj",
+    ) -> DeleteResponseResult:
+        deleted_id: Final = unquote(raw_response.request.url.path.rsplit("/", 1)[-1])
+        return DeleteResponseResult(id=deleted_id, object="response", deleted=True)
+
     def supports_native_websocket(self) -> bool:
         return False
diff --git a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
index 207fd518f89..997949bd03a 100644
--- a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
+++ b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
@@ -3,6 +3,7 @@ from collections.abc import Mapping
 from types import MappingProxyType
 from typing import Final, TypedDict
 from unittest.mock import MagicMock, patch
+from urllib.parse import quote
 
 import httpx
 import pytest
@@ -261,3 +262,16 @@ def test_validate_environment_without_any_key_raises() -> None:
         FireworksAIResponsesAPIConfig().validate_environment(
             headers=NO_HEADERS, model="accounts/fireworks/models/kimi-k3", litellm_params=None
         )
+
+
+def test_delete_responses_maps_fireworks_message_only_body_to_deleted_result() -> None:
+    response_id: Final = "resp_xFaIJR9Nc_OXmqKRqL78UuAGj2Te5GY5BT_knpZiMrYoNOVmu5oc2mQW1HI7hCtEYB4mcx2lEYS0DYP1U5yEQskHunuB4=="
+    request: Final = httpx.Request("DELETE", f"{FIREWORKS_RESPONSES_URL}/{quote(response_id, safe='')}")
+    client: Final = MagicMock()
+    client.delete.return_value = httpx.Response(200, json={"message": "Response deleted successfully"}, request=request)
+    with patch(HTTPX_CLIENT_FACTORY, return_value=client):
+        result: Final = litellm.delete_responses(
+            response_id=response_id, custom_llm_provider="fireworks_ai", api_key="fw-test-key"
+        )
+    assert client.delete.call_args.kwargs["url"] == str(request.url)
+    assert (result.id, result.object, result.deleted) == (response_id, "response", True)

From a93ba9229272100ddec777dedd77f02356b39954 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 17:28:51 -0700
Subject: [PATCH 037/164] test(fireworks_ai): cover native Responses API
 streaming end to end

---
 ...t_fireworks_ai_responses_transformation.py | 112 +++++++++++++++++-
 1 file changed, 111 insertions(+), 1 deletion(-)

diff --git a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
index 997949bd03a..8fac91ee475 100644
--- a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
+++ b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
@@ -265,7 +265,9 @@ def test_validate_environment_without_any_key_raises() -> None:
 
 
 def test_delete_responses_maps_fireworks_message_only_body_to_deleted_result() -> None:
-    response_id: Final = "resp_xFaIJR9Nc_OXmqKRqL78UuAGj2Te5GY5BT_knpZiMrYoNOVmu5oc2mQW1HI7hCtEYB4mcx2lEYS0DYP1U5yEQskHunuB4=="
+    response_id: Final = (
+        "resp_xFaIJR9Nc_OXmqKRqL78UuAGj2Te5GY5BT_knpZiMrYoNOVmu5oc2mQW1HI7hCtEYB4mcx2lEYS0DYP1U5yEQskHunuB4=="
+    )
     request: Final = httpx.Request("DELETE", f"{FIREWORKS_RESPONSES_URL}/{quote(response_id, safe='')}")
     client: Final = MagicMock()
     client.delete.return_value = httpx.Response(200, json={"message": "Response deleted successfully"}, request=request)
@@ -275,3 +277,111 @@ def test_delete_responses_maps_fireworks_message_only_body_to_deleted_result() -
         )
     assert client.delete.call_args.kwargs["url"] == str(request.url)
     assert (result.id, result.object, result.deleted) == (response_id, "response", True)
+
+
+def _fireworks_stream_response(status: str, output: tuple[Mapping[str, object], ...]) -> Mapping[str, object]:
+    return {
+        "id": "resp_htnkJ8piNKeOHkn9LfAusC38O2OgcDQs4S8trSOJ6anLeqjUDGqu2PkWmg5N",
+        "object": "response",
+        "created_at": 1788567245,
+        "model": "accounts/fireworks/models/kimi-k3",
+        "status": status,
+        "output": output,
+        "usage": None
+        if status == "in_progress"
+        else {
+            "input_tokens": 95,
+            "output_tokens": 89,
+            "total_tokens": 184,
+            "input_tokens_details": {"cached_tokens": 94},
+        },
+    }
+
+
+FIREWORKS_SSE_EVENTS: Final[tuple[Mapping[str, object], ...]] = (
+    {"type": "response.created", "sequence_number": 0, "response": _fireworks_stream_response("in_progress", ())},
+    {
+        "type": "response.output_item.added",
+        "sequence_number": 1,
+        "output_index": 0,
+        "item": {"id": "rs_1", "type": "reasoning", "summary": []},
+    },
+    {
+        "type": "response.reasoning_summary_text.delta",
+        "sequence_number": 2,
+        "item_id": "rs_1",
+        "output_index": 0,
+        "summary_index": 0,
+        "delta": "pong",
+    },
+    {
+        "type": "response.output_item.added",
+        "sequence_number": 3,
+        "output_index": 1,
+        "item": {"id": "msg_1", "type": "message", "role": "assistant", "status": "in_progress", "content": []},
+    },
+    {
+        "type": "response.output_text.delta",
+        "sequence_number": 4,
+        "item_id": "msg_1",
+        "output_index": 1,
+        "content_index": 0,
+        "delta": "po",
+    },
+    {
+        "type": "response.output_text.delta",
+        "sequence_number": 5,
+        "item_id": "msg_1",
+        "output_index": 1,
+        "content_index": 0,
+        "delta": "ng",
+    },
+    {
+        "type": "response.completed",
+        "sequence_number": 6,
+        "response": _fireworks_stream_response(
+            "completed",
+            (
+                {"id": "rs_1", "type": "reasoning", "summary": [{"type": "summary_text", "text": "pong"}]},
+                {
+                    "id": "msg_1",
+                    "type": "message",
+                    "role": "assistant",
+                    "status": "completed",
+                    "content": [{"type": "output_text", "text": "pong", "annotations": []}],
+                },
+            ),
+        ),
+    },
+)
+
+
+def _sse_body(events: tuple[Mapping[str, object], ...]) -> bytes:
+    return b"".join(f"data: {json.dumps(dict(event))}\n\n".encode() for event in events) + b"data: [DONE]\n\n"
+
+
+def test_streaming_responses_call_hits_native_endpoint_and_yields_every_fireworks_event() -> None:
+    request: Final = httpx.Request("POST", FIREWORKS_RESPONSES_URL)
+    client: Final = MagicMock()
+    client.post.return_value = httpx.Response(
+        200, content=_sse_body(FIREWORKS_SSE_EVENTS), headers={"content-type": "text/event-stream"}, request=request
+    )
+    with patch(HTTPX_CLIENT_FACTORY, return_value=client):
+        received: Final = tuple(
+            litellm.responses(
+                model="fireworks_ai/kimi-k3",
+                input="Reply with the single word pong.",
+                stream=True,
+                api_key="fw-test-key",
+            )
+        )
+    url, _, body = _sent_request(client)
+    assert (url, body["model"], body["stream"], client.post.call_args.kwargs["stream"]) == (
+        FIREWORKS_RESPONSES_URL,
+        "accounts/fireworks/models/kimi-k3",
+        True,
+        True,
+    )
+    assert tuple(event.type for event in received) == tuple(event["type"] for event in FIREWORKS_SSE_EVENTS)
+    assert "".join(event.delta for event in received if event.type == "response.output_text.delta") == "pong"
+    assert received[-1].response.usage.output_tokens == 89

From 6463994f78161ee4448380a5e9fb2ad28729b98f Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 17:29:10 -0700
Subject: [PATCH 038/164] fix(ai-gateway): dial upstream WebSockets over an
 explicit rustls provider

Build one ClientConfig that names ring and loads the native roots once, and hand it to every tokio-tungstenite dial as its connector instead of installing a process-wide default from the dial path. Each of the three dial sites gets a wss:// test that reproduces the panic if its connector is dropped.
---
 litellm-rust/Cargo.lock                       |   1 +
 litellm-rust/Cargo.toml                       |   1 +
 litellm-rust/crates/ai-gateway/Cargo.toml     |   5 +-
 .../crates/ai-gateway/src/io/realtime.rs      |  27 +++++
 .../crates/ai-gateway/src/io/responses_ws.rs  |  23 ++++
 litellm-rust/crates/ai-gateway/src/io/tls.rs  | 113 ++++++++++--------
 .../tests/crypto_provider_wiring.rs           |  19 ++-
 7 files changed, 134 insertions(+), 55 deletions(-)

diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock
index 2ed998174e4..d214e80d818 100644
--- a/litellm-rust/Cargo.lock
+++ b/litellm-rust/Cargo.lock
@@ -1416,6 +1416,7 @@ dependencies = [
  "pyo3",
  "reqwest",
  "rustls 0.23.42",
+ "rustls-native-certs",
  "serde",
  "serde_json",
  "sha2 0.10.9",
diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml
index 2e2e8809b7f..d3d25cbda8d 100644
--- a/litellm-rust/Cargo.toml
+++ b/litellm-rust/Cargo.toml
@@ -27,6 +27,7 @@ rand = "0.8"
 reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
 rstest = "0.26.1"
 rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
+rustls-native-certs = "0.8"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 sha2 = "0.10"
diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml
index 414abc2356d..82bedd0c8a0 100644
--- a/litellm-rust/crates/ai-gateway/Cargo.toml
+++ b/litellm-rust/crates/ai-gateway/Cargo.toml
@@ -19,9 +19,10 @@ litellm-core = { workspace = true, features = ["bedrock-auth"] }
 # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
 # Python proxy callbacks API.
 reqwest.workspace = true
-# rustls is a direct dependency so `io::tls` can install a process-level
-# crypto provider; see that module for why the graph needs one.
+# rustls and its root store are direct dependencies so `io::tls` can build the
+# one TLS config the outbound dials use; see that module for why it has to.
 rustls.workspace = true
+rustls-native-certs.workspace = true
 # `sync` powers the bounded mpsc channel the realtime logger drains.
 tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] }
 tokio-tungstenite.workspace = true
diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs
index 53d87848342..207c31dffa0 100644
--- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs
+++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs
@@ -286,6 +286,33 @@ mod tests {
         serde_json::from_str(raw).expect("valid event json")
     }
 
+    /// The realtime dial has to reach a `wss://` upstream without a process-wide
+    /// crypto provider installed, which is what dialing through `io::tls` buys.
+    #[tokio::test]
+    async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() {
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("bind a loopback port");
+        let port = listener
+            .local_addr()
+            .expect("read the bound address")
+            .port();
+        tokio::spawn(async move {
+            while let Ok((stream, _peer)) = listener.accept().await {
+                drop(stream);
+            }
+        });
+
+        let result = dial_upstream(
+            "gpt-realtime",
+            "sk-test",
+            Some(&format!("wss://127.0.0.1:{port}")),
+        )
+        .await;
+
+        assert!(matches!(result, Err(Error::Network(_))));
+    }
+
     #[test]
     fn resolve_api_key_prefers_param_then_blank_falls_through() {
         assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test");
diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
index ee181791413..9df3d0c6cc5 100644
--- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
+++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
@@ -326,6 +326,29 @@ mod tests {
     use tokio::net::TcpListener;
     use tokio_tungstenite::accept_async;
 
+    /// The Responses dial has to reach a `wss://` upstream without a process-wide
+    /// crypto provider installed, which is what dialing through `io::tls` buys.
+    #[tokio::test]
+    async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() {
+        let listener = TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("bind a loopback port");
+        let port = listener
+            .local_addr()
+            .expect("read the bound address")
+            .port();
+        tokio::spawn(async move {
+            while let Ok((stream, _peer)) = listener.accept().await {
+                drop(stream);
+            }
+        });
+
+        let result =
+            dial_upstream("gpt-5", "sk-test", Some(&format!("wss://127.0.0.1:{port}"))).await;
+
+        assert!(matches!(result, Err(Error::Network(_))));
+    }
+
     async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) {
         let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
         let address = listener.local_addr().expect("local address");
diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs
index 96544adffb7..16fd11e2e79 100644
--- a/litellm-rust/crates/ai-gateway/src/io/tls.rs
+++ b/litellm-rust/crates/ai-gateway/src/io/tls.rs
@@ -1,29 +1,54 @@
-//! Outbound WebSocket dials, with the rustls crypto provider settled first.
+//! Outbound WebSocket dials over a TLS config this crate builds once and owns.
 //!
 //! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth`
-//! enables `rustls/aws-lc-rs`, so `ClientConfig::builder()` — which is how
-//! `tokio-tungstenite` builds its TLS config — panics rather than guess between
-//! them. reqwest and the AWS SDK pick a provider explicitly and never panic.
-//!
-//! Installing from the dial rather than from a `main` also covers the `cdylib`
-//! the Python bridge loads, the tests, and the benches, none of which have one.
-//! ring is what reqwest already falls back to, so installing it changes no
-//! working path, and whoever installs into this rustls build first still wins.
+//! enables `rustls/aws-lc-rs`, so the bare `ClientConfig::builder()` that
+//! `tokio-tungstenite` uses when handed no connector panics rather than guess
+//! between them. Naming ring on a connector of our own settles that for these
+//! dials without touching the process-wide default, and building the config
+//! once keeps the platform trust store, which `tokio-tungstenite` would
+//! otherwise re-read on every dial, off the dial path.
 
-use std::sync::Once;
+use std::io;
+use std::sync::{Arc, OnceLock};
 
+use rustls::{ClientConfig, RootCertStore};
 use tokio::net::TcpStream;
 use tokio_tungstenite::tungstenite::Error;
 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+use tokio_tungstenite::tungstenite::error::TlsError;
 use tokio_tungstenite::tungstenite::handshake::client::Response;
-use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
+use tokio_tungstenite::{
+    Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
+};
 
-static INSTALL_CRYPTO_PROVIDER: Once = Once::new();
+static TLS_CONFIG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
 
-pub(crate) fn ensure_crypto_provider() {
-    INSTALL_CRYPTO_PROVIDER.call_once(|| {
-        let _ = rustls::crypto::ring::default_provider().install_default();
-    });
+fn build_config() -> Result<ClientConfig, Box<Error>> {
+    let native = rustls_native_certs::load_native_certs();
+    let roots = {
+        let mut store = RootCertStore::empty();
+        let (added, _ignored) = store.add_parsable_certificates(native.certs);
+        if added == 0 {
+            return Err(Box::new(Error::Io(io::Error::other(format!(
+                "no usable native root certificates: {:?}",
+                native.errors
+            )))));
+        }
+        store
+    };
+
+    ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
+        .with_safe_default_protocol_versions()
+        .map(|builder| builder.with_root_certificates(roots).with_no_client_auth())
+        .map_err(|error| Box::new(Error::Tls(TlsError::Rustls(error))))
+}
+
+fn tls_config() -> Result<Arc<ClientConfig>, Box<Error>> {
+    if let Some(config) = TLS_CONFIG.get() {
+        return Ok(Arc::clone(config));
+    }
+    let built = Arc::new(build_config()?);
+    Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built)))
 }
 
 pub(crate) async fn connect_upstream<R>(
@@ -32,19 +57,25 @@ pub(crate) async fn connect_upstream<R>(
 where
     R: IntoClientRequest + Unpin,
 {
-    ensure_crypto_provider();
-    connect_async(request).await.map_err(Box::new)
+    let request = request.into_client_request().map_err(Box::new)?;
+    let connector = match request.uri().scheme_str() {
+        Some("wss") => Some(Connector::Rustls(tls_config()?)),
+        _ => None,
+    };
+    connect_async_tls_with_config(request, None, false, connector)
+        .await
+        .map_err(Box::new)
 }
 
 #[cfg(test)]
 mod tests {
-    use rustls::crypto::CryptoProvider;
+    use rustls::CipherSuite;
+    use rustls::NamedGroup;
+    use rustls::crypto::{CryptoProvider, aws_lc_rs, ring};
 
-    use super::ensure_crypto_provider;
+    use super::{Arc, build_config, tls_config};
 
-    fn fingerprint(
-        provider: &CryptoProvider,
-    ) -> (Vec<rustls::CipherSuite>, Vec<rustls::NamedGroup>) {
+    fn fingerprint(provider: &CryptoProvider) -> (Vec<CipherSuite>, Vec<NamedGroup>) {
         (
             provider
                 .cipher_suites
@@ -60,43 +91,31 @@ mod tests {
     }
 
     #[test]
-    fn client_config_builder_works_with_both_provider_features_enabled() {
-        ensure_crypto_provider();
-
-        assert!(CryptoProvider::get_default().is_some());
-
-        let config = rustls::ClientConfig::builder()
-            .with_root_certificates(rustls::RootCertStore::empty())
-            .with_no_client_auth();
+    fn builds_a_usable_config_with_both_provider_features_enabled() {
+        let config = build_config().expect("a client config");
 
         assert!(!config.crypto_provider().cipher_suites.is_empty());
     }
 
     #[test]
-    fn installs_ring_rather_than_aws_lc_rs() {
-        ensure_crypto_provider();
-
-        let installed = CryptoProvider::get_default().expect("a provider is installed");
+    fn dials_with_ring_rather_than_aws_lc_rs() {
+        let config = build_config().expect("a client config");
 
         assert_eq!(
-            fingerprint(installed),
-            fingerprint(&rustls::crypto::ring::default_provider())
+            fingerprint(config.crypto_provider()),
+            fingerprint(&ring::default_provider())
         );
         assert_ne!(
-            fingerprint(installed),
-            fingerprint(&rustls::crypto::aws_lc_rs::default_provider())
+            fingerprint(config.crypto_provider()),
+            fingerprint(&aws_lc_rs::default_provider())
         );
     }
 
     #[test]
-    fn ensure_crypto_provider_is_idempotent() {
-        ensure_crypto_provider();
-        let first = CryptoProvider::get_default().cloned();
+    fn the_trust_store_is_loaded_once_and_shared() {
+        let first = tls_config().expect("a client config");
+        let second = tls_config().expect("a client config");
 
-        ensure_crypto_provider();
-        let second = CryptoProvider::get_default().cloned();
-
-        assert!(first.is_some());
-        assert!(std::sync::Arc::ptr_eq(&first.unwrap(), &second.unwrap()));
+        assert!(Arc::ptr_eq(&first, &second));
     }
 }
diff --git a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs
index db5698f8460..05f7d9610d5 100644
--- a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs
+++ b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs
@@ -1,5 +1,6 @@
-//! Guards the wiring, not just the helper: the dial itself has to install the
-//! rustls provider, in a test binary where nothing else has installed one.
+//! Guards the wiring, not just the helper: a `wss://` dial through the public
+//! API has to resolve its own crypto provider, in a test binary where nothing
+//! has installed a process-wide one, and has to leave it uninstalled.
 
 use std::collections::HashMap;
 use std::time::Duration;
@@ -7,8 +8,7 @@ use std::time::Duration;
 use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection;
 use tokio::net::TcpListener;
 
-#[tokio::test]
-async fn dialing_wss_returns_an_error_instead_of_panicking() {
+async fn dead_tls_server() -> u16 {
     let listener = TcpListener::bind("127.0.0.1:0")
         .await
         .expect("bind a loopback port");
@@ -23,6 +23,13 @@ async fn dialing_wss_returns_an_error_instead_of_panicking() {
         }
     });
 
+    port
+}
+
+#[tokio::test]
+async fn dialing_wss_returns_an_error_instead_of_panicking() {
+    let port = dead_tls_server().await;
+
     let result = ResponsesWebSocketConnection::connect_url(
         &format!("wss://127.0.0.1:{port}/"),
         &HashMap::new(),
@@ -35,7 +42,7 @@ async fn dialing_wss_returns_an_error_instead_of_panicking() {
         "a plain TCP server cannot finish a TLS handshake"
     );
     assert!(
-        rustls::crypto::CryptoProvider::get_default().is_some(),
-        "the dial is what installs the process-wide provider"
+        rustls::crypto::CryptoProvider::get_default().is_none(),
+        "the dial settles its provider on its own connector, not process-wide"
     );
 }

From 199b44a475719500915dbd9059c0ea5dfeeae782 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 17:48:56 -0700
Subject: [PATCH 039/164] fix(async): move remote image fetches off the event
 loop for Snowflake, Bedrock invoke Claude, Mantle and Gemini

---
 .../prompt_templates/image_handling.py        | 103 ++++++++++
 litellm/llms/base_llm/chat/transformation.py  |   4 +
 .../anthropic_claude3_transformation.py       |  61 +-----
 .../bedrock/chat/mantle/transformation.py     |  15 +-
 litellm/llms/custom_httpx/llm_http_handler.py | 194 ++++++++++--------
 litellm/llms/snowflake/chat/transformation.py |  16 ++
 .../llms/vertex_ai/gemini/transformation.py   |   9 +-
 litellm/main.py                               |   8 +-
 tests/test_litellm/conftest.py                |  35 ++++
 .../litellm_core_utils/test_image_handling.py |  62 ++++++
 ...ations_anthropic_claude3_transformation.py |  49 +++++
 ...test_bedrock_chat_mantle_transformation.py |  51 +++++
 .../custom_httpx/test_llm_http_handler.py     | 106 +++++++++-
 .../test_snowflake_chat_transformation.py     |  46 +++++
 .../vertex_ai/gemini/test_transformation.py   |  42 ++++
 15 files changed, 644 insertions(+), 157 deletions(-)
 create mode 100644 tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py

diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py
index 81132fa89a5..1890d4eb682 100644
--- a/litellm/litellm_core_utils/prompt_templates/image_handling.py
+++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py
@@ -2,7 +2,11 @@
 Helper functions to handle images passed in messages
 """
 
+import asyncio
 import base64
+from collections.abc import Mapping
+from dataclasses import dataclass
+from types import MappingProxyType
 from typing import Final
 
 from httpx import Response
@@ -12,6 +16,7 @@ from litellm import verbose_logger
 from litellm.caching.caching import InMemoryCache
 from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB
 from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
+from litellm.types.llms.openai import AllMessageValues
 
 MAX_IMGS_IN_MEMORY: Final = 10
 
@@ -124,3 +129,101 @@ def convert_url_to_base64(url: str) -> str:
     raise litellm.ImageFetchError(
         f"Error: Unable to fetch image from URL after 3 attempts. url={url}",
     )
+
+
+_REMOTE_URL_PREFIXES: Final = ("http://", "https://")
+
+
+@dataclass(frozen=True, slots=True)
+class _RemoteImage:
+    part: Mapping[str, object]
+    image_url: Mapping[str, object] | None
+    url: str
+
+
+@dataclass(frozen=True, slots=True)
+class _RemoteFile:
+    part: Mapping[str, object]
+    file: Mapping[str, object]
+    url: str
+
+
+def _as_mapping(value: object) -> Mapping[str, object] | None:
+    return value if isinstance(value, Mapping) else None  # pyright: ignore[reportUnknownVariableType]  # fields are parsed one by one
+
+
+def _remote_url(candidate: object) -> str | None:
+    return candidate if isinstance(candidate, str) and candidate.startswith(_REMOTE_URL_PREFIXES) else None
+
+
+def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | None:
+    fields: Final = _as_mapping(part)
+    if fields is None:
+        return None
+    if fields.get("type") == "image_url":
+        image_url: Final = fields.get("image_url")
+        image_url_fields: Final = _as_mapping(image_url)
+        url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url)
+        return _RemoteImage(fields, image_url_fields, url) if url is not None else None
+    file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None
+    file_url: Final = _remote_url(file.get("file_id")) if file is not None else None
+    return _RemoteFile(fields, file, file_url) if file is not None and file_url is not None else None
+
+
+_PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"})
+
+
+def _inferred_format(file: Mapping[str, object], url: str) -> Mapping[str, str]:
+    return _PDF_FORMAT if "format" not in file and url.lower().endswith(".pdf") else MappingProxyType({})
+
+
+def _inlined_image_url(image_url: Mapping[str, object] | None, data_url: str) -> Mapping[str, object] | str:
+    return {**image_url, "url": data_url} if image_url is not None else data_url  # mutable-ok: json-serialized part
+
+
+def _inlined_file(file: Mapping[str, object], url: str, data_url: str) -> Mapping[str, object]:
+    kept: Final = {k: v for k, v in file.items() if k != "file_id"}  # mutable-ok: json-serialized message part
+    return {**kept, **_inferred_format(file, url), "file_data": data_url}  # mutable-ok: json-serialized part
+
+
+def _inline(remote: _RemoteImage | _RemoteFile, data_url: str) -> Mapping[str, object]:
+    match remote:
+        case _RemoteImage(part, image_url, _):
+            return {**part, "image_url": _inlined_image_url(image_url, data_url)}  # mutable-ok: json-serialized part
+        case _RemoteFile(part, file, url):
+            return {**part, "file": _inlined_file(file, url, data_url)}  # mutable-ok: json-serialized message part
+
+
+def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]:
+    content: Final = message.get("content")
+    return tuple(content) if isinstance(content, list) else ()  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]  # parts are parsed one by one
+
+
+def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> AllMessageValues:
+    parts: Final = _content_parts(message)
+    if not parts:
+        return message
+    inlined_parts: Final = [  # mutable-ok: content must stay a list for the transforms' isinstance checks
+        _inline(remote, data_urls[remote.url]) if (remote := _parse_remote_part(part)) is not None else part
+        for part in parts
+    ]
+    inlined_message: Final = {**message, "content": inlined_parts}  # mutable-ok: json-serialized message
+    return inlined_message  # pyright: ignore[reportReturnType]  # the same message with its remote parts inlined
+
+
+async def async_inline_remote_media(
+    messages: list[AllMessageValues],  # mutable-ok: every transform_request takes list[AllMessageValues]
+) -> list[AllMessageValues]:  # mutable-ok: every transform_request takes list[AllMessageValues]
+    remote_urls: Final = tuple(
+        dict.fromkeys(
+            remote.url
+            for message in messages
+            for part in _content_parts(message)
+            if (remote := _parse_remote_part(part)) is not None
+        )
+    )
+    if not remote_urls:
+        return messages
+    data_urls: Final = await asyncio.gather(*(async_convert_url_to_base64(url) for url in remote_urls))
+    inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True)))
+    return [_inline_message(message, inlined) for message in messages]  # mutable-ok: transform_request takes a list
diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py
index bbe1cc85df1..7bfc87a30d6 100644
--- a/litellm/llms/base_llm/chat/transformation.py
+++ b/litellm/llms/base_llm/chat/transformation.py
@@ -411,6 +411,10 @@ class BaseConfig(ABC):
     def has_custom_stream_wrapper(self) -> bool:
         return False
 
+    @property
+    def uses_async_transform_request(self) -> bool:
+        return False
+
     @property
     def supports_stream_param_in_request_body(self) -> bool:
         """
diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py
index 67720451c00..38f280eef03 100644
--- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py
+++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py
@@ -8,7 +8,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
     convert_to_anthropic_image_obj,
 )
 from litellm.litellm_core_utils.prompt_templates.image_handling import (
-    async_convert_url_to_base64,
+    async_inline_remote_media,
     convert_url_to_base64,
 )
 from litellm.llms.anthropic.chat.transformation import AnthropicConfig
@@ -172,6 +172,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
 
         return _anthropic_request
 
+    @property
+    def uses_async_transform_request(self) -> bool:
+        return True
+
     async def async_transform_request(
         self,
         model: str,
@@ -180,26 +184,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
         litellm_params: dict,
         headers: dict,
     ) -> dict:
-        _anthropic_request: Final = self._build_bedrock_anthropic_request_base(
+        return self.transform_request(
             model=model,
-            messages=messages,
+            messages=await async_inline_remote_media(messages),
             optional_params=optional_params,
             litellm_params=litellm_params,
             headers=headers,
         )
 
-        await self._async_convert_document_url_sources_to_base64(_anthropic_request)
-        beta_list: Final = self._compute_bedrock_invoke_beta_headers(
-            model=model,
-            messages=messages,
-            optional_params=optional_params,
-            headers=headers,
-        )
-        if beta_list:
-            _anthropic_request["anthropic_beta"] = beta_list
-
-        return _anthropic_request
-
     def _build_bedrock_anthropic_request_base(
         self,
         model: str,
@@ -321,45 +313,6 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
                     "data": image_chunk["data"],
                 }
 
-    async def _async_convert_document_url_sources_to_base64(self, anthropic_request: dict) -> None:
-        """
-        Async version of document URL conversion for async completion paths.
-        """
-        messages: Final = anthropic_request.get("messages")
-        if not isinstance(messages, list):
-            return
-
-        for message in messages:
-            if not isinstance(message, dict):
-                continue
-            content = message.get("content")
-            if not isinstance(content, list):
-                continue
-
-            for block in content:
-                if not isinstance(block, dict) or block.get("type") != "document":
-                    continue
-                source = block.get("source")
-                if not isinstance(source, dict) or source.get("type") != "url":
-                    continue
-                source_url = source.get("url")
-                if not isinstance(source_url, str):
-                    continue
-
-                inferred_format: str | None = None
-                if source_url.lower().endswith(".pdf"):
-                    inferred_format = "application/pdf"
-                base64_url = await async_convert_url_to_base64(url=source_url)
-                image_chunk = convert_to_anthropic_image_obj(
-                    openai_image_url=base64_url,
-                    format=inferred_format,
-                )
-                block["source"] = {
-                    "type": "base64",
-                    "media_type": image_chunk["media_type"],
-                    "data": image_chunk["data"],
-                }
-
     def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict:
         """
         Convert tool search entries to the format supported by the Bedrock Invoke API.
diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py
index 31fc079c0c9..7e2037c33f1 100644
--- a/litellm/llms/bedrock/chat/mantle/transformation.py
+++ b/litellm/llms/bedrock/chat/mantle/transformation.py
@@ -10,6 +10,7 @@ at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth.
 from collections.abc import AsyncIterator, Iterator
 from typing import TYPE_CHECKING, Any, Final
 
+from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
 from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
     AmazonAnthropicClaudeConfig,
 )
@@ -110,21 +111,13 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig):
         litellm_params: dict,
         headers: dict,
     ) -> dict:
-        model_id: Final = model.replace("mantle/", "", 1)
-
-        request: Final = self._build_bedrock_anthropic_request_base(
-            model=model_id,
-            messages=messages,
+        return self.transform_request(
+            model=model,
+            messages=await async_inline_remote_media(messages),
             optional_params=optional_params,
             litellm_params=litellm_params,
             headers=headers,
         )
-        await self._async_convert_document_url_sources_to_base64(request)
-        return self._restore_mantle_body_fields(
-            request=request,
-            model_id=model_id,
-            optional_params=optional_params,
-        )
 
     @staticmethod
     def _restore_mantle_body_fields(request: dict, model_id: str, optional_params: dict) -> dict:
diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index f281c249c72..54c21f8d5b3 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -116,6 +116,7 @@ from litellm.types.llms.anthropic_skills import (
     Skill,
 )
 from litellm.types.llms.openai import (
+    AllMessageValues,
     CreateBatchRequest,
     CreateFileRequest,
     FileContentRequest,
@@ -488,7 +489,7 @@ class BaseLLMHTTPHandler:
     def completion(
         self,
         model: str,
-        messages: list,
+        messages: list[AllMessageValues],
         api_base: str | None,
         custom_llm_provider: str,
         model_response: ModelResponse,
@@ -507,7 +508,7 @@ class BaseLLMHTTPHandler:
         shared_session: Optional["ClientSession"] = None,
     ):
         json_mode: Final[bool] = optional_params.pop("json_mode", False)
-        extra_body: Final[dict | None] = optional_params.pop("extra_body", None)
+        extra_body: Final[Mapping[str, object] | None] = optional_params.pop("extra_body", None)
 
         provider_config = provider_config or ProviderConfigManager.get_provider_chat_config(
             model=model, provider=litellm.LlmProviders(custom_llm_provider)
@@ -522,14 +523,17 @@ class BaseLLMHTTPHandler:
         )
 
         # get config from model, custom llm provider
-        headers = provider_config.validate_environment(
-            api_key=api_key,
-            headers=headers or {},
-            model=model,
-            messages=messages,
-            optional_params=optional_params,
-            api_base=api_base,
-            litellm_params=litellm_params,
+        request_headers: Final = cast(  # cast-ok: validate_environment is declared as a bare dict
+            "dict[str, object]",
+            provider_config.validate_environment(
+                api_key=api_key,
+                headers=headers or {},
+                model=model,
+                messages=messages,
+                optional_params=optional_params,
+                api_base=api_base,
+                litellm_params=litellm_params,
+            ),
         )
 
         api_base = provider_config.get_complete_url(
@@ -541,93 +545,117 @@ class BaseLLMHTTPHandler:
             litellm_params=litellm_params,
         )
 
-        data: dict[str, object] = provider_config.transform_request(
-            model=model,
-            messages=messages,
-            optional_params=optional_params,
-            litellm_params=litellm_params,
-            headers=headers,
-        )
-
-        if extra_body is not None:
-            data = {**data, **extra_body}
-
-        headers, signed_json_body = provider_config.sign_request(
-            headers=headers,
-            optional_params={
-                **optional_params,
-                **_aws_signing_overrides(optional_params, litellm_params),
-            },
-            request_data=data,
-            api_base=api_base,
-            api_key=api_key,
-            stream=stream,
-            fake_stream=fake_stream,
-            model=model,
-        )
-
-        ## LOGGING
-        logging_obj.pre_call(
-            input=messages,
-            api_key=api_key,
-            additional_args={
-                "complete_input_dict": data,
-                "api_base": api_base,
-                "headers": headers,
-            },
-        )
-
-        # Check if stream was converted for WebSearch interception
-        # This is set by the async_pre_request_hook in WebSearchInterceptionLogger
-        if litellm_params.get("_websearch_interception_converted_stream", False):
-            logging_obj.model_call_details["websearch_interception_converted_stream"] = True
-
-        if acompletion is True:
-            if stream is True:
-                data = self._add_stream_param_to_request_body(
-                    data=data,
-                    provider_config=provider_config,
+        def sign_and_log(
+            transformed: dict[str, object],  # mutable-ok: async_completion takes dict
+        ) -> tuple[dict[str, object], dict[str, object], bytes | None]:  # mutable-ok: async_completion takes dict
+            data: Final = {**transformed, **extra_body} if extra_body is not None else transformed
+            signed: Final = cast(  # cast-ok: sign_request is declared as a bare dict
+                "tuple[dict[str, object], bytes | None]",
+                provider_config.sign_request(
+                    headers=request_headers,
+                    optional_params={
+                        **optional_params,
+                        **_aws_signing_overrides(optional_params, litellm_params),
+                    },
+                    request_data=data,
+                    api_base=api_base,
+                    api_key=api_key,
+                    stream=stream,
                     fake_stream=fake_stream,
-                )
+                    model=model,
+                ),
+            )
+            logging_obj.pre_call(
+                input=messages,
+                api_key=api_key,
+                additional_args={
+                    "complete_input_dict": data,
+                    "api_base": api_base,
+                    "headers": signed[0],
+                },
+            )
+            if litellm_params.get("_websearch_interception_converted_stream", False):
+                logging_obj.model_call_details["websearch_interception_converted_stream"] = True
+            return data, signed[0], signed[1]
+
+        def dispatch_async(
+            data: dict[str, object],  # mutable-ok: async_completion takes dict
+            signed_headers: dict[str, object],  # mutable-ok: async_completion takes dict
+            signed_json_body: bytes | None,
+        ):
+            async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None
+            if stream is True:
                 return self.acompletion_stream_function(
                     model=model,
                     messages=messages,
                     api_base=api_base,
-                    headers=headers,
+                    headers=signed_headers,
                     custom_llm_provider=custom_llm_provider,
                     provider_config=provider_config,
                     timeout=timeout,
                     logging_obj=logging_obj,
-                    data=data,
+                    data=self._add_stream_param_to_request_body(
+                        data=data,
+                        provider_config=provider_config,
+                        fake_stream=fake_stream,
+                    ),
                     fake_stream=fake_stream,
-                    client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
+                    client=async_client,
                     litellm_params=litellm_params,
                     json_mode=json_mode,
                     optional_params=optional_params,
                     signed_json_body=signed_json_body,
                 )
+            return self.async_completion(
+                custom_llm_provider=custom_llm_provider,
+                provider_config=provider_config,
+                api_base=api_base,
+                headers=signed_headers,
+                data=data,
+                timeout=timeout,
+                model=model,
+                model_response=model_response,
+                logging_obj=logging_obj,
+                api_key=api_key,
+                messages=messages,
+                optional_params=optional_params,
+                litellm_params=litellm_params,
+                encoding=encoding,
+                client=async_client,
+                json_mode=json_mode,
+                signed_json_body=signed_json_body,
+                shared_session=shared_session,
+            )
 
-            else:
-                return self.async_completion(
-                    custom_llm_provider=custom_llm_provider,
-                    provider_config=provider_config,
-                    api_base=api_base,
-                    headers=headers,
-                    data=data,
-                    timeout=timeout,
-                    model=model,
-                    model_response=model_response,
-                    logging_obj=logging_obj,
-                    api_key=api_key,
-                    messages=messages,
-                    optional_params=optional_params,
-                    litellm_params=litellm_params,
-                    encoding=encoding,
-                    client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
-                    json_mode=json_mode,
-                    signed_json_body=signed_json_body,
-                    shared_session=shared_session,
+        if acompletion is True and provider_config.uses_async_transform_request:
+
+            async def transform_then_dispatch():
+                transformed: Final = cast(  # cast-ok: async_transform_request is declared as a bare dict
+                    "dict[str, object]",
+                    await provider_config.async_transform_request(
+                        model=model,
+                        messages=messages,
+                        optional_params=optional_params,
+                        litellm_params=litellm_params,
+                        headers=request_headers,
+                    ),
                 )
+                return await dispatch_async(*sign_and_log(transformed))
+
+            return transform_then_dispatch()
+
+        data, signed_headers, signed_json_body = sign_and_log(
+            provider_config.transform_request(
+                model=model,
+                messages=messages,
+                optional_params=optional_params,
+                litellm_params=litellm_params,
+                headers=request_headers,
+            )
+        )
+
+        if acompletion is True:
+            return dispatch_async(data, signed_headers, signed_json_body)
 
         if stream is True:
             data = self._add_stream_param_to_request_body(
@@ -641,7 +669,7 @@ class BaseLLMHTTPHandler:
                     custom_llm_provider=custom_llm_provider,
                     logging_obj=logging_obj,
                     api_base=api_base,
-                    headers=headers,
+                    headers=signed_headers,
                     data=data,
                     signed_json_body=signed_json_body,
                     messages=messages,
@@ -651,7 +679,7 @@ class BaseLLMHTTPHandler:
             completion_stream, headers = self.make_sync_call(
                 provider_config=provider_config,
                 api_base=api_base,
-                headers=headers,
+                headers=signed_headers,
                 data=data,
                 signed_json_body=signed_json_body,
                 original_data=data,
@@ -684,7 +712,7 @@ class BaseLLMHTTPHandler:
             sync_httpx_client=sync_httpx_client,
             provider_config=provider_config,
             api_base=api_base,
-            headers=headers,
+            headers=signed_headers,
             data=data,
             signed_json_body=signed_json_body,
             timeout=timeout,
diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py
index c64fc583edc..f65b0876202 100644
--- a/litellm/llms/snowflake/chat/transformation.py
+++ b/litellm/llms/snowflake/chat/transformation.py
@@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
     create_anthropic_image_param,
     select_anthropic_content_block_type_for_file,
 )
+from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
 from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser
 from litellm.llms.anthropic.chat.transformation import AnthropicConfig
 from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload
@@ -421,6 +422,21 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
             return self._transform_request_anthropic(model, messages, optional_params, stream, extra_body)
         return self._transform_request_openai(model, messages, optional_params, stream, extra_body)
 
+    @property
+    def uses_async_transform_request(self) -> bool:
+        return True
+
+    async def async_transform_request(
+        self,
+        model: str,
+        messages: list[AllMessageValues],  # mutable-ok: BaseConfig signature
+        optional_params: dict[str, object],  # mutable-ok: BaseConfig signature
+        litellm_params: dict[str, object],  # mutable-ok: BaseConfig signature
+        headers: dict[str, object],  # mutable-ok: BaseConfig signature
+    ) -> dict[str, object]:  # mutable-ok: BaseConfig signature
+        inlined_messages: Final = await async_inline_remote_media(messages) if _is_claude_model(model) else messages
+        return self.transform_request(model, inlined_messages, optional_params, litellm_params, headers)
+
     def _transform_request_openai(
         self,
         model: str,
diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py
index e2d62be6a69..c9480f07150 100644
--- a/litellm/llms/vertex_ai/gemini/transformation.py
+++ b/litellm/llms/vertex_ai/gemini/transformation.py
@@ -27,6 +27,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
     convert_to_gemini_tool_call_result,
     response_schema_prompt,
 )
+from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
 from litellm.types.files import (
@@ -1348,13 +1349,15 @@ async def async_transform_request_body(
         vertex_auth_header=vertex_auth_header,
     )
 
-    if _openai_messages_may_need_sync_gcs_metadata_fetch(messages):
+    inlined_messages: Final = await async_inline_remote_media(messages) if custom_llm_provider == "gemini" else messages
+
+    if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages):
         # _transform_request_body may issue a sync httpx.get (up to 5s timeout)
         # via _get_gcs_object_content_type to fetch GCS object metadata. Run the
         # whole sync transformation on a worker thread so it does not block the
         # async event loop.
         return await asyncify(_transform_request_body)(
-            messages=messages,
+            messages=inlined_messages,
             model=model,
             custom_llm_provider=custom_llm_provider,
             litellm_params=litellm_params,
@@ -1363,7 +1366,7 @@ async def async_transform_request_body(
         )
 
     return _transform_request_body(
-        messages=messages,
+        messages=inlined_messages,
         model=model,
         custom_llm_provider=custom_llm_provider,
         litellm_params=litellm_params,
diff --git a/litellm/main.py b/litellm/main.py
index 2929790f2bd..9970c203e03 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -4474,7 +4474,7 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR
     acompletion: Final = ctx.acompletion
     api_base: Final = ctx.api_base
     api_key: Final = ctx.api_key
-    client = _dispatch_client_http(ctx)
+    injected_client: Final = _dispatch_client_http(ctx)
     custom_llm_provider: Final = ctx.custom_llm_provider
     headers: Final = ctx.headers
     litellm_params: Final = ctx.litellm_params
@@ -4486,11 +4486,11 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR
     shared_session: Final = ctx.shared_session
     stream: Final = ctx.stream
     timeout: Final = ctx.timeout
+    client: Final = (
+        injected_client if injected_client is not None else (HTTPHandler(timeout=timeout) if stream is False else None)
+    )  # Keep this here, otherwise, the httpx.client closes and streaming is impossible
 
     try:
-        client = (
-            HTTPHandler(timeout=timeout) if stream is False else None
-        )  # Keep this here, otherwise, the httpx.client closes and streaming is impossible
         response: Final = base_llm_http_handler.completion(
             model=model,
             messages=messages,
diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py
index 62c95cb100b..aa8ba168ecc 100644
--- a/tests/test_litellm/conftest.py
+++ b/tests/test_litellm/conftest.py
@@ -7,9 +7,12 @@
 # 4. Added proper cleanup in fixtures
 # 5. Added worker-specific isolation for parallel execution
 
+import base64
 import importlib
 import os
 from pathlib import Path
+from types import SimpleNamespace
+import httpx
 import pytest
 
 import asyncio
@@ -595,3 +598,35 @@ def pytest_sessionfinish(session, exitstatus):
     _close_handler_if_needed(getattr(litellm, "aclient", None))
     _close_handler_if_needed(getattr(litellm, "client", None))
     _run_coroutine_if_needed(close_litellm_async_clients())
+
+
+ONE_PIXEL_PNG = base64.b64decode(
+    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
+)
+
+
+@pytest.fixture
+def async_only_image_fetch(monkeypatch):
+    from litellm.litellm_core_utils.prompt_templates import image_handling
+
+    fetch = SimpleNamespace(
+        fetched=[],
+        base64_png=base64.b64encode(ONE_PIXEL_PNG).decode(),
+        data_url="data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode(),
+    )
+
+    def forbid_sync_fetch(client, url, **kwargs):
+        raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}")
+
+    async def serve_png(client, url, **kwargs):
+        fetch.fetched.append(url)
+        return httpx.Response(
+            200,
+            content=ONE_PIXEL_PNG,
+            headers={"content-type": "image/png"},
+            request=httpx.Request("GET", url),
+        )
+
+    monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch)
+    monkeypatch.setattr(image_handling, "async_safe_get", serve_png)
+    return fetch
diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py
index 893472d63ae..5b6d403fa21 100644
--- a/tests/test_litellm/litellm_core_utils/test_image_handling.py
+++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py
@@ -1,3 +1,5 @@
+import copy
+import uuid
 from unittest.mock import patch
 
 import pytest
@@ -8,6 +10,7 @@ from litellm import constants
 from litellm.litellm_core_utils.prompt_templates import image_handling
 from litellm.litellm_core_utils.prompt_templates.image_handling import (
     async_convert_url_to_base64,
+    async_inline_remote_media,
     convert_url_to_base64,
 )
 
@@ -268,3 +271,62 @@ def test_image_size_limit_disabled(monkeypatch):
 
     assert "Image URL download is disabled" in str(excinfo.value)
     assert "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0" in str(excinfo.value)
+
+
+async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_only_image_fetch):
+    image_url = f"http://img.example/{uuid.uuid4()}.png"
+    pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf"
+    messages = [
+        {"role": "system", "content": "be terse"},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "what is this?"},
+                {"type": "image_url", "image_url": {"url": image_url, "detail": "low"}},
+                {"type": "image_url", "image_url": image_url},
+                {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
+                {"type": "file", "file": {"file_id": pdf_url}},
+                {"type": "file", "file": {"file_id": image_url, "format": "image/png"}},
+            ],
+        },
+    ]
+    snapshot = copy.deepcopy(messages)
+
+    inlined = await async_inline_remote_media(messages)
+
+    data_url = async_only_image_fetch.data_url
+    assert inlined[0] == {"role": "system", "content": "be terse"}
+    assert inlined[1]["content"] == [
+        {"type": "text", "text": "what is this?"},
+        {"type": "image_url", "image_url": {"url": data_url, "detail": "low"}},
+        {"type": "image_url", "image_url": data_url},
+        {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
+        {"type": "file", "file": {"format": "application/pdf", "file_data": data_url}},
+        {"type": "file", "file": {"format": "image/png", "file_data": data_url}},
+    ]
+    assert sorted(async_only_image_fetch.fetched) == sorted([image_url, pdf_url])
+    assert messages == snapshot
+
+
+async def test_async_inline_remote_media_leaves_messages_without_remote_parts_alone(async_only_image_fetch):
+    messages = [
+        {"role": "user", "content": [{"type": "text", "text": "hi"}]},
+        {
+            "role": "user",
+            "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}],
+        },
+    ]
+
+    assert await async_inline_remote_media(messages) is messages
+    assert async_only_image_fetch.fetched == []
+
+
+async def test_async_inline_remote_media_raises_image_fetch_error_when_the_fetch_fails(monkeypatch):
+    async def serve_404(client, url, **kwargs):
+        return Response(404, request=Request("GET", url))
+
+    monkeypatch.setattr(image_handling, "async_safe_get", serve_404)
+    url = f"http://img.example/{uuid.uuid4()}.png"
+
+    with pytest.raises(litellm.ImageFetchError, match="Status code: 404"):
+        await async_inline_remote_media([{"role": "user", "content": [{"type": "image_url", "image_url": url}]}])
diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py
index 0d7573a2536..e808a087e3d 100644
--- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py
+++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py
@@ -1,15 +1,19 @@
 import asyncio
 import json
+import uuid
 from unittest.mock import patch
 
+import httpx
 import pytest
 
 # Ensure the project root is on the import path so `litellm` can be imported when
 # tests are executed from any working directory.
 
+import litellm
 from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
     AmazonAnthropicClaudeConfig,
 )
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 
 
 def test_get_supported_params_thinking():
@@ -714,3 +718,48 @@ def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking
 
     assert result["thinking"] == {"type": "adaptive"}
     assert result["output_config"] == {"effort": "high"}
+
+
+async def test_bedrock_invoke_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch):
+    image_url = f"http://img.example/{uuid.uuid4()}.png"
+    captured = {}
+
+    def handle(request):
+        captured["body"] = request.content.decode()
+        return httpx.Response(
+            200,
+            json={
+                "id": "msg_1",
+                "type": "message",
+                "role": "assistant",
+                "model": "us.anthropic.claude-sonnet-5",
+                "content": [{"type": "text", "text": "Green"}],
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 1, "output_tokens": 1},
+            },
+        )
+
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+    response = await litellm.acompletion(
+        model="bedrock/invoke/us.anthropic.claude-sonnet-5",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What colour is this?"},
+                    {"type": "image_url", "image_url": {"url": image_url}},
+                ],
+            }
+        ],
+        aws_access_key_id="AKIAEXAMPLE",
+        aws_secret_access_key="fake-secret",
+        aws_region_name="us-east-1",
+        client=client,
+    )
+
+    assert response.choices[0].message.content == "Green"
+    assert async_only_image_fetch.fetched == [image_url]
+    assert image_url not in captured["body"]
+    assert async_only_image_fetch.base64_png in captured["body"]
diff --git a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py
new file mode 100644
index 00000000000..9b491d305c7
--- /dev/null
+++ b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py
@@ -0,0 +1,51 @@
+import uuid
+
+import httpx
+
+import litellm
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+
+
+async def test_bedrock_mantle_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch):
+    image_url = f"http://img.example/{uuid.uuid4()}.png"
+    captured = {}
+
+    def handle(request):
+        captured["body"] = request.content.decode()
+        return httpx.Response(
+            200,
+            json={
+                "id": "msg_1",
+                "type": "message",
+                "role": "assistant",
+                "model": "us.anthropic.claude-sonnet-5",
+                "content": [{"type": "text", "text": "Green"}],
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 1, "output_tokens": 1},
+            },
+        )
+
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+    response = await litellm.acompletion(
+        model="bedrock/mantle/us.anthropic.claude-sonnet-5",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What colour is this?"},
+                    {"type": "image_url", "image_url": {"url": image_url}},
+                ],
+            }
+        ],
+        aws_access_key_id="AKIAEXAMPLE",
+        aws_secret_access_key="fake-secret",
+        aws_region_name="us-east-1",
+        client=client,
+    )
+
+    assert response.choices[0].message.content == "Green"
+    assert async_only_image_fetch.fetched == [image_url]
+    assert image_url not in captured["body"]
+    assert async_only_image_fetch.base64_png in captured["body"]
diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
index 023e2d8843f..e86e671b939 100644
--- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
+++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
@@ -17,7 +17,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import (
     AudioTranscriptionRequestData,
     BaseAudioTranscriptionConfig,
 )
-from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 from litellm.llms.custom_httpx.llm_http_handler import (
     BaseLLMHTTPHandler,
@@ -30,7 +30,7 @@ from litellm.llms.azure.videos.transformation import AzureVideoConfig
 from litellm.llms.openai.videos.transformation import OpenAIVideoConfig
 from litellm.types.llms.openai import ResponsesAPIResponse
 from litellm.types.router import GenericLiteLLMParams
-from litellm.types.utils import TranscriptionResponse
+from litellm.types.utils import ModelResponse, TranscriptionResponse
 
 _ACTIVE_KEY = "_code_interpreter_interception_active"
 _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key"
@@ -3186,3 +3186,105 @@ async def test_async_container_list_handler_transforms_success_response():
 
     assert [container.id for container in response.data] == ["cntr_a"]
     assert response.has_more is True
+
+
+class _TransformRecordingConfig(BaseConfig):
+    def __init__(self, transform_async: bool):
+        self.transform_async = transform_async
+        self.transform_calls = []
+
+    @property
+    def uses_async_transform_request(self) -> bool:
+        return self.transform_async
+
+    def get_supported_openai_params(self, model):
+        return []
+
+    def map_openai_params(self, non_default_params, optional_params, model, drop_params):
+        return optional_params
+
+    def validate_environment(
+        self, headers, model, messages, optional_params, litellm_params, api_key=None, api_base=None
+    ):
+        return {}
+
+    def transform_request(self, model, messages, optional_params, litellm_params, headers):
+        self.transform_calls.append("sync")
+        return {"transformed_by": "sync"}
+
+    async def async_transform_request(self, model, messages, optional_params, litellm_params, headers):
+        self.transform_calls.append("async")
+        return {"transformed_by": "async"}
+
+    def transform_response(
+        self,
+        model,
+        raw_response,
+        model_response,
+        logging_obj,
+        request_data,
+        messages,
+        optional_params,
+        litellm_params,
+        encoding,
+        api_key=None,
+        json_mode=None,
+    ):
+        model_response.choices[0].message.content = raw_response.json()["transformed_by"]
+        return model_response
+
+    def get_error_class(self, error_message, status_code, headers):
+        return BaseLLMException(status_code=status_code, message=error_message, headers=headers)
+
+
+def _start_async_completion(config):
+    captured = {}
+
+    def handle(request):
+        captured["body"] = json.loads(request.content)
+        return httpx.Response(200, json=captured["body"])
+
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+    pending = BaseLLMHTTPHandler().completion(
+        model="stub-model",
+        messages=[{"role": "user", "content": "hi"}],
+        api_base="https://llm.example/v1/chat",
+        custom_llm_provider="openai",
+        model_response=ModelResponse(),
+        encoding=None,
+        logging_obj=Mock(dynamic_success_callbacks=None, model_call_details={}),
+        optional_params={},
+        timeout=10.0,
+        litellm_params={},
+        acompletion=True,
+        client=client,
+        provider_config=config,
+    )
+    return pending, captured
+
+
+async def test_completion_awaits_async_transform_request_when_config_opts_in():
+    config = _TransformRecordingConfig(transform_async=True)
+
+    pending, captured = _start_async_completion(config)
+    assert config.transform_calls == []
+
+    response = await pending
+
+    assert config.transform_calls == ["async"]
+    assert captured["body"] == {"transformed_by": "async"}
+    assert response.choices[0].message.content == "async"
+
+
+async def test_completion_keeps_sync_transform_request_before_returning_by_default():
+    config = _TransformRecordingConfig(transform_async=False)
+
+    pending, captured = _start_async_completion(config)
+    assert config.transform_calls == ["sync"]
+
+    response = await pending
+
+    assert config.transform_calls == ["sync"]
+    assert captured["body"] == {"transformed_by": "sync"}
+    assert response.choices[0].message.content == "sync"
diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py
index 25a961c3413..5687a319f06 100644
--- a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py
+++ b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py
@@ -6,6 +6,7 @@ Tests tool calling request/response transformations and chat completions
 import asyncio
 import os
 import copy
+import uuid
 import json
 from typing import Any, Dict, List
 
@@ -945,3 +946,48 @@ class TestSnowflakeChatCompletion:
 
         assert len(chunks_received) > 0
         content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content)
+
+
+async def test_snowflake_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch):
+    image_url = f"http://img.example/{uuid.uuid4()}.png"
+    captured = {}
+
+    def handle(request):
+        captured["body"] = request.content.decode()
+        return httpx.Response(
+            200,
+            json={
+                "id": "msg_1",
+                "type": "message",
+                "role": "assistant",
+                "model": "claude-sonnet-4-6",
+                "content": [{"type": "text", "text": "Green"}],
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 1, "output_tokens": 1},
+            },
+        )
+
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+    response = await litellm.acompletion(
+        model="snowflake/claude-sonnet-4-6",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What colour is this?"},
+                    {"type": "image_url", "image_url": {"url": image_url}},
+                ],
+            }
+        ],
+        api_key="fake-jwt",
+        account_id="FAKE-ACCOUNT",
+        api_base=FAKE_API_BASE,
+        client=client,
+    )
+
+    assert response.choices[0].message.content == "Green"
+    assert async_only_image_fetch.fetched == [image_url]
+    assert image_url not in captured["body"]
+    assert async_only_image_fetch.base64_png in captured["body"]
diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py
index fad310fc5c0..ec445342523 100644
--- a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py
@@ -1,6 +1,10 @@
 
+import uuid
+import httpx
 import pytest
 
+import litellm
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.llms.vertex_ai.gemini import transformation
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
     VertexGeminiConfig,
@@ -338,3 +342,41 @@ def test_map_function_enterprise_web_search_snake_case():
 
     assert len(result) == 1
     assert "enterpriseWebSearch" in result[0]
+
+
+async def test_gemini_ai_studio_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch):
+    image_url = f"http://img.example/{uuid.uuid4()}.png"
+    captured = {}
+
+    def handle(request):
+        captured["body"] = request.content.decode()
+        return httpx.Response(
+            200,
+            json={
+                "candidates": [{"content": {"parts": [{"text": "Green"}], "role": "model"}, "finishReason": "STOP"}],
+                "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
+            },
+        )
+
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+    response = await litellm.acompletion(
+        model="gemini/gemini-3.8-flash",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What colour is this?"},
+                    {"type": "image_url", "image_url": {"url": image_url}},
+                ],
+            }
+        ],
+        api_key="fake-gemini-key",
+        client=client,
+    )
+
+    assert response.choices[0].message.content == "Green"
+    assert async_only_image_fetch.fetched == [image_url]
+    assert image_url not in captured["body"]
+    assert async_only_image_fetch.base64_png in captured["body"]

From e355203014326462c3e70e6b6c8d37fcc0abe656 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 18:11:41 -0700
Subject: [PATCH 040/164] fix(gemini): keep Files API refs out of the async
 remote media walker

Gemini AI Studio file URIs under generativelanguage.googleapis.com/v1beta/files/
answer 403 when fetched and must pass through as file_data.file_uri, which the
sync transform already did. Give the shared walker a skip_url_prefixes parameter,
pass that prefix from the Gemini body builder, and skip it in the AI Studio
message transform for both image_url and file parts
---
 .../prompt_templates/image_handling.py        | 12 +++--
 litellm/llms/gemini/chat/transformation.py    | 14 ++++--
 .../llms/vertex_ai/gemini/transformation.py   |  9 +++-
 .../litellm_core_utils/test_image_handling.py | 28 ++++++++++++
 .../vertex_ai/gemini/test_transformation.py   | 44 +++++++++++++++++++
 5 files changed, 99 insertions(+), 8 deletions(-)

diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py
index 1890d4eb682..5ae8d224556 100644
--- a/litellm/litellm_core_utils/prompt_templates/image_handling.py
+++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py
@@ -199,13 +199,18 @@ def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]:
     return tuple(content) if isinstance(content, list) else ()  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]  # parts are parsed one by one
 
 
+def _inline_part(part: object, data_urls: Mapping[str, str]) -> object:
+    remote: Final = _parse_remote_part(part)
+    data_url: Final = data_urls.get(remote.url) if remote is not None else None
+    return _inline(remote, data_url) if remote is not None and data_url is not None else part
+
+
 def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> AllMessageValues:
     parts: Final = _content_parts(message)
     if not parts:
         return message
     inlined_parts: Final = [  # mutable-ok: content must stay a list for the transforms' isinstance checks
-        _inline(remote, data_urls[remote.url]) if (remote := _parse_remote_part(part)) is not None else part
-        for part in parts
+        _inline_part(part, data_urls) for part in parts
     ]
     inlined_message: Final = {**message, "content": inlined_parts}  # mutable-ok: json-serialized message
     return inlined_message  # pyright: ignore[reportReturnType]  # the same message with its remote parts inlined
@@ -213,13 +218,14 @@ def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) ->
 
 async def async_inline_remote_media(
     messages: list[AllMessageValues],  # mutable-ok: every transform_request takes list[AllMessageValues]
+    skip_url_prefixes: tuple[str, ...] = (),
 ) -> list[AllMessageValues]:  # mutable-ok: every transform_request takes list[AllMessageValues]
     remote_urls: Final = tuple(
         dict.fromkeys(
             remote.url
             for message in messages
             for part in _content_parts(message)
-            if (remote := _parse_remote_part(part)) is not None
+            if (remote := _parse_remote_part(part)) is not None and not remote.url.startswith(skip_url_prefixes)
         )
     )
     if not remote_urls:
diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py
index 1a67b33665b..42c9ef13730 100644
--- a/litellm/llms/gemini/chat/transformation.py
+++ b/litellm/llms/gemini/chat/transformation.py
@@ -12,7 +12,7 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionFileObject
 from litellm.types.llms.vertex_ai import ContentType, PartType
 from litellm.utils import supports_reasoning
 
-from ...vertex_ai.gemini.transformation import _gemini_convert_messages_with_history
+from ...vertex_ai.gemini.transformation import GEMINI_FILES_API_URI_PREFIX, _gemini_convert_messages_with_history
 from ...vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
 
 
@@ -127,7 +127,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
                     if element.get("type") == "image_url":
                         img_element = cast(ChatCompletionImageObject, element)  # cast-ok: runtime type tag checked
                         _image_url, format, detail = _image_url_fields(img_element)
-                        if _image_url and "https://" in _image_url:
+                        if (
+                            _image_url
+                            and "https://" in _image_url
+                            and not _image_url.startswith(GEMINI_FILES_API_URI_PREFIX)
+                        ):
                             image_obj = convert_to_anthropic_image_obj(_image_url, format=format)
                             converted_image_url = convert_generic_image_chunk_to_openai_image_obj(image_obj)
                             if detail is not None:
@@ -147,7 +151,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
                                 llm_provider="gemini",
                             )
                         file_id = _file_field.get("file_id")
-                        if file_id and ("http://" in file_id or "https://" in file_id):
+                        if (
+                            file_id
+                            and ("http://" in file_id or "https://" in file_id)
+                            and not file_id.startswith(GEMINI_FILES_API_URI_PREFIX)
+                        ):
                             # Convert HTTP/HTTPS file URL to base64 data
                             try:
                                 base64_data = convert_url_to_base64(file_id)
diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py
index c9480f07150..6d100143e52 100644
--- a/litellm/llms/vertex_ai/gemini/transformation.py
+++ b/litellm/llms/vertex_ai/gemini/transformation.py
@@ -69,6 +69,7 @@ _GCS_METADATA_VERTEX_BASE: Any | None = None
 # Shared sync client for GCS JSON API metadata reads so proxy/SSL settings
 # from litellm's HTTP stack apply (see Greptile review on PR #27278).
 _GCS_METADATA_HTTP_HANDLER: HTTPHandler | None = None
+GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/"
 _GEMINI_MIME_TYPE_ALIASES: Final[dict[str, str]] = {
     "image/jpg": "image/jpeg",
 }
@@ -557,7 +558,7 @@ def _process_gemini_media(
             file_data = FileDataType(mime_type=mime_type, file_uri=image_url)
             part: PartType = {"file_data": file_data}
             return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata)
-        elif image_url.startswith("https://generativelanguage.googleapis.com/v1beta/files/"):
+        elif image_url.startswith(GEMINI_FILES_API_URI_PREFIX):
             # Gemini Files API URIs — the file is already uploaded to Google's
             # servers; pass the URI through as file_data without fetching it.
             # These URLs return 403 when accessed directly, so we must not try
@@ -1349,7 +1350,11 @@ async def async_transform_request_body(
         vertex_auth_header=vertex_auth_header,
     )
 
-    inlined_messages: Final = await async_inline_remote_media(messages) if custom_llm_provider == "gemini" else messages
+    inlined_messages: Final = (
+        await async_inline_remote_media(messages, skip_url_prefixes=(GEMINI_FILES_API_URI_PREFIX,))
+        if custom_llm_provider == "gemini"
+        else messages
+    )
 
     if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages):
         # _transform_request_body may issue a sync httpx.get (up to 5s timeout)
diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py
index 5b6d403fa21..bcf4c7cb6ff 100644
--- a/tests/test_litellm/litellm_core_utils/test_image_handling.py
+++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py
@@ -308,6 +308,34 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o
     assert messages == snapshot
 
 
+async def test_async_inline_remote_media_leaves_skipped_url_prefixes_untouched(async_only_image_fetch):
+    skipped_prefix = "https://generativelanguage.googleapis.com/v1beta/files/"
+    skipped_file = f"{skipped_prefix}{uuid.uuid4().hex}"
+    skipped_image = f"{skipped_prefix}{uuid.uuid4().hex}"
+    fetched_image = f"https://img.example/{uuid.uuid4()}.png"
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}},
+                {"type": "image_url", "image_url": {"url": skipped_image}},
+                {"type": "image_url", "image_url": fetched_image},
+            ],
+        }
+    ]
+    snapshot = copy.deepcopy(messages)
+
+    inlined = await async_inline_remote_media(messages, skip_url_prefixes=(skipped_prefix,))
+
+    assert inlined[0]["content"] == [
+        {"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}},
+        {"type": "image_url", "image_url": {"url": skipped_image}},
+        {"type": "image_url", "image_url": async_only_image_fetch.data_url},
+    ]
+    assert async_only_image_fetch.fetched == [fetched_image]
+    assert messages == snapshot
+
+
 async def test_async_inline_remote_media_leaves_messages_without_remote_parts_alone(async_only_image_fetch):
     messages = [
         {"role": "user", "content": [{"type": "text", "text": "hi"}]},
diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py
index ec445342523..d31254746d4 100644
--- a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py
@@ -1,4 +1,5 @@
 
+import json
 import uuid
 import httpx
 import pytest
@@ -380,3 +381,46 @@ async def test_gemini_ai_studio_async_completion_inlines_remote_images_off_the_e
     assert async_only_image_fetch.fetched == [image_url]
     assert image_url not in captured["body"]
     assert async_only_image_fetch.base64_png in captured["body"]
+
+
+async def test_gemini_ai_studio_async_completion_passes_files_api_uris_through_unfetched(async_only_image_fetch):
+    files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}"
+    files_api_image = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}"
+    captured = {}
+
+    def handle(request):
+        captured["body"] = json.loads(request.content)
+        return httpx.Response(
+            200,
+            json={
+                "candidates": [{"content": {"parts": [{"text": "A report"}], "role": "model"}, "finishReason": "STOP"}],
+                "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
+            },
+        )
+
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+    response = await litellm.acompletion(
+        model="gemini/gemini-3.8-flash",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "Summarize these"},
+                    {"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}},
+                    {"type": "image_url", "image_url": {"url": files_api_image, "format": "image/png"}},
+                ],
+            }
+        ],
+        api_key="fake-gemini-key",
+        client=client,
+    )
+
+    assert response.choices[0].message.content == "A report"
+    assert async_only_image_fetch.fetched == []
+    file_parts = [part["file_data"] for part in captured["body"]["contents"][0]["parts"] if "file_data" in part]
+    assert file_parts == [
+        {"mime_type": "application/pdf", "file_uri": files_api_pdf},
+        {"mime_type": "image/png", "file_uri": files_api_image},
+    ]

From 08bb7de868f0d740c937585835333ac226faccad Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 18:28:17 -0700
Subject: [PATCH 041/164] fix(image_handling): cap in-flight remote media
 fetches per request

---
 .../prompt_templates/image_handling.py        |  9 +++++++-
 .../litellm_core_utils/test_image_handling.py | 22 +++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py
index 5ae8d224556..4beaabc8b24 100644
--- a/litellm/litellm_core_utils/prompt_templates/image_handling.py
+++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py
@@ -19,6 +19,7 @@ from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
 from litellm.types.llms.openai import AllMessageValues
 
 MAX_IMGS_IN_MEMORY: Final = 10
+MAX_CONCURRENT_REMOTE_MEDIA_FETCHES: Final = 20
 
 in_memory_cache: Final = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY)
 
@@ -216,6 +217,11 @@ def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) ->
     return inlined_message  # pyright: ignore[reportReturnType]  # the same message with its remote parts inlined
 
 
+async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore) -> str:
+    async with in_flight:
+        return await async_convert_url_to_base64(url)
+
+
 async def async_inline_remote_media(
     messages: list[AllMessageValues],  # mutable-ok: every transform_request takes list[AllMessageValues]
     skip_url_prefixes: tuple[str, ...] = (),
@@ -230,6 +236,7 @@ async def async_inline_remote_media(
     )
     if not remote_urls:
         return messages
-    data_urls: Final = await asyncio.gather(*(async_convert_url_to_base64(url) for url in remote_urls))
+    in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES)
+    data_urls: Final = await asyncio.gather(*(_fetch_data_url(url, in_flight) for url in remote_urls))
     inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True)))
     return [_inline_message(message, inlined) for message in messages]  # mutable-ok: transform_request takes a list
diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py
index bcf4c7cb6ff..ae9016f2f9e 100644
--- a/tests/test_litellm/litellm_core_utils/test_image_handling.py
+++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py
@@ -1,3 +1,4 @@
+import asyncio
 import copy
 import uuid
 from unittest.mock import patch
@@ -9,6 +10,7 @@ import litellm
 from litellm import constants
 from litellm.litellm_core_utils.prompt_templates import image_handling
 from litellm.litellm_core_utils.prompt_templates.image_handling import (
+    MAX_CONCURRENT_REMOTE_MEDIA_FETCHES,
     async_convert_url_to_base64,
     async_inline_remote_media,
     convert_url_to_base64,
@@ -336,6 +338,26 @@ async def test_async_inline_remote_media_leaves_skipped_url_prefixes_untouched(a
     assert messages == snapshot
 
 
+async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch):
+    in_flight = {"now": 0, "peak": 0}
+
+    async def serve_png_slowly(client, url, **kwargs):
+        in_flight["now"] += 1
+        in_flight["peak"] = max(in_flight["peak"], in_flight["now"])
+        await asyncio.sleep(0.01)
+        in_flight["now"] -= 1
+        return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url))
+
+    monkeypatch.setattr(image_handling, "async_safe_get", serve_png_slowly)
+    urls = [f"https://img.example/{uuid.uuid4()}.png" for _ in range(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES + 5)]
+    messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}} for url in urls]}]
+
+    inlined = await async_inline_remote_media(messages)
+
+    assert in_flight["peak"] == MAX_CONCURRENT_REMOTE_MEDIA_FETCHES
+    assert all(part["image_url"]["url"].startswith("data:image/png;base64,") for part in inlined[0]["content"])
+
+
 async def test_async_inline_remote_media_leaves_messages_without_remote_parts_alone(async_only_image_fetch):
     messages = [
         {"role": "user", "content": [{"type": "text", "text": "hi"}]},

From a534b9fac595a743146a02dc2d1ffa1cab94e533 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 18:35:23 -0700
Subject: [PATCH 042/164] fix(fireworks_ai): send developer input items as
 system messages on the native responses path

---
 .../fireworks_ai/responses/transformation.py  | 15 +++++++++++++-
 ...t_fireworks_ai_responses_transformation.py | 20 +++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py
index 7f29e03ff0f..9265d12e75e 100644
--- a/litellm/llms/fireworks_ai/responses/transformation.py
+++ b/litellm/llms/fireworks_ai/responses/transformation.py
@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Final
 from urllib.parse import unquote
 
 import httpx
+from openai.types.responses import EasyInputMessageParam, ResponseInputItemParam
 
 from litellm.llms.fireworks_ai.common_utils import (
     resolve_fireworks_api_key,
@@ -30,6 +31,18 @@ def _session_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, object
     )
 
 
+def _developer_item_as_system(item: ResponseInputItemParam) -> ResponseInputItemParam:
+    if "role" not in item or item["role"] != "developer":
+        return item
+    return EasyInputMessageParam(role="system", content=item["content"], type="message")
+
+
+def _developer_items_as_system(input: str | ResponseInputParam) -> str | ResponseInputParam:
+    if isinstance(input, str):
+        return input
+    return [_developer_item_as_system(item) for item in input]
+
+
 class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
     @property
     def custom_llm_provider(self) -> LlmProviders:
@@ -65,7 +78,7 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
     ) -> dict:  # mutable-ok: overrides the base class signature
         return super().transform_responses_api_request(
             model=resolve_fireworks_resource_name(model),
-            input=input,
+            input=_developer_items_as_system(input),
             response_api_optional_request_params=response_api_optional_request_params,
             litellm_params=litellm_params,
             headers=headers,
diff --git a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
index 8fac91ee475..e6e92824ee1 100644
--- a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
+++ b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
@@ -160,6 +160,26 @@ def test_responses_call_forwards_previous_response_id_and_store() -> None:
     assert body["input"][0]["call_id"] == "call_abc123"
 
 
+def test_responses_call_sends_developer_items_as_system_messages() -> None:
+    client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
+    with patch(HTTPX_CLIENT_FACTORY, return_value=client):
+        litellm.responses(
+            model="fireworks_ai/accounts/fireworks/models/kimi-k3",
+            input=[  # mutable-ok: the Responses API takes input as a JSON list
+                {"role": "user", "content": "Hi there"},
+                {"role": "developer", "content": "Answer with exactly one word."},
+                {"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]},
+            ],
+            api_key="fw-test-key",
+        )
+    _, _, body = _sent_request(client)
+    assert tuple(body["input"]) == (
+        {"role": "user", "content": "Hi there"},
+        {"role": "system", "content": "Answer with exactly one word.", "type": "message"},
+        {"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]},
+    )
+
+
 def test_responses_call_sends_session_affinity_for_caller_session_id() -> None:
     client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
     with patch(HTTPX_CLIENT_FACTORY, return_value=client):

From 567915aeee82953c4127456654f3a63ec9862970 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 18:51:14 -0700
Subject: [PATCH 043/164] fix(image_handling): inline url-sourced Anthropic
 document and image blocks in the async walker

---
 .../prompt_templates/image_handling.py        | 51 +++++++++++++++----
 .../litellm_core_utils/test_image_handling.py | 11 ++++
 ...ations_anthropic_claude3_transformation.py | 47 +++++++++++++++++
 ...test_bedrock_chat_mantle_transformation.py | 48 +++++++++++++++++
 4 files changed, 147 insertions(+), 10 deletions(-)

diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py
index 4beaabc8b24..55c01b3b1cb 100644
--- a/litellm/litellm_core_utils/prompt_templates/image_handling.py
+++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py
@@ -157,18 +157,41 @@ def _remote_url(candidate: object) -> str | None:
     return candidate if isinstance(candidate, str) and candidate.startswith(_REMOTE_URL_PREFIXES) else None
 
 
-def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | None:
+_ANTHROPIC_MEDIA_BLOCK_TYPES: Final = frozenset({"document", "image"})
+
+
+@dataclass(frozen=True, slots=True)
+class _RemoteSource:
+    part: Mapping[str, object]
+    url: str
+
+
+def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None:
+    if fields.get("type") != "image_url":
+        return None
+    image_url: Final = fields.get("image_url")
+    image_url_fields: Final = _as_mapping(image_url)
+    url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url)
+    return _RemoteImage(fields, image_url_fields, url) if url is not None else None
+
+
+def _parse_remote_file(fields: Mapping[str, object]) -> _RemoteFile | None:
+    file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None
+    url: Final = _remote_url(file.get("file_id")) if file is not None else None
+    return _RemoteFile(fields, file, url) if file is not None and url is not None else None
+
+
+def _parse_remote_source(fields: Mapping[str, object]) -> _RemoteSource | None:
+    source: Final = _as_mapping(fields.get("source")) if fields.get("type") in _ANTHROPIC_MEDIA_BLOCK_TYPES else None
+    url: Final = _remote_url(source.get("url")) if source is not None and source.get("type") == "url" else None
+    return _RemoteSource(fields, url) if url is not None else None
+
+
+def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSource | None:
     fields: Final = _as_mapping(part)
     if fields is None:
         return None
-    if fields.get("type") == "image_url":
-        image_url: Final = fields.get("image_url")
-        image_url_fields: Final = _as_mapping(image_url)
-        url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url)
-        return _RemoteImage(fields, image_url_fields, url) if url is not None else None
-    file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None
-    file_url: Final = _remote_url(file.get("file_id")) if file is not None else None
-    return _RemoteFile(fields, file, file_url) if file is not None and file_url is not None else None
+    return _parse_remote_image(fields) or _parse_remote_file(fields) or _parse_remote_source(fields)
 
 
 _PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"})
@@ -187,12 +210,20 @@ def _inlined_file(file: Mapping[str, object], url: str, data_url: str) -> Mappin
     return {**kept, **_inferred_format(file, url), "file_data": data_url}  # mutable-ok: json-serialized part
 
 
-def _inline(remote: _RemoteImage | _RemoteFile, data_url: str) -> Mapping[str, object]:
+def _base64_source(url: str, data_url: str) -> Mapping[str, str]:
+    fetched_media_type, data = data_url.removeprefix("data:").split(";base64,", 1)
+    media_type: Final = "application/pdf" if url.lower().endswith(".pdf") else fetched_media_type
+    return {"type": "base64", "media_type": media_type, "data": data}  # mutable-ok: json-serialized message part
+
+
+def _inline(remote: _RemoteImage | _RemoteFile | _RemoteSource, data_url: str) -> Mapping[str, object]:
     match remote:
         case _RemoteImage(part, image_url, _):
             return {**part, "image_url": _inlined_image_url(image_url, data_url)}  # mutable-ok: json-serialized part
         case _RemoteFile(part, file, url):
             return {**part, "file": _inlined_file(file, url, data_url)}  # mutable-ok: json-serialized message part
+        case _RemoteSource(part, url):
+            return {**part, "source": _base64_source(url, data_url)}  # mutable-ok: json-serialized message part
 
 
 def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]:
diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py
index ae9016f2f9e..106bf90213b 100644
--- a/tests/test_litellm/litellm_core_utils/test_image_handling.py
+++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py
@@ -289,6 +289,9 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o
                 {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
                 {"type": "file", "file": {"file_id": pdf_url}},
                 {"type": "file", "file": {"file_id": image_url, "format": "image/png"}},
+                {"type": "document", "source": {"type": "url", "url": pdf_url}, "title": "the doc"},
+                {"type": "image", "source": {"type": "url", "url": image_url}},
+                {"type": "document", "source": {"type": "file", "file_id": "file_abc"}},
             ],
         },
     ]
@@ -297,6 +300,7 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o
     inlined = await async_inline_remote_media(messages)
 
     data_url = async_only_image_fetch.data_url
+    base64_png = async_only_image_fetch.base64_png
     assert inlined[0] == {"role": "system", "content": "be terse"}
     assert inlined[1]["content"] == [
         {"type": "text", "text": "what is this?"},
@@ -305,6 +309,13 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o
         {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
         {"type": "file", "file": {"format": "application/pdf", "file_data": data_url}},
         {"type": "file", "file": {"format": "image/png", "file_data": data_url}},
+        {
+            "type": "document",
+            "source": {"type": "base64", "media_type": "application/pdf", "data": base64_png},
+            "title": "the doc",
+        },
+        {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": base64_png}},
+        {"type": "document", "source": {"type": "file", "file_id": "file_abc"}},
     ]
     assert sorted(async_only_image_fetch.fetched) == sorted([image_url, pdf_url])
     assert messages == snapshot
diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py
index e808a087e3d..cbf160c451f 100644
--- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py
+++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py
@@ -763,3 +763,50 @@ async def test_bedrock_invoke_claude_async_completion_inlines_remote_images_off_
     assert async_only_image_fetch.fetched == [image_url]
     assert image_url not in captured["body"]
     assert async_only_image_fetch.base64_png in captured["body"]
+
+
+async def test_bedrock_invoke_claude_async_completion_inlines_document_url_sources_off_the_event_loop(async_only_image_fetch):
+    pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf"
+    captured = {}
+
+    def handle(request):
+        captured["body"] = json.loads(request.content)
+        return httpx.Response(
+            200,
+            json={
+                "id": "msg_1",
+                "type": "message",
+                "role": "assistant",
+                "model": "us.anthropic.claude-sonnet-5",
+                "content": [{"type": "text", "text": "A lease"}],
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 1, "output_tokens": 1},
+            },
+        )
+
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+    response = await litellm.acompletion(
+        model="bedrock/invoke/us.anthropic.claude-sonnet-5",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What is this document?"},
+                    {"type": "document", "source": {"type": "url", "url": pdf_url}},
+                ],
+            }
+        ],
+        aws_access_key_id="AKIAEXAMPLE",
+        aws_secret_access_key="fake-secret",
+        aws_region_name="us-east-1",
+        client=client,
+    )
+
+    assert response.choices[0].message.content == "A lease"
+    assert async_only_image_fetch.fetched == [pdf_url]
+    assert {
+        "type": "document",
+        "source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png},
+    } in captured["body"]["messages"][0]["content"]
diff --git a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py
index 9b491d305c7..a8448f5fa7a 100644
--- a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py
+++ b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py
@@ -1,3 +1,4 @@
+import json
 import uuid
 
 import httpx
@@ -49,3 +50,50 @@ async def test_bedrock_mantle_claude_async_completion_inlines_remote_images_off_
     assert async_only_image_fetch.fetched == [image_url]
     assert image_url not in captured["body"]
     assert async_only_image_fetch.base64_png in captured["body"]
+
+
+async def test_bedrock_mantle_claude_async_completion_inlines_document_url_sources_off_the_event_loop(async_only_image_fetch):
+    pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf"
+    captured = {}
+
+    def handle(request):
+        captured["body"] = json.loads(request.content)
+        return httpx.Response(
+            200,
+            json={
+                "id": "msg_1",
+                "type": "message",
+                "role": "assistant",
+                "model": "us.anthropic.claude-sonnet-5",
+                "content": [{"type": "text", "text": "A lease"}],
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 1, "output_tokens": 1},
+            },
+        )
+
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+    response = await litellm.acompletion(
+        model="bedrock/mantle/us.anthropic.claude-sonnet-5",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What is this document?"},
+                    {"type": "document", "source": {"type": "url", "url": pdf_url}},
+                ],
+            }
+        ],
+        aws_access_key_id="AKIAEXAMPLE",
+        aws_secret_access_key="fake-secret",
+        aws_region_name="us-east-1",
+        client=client,
+    )
+
+    assert response.choices[0].message.content == "A lease"
+    assert async_only_image_fetch.fetched == [pdf_url]
+    assert {
+        "type": "document",
+        "source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png},
+    } in captured["body"]["messages"][0]["content"]

From 0157808187bc8dec07205b78cd67dc9e2cea40d8 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 18:58:09 -0700
Subject: [PATCH 044/164] fix(fireworks_ai): keep file_search on LiteLLM's
 emulated search for the Responses API

---
 litellm/llms/fireworks_ai/responses/transformation.py    | 3 +++
 .../test_fireworks_ai_responses_transformation.py        | 9 +++++++++
 2 files changed, 12 insertions(+)

diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py
index 9265d12e75e..121d3b5d9a9 100644
--- a/litellm/llms/fireworks_ai/responses/transformation.py
+++ b/litellm/llms/fireworks_ai/responses/transformation.py
@@ -94,3 +94,6 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
 
     def supports_native_websocket(self) -> bool:
         return False
+
+    def supports_native_file_search(self) -> bool:
+        return False
diff --git a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
index e6e92824ee1..f948acb7bf3 100644
--- a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
+++ b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
@@ -19,6 +19,7 @@ from typing_extensions import ReadOnly
 
 import litellm
 from litellm.llms.fireworks_ai.responses.transformation import FireworksAIResponsesAPIConfig
+from litellm.responses.file_search.emulated_handler import should_use_emulated_file_search
 from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponsesAPIResponse
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import LlmProviders
@@ -180,6 +181,14 @@ def test_responses_call_sends_developer_items_as_system_messages() -> None:
     )
 
 
+def test_file_search_tools_take_litellm_emulated_search_not_fireworks() -> None:
+    config: Final = FireworksAIResponsesAPIConfig()
+    file_search: Final = ({"type": "file_search", "vector_store_ids": ("vs_kb",)},)
+    function_tool: Final = ({"type": "function", "name": "get_weather", "parameters": {"type": "object"}},)
+    assert should_use_emulated_file_search(tools=file_search, provider_config=config)
+    assert not should_use_emulated_file_search(tools=function_tool, provider_config=config)
+
+
 def test_responses_call_sends_session_affinity_for_caller_session_id() -> None:
     client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
     with patch(HTTPX_CLIENT_FACTORY, return_value=client):

From 001b531ae960aad3fa38447508d4ce0d05c118c4 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 19:11:27 -0700
Subject: [PATCH 045/164] fix(cost): bill realtime reasoning tokens nested in
 text_tokens once

OpenAI and Azure realtime usage reports output_tokens == text_tokens + audio_tokens
with reasoning_tokens counted inside text_tokens, so generic_cost_per_token billed
the reasoning share twice. When the output token details sum past completion_tokens,
the nested reasoning overlap is now subtracted from text_tokens before pricing;
shapes where text_tokens already excludes reasoning are unchanged.
---
 .../litellm_core_utils/llm_cost_calc/utils.py | 19 ++++++-
 .../llm_cost_calc/test_llm_cost_calc_utils.py | 51 +++++++++++++++++++
 tests/test_litellm/test_cost_calculator.py    | 42 +++++++++++++++
 3 files changed, 111 insertions(+), 1 deletion(-)

diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py
index 8e24302b440..1b8980abd35 100644
--- a/litellm/litellm_core_utils/llm_cost_calc/utils.py
+++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py
@@ -852,6 +852,17 @@ class CompletionTokensDetailsResult(TypedDict):
     video_tokens: int
 
 
+def _text_tokens_without_nested_reasoning(
+    completion_tokens: int,
+    text_tokens: int,
+    reasoning_tokens: int,
+    other_modality_tokens: int,
+) -> int:
+    reported_total: Final = text_tokens + reasoning_tokens + other_modality_tokens
+    nested_reasoning_tokens: Final = min(reasoning_tokens, text_tokens, max(reported_total - completion_tokens, 0))
+    return text_tokens - nested_reasoning_tokens
+
+
 def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult:
     audio_tokens: Final = (
         cast(
@@ -860,7 +871,7 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu
         )
         or 0
     )
-    text_tokens: Final = (
+    reported_text_tokens: Final = (
         cast(
             int | None,
             getattr(usage.completion_tokens_details, "text_tokens", None),
@@ -882,6 +893,12 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu
         or 0
     )
     video_tokens: Final = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0))
+    text_tokens: Final = _text_tokens_without_nested_reasoning(
+        completion_tokens=usage.completion_tokens,
+        text_tokens=reported_text_tokens,
+        reasoning_tokens=reasoning_tokens,
+        other_modality_tokens=audio_tokens + image_tokens + video_tokens,
+    )
 
     return CompletionTokensDetailsResult(
         audio_tokens=audio_tokens,
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 0f8084643ea..f7ca2a89048 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
@@ -4723,3 +4723,54 @@ def test_route_image_generation_cost_falls_back_to_requested_size(monkeypatch, r
     )
 
     assert cost == expected_cost
+
+
+def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_local_model_cost_map):
+    """
+    Realtime usage (OpenAI and Azure) reports output_tokens == text_tokens + audio_tokens with
+    reasoning_tokens already counted inside text_tokens, so reasoning must not be billed on top.
+    """
+
+    model = "gpt-realtime-2.1-mini"
+    usage = Usage(
+        prompt_tokens=346,
+        completion_tokens=29,
+        total_tokens=375,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            text_tokens=152, image_tokens=194, audio_tokens=0, cached_tokens=128
+        ),
+        completion_tokens_details=CompletionTokensDetailsWrapper(
+            text_tokens=29, audio_tokens=0, reasoning_tokens=19
+        ),
+    )
+
+    prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai")
+    breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage)
+
+    info = litellm.get_model_info(model=model, custom_llm_provider="openai")
+    assert completion_cost == pytest.approx(29 * info["output_cost_per_token"])
+    assert completion_cost - breakdown.reasoning_cost == pytest.approx(10 * info["output_cost_per_token"])
+    assert prompt_cost == pytest.approx(
+        24 * info["input_cost_per_token"]
+        + 128 * info["cache_read_input_token_cost"]
+        + 194 * info["input_cost_per_image_token"]
+    )
+
+
+def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens(_local_model_cost_map):
+    """Providers whose text_tokens exclude reasoning (text + reasoning == completion) stay billed in full."""
+
+    model = "gpt-realtime-2.1-mini"
+    usage = Usage(
+        prompt_tokens=100,
+        completion_tokens=44,
+        total_tokens=144,
+        completion_tokens_details=CompletionTokensDetailsWrapper(
+            text_tokens=25, audio_tokens=0, reasoning_tokens=19
+        ),
+    )
+
+    _, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai")
+
+    info = litellm.get_model_info(model=model, custom_llm_provider="openai")
+    assert completion_cost == pytest.approx(44 * info["output_cost_per_token"])
diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index 2046695f151..203f0a9e840 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -4492,3 +4492,45 @@ def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_m
 
     assert prompt_cost == pytest.approx(1000 * 5e-6)
     assert completion_cost == pytest.approx(500 * 2.5e-5)
+
+
+def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once(_local_model_cost_map):
+    """Realtime response.done nests reasoning_tokens inside text_tokens, so they are billed once."""
+    results: OpenAIRealtimeStreamList = [
+        {"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}},
+        {
+            "type": "response.done",
+            "response": {
+                "usage": {
+                    "total_tokens": 260,
+                    "input_tokens": 237,
+                    "output_tokens": 23,
+                    "input_token_details": {
+                        "text_tokens": 43,
+                        "audio_tokens": 0,
+                        "image_tokens": 194,
+                        "cached_tokens": 0,
+                        "cached_tokens_details": {"text_tokens": 0, "audio_tokens": 0, "image_tokens": 0},
+                    },
+                    "output_token_details": {"text_tokens": 23, "audio_tokens": 0, "reasoning_tokens": 18},
+                }
+            },
+        },
+    ]
+    combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(
+        results=results,
+    )
+
+    total_cost = handle_realtime_stream_cost_calculation(
+        results=results,
+        combined_usage_object=combined_usage_object,
+        custom_llm_provider="azure",
+        litellm_model_name="azure/gpt-realtime-2.1-mini",
+    )
+
+    info = litellm.get_model_info(model="azure/gpt-realtime-2.1-mini", custom_llm_provider="azure")
+    expected = (
+        43 * info["input_cost_per_token"] + 194 * info["input_cost_per_image_token"] + 23 * info["output_cost_per_token"]
+    )
+    assert total_cost == pytest.approx(expected)
+    assert total_cost == pytest.approx(0.0002362)

From 5a7e938f371e787d1e60af1eddea70bc10978813 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 19:40:46 -0700
Subject: [PATCH 046/164] test(cost): type the cost map fixture parameter on
 the new tests

---
 .../llm_cost_calc/test_llm_cost_calc_utils.py                 | 4 ++--
 tests/test_litellm/test_cost_calculator.py                    | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

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 f7ca2a89048..19464dfd2a0 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
@@ -4725,7 +4725,7 @@ def test_route_image_generation_cost_falls_back_to_requested_size(monkeypatch, r
     assert cost == expected_cost
 
 
-def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_local_model_cost_map):
+def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_local_model_cost_map: None) -> None:
     """
     Realtime usage (OpenAI and Azure) reports output_tokens == text_tokens + audio_tokens with
     reasoning_tokens already counted inside text_tokens, so reasoning must not be billed on top.
@@ -4757,7 +4757,7 @@ def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_loca
     )
 
 
-def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens(_local_model_cost_map):
+def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens(_local_model_cost_map: None) -> None:
     """Providers whose text_tokens exclude reasoning (text + reasoning == completion) stay billed in full."""
 
     model = "gpt-realtime-2.1-mini"
diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index 203f0a9e840..3eb051982e4 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -4494,7 +4494,7 @@ def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_m
     assert completion_cost == pytest.approx(500 * 2.5e-5)
 
 
-def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once(_local_model_cost_map):
+def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once(_local_model_cost_map: None) -> None:
     """Realtime response.done nests reasoning_tokens inside text_tokens, so they are billed once."""
     results: OpenAIRealtimeStreamList = [
         {"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}},

From 92e449d4d0c486f24c7a09bfce988a6be5b1dd70 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 20:35:21 -0700
Subject: [PATCH 047/164] test(cost): cover reasoning nested in text_tokens
 beside audio output

---
 .../llm_cost_calc/test_llm_cost_calc_utils.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

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 19464dfd2a0..29874c5cca7 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
@@ -4774,3 +4774,22 @@ def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tok
 
     info = litellm.get_model_info(model=model, custom_llm_provider="openai")
     assert completion_cost == pytest.approx(44 * info["output_cost_per_token"])
+
+
+def test_generic_cost_per_token_bills_nested_reasoning_once_beside_audio_output(_local_model_cost_map: None) -> None:
+    """Audio-output realtime usage nests reasoning inside text_tokens next to audio_tokens; text is billed net of it."""
+
+    model = "gpt-realtime-2.1-mini"
+    usage = Usage(
+        prompt_tokens=120,
+        completion_tokens=100,
+        total_tokens=220,
+        completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=30, audio_tokens=70, reasoning_tokens=20),
+    )
+
+    _, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai")
+    breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage)
+
+    info = litellm.get_model_info(model=model, custom_llm_provider="openai")
+    assert breakdown.reasoning_cost == pytest.approx(20 * info["output_cost_per_token"])
+    assert completion_cost == pytest.approx(30 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"])

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 048/164] 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 f09992b73082e37d7337402954387def48fc7420 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 20:54:40 -0700
Subject: [PATCH 049/164] fix(image_handling): keep DNS, signing, and Vertex
 Gemini fetches off the event loop

The SSRF check in async_safe_get resolved DNS on the event loop and a blocked
address was retried three times; validate_url now runs in a thread and an
SSRFError fails the fetch on the first attempt in both fetchers. The shared
HTTP handler signed the request and ran pre_call logging on the loop after an
async transform; both now run in a thread. Vertex AI Gemini still fetched
http:// images and https images without an inferrable mime type with the sync
converter inside its async body builder; the walker takes a should_inline
predicate and Vertex AI inlines exactly those URLs, leaving https images with a
known mime type and Files API refs to Google. When one download fails the
other in-flight downloads for that request are now cancelled instead of
finishing in the background
---
 .../prompt_templates/image_handling.py        |  73 ++++++++--
 litellm/litellm_core_utils/url_utils.py       |   3 +-
 litellm/llms/custom_httpx/llm_http_handler.py |   2 +-
 .../llms/vertex_ai/gemini/transformation.py   |  26 +++-
 .../litellm_core_utils/test_image_handling.py | 135 ++++++++++++++++--
 .../litellm_core_utils/test_url_utils.py      |  35 +++++
 .../custom_httpx/test_llm_http_handler.py     |  28 +++-
 .../vertex_ai/gemini/test_transformation.py   |  51 +++++++
 8 files changed, 318 insertions(+), 35 deletions(-)

diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py
index 55c01b3b1cb..dcad7776b58 100644
--- a/litellm/litellm_core_utils/prompt_templates/image_handling.py
+++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py
@@ -4,7 +4,7 @@ Helper functions to handle images passed in messages
 
 import asyncio
 import base64
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass
 from types import MappingProxyType
 from typing import Final
@@ -15,7 +15,7 @@ import litellm
 from litellm import verbose_logger
 from litellm.caching.caching import InMemoryCache
 from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB
-from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
+from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get
 from litellm.types.llms.openai import AllMessageValues
 
 MAX_IMGS_IN_MEMORY: Final = 10
@@ -99,6 +99,8 @@ async def async_convert_url_to_base64(url: str) -> str:
             return _process_image_response(response, url)
         except litellm.ImageFetchError:
             raise
+        except SSRFError as e:
+            raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e
         except Exception:
             pass
     raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}")
@@ -125,6 +127,8 @@ def convert_url_to_base64(url: str) -> str:
             return _process_image_response(response, url)
         except litellm.ImageFetchError:
             raise
+        except SSRFError as e:
+            raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e
         except Exception as e:
             verbose_logger.exception(e)
     raise litellm.ImageFetchError(
@@ -163,9 +167,23 @@ _ANTHROPIC_MEDIA_BLOCK_TYPES: Final = frozenset({"document", "image"})
 @dataclass(frozen=True, slots=True)
 class _RemoteSource:
     part: Mapping[str, object]
+    source: Mapping[str, object]
     url: str
 
 
+@dataclass(frozen=True, slots=True)
+class RemoteMedia:
+    url: str
+    fields: Mapping[str, object]
+
+
+_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
+
+
+def inline_every_remote_url(_media: RemoteMedia) -> bool:
+    return True
+
+
 def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None:
     if fields.get("type") != "image_url":
         return None
@@ -184,7 +202,7 @@ def _parse_remote_file(fields: Mapping[str, object]) -> _RemoteFile | None:
 def _parse_remote_source(fields: Mapping[str, object]) -> _RemoteSource | None:
     source: Final = _as_mapping(fields.get("source")) if fields.get("type") in _ANTHROPIC_MEDIA_BLOCK_TYPES else None
     url: Final = _remote_url(source.get("url")) if source is not None and source.get("type") == "url" else None
-    return _RemoteSource(fields, url) if url is not None else None
+    return _RemoteSource(fields, source, url) if source is not None and url is not None else None
 
 
 def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSource | None:
@@ -194,6 +212,16 @@ def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSour
     return _parse_remote_image(fields) or _parse_remote_file(fields) or _parse_remote_source(fields)
 
 
+def _remote_media(remote: _RemoteImage | _RemoteFile | _RemoteSource) -> RemoteMedia:
+    match remote:
+        case _RemoteImage(_, image_url, url):
+            return RemoteMedia(url, image_url if image_url is not None else _NO_FIELDS)
+        case _RemoteFile(_, file, url):
+            return RemoteMedia(url, file)
+        case _RemoteSource(_, source, url):
+            return RemoteMedia(url, source)
+
+
 _PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"})
 
 
@@ -222,7 +250,7 @@ def _inline(remote: _RemoteImage | _RemoteFile | _RemoteSource, data_url: str) -
             return {**part, "image_url": _inlined_image_url(image_url, data_url)}  # mutable-ok: json-serialized part
         case _RemoteFile(part, file, url):
             return {**part, "file": _inlined_file(file, url, data_url)}  # mutable-ok: json-serialized message part
-        case _RemoteSource(part, url):
+        case _RemoteSource(part, _, url):
             return {**part, "source": _base64_source(url, data_url)}  # mutable-ok: json-serialized message part
 
 
@@ -231,18 +259,22 @@ def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]:
     return tuple(content) if isinstance(content, list) else ()  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]  # parts are parsed one by one
 
 
-def _inline_part(part: object, data_urls: Mapping[str, str]) -> object:
+def _inline_part(part: object, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]) -> object:
     remote: Final = _parse_remote_part(part)
-    data_url: Final = data_urls.get(remote.url) if remote is not None else None
-    return _inline(remote, data_url) if remote is not None and data_url is not None else part
+    if remote is None or not should_inline(_remote_media(remote)):
+        return part
+    data_url: Final = data_urls.get(remote.url)
+    return _inline(remote, data_url) if data_url is not None else part
 
 
-def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> AllMessageValues:
+def _inline_message(
+    message: AllMessageValues, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]
+) -> AllMessageValues:
     parts: Final = _content_parts(message)
     if not parts:
         return message
     inlined_parts: Final = [  # mutable-ok: content must stay a list for the transforms' isinstance checks
-        _inline_part(part, data_urls) for part in parts
+        _inline_part(part, data_urls, should_inline) for part in parts
     ]
     inlined_message: Final = {**message, "content": inlined_parts}  # mutable-ok: json-serialized message
     return inlined_message  # pyright: ignore[reportReturnType]  # the same message with its remote parts inlined
@@ -253,21 +285,34 @@ async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore) -> str:
         return await async_convert_url_to_base64(url)
 
 
+async def _fetch_data_urls(remote_urls: tuple[str, ...]) -> tuple[str, ...]:
+    in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES)
+    fetches: Final = tuple(asyncio.create_task(_fetch_data_url(url, in_flight)) for url in remote_urls)
+    try:
+        return tuple(await asyncio.gather(*fetches))
+    except BaseException:
+        for fetch in fetches:
+            fetch.cancel()
+        await asyncio.gather(*fetches, return_exceptions=True)
+        raise
+
+
 async def async_inline_remote_media(
     messages: list[AllMessageValues],  # mutable-ok: every transform_request takes list[AllMessageValues]
-    skip_url_prefixes: tuple[str, ...] = (),
+    should_inline: Callable[[RemoteMedia], bool] = inline_every_remote_url,
 ) -> list[AllMessageValues]:  # mutable-ok: every transform_request takes list[AllMessageValues]
     remote_urls: Final = tuple(
         dict.fromkeys(
             remote.url
             for message in messages
             for part in _content_parts(message)
-            if (remote := _parse_remote_part(part)) is not None and not remote.url.startswith(skip_url_prefixes)
+            if (remote := _parse_remote_part(part)) is not None and should_inline(_remote_media(remote))
         )
     )
     if not remote_urls:
         return messages
-    in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES)
-    data_urls: Final = await asyncio.gather(*(_fetch_data_url(url, in_flight) for url in remote_urls))
+    data_urls: Final = await _fetch_data_urls(remote_urls)
     inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True)))
-    return [_inline_message(message, inlined) for message in messages]  # mutable-ok: transform_request takes a list
+    return [  # mutable-ok: transform_request takes a list
+        _inline_message(message, inlined, should_inline) for message in messages
+    ]
diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py
index 1e43117933d..fa070a648f5 100644
--- a/litellm/litellm_core_utils/url_utils.py
+++ b/litellm/litellm_core_utils/url_utils.py
@@ -19,6 +19,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config):
   check but still resolve DNS and still rewrite HTTP to the resolved IP.
 """
 
+import asyncio
 import socket
 from ipaddress import ip_address, ip_network
 from typing import Any, Final, Protocol
@@ -471,7 +472,7 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response
     kwargs.pop("follow_redirects", None)
     headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})}
     for _ in range(_MAX_REDIRECTS):
-        validated_url, original_host = validate_url(url)
+        validated_url, original_host = await asyncio.to_thread(validate_url, url)
         response = await fetcher.get(
             validated_url,
             headers={**headers_view["headers"], "Host": original_host},
diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index 54c21f8d5b3..d57aee025d5 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -640,7 +640,7 @@ class BaseLLMHTTPHandler:
                         headers=request_headers,
                     ),
                 )
-                return await dispatch_async(*sign_and_log(transformed))
+                return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed))
 
             return transform_then_dispatch()
 
diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py
index 6d100143e52..13e2238fdf6 100644
--- a/litellm/llms/vertex_ai/gemini/transformation.py
+++ b/litellm/llms/vertex_ai/gemini/transformation.py
@@ -7,6 +7,7 @@ Why separate file? Make it easy to see how transformation works
 import json
 import os
 import re
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, Final, Literal, cast
 from urllib.parse import quote
 
@@ -27,7 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
     convert_to_gemini_tool_call_result,
     response_schema_prompt,
 )
-from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
+from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, async_inline_remote_media
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
 from litellm.types.files import (
@@ -1309,6 +1310,23 @@ def sync_transform_request_body(
     )
 
 
+def _explicit_mime_type(fields: Mapping[str, object]) -> str | None:
+    hint: Final = fields.get("format") or fields.get("mime_type") or fields.get("content_type")
+    return hint if isinstance(hint, str) else None
+
+
+def _ai_studio_inlines(media: RemoteMedia) -> bool:
+    return not media.url.startswith(GEMINI_FILES_API_URI_PREFIX)
+
+
+def _vertex_inlines(media: RemoteMedia) -> bool:
+    if media.url.startswith(GEMINI_FILES_API_URI_PREFIX):
+        return False
+    return media.url.startswith("http://") or (
+        _explicit_mime_type(media.fields) is None and _get_image_mime_type_from_url(media.url) is None
+    )
+
+
 async def async_transform_request_body(
     gemini_api_key: str | None,
     messages: list[AllMessageValues],
@@ -1350,10 +1368,8 @@ async def async_transform_request_body(
         vertex_auth_header=vertex_auth_header,
     )
 
-    inlined_messages: Final = (
-        await async_inline_remote_media(messages, skip_url_prefixes=(GEMINI_FILES_API_URI_PREFIX,))
-        if custom_llm_provider == "gemini"
-        else messages
+    inlined_messages: Final = await async_inline_remote_media(
+        messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines
     )
 
     if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages):
diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py
index 106bf90213b..e7b28d0d3a2 100644
--- a/tests/test_litellm/litellm_core_utils/test_image_handling.py
+++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py
@@ -1,5 +1,6 @@
 import asyncio
 import copy
+import time
 import uuid
 from unittest.mock import patch
 
@@ -11,10 +12,12 @@ from litellm import constants
 from litellm.litellm_core_utils.prompt_templates import image_handling
 from litellm.litellm_core_utils.prompt_templates.image_handling import (
     MAX_CONCURRENT_REMOTE_MEDIA_FETCHES,
+    RemoteMedia,
     async_convert_url_to_base64,
     async_inline_remote_media,
     convert_url_to_base64,
 )
+from litellm.litellm_core_utils.url_utils import SSRFError
 
 
 @pytest.fixture(autouse=True)
@@ -321,34 +324,142 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o
     assert messages == snapshot
 
 
-async def test_async_inline_remote_media_leaves_skipped_url_prefixes_untouched(async_only_image_fetch):
-    skipped_prefix = "https://generativelanguage.googleapis.com/v1beta/files/"
-    skipped_file = f"{skipped_prefix}{uuid.uuid4().hex}"
-    skipped_image = f"{skipped_prefix}{uuid.uuid4().hex}"
-    fetched_image = f"https://img.example/{uuid.uuid4()}.png"
+async def test_async_inline_remote_media_inlines_only_the_parts_the_predicate_accepts(async_only_image_fetch):
+    files_api_prefix = "https://generativelanguage.googleapis.com/v1beta/files/"
+    files_api_pdf = f"{files_api_prefix}{uuid.uuid4().hex}"
+    hinted_image = f"https://img.example/{uuid.uuid4()}.png"
+    plain_image = f"https://img.example/{uuid.uuid4()}.png"
+    hinted_document = f"https://docs.example/{uuid.uuid4()}.pdf"
+    seen = []
+
+    def inline_unhinted_outside_files_api(media: RemoteMedia) -> bool:
+        seen.append(media)
+        return not media.url.startswith(files_api_prefix) and "format" not in media.fields
+
     messages = [
         {
             "role": "user",
             "content": [
-                {"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}},
-                {"type": "image_url", "image_url": {"url": skipped_image}},
-                {"type": "image_url", "image_url": fetched_image},
+                {"type": "file", "file": {"file_id": files_api_pdf}},
+                {"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}},
+                {"type": "image_url", "image_url": {"url": plain_image}},
+                {"type": "image_url", "image_url": plain_image},
+                {"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}},
             ],
         }
     ]
     snapshot = copy.deepcopy(messages)
 
-    inlined = await async_inline_remote_media(messages, skip_url_prefixes=(skipped_prefix,))
+    inlined = await async_inline_remote_media(messages, should_inline=inline_unhinted_outside_files_api)
 
     assert inlined[0]["content"] == [
-        {"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}},
-        {"type": "image_url", "image_url": {"url": skipped_image}},
+        {"type": "file", "file": {"file_id": files_api_pdf}},
+        {"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}},
+        {"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}},
         {"type": "image_url", "image_url": async_only_image_fetch.data_url},
+        {"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}},
+    ]
+    assert async_only_image_fetch.fetched == [plain_image]
+    assert [(media.url, dict(media.fields)) for media in seen[:5]] == [
+        (files_api_pdf, {"file_id": files_api_pdf}),
+        (hinted_image, {"url": hinted_image, "format": "image/png"}),
+        (plain_image, {"url": plain_image}),
+        (plain_image, {}),
+        (hinted_document, {"type": "url", "url": hinted_document, "format": "application/pdf"}),
     ]
-    assert async_only_image_fetch.fetched == [fetched_image]
     assert messages == snapshot
 
 
+async def test_async_inline_remote_media_inlines_a_shared_url_only_where_the_predicate_accepts_it(
+    async_only_image_fetch,
+):
+    shared = f"https://img.example/{uuid.uuid4()}.png"
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image_url", "image_url": {"url": shared, "format": "image/png"}},
+                {"type": "image_url", "image_url": {"url": shared}},
+            ],
+        }
+    ]
+
+    inlined = await async_inline_remote_media(messages, should_inline=lambda media: "format" not in media.fields)
+
+    assert inlined[0]["content"] == [
+        {"type": "image_url", "image_url": {"url": shared, "format": "image/png"}},
+        {"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}},
+    ]
+    assert async_only_image_fetch.fetched == [shared]
+
+
+async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fails(monkeypatch):
+    missing = f"http://img.example/{uuid.uuid4()}-missing.png"
+    slow = f"http://img.example/{uuid.uuid4()}-slow.png"
+    slow_fetch_outcomes = []
+
+    async def serve(client, url, **kwargs):
+        if url == missing:
+            return Response(404, request=Request("GET", url))
+        try:
+            await asyncio.sleep(5)
+        except asyncio.CancelledError:
+            slow_fetch_outcomes.append("cancelled")
+            raise
+        slow_fetch_outcomes.append("finished")
+        return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url))
+
+    monkeypatch.setattr(image_handling, "async_safe_get", serve)
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image_url", "image_url": {"url": missing}},
+                {"type": "image_url", "image_url": {"url": slow}},
+            ],
+        }
+    ]
+    started = time.perf_counter()
+
+    with pytest.raises(litellm.ImageFetchError, match="Status code: 404"):
+        await async_inline_remote_media(messages)
+
+    assert slow_fetch_outcomes == ["cancelled"]
+    assert time.perf_counter() - started < 1
+
+
+async def test_async_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch):
+    attempts = []
+
+    async def block(client, url, **kwargs):
+        attempts.append(url)
+        raise SSRFError("URL targets a blocked address (10.0.0.8)")
+
+    monkeypatch.setattr(image_handling, "async_safe_get", block)
+    url = f"http://internal.example/{uuid.uuid4()}.png"
+
+    with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"):
+        await async_convert_url_to_base64(url)
+
+    assert attempts == [url]
+
+
+def test_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch):
+    attempts = []
+
+    def block(client, url, **kwargs):
+        attempts.append(url)
+        raise SSRFError("URL targets a blocked address (10.0.0.8)")
+
+    monkeypatch.setattr(image_handling, "safe_get", block)
+    url = f"http://internal.example/{uuid.uuid4()}.png"
+
+    with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"):
+        convert_url_to_base64(url)
+
+    assert attempts == [url]
+
+
 async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch):
     in_flight = {"now": 0, "peak": 0}
 
diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py
index aaaa43a0dc4..fccdc1a2a0a 100644
--- a/tests/test_litellm/litellm_core_utils/test_url_utils.py
+++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py
@@ -1,5 +1,9 @@
+import asyncio
 import socket
+import threading
+import time
 
+import httpx
 import pytest
 
 import litellm
@@ -535,3 +539,34 @@ def test_assert_same_origin_error_message_does_not_leak_hostnames():
     detail = str(exc.value)
     assert "attacker.example.com" not in detail
     assert "api.internal-corp.example" not in detail
+
+
+async def test_async_safe_get_resolves_dns_off_the_event_loop(monkeypatch):
+    loop_thread = threading.current_thread()
+    resolver_threads = []
+
+    def slow_getaddrinfo(host, port, *args, **kwargs):
+        resolver_threads.append(threading.current_thread())
+        time.sleep(0.4)
+        return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port or 443))]
+
+    monkeypatch.setattr(url_utils.socket, "getaddrinfo", slow_getaddrinfo)
+
+    class FakeClient:
+        async def get(self, url, **kwargs):
+            return httpx.Response(200, request=httpx.Request("GET", url))
+
+    ticks = [time.perf_counter()]
+
+    async def heartbeat():
+        while True:
+            await asyncio.sleep(0.01)
+            ticks.append(time.perf_counter())
+
+    beating = asyncio.create_task(heartbeat())
+    response = await url_utils.async_safe_get(FakeClient(), "https://img.example/a.png")
+    beating.cancel()
+
+    assert response.status_code == 200
+    assert resolver_threads and all(thread is not loop_thread for thread in resolver_threads)
+    assert max(b - a for a, b in zip(ticks, ticks[1:])) < 0.2
diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
index e86e671b939..f1614654ffb 100644
--- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
+++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
@@ -1,6 +1,7 @@
 import asyncio
 import json
 import logging
+import threading
 import time
 from unittest.mock import AsyncMock, Mock, patch
 
@@ -3192,6 +3193,7 @@ class _TransformRecordingConfig(BaseConfig):
     def __init__(self, transform_async: bool):
         self.transform_async = transform_async
         self.transform_calls = []
+        self.sign_threads = []
 
     @property
     def uses_async_transform_request(self) -> bool:
@@ -3216,6 +3218,12 @@ class _TransformRecordingConfig(BaseConfig):
         self.transform_calls.append("async")
         return {"transformed_by": "async"}
 
+    def sign_request(
+        self, headers, optional_params, request_data, api_base, api_key=None, model=None, stream=None, fake_stream=None
+    ):
+        self.sign_threads.append(threading.current_thread())
+        return headers, None
+
     def transform_response(
         self,
         model,
@@ -3237,7 +3245,7 @@ class _TransformRecordingConfig(BaseConfig):
         return BaseLLMException(status_code=status_code, message=error_message, headers=headers)
 
 
-def _start_async_completion(config):
+def _start_async_completion(config, logging_obj=None):
     captured = {}
 
     def handle(request):
@@ -3253,7 +3261,7 @@ def _start_async_completion(config):
         custom_llm_provider="openai",
         model_response=ModelResponse(),
         encoding=None,
-        logging_obj=Mock(dynamic_success_callbacks=None, model_call_details={}),
+        logging_obj=logging_obj if logging_obj is not None else Mock(dynamic_success_callbacks=None, model_call_details={}),
         optional_params={},
         timeout=10.0,
         litellm_params={},
@@ -3277,6 +3285,22 @@ async def test_completion_awaits_async_transform_request_when_config_opts_in():
     assert response.choices[0].message.content == "async"
 
 
+async def test_completion_signs_and_logs_off_the_event_loop_after_the_async_transform():
+    config = _TransformRecordingConfig(transform_async=True)
+    loop_thread = threading.current_thread()
+    pre_call_threads = []
+    logging_obj = Mock(dynamic_success_callbacks=None, model_call_details={})
+    logging_obj.pre_call.side_effect = lambda **kwargs: pre_call_threads.append(threading.current_thread())
+
+    pending, captured = _start_async_completion(config, logging_obj)
+    response = await pending
+
+    assert response.choices[0].message.content == "async"
+    assert captured["body"] == {"transformed_by": "async"}
+    assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads)
+    assert pre_call_threads and all(thread is not loop_thread for thread in pre_call_threads)
+
+
 async def test_completion_keeps_sync_transform_request_before_returning_by_default():
     config = _TransformRecordingConfig(transform_async=False)
 
diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py
index d31254746d4..f135acd094f 100644
--- a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py
@@ -1,6 +1,8 @@
 
 import json
 import uuid
+from unittest.mock import Mock
+
 import httpx
 import pytest
 
@@ -424,3 +426,52 @@ async def test_gemini_ai_studio_async_completion_passes_files_api_uris_through_u
         {"mime_type": "application/pdf", "file_uri": files_api_pdf},
         {"mime_type": "image/png", "file_uri": files_api_image},
     ]
+
+
+async def test_vertex_ai_async_transform_inlines_only_the_urls_gemini_cannot_fetch_itself(async_only_image_fetch):
+    plain_http_png = f"http://img.example/{uuid.uuid4()}.png"
+    extensionless_https = f"https://cdn.example/files/{uuid.uuid4().hex}"
+    https_png = f"https://img.example/{uuid.uuid4()}.png"
+    hinted_extensionless = f"https://cdn.example/files/{uuid.uuid4().hex}"
+    files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}"
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe these"},
+                {"type": "image_url", "image_url": {"url": plain_http_png}},
+                {"type": "image_url", "image_url": {"url": extensionless_https}},
+                {"type": "image_url", "image_url": {"url": https_png}},
+                {"type": "image_url", "image_url": {"url": hinted_extensionless, "mime_type": "image/webp"}},
+                {"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}},
+            ],
+        }
+    ]
+
+    body = await transformation.async_transform_request_body(
+        gemini_api_key=None,
+        messages=messages,
+        api_base=None,
+        model="gemini-3.8-flash",
+        client=None,
+        timeout=None,
+        extra_headers=None,
+        optional_params={},
+        logging_obj=Mock(),
+        custom_llm_provider="vertex_ai",
+        litellm_params={},
+        vertex_project="qa-project",
+        vertex_location="us-central1",
+        vertex_auth_header=None,
+    )
+
+    inlined = {"inline_data": {"mime_type": "image/png", "data": async_only_image_fetch.base64_png}}
+    assert body["contents"][0]["parts"] == [
+        {"text": "Describe these"},
+        inlined,
+        inlined,
+        {"file_data": {"mime_type": "image/png", "file_uri": https_png}},
+        {"file_data": {"mime_type": "image/webp", "file_uri": hinted_extensionless}},
+        {"file_data": {"mime_type": "application/pdf", "file_uri": files_api_pdf}},
+    ]
+    assert sorted(async_only_image_fetch.fetched) == sorted([plain_http_png, extensionless_https])

From e9a40ad4d2d6712812e9b9210b28d4c291cdfc0b Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 22:16:14 -0700
Subject: [PATCH 050/164] fix(fireworks_ai): map developer items after pydantic
 input items are dumped

---
 .../fireworks_ai/responses/transformation.py  |  5 ++-
 ...t_fireworks_ai_responses_transformation.py | 43 ++++++++++++++++++-
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py
index 121d3b5d9a9..fb0587553d4 100644
--- a/litellm/llms/fireworks_ai/responses/transformation.py
+++ b/litellm/llms/fireworks_ai/responses/transformation.py
@@ -68,6 +68,9 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
         base: Final = (api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE).rstrip("/")
         return f"{base}/responses"
 
+    def _validate_input_param(self, input: str | ResponseInputParam) -> str | ResponseInputParam:
+        return _developer_items_as_system(super()._validate_input_param(input))
+
     def transform_responses_api_request(
         self,
         model: str,
@@ -78,7 +81,7 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
     ) -> dict:  # mutable-ok: overrides the base class signature
         return super().transform_responses_api_request(
             model=resolve_fireworks_resource_name(model),
-            input=_developer_items_as_system(input),
+            input=input,
             response_api_optional_request_params=response_api_optional_request_params,
             litellm_params=litellm_params,
             headers=headers,
diff --git a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
index f948acb7bf3..b9408d44e9a 100644
--- a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
+++ b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
@@ -1,13 +1,14 @@
 import json
 from collections.abc import Mapping
 from types import MappingProxyType
-from typing import Final, TypedDict
+from typing import Final, TypedDict, cast
 from unittest.mock import MagicMock, patch
 from urllib.parse import quote
 
 import httpx
 import pytest
 from openai.types.responses import (
+    EasyInputMessage,
     ResponseFunctionToolCall,
     ResponseOutputMessage,
     ResponseOutputText,
@@ -20,7 +21,7 @@ from typing_extensions import ReadOnly
 import litellm
 from litellm.llms.fireworks_ai.responses.transformation import FireworksAIResponsesAPIConfig
 from litellm.responses.file_search.emulated_handler import should_use_emulated_file_search
-from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponsesAPIResponse
+from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponseInputParam, ResponsesAPIResponse
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import LlmProviders
 from litellm.utils import ProviderConfigManager
@@ -181,6 +182,44 @@ def test_responses_call_sends_developer_items_as_system_messages() -> None:
     )
 
 
+def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_output_items() -> None:
+    client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
+    pydantic_input: Final = cast(
+        ResponseInputParam,
+        [  # mutable-ok: the Responses API takes input as a JSON list
+            EasyInputMessage(role="developer", content="Answer with exactly one word.", type="message"),
+            ResponseReasoningItem(id="rs_1", summary=(), type="reasoning"),
+            ResponseFunctionToolCall(
+                id="fc_1",
+                call_id="call_abc123",
+                name="get_weather",
+                arguments='{"city": "Paris"}',
+                status="completed",
+                type="function_call",
+            ),
+            FunctionCallOutput(type="function_call_output", call_id="call_abc123", output="21C"),
+        ],
+    )
+    with patch(HTTPX_CLIENT_FACTORY, return_value=client):
+        litellm.responses(
+            model="fireworks_ai/accounts/fireworks/models/kimi-k3", input=pydantic_input, api_key="fw-test-key"
+        )
+    _, _, body = _sent_request(client)
+    assert tuple(body["input"]) == (
+        {"role": "system", "content": "Answer with exactly one word.", "type": "message"},
+        {"id": "rs_1", "summary": [], "type": "reasoning"},
+        {
+            "id": "fc_1",
+            "call_id": "call_abc123",
+            "name": "get_weather",
+            "arguments": '{"city": "Paris"}',
+            "status": "completed",
+            "type": "function_call",
+        },
+        {"type": "function_call_output", "call_id": "call_abc123", "output": "21C"},
+    )
+
+
 def test_file_search_tools_take_litellm_emulated_search_not_fireworks() -> None:
     config: Final = FireworksAIResponsesAPIConfig()
     file_search: Final = ({"type": "file_search", "vector_store_ids": ("vs_kb",)},)

From 836c20a4b6dd931d119ac2cce92977c3c947e062 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 22:31:06 -0700
Subject: [PATCH 051/164] test(ai-gateway): drop provider fingerprint and cache
 identity assertions

---
 litellm-rust/crates/ai-gateway/src/io/tls.rs | 43 +-------------------
 1 file changed, 1 insertion(+), 42 deletions(-)

diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs
index 16fd11e2e79..a2562f60345 100644
--- a/litellm-rust/crates/ai-gateway/src/io/tls.rs
+++ b/litellm-rust/crates/ai-gateway/src/io/tls.rs
@@ -69,26 +69,7 @@ where
 
 #[cfg(test)]
 mod tests {
-    use rustls::CipherSuite;
-    use rustls::NamedGroup;
-    use rustls::crypto::{CryptoProvider, aws_lc_rs, ring};
-
-    use super::{Arc, build_config, tls_config};
-
-    fn fingerprint(provider: &CryptoProvider) -> (Vec<CipherSuite>, Vec<NamedGroup>) {
-        (
-            provider
-                .cipher_suites
-                .iter()
-                .map(|suite| suite.suite())
-                .collect(),
-            provider
-                .kx_groups
-                .iter()
-                .map(|group| group.name())
-                .collect(),
-        )
-    }
+    use super::build_config;
 
     #[test]
     fn builds_a_usable_config_with_both_provider_features_enabled() {
@@ -96,26 +77,4 @@ mod tests {
 
         assert!(!config.crypto_provider().cipher_suites.is_empty());
     }
-
-    #[test]
-    fn dials_with_ring_rather_than_aws_lc_rs() {
-        let config = build_config().expect("a client config");
-
-        assert_eq!(
-            fingerprint(config.crypto_provider()),
-            fingerprint(&ring::default_provider())
-        );
-        assert_ne!(
-            fingerprint(config.crypto_provider()),
-            fingerprint(&aws_lc_rs::default_provider())
-        );
-    }
-
-    #[test]
-    fn the_trust_store_is_loaded_once_and_shared() {
-        let first = tls_config().expect("a client config");
-        let second = tls_config().expect("a client config");
-
-        assert!(Arc::ptr_eq(&first, &second));
-    }
 }

From e1d900d1c29e7f5dbd5db2ca09ab3592312bb1fe Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 22:47:09 -0700
Subject: [PATCH 052/164] fix(realtime): store text_tokens without the nested
 reasoning share

The realtime usage writer passed the provider's output_token_details through as sent, so spend logs and callbacks kept a text_tokens that still contained reasoning_tokens while every other completion_tokens_details producer stores the partitioned share. The writer now applies the same rule the cost calculator uses, moved to litellm/types/utils.py so both read one definition, and the calculator keeps it for usage objects that arrive nested from elsewhere
---
 .../litellm_core_utils/llm_cost_calc/utils.py | 14 +-----
 litellm/responses/utils.py                    | 25 ++++++++--
 litellm/types/utils.py                        | 11 ++++
 .../responses/test_responses_utils.py         | 50 +++++++++++++++++++
 tests/test_litellm/test_cost_calculator.py    | 39 +++++++++++++++
 5 files changed, 123 insertions(+), 16 deletions(-)

diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py
index 1b8980abd35..9432fefc368 100644
--- a/litellm/litellm_core_utils/llm_cost_calc/utils.py
+++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py
@@ -26,6 +26,7 @@ from litellm.types.utils import (
     PromptTokensDetailsWrapper,
     ServiceTier,
     Usage,
+    text_tokens_without_nested_reasoning,
 )
 from litellm.utils import get_model_info
 
@@ -852,17 +853,6 @@ class CompletionTokensDetailsResult(TypedDict):
     video_tokens: int
 
 
-def _text_tokens_without_nested_reasoning(
-    completion_tokens: int,
-    text_tokens: int,
-    reasoning_tokens: int,
-    other_modality_tokens: int,
-) -> int:
-    reported_total: Final = text_tokens + reasoning_tokens + other_modality_tokens
-    nested_reasoning_tokens: Final = min(reasoning_tokens, text_tokens, max(reported_total - completion_tokens, 0))
-    return text_tokens - nested_reasoning_tokens
-
-
 def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult:
     audio_tokens: Final = (
         cast(
@@ -893,7 +883,7 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu
         or 0
     )
     video_tokens: Final = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0))
-    text_tokens: Final = _text_tokens_without_nested_reasoning(
+    text_tokens: Final = text_tokens_without_nested_reasoning(
         completion_tokens=usage.completion_tokens,
         text_tokens=reported_text_tokens,
         reasoning_tokens=reasoning_tokens,
diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py
index 39675faf735..28f590dcb66 100644
--- a/litellm/responses/utils.py
+++ b/litellm/responses/utils.py
@@ -25,9 +25,15 @@ from litellm.types.utils import (
     PromptTokensDetailsWrapper,
     SpecialEnums,
     Usage,
+    text_tokens_without_nested_reasoning,
 )
 
 
+def _output_token_detail(details: object, field: str) -> int | None:
+    value: Final = getattr(details, field, None)
+    return value if isinstance(value, int) else None
+
+
 def _is_object_sequence(value: object) -> TypeIs[Sequence[object]]:  # guard-ok: a list is a Sequence of anything
     return isinstance(value, list)
 
@@ -1134,11 +1140,22 @@ class ResponseAPILoggingUtils:
             response_api_usage, "output_tokens_details", None
         )
         if output_tokens_details:
+            reasoning_tokens: Final = _output_token_detail(output_tokens_details, "reasoning_tokens")
+            image_tokens: Final = _output_token_detail(output_tokens_details, "image_tokens")
+            audio_tokens: Final = _output_token_detail(output_tokens_details, "audio_tokens")
+            reported_text_tokens: Final = _output_token_detail(output_tokens_details, "text_tokens")
             completion_tokens_details = CompletionTokensDetailsWrapper(
-                reasoning_tokens=getattr(output_tokens_details, "reasoning_tokens", None),
-                image_tokens=getattr(output_tokens_details, "image_tokens", None),
-                text_tokens=getattr(output_tokens_details, "text_tokens", None),
-                audio_tokens=getattr(output_tokens_details, "audio_tokens", None),
+                reasoning_tokens=reasoning_tokens,
+                image_tokens=image_tokens,
+                text_tokens=None
+                if reported_text_tokens is None
+                else text_tokens_without_nested_reasoning(
+                    completion_tokens=completion_tokens,
+                    text_tokens=reported_text_tokens,
+                    reasoning_tokens=reasoning_tokens or 0,
+                    other_modality_tokens=(audio_tokens or 0) + (image_tokens or 0),
+                ),
+                audio_tokens=audio_tokens,
             )
 
         extra_usage_fields: Final = {
diff --git a/litellm/types/utils.py b/litellm/types/utils.py
index 5052cd6ef48..9b24f1d837b 100644
--- a/litellm/types/utils.py
+++ b/litellm/types/utils.py
@@ -1622,6 +1622,17 @@ class Choices(SafeAttributeModel, OpenAIObject):
         setattr(self, key, value)
 
 
+def text_tokens_without_nested_reasoning(
+    completion_tokens: int,
+    text_tokens: int,
+    reasoning_tokens: int,
+    other_modality_tokens: int,
+) -> int:
+    reported_total: Final = text_tokens + reasoning_tokens + other_modality_tokens
+    nested_reasoning_tokens: Final = min(reasoning_tokens, text_tokens, max(reported_total - completion_tokens, 0))
+    return text_tokens - nested_reasoning_tokens
+
+
 class CompletionTokensDetailsWrapper(CompletionTokensDetails):  # wrapper for older openai versions
     text_tokens: int | None = None
     """Text tokens generated by the model."""
diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py
index cb6efa21036..9d9eefdceb3 100644
--- a/tests/test_litellm/responses/test_responses_utils.py
+++ b/tests/test_litellm/responses/test_responses_utils.py
@@ -440,6 +440,56 @@ class TestResponseAPILoggingUtils:
         assert result.completion_tokens_details.text_tokens == 20
         assert result.completion_tokens_details.audio_tokens is None
 
+    def test_transform_realtime_usage_partitions_reasoning_out_of_text_tokens(self):
+        """Realtime nests reasoning_tokens inside text_tokens; the stored text share excludes them."""
+        usage = {
+            "input_tokens": 237,
+            "output_tokens": 70,
+            "total_tokens": 307,
+            "input_token_details": {"text_tokens": 43, "audio_tokens": 0, "image_tokens": 194, "cached_tokens": 0},
+            "output_token_details": {"text_tokens": 70, "audio_tokens": 0, "reasoning_tokens": 52},
+        }
+
+        result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage)
+
+        assert result.completion_tokens == 70
+        assert result.completion_tokens_details is not None
+        assert result.completion_tokens_details.text_tokens == 18
+        assert result.completion_tokens_details.reasoning_tokens == 52
+        assert result.completion_tokens_details.audio_tokens == 0
+
+    def test_transform_realtime_usage_partitions_reasoning_beside_audio_output(self):
+        """Audio output stays as reported; only the text share sheds the nested reasoning tokens."""
+        usage = {
+            "input_tokens": 100,
+            "output_tokens": 70,
+            "total_tokens": 170,
+            "input_token_details": {"text_tokens": 100, "audio_tokens": 0, "cached_tokens": 0},
+            "output_token_details": {"text_tokens": 39, "audio_tokens": 31, "reasoning_tokens": 23},
+        }
+
+        result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage)
+
+        assert result.completion_tokens_details is not None
+        assert result.completion_tokens_details.text_tokens == 16
+        assert result.completion_tokens_details.audio_tokens == 31
+        assert result.completion_tokens_details.reasoning_tokens == 23
+
+    def test_transform_response_api_usage_keeps_partitioned_text_tokens(self):
+        """A provider already reporting text_tokens beside reasoning_tokens is stored as sent."""
+        usage = {
+            "input_tokens": 10,
+            "output_tokens": 20,
+            "total_tokens": 30,
+            "output_tokens_details": {"text_tokens": 12, "reasoning_tokens": 5},
+        }
+
+        result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage)
+
+        assert result.completion_tokens_details is not None
+        assert result.completion_tokens_details.text_tokens == 12
+        assert result.completion_tokens_details.reasoning_tokens == 5
+
     def test_transform_response_api_usage_carries_extra_provider_fields(self):
         """Non-standard usage fields (e.g. xAI tool details) must survive chat normalization."""
         details = {"web_search_calls": 2, "x_search_calls": 0}
diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index 3eb051982e4..1f14fdc2c67 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -4534,3 +4534,42 @@ def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_o
     )
     assert total_cost == pytest.approx(expected)
     assert total_cost == pytest.approx(0.0002362)
+
+
+def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() -> None:
+    """The combined usage that lands in spend logs keeps reasoning out of text_tokens for every turn."""
+    results: OpenAIRealtimeStreamList = [
+        {"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}},
+        {
+            "type": "response.done",
+            "response": {
+                "usage": {
+                    "total_tokens": 307,
+                    "input_tokens": 237,
+                    "output_tokens": 70,
+                    "input_token_details": {"text_tokens": 43, "audio_tokens": 0, "image_tokens": 194, "cached_tokens": 0},
+                    "output_token_details": {"text_tokens": 70, "audio_tokens": 0, "reasoning_tokens": 52},
+                }
+            },
+        },
+        {
+            "type": "response.done",
+            "response": {
+                "usage": {
+                    "total_tokens": 363,
+                    "input_tokens": 300,
+                    "output_tokens": 63,
+                    "input_token_details": {"text_tokens": 106, "audio_tokens": 0, "image_tokens": 194, "cached_tokens": 0},
+                    "output_token_details": {"text_tokens": 63, "audio_tokens": 0, "reasoning_tokens": 43},
+                }
+            },
+        },
+    ]
+
+    combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results)
+
+    assert combined.completion_tokens == 133
+    assert combined.completion_tokens_details is not None
+    assert combined.completion_tokens_details.reasoning_tokens == 95
+    assert combined.completion_tokens_details.text_tokens == 38
+    assert combined.completion_tokens_details.audio_tokens == 0

From 53178486a0c919481bd40da62bfa4ff98c5c2e5a Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 22:59:22 -0700
Subject: [PATCH 053/164] style(tests): wrap realtime cost test lines to the
 120-column limit

---
 .../llm_cost_calc/test_llm_cost_calc_utils.py |  8 +++++--
 tests/test_litellm/test_cost_calculator.py    | 22 +++++++++++++++----
 2 files changed, 24 insertions(+), 6 deletions(-)

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 e7a92a96f34..e8eae73df3f 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
@@ -4800,7 +4800,9 @@ def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_loca
     )
 
 
-def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens(_local_model_cost_map: None) -> None:
+def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens(
+    _local_model_cost_map: None,
+) -> None:
     """Providers whose text_tokens exclude reasoning (text + reasoning == completion) stay billed in full."""
 
     model = "gpt-realtime-2.1-mini"
@@ -4835,4 +4837,6 @@ def test_generic_cost_per_token_bills_nested_reasoning_once_beside_audio_output(
 
     info = litellm.get_model_info(model=model, custom_llm_provider="openai")
     assert breakdown.reasoning_cost == pytest.approx(20 * info["output_cost_per_token"])
-    assert completion_cost == pytest.approx(30 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"])
+    assert completion_cost == pytest.approx(
+        30 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"]
+    )
diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index 1f14fdc2c67..57246ec08a5 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -4494,7 +4494,9 @@ def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_m
     assert completion_cost == pytest.approx(500 * 2.5e-5)
 
 
-def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once(_local_model_cost_map: None) -> None:
+def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once(
+    _local_model_cost_map: None,
+) -> None:
     """Realtime response.done nests reasoning_tokens inside text_tokens, so they are billed once."""
     results: OpenAIRealtimeStreamList = [
         {"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}},
@@ -4530,7 +4532,9 @@ def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_o
 
     info = litellm.get_model_info(model="azure/gpt-realtime-2.1-mini", custom_llm_provider="azure")
     expected = (
-        43 * info["input_cost_per_token"] + 194 * info["input_cost_per_image_token"] + 23 * info["output_cost_per_token"]
+        43 * info["input_cost_per_token"]
+        + 194 * info["input_cost_per_image_token"]
+        + 23 * info["output_cost_per_token"]
     )
     assert total_cost == pytest.approx(expected)
     assert total_cost == pytest.approx(0.0002362)
@@ -4547,7 +4551,12 @@ def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() ->
                     "total_tokens": 307,
                     "input_tokens": 237,
                     "output_tokens": 70,
-                    "input_token_details": {"text_tokens": 43, "audio_tokens": 0, "image_tokens": 194, "cached_tokens": 0},
+                    "input_token_details": {
+                        "text_tokens": 43,
+                        "audio_tokens": 0,
+                        "image_tokens": 194,
+                        "cached_tokens": 0,
+                    },
                     "output_token_details": {"text_tokens": 70, "audio_tokens": 0, "reasoning_tokens": 52},
                 }
             },
@@ -4559,7 +4568,12 @@ def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() ->
                     "total_tokens": 363,
                     "input_tokens": 300,
                     "output_tokens": 63,
-                    "input_token_details": {"text_tokens": 106, "audio_tokens": 0, "image_tokens": 194, "cached_tokens": 0},
+                    "input_token_details": {
+                        "text_tokens": 106,
+                        "audio_tokens": 0,
+                        "image_tokens": 194,
+                        "cached_tokens": 0,
+                    },
                     "output_token_details": {"text_tokens": 63, "audio_tokens": 0, "reasoning_tokens": 43},
                 }
             },

From 0ee3bec046bf0f15eebdddc1d52610d389ca220f Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 23:20:21 -0700
Subject: [PATCH 054/164] fix(image_edit): await an async transform hook and
 hide the URL policy verdict from callers

The image edit handler now awaits BaseImageEditConfig.async_transform_image_edit_request, and
Black Forest Labs overrides it so URL images and masks download through async_safe_get instead of
the blocking safe_get on the event loop. Rejected image fetches raise a fixed policy message with
the user_url_allowed_hosts hint rather than echoing the resolver's verdict (resolved IP, DNS
failure) back to the caller. The test fixture also fails any request-path call of the sync
convert_url_to_base64 so a regression cannot pass unnoticed.
---
 .../prompt_templates/image_handling.py        |  12 +-
 .../base_llm/image_edit/transformation.py     |  19 ++
 .../image_edit/transformation.py              | 102 ++++++-----
 litellm/llms/custom_httpx/llm_http_handler.py |   2 +-
 tests/test_litellm/conftest.py                |  10 +-
 .../litellm_core_utils/test_image_handling.py |  65 ++++---
 .../test_bfl_image_edit_transformation.py     |  78 ++++++++-
 .../custom_httpx/test_llm_http_handler.py     | 163 +++++++++++++++++-
 8 files changed, 379 insertions(+), 72 deletions(-)

diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py
index dcad7776b58..4e7dfdb4017 100644
--- a/litellm/litellm_core_utils/prompt_templates/image_handling.py
+++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py
@@ -78,6 +78,14 @@ def _process_image_response(response: Response, url: str) -> str:
     return result
 
 
+def _url_policy_rejection(url: str, verdict: SSRFError) -> "litellm.ImageFetchError":
+    verbose_logger.warning("Image fetch of %s rejected by the URL policy: %s", url, verdict)
+    return litellm.ImageFetchError(
+        "Error: Unable to fetch image from URL. The proxy's URL policy rejected this host; "
+        f"an admin can allow it with `user_url_allowed_hosts` in general_settings. url={url}"
+    )
+
+
 async def async_convert_url_to_base64(url: str) -> str:
     if url.startswith("data:") and ";base64," in url:
         return url
@@ -100,7 +108,7 @@ async def async_convert_url_to_base64(url: str) -> str:
         except litellm.ImageFetchError:
             raise
         except SSRFError as e:
-            raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e
+            raise _url_policy_rejection(url, e) from e
         except Exception:
             pass
     raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}")
@@ -128,7 +136,7 @@ def convert_url_to_base64(url: str) -> str:
         except litellm.ImageFetchError:
             raise
         except SSRFError as e:
-            raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e
+            raise _url_policy_rejection(url, e) from e
         except Exception as e:
             verbose_logger.exception(e)
     raise litellm.ImageFetchError(
diff --git a/litellm/llms/base_llm/image_edit/transformation.py b/litellm/llms/base_llm/image_edit/transformation.py
index 9a25d3294e0..4faf0aaaf30 100644
--- a/litellm/llms/base_llm/image_edit/transformation.py
+++ b/litellm/llms/base_llm/image_edit/transformation.py
@@ -1,5 +1,6 @@
 import types
 from abc import ABC, abstractmethod
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any
 
 import httpx
@@ -102,6 +103,24 @@ class BaseImageEditConfig(ABC):
     ) -> tuple[dict, RequestFiles]:
         pass
 
+    async def async_transform_image_edit_request(
+        self,
+        model: str,
+        prompt: str | None,
+        image: FileTypes | None,
+        image_edit_optional_request_params: Mapping[str, object],
+        litellm_params: GenericLiteLLMParams,
+        headers: Mapping[str, str],
+    ) -> tuple[dict, RequestFiles]:
+        return self.transform_image_edit_request(
+            model=model,
+            prompt=prompt,
+            image=image,
+            image_edit_optional_request_params=dict(image_edit_optional_request_params),
+            litellm_params=litellm_params,
+            headers=dict(headers),
+        )
+
     def finalize_image_edit_request_data(self, data: dict, resolved_request_url: str) -> dict:
         """
         Last pass on the request dict after ``transform_image_edit_request``, using the
diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py
index 013053e5bd5..c6de9f9ba65 100644
--- a/litellm/llms/black_forest_labs/image_edit/transformation.py
+++ b/litellm/llms/black_forest_labs/image_edit/transformation.py
@@ -9,6 +9,7 @@ API Reference: https://docs.bfl.ai/
 
 import base64
 import time
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, Final
 
 import httpx
@@ -16,7 +17,7 @@ from httpx._types import RequestFiles
 
 import litellm
 from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
-from litellm.litellm_core_utils.url_utils import safe_get
+from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
 from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.images.main import ImageEditOptionalRequestParams
@@ -37,6 +38,22 @@ else:
     LiteLLMLoggingObj = Any
 
 
+_BFL_REQUEST_PARAMS: Final = (
+    "seed",
+    "output_format",
+    "safety_tolerance",
+    "prompt_upsampling",
+    "aspect_ratio",
+    "steps",
+    "guidance",
+    "grow_mask",
+    "top",
+    "bottom",
+    "left",
+    "right",
+)
+
+
 class BlackForestLabsImageEditConfig(BaseImageEditConfig):
     """
     Configuration for Black Forest Labs image editing.
@@ -85,34 +102,10 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
         BFL-specific params are passed through directly.
         """
         optional_params: Final[dict[str, Any]] = {}
-
-        # Pass through BFL-specific params
-        bfl_params: Final = [
-            "seed",
-            "output_format",
-            "safety_tolerance",
-            "prompt_upsampling",
-            # Kontext-specific
-            "aspect_ratio",
-            # Fill/Inpaint-specific
-            "steps",
-            "guidance",
-            "grow_mask",
-            # Expand-specific
-            "top",
-            "bottom",
-            "left",
-            "right",
-        ]
-
-        # Convert TypedDict to regular dict for access
-        params_dict: Final = dict(image_edit_optional_params)
-
-        for param in bfl_params:
-            if param in params_dict:
-                value = params_dict[param]
-                if value is not None:
-                    optional_params[param] = value
+        params: Final[Mapping[str, object]] = image_edit_optional_params
+        for param in _BFL_REQUEST_PARAMS:
+            if (value := params.get(param)) is not None:
+                optional_params[param] = value
 
         # Set default output format
         if "output_format" not in optional_params:
@@ -251,23 +244,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
             "input_image": b64_image,
         }
 
-        # Add optional params (only BFL-recognized parameters)
-        bfl_request_params: Final = [
-            "seed",
-            "output_format",
-            "safety_tolerance",
-            "prompt_upsampling",
-            "aspect_ratio",
-            "steps",
-            "guidance",
-            "grow_mask",
-            "top",
-            "bottom",
-            "left",
-            "right",
-        ]
         for key, value in image_edit_optional_request_params.items():
-            if key in bfl_request_params and value is not None:
+            if key in _BFL_REQUEST_PARAMS and value is not None:
                 request_body[key] = value
 
         # Handle mask if provided (for inpainting)
@@ -277,7 +255,39 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
             request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8")
 
         # BFL uses JSON, not multipart - return empty files
-        return request_body, []
+        return request_body, ()
+
+    async def async_transform_image_edit_request(
+        self,
+        model: str,
+        prompt: str | None,
+        image: FileTypes | None,
+        image_edit_optional_request_params: Mapping[str, object],
+        litellm_params: GenericLiteLLMParams,
+        headers: Mapping[str, str],
+    ) -> tuple[dict, RequestFiles]:
+        downloaded_image: Final = await self._fetch_remote_image(image)
+        downloaded_mask: Final = await self._fetch_remote_image(image_edit_optional_request_params.get("mask"))
+        return self.transform_image_edit_request(
+            model=model,
+            prompt=prompt,
+            image=image if downloaded_image is None else downloaded_image,
+            image_edit_optional_request_params=(
+                dict(image_edit_optional_request_params)
+                if downloaded_mask is None
+                else {**image_edit_optional_request_params, "mask": downloaded_mask}
+            ),
+            litellm_params=litellm_params,
+            headers=dict(headers),
+        )
+
+    async def _fetch_remote_image(self, image: object) -> bytes | None:
+        candidate: Final = image[0] if isinstance(image, list) and image else image
+        if not isinstance(candidate, str) or not candidate.startswith(("http://", "https://")):
+            return None
+        response: Final = await async_safe_get(litellm.module_level_aclient, candidate, timeout=60.0)
+        response.raise_for_status()
+        return response.content
 
     def transform_image_edit_response(
         self,
diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index d57aee025d5..6a17841faa2 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -6787,7 +6787,7 @@ class BaseLLMHTTPHandler:
             litellm_params=dict(litellm_params),
         )
 
-        data, files = image_edit_provider_config.transform_image_edit_request(
+        data, files = await image_edit_provider_config.async_transform_image_edit_request(
             model=model,
             image=image,
             prompt=prompt,
diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py
index aa8ba168ecc..a4f32df46ae 100644
--- a/tests/test_litellm/conftest.py
+++ b/tests/test_litellm/conftest.py
@@ -607,7 +607,8 @@ ONE_PIXEL_PNG = base64.b64decode(
 
 @pytest.fixture
 def async_only_image_fetch(monkeypatch):
-    from litellm.litellm_core_utils.prompt_templates import image_handling
+    from litellm.litellm_core_utils.prompt_templates import factory, image_handling
+    from litellm.llms.gemini.chat import transformation as gemini_chat_transformation
 
     fetch = SimpleNamespace(
         fetched=[],
@@ -627,6 +628,13 @@ def async_only_image_fetch(monkeypatch):
             request=httpx.Request("GET", url),
         )
 
+    def forbid_sync_convert(url, *args, **kwargs):
+        if url.startswith(("http://", "https://")):
+            raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}")
+        return url
+
     monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch)
     monkeypatch.setattr(image_handling, "async_safe_get", serve_png)
+    for module in (image_handling, factory, gemini_chat_transformation):
+        monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert)
     return fetch
diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py
index e7b28d0d3a2..d8ccc9251dd 100644
--- a/tests/test_litellm/litellm_core_utils/test_image_handling.py
+++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py
@@ -428,36 +428,61 @@ async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fail
     assert time.perf_counter() - started < 1
 
 
-async def test_async_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch):
+_SSRF_VERDICTS = (
+    "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, "
+    "add the host to `user_url_allowed_hosts` in general_settings.",
+    "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known",
+    "No addresses found for 'internal.example'",
+)
+
+
+def _assert_one_verdict_free_message(messages, url):
+    assert len(set(messages)) == 1
+    assert "10.0.0.8" not in messages[0]
+    assert "DNS" not in messages[0]
+    assert "No addresses" not in messages[0]
+    assert "user_url_allowed_hosts" in messages[0]
+    assert url in messages[0]
+
+
+async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch):
     attempts = []
-
-    async def block(client, url, **kwargs):
-        attempts.append(url)
-        raise SSRFError("URL targets a blocked address (10.0.0.8)")
-
-    monkeypatch.setattr(image_handling, "async_safe_get", block)
+    messages = []
     url = f"http://internal.example/{uuid.uuid4()}.png"
 
-    with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"):
-        await async_convert_url_to_base64(url)
+    for verdict in _SSRF_VERDICTS:
 
-    assert attempts == [url]
+        async def block(client, fetched_url, verdict=verdict, **kwargs):
+            attempts.append(fetched_url)
+            raise SSRFError(verdict)
+
+        monkeypatch.setattr(image_handling, "async_safe_get", block)
+        with pytest.raises(litellm.ImageFetchError) as raised:
+            await async_convert_url_to_base64(url)
+        messages.append(raised.value.message)
+
+    assert attempts == [url] * len(_SSRF_VERDICTS)
+    _assert_one_verdict_free_message(messages, url)
 
 
-def test_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch):
+def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch):
     attempts = []
-
-    def block(client, url, **kwargs):
-        attempts.append(url)
-        raise SSRFError("URL targets a blocked address (10.0.0.8)")
-
-    monkeypatch.setattr(image_handling, "safe_get", block)
+    messages = []
     url = f"http://internal.example/{uuid.uuid4()}.png"
 
-    with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"):
-        convert_url_to_base64(url)
+    for verdict in _SSRF_VERDICTS:
 
-    assert attempts == [url]
+        def block(client, fetched_url, verdict=verdict, **kwargs):
+            attempts.append(fetched_url)
+            raise SSRFError(verdict)
+
+        monkeypatch.setattr(image_handling, "safe_get", block)
+        with pytest.raises(litellm.ImageFetchError) as raised:
+            convert_url_to_base64(url)
+        messages.append(raised.value.message)
+
+    assert attempts == [url] * len(_SSRF_VERDICTS)
+    _assert_one_verdict_free_message(messages, url)
 
 
 async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch):
diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py
index ec243b7058d..d9d6e813d86 100644
--- a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py
+++ b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py
@@ -16,6 +16,9 @@ import httpx
 import pytest
 
 
+from litellm.llms.black_forest_labs.image_edit import (
+    transformation as bfl_transformation,
+)
 from litellm.llms.black_forest_labs.image_edit.transformation import (
     BlackForestLabsImageEditConfig,
 )
@@ -186,7 +189,7 @@ class TestBlackForestLabsImageEditTransformation:
         assert data["output_format"] == "jpeg"
 
         # BFL uses JSON, not multipart - files should be empty
-        assert files == []
+        assert files == ()
 
     def test_transform_image_edit_request_with_mask(self):
         """Test request transformation with mask for inpainting."""
@@ -299,3 +302,76 @@ class TestBlackForestLabsImageEditTransformation:
     def test_use_multipart_form_data_returns_false(self):
         """Test that use_multipart_form_data returns False for BFL."""
         assert self.config.use_multipart_form_data() is False
+
+
+async def test_async_transform_image_edit_request_downloads_url_images_with_the_async_fetcher(monkeypatch):
+    served = b"png-bytes-from-cdn"
+    fetched = []
+
+    def forbid_sync_fetch(client, url, **kwargs):
+        raise AssertionError(f"sync image fetch ran on the event loop: {url}")
+
+    async def serve(client, url, **kwargs):
+        fetched.append((url, kwargs.get("timeout")))
+        return httpx.Response(200, content=served, request=httpx.Request("GET", url))
+
+    monkeypatch.setattr(bfl_transformation, "safe_get", forbid_sync_fetch)
+    monkeypatch.setattr(bfl_transformation, "async_safe_get", serve)
+
+    data, files = await BlackForestLabsImageEditConfig().async_transform_image_edit_request(
+        model="flux-kontext-pro",
+        prompt="Add a red hat",
+        image="https://cdn.example/photo.png",
+        image_edit_optional_request_params={"mask": "https://cdn.example/mask.png", "seed": 7},
+        litellm_params=GenericLiteLLMParams(),
+        headers={},
+    )
+
+    assert base64.b64decode(data["input_image"]) == served
+    assert base64.b64decode(data["mask"]) == served
+    assert data["seed"] == 7
+    assert files == ()
+    assert fetched == [("https://cdn.example/photo.png", 60.0), ("https://cdn.example/mask.png", 60.0)]
+
+
+async def test_async_transform_image_edit_request_never_fetches_for_local_images(monkeypatch):
+    def refuse(*args, **kwargs):
+        raise AssertionError("no network fetch expected for local image bytes")
+
+    monkeypatch.setattr(bfl_transformation, "safe_get", refuse)
+    monkeypatch.setattr(bfl_transformation, "async_safe_get", refuse)
+
+    data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request(
+        model="flux-kontext-pro",
+        prompt="Add a red hat",
+        image=[BytesIO(b"first"), BytesIO(b"other")],
+        image_edit_optional_request_params={"mask": b"mask-bytes"},
+        litellm_params=GenericLiteLLMParams(),
+        headers={},
+    )
+
+    assert base64.b64decode(data["input_image"]) == b"first"
+    assert base64.b64decode(data["mask"]) == b"mask-bytes"
+
+
+async def test_async_transform_image_edit_request_downloads_only_the_first_url_of_a_list(monkeypatch):
+    fetched = []
+
+    async def serve(client, url, **kwargs):
+        fetched.append(url)
+        return httpx.Response(200, content=b"first-bytes", request=httpx.Request("GET", url))
+
+    monkeypatch.setattr(bfl_transformation, "safe_get", lambda *args, **kwargs: pytest.fail("sync fetch ran"))
+    monkeypatch.setattr(bfl_transformation, "async_safe_get", serve)
+
+    data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request(
+        model="flux-kontext-pro",
+        prompt="Add a red hat",
+        image=["https://cdn.example/a.png", "https://cdn.example/b.png"],
+        image_edit_optional_request_params={},
+        litellm_params=GenericLiteLLMParams(),
+        headers={},
+    )
+
+    assert fetched == ["https://cdn.example/a.png"]
+    assert base64.b64decode(data["input_image"]) == b"first-bytes"
diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
index f1614654ffb..3afd57bbf34 100644
--- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
+++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
@@ -19,6 +19,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import (
     BaseAudioTranscriptionConfig,
 )
 from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
+from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 from litellm.llms.custom_httpx.llm_http_handler import (
     BaseLLMHTTPHandler,
@@ -31,7 +32,7 @@ from litellm.llms.azure.videos.transformation import AzureVideoConfig
 from litellm.llms.openai.videos.transformation import OpenAIVideoConfig
 from litellm.types.llms.openai import ResponsesAPIResponse
 from litellm.types.router import GenericLiteLLMParams
-from litellm.types.utils import ModelResponse, TranscriptionResponse
+from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse
 
 _ACTIVE_KEY = "_code_interpreter_interception_active"
 _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key"
@@ -3244,6 +3245,11 @@ class _TransformRecordingConfig(BaseConfig):
     def get_error_class(self, error_message, status_code, headers):
         return BaseLLMException(status_code=status_code, message=error_message, headers=headers)
 
+    def get_model_response_iterator(self, streaming_response, sync_stream, json_mode=False):
+        return litellm.OpenAIGPTConfig().get_model_response_iterator(
+            streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode
+        )
+
 
 def _start_async_completion(config, logging_obj=None):
     captured = {}
@@ -3312,3 +3318,158 @@ async def test_completion_keeps_sync_transform_request_before_returning_by_defau
     assert config.transform_calls == ["sync"]
     assert captured["body"] == {"transformed_by": "sync"}
     assert response.choices[0].message.content == "sync"
+
+
+def _sse_echoing_transformed_by(request):
+    transformed_by = json.loads(request.content)["transformed_by"]
+    chunk = {
+        "id": "chatcmpl-1",
+        "object": "chat.completion.chunk",
+        "created": 1,
+        "model": "stub-model",
+        "choices": [{"index": 0, "delta": {"content": transformed_by}, "finish_reason": None}],
+    }
+    return httpx.Response(
+        200,
+        content=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode(),
+        headers={"content-type": "text/event-stream"},
+        request=request,
+    )
+
+
+def _streaming_logging_obj():
+    from litellm.litellm_core_utils.litellm_logging import Logging
+
+    logging_obj = Logging(
+        model="stub-model",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=True,
+        call_type="acompletion",
+        start_time=time.time(),
+        litellm_call_id="async-transform-stream",
+        function_id="f",
+    )
+    logging_obj.update_environment_variables(
+        model="stub-model", user="", optional_params={}, litellm_params={}, custom_llm_provider="openai"
+    )
+    return logging_obj
+
+
+async def test_completion_streams_after_the_async_transform_request():
+    config = _TransformRecordingConfig(transform_async=True)
+    loop_thread = threading.current_thread()
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(_sse_echoing_transformed_by))
+
+    stream = await BaseLLMHTTPHandler().completion(
+        model="stub-model",
+        messages=[{"role": "user", "content": "hi"}],
+        api_base="https://llm.example/v1/chat",
+        custom_llm_provider="openai",
+        model_response=ModelResponse(),
+        encoding=None,
+        logging_obj=_streaming_logging_obj(),
+        optional_params={},
+        timeout=10.0,
+        litellm_params={},
+        acompletion=True,
+        stream=True,
+        client=client,
+        provider_config=config,
+    )
+    collected = [chunk async for chunk in stream]
+
+    assert config.transform_calls == ["async"]
+    assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads)
+    assert "".join(chunk.choices[0].delta.content or "" for chunk in collected) == "async"
+
+
+class _ImageEditRecordingConfig(BaseImageEditConfig):
+    def __init__(self):
+        self.transform_calls = []
+
+    def get_supported_openai_params(self, model):
+        return []
+
+    def map_openai_params(self, image_edit_optional_params, model, drop_params):
+        return dict(image_edit_optional_params)
+
+    def validate_environment(self, headers, model, api_key=None, litellm_params=None, api_base=None):
+        return {}
+
+    def get_complete_url(self, model, api_base, litellm_params):
+        return "https://images.example/v1/edits"
+
+    def use_multipart_form_data(self):
+        return False
+
+    def transform_image_edit_request(
+        self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers
+    ):
+        self.transform_calls.append("sync")
+        return {"transformed_by": "sync"}, []
+
+    async def async_transform_image_edit_request(
+        self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers
+    ):
+        self.transform_calls.append("async")
+        return {"transformed_by": "async"}, []
+
+    def transform_image_edit_response(self, model, raw_response, logging_obj):
+        return ImageResponse(data=[ImageObject(b64_json=raw_response.json()["transformed_by"])])
+
+
+def _echo_json_transport(captured):
+    def handle(request):
+        captured["body"] = json.loads(request.content)
+        return httpx.Response(200, json=captured["body"])
+
+    return httpx.MockTransport(handle)
+
+
+async def test_async_image_edit_handler_awaits_the_async_transform():
+    config = _ImageEditRecordingConfig()
+    captured = {}
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=_echo_json_transport(captured))
+
+    response = await BaseLLMHTTPHandler().async_image_edit_handler(
+        model="edit-model",
+        image=b"raw-image",
+        prompt="add a hat",
+        image_edit_provider_config=config,
+        image_edit_optional_request_params={},
+        custom_llm_provider="openai",
+        litellm_params=GenericLiteLLMParams(),
+        logging_obj=Mock(),
+        timeout=10.0,
+        client=client,
+    )
+
+    assert config.transform_calls == ["async"]
+    assert captured["body"] == {"transformed_by": "async"}
+    assert response.data[0].b64_json == "async"
+
+
+def test_image_edit_handler_keeps_the_sync_transform():
+    config = _ImageEditRecordingConfig()
+    captured = {}
+    client = HTTPHandler()
+    client.client = httpx.Client(transport=_echo_json_transport(captured))
+
+    response = BaseLLMHTTPHandler().image_edit_handler(
+        model="edit-model",
+        image=b"raw-image",
+        prompt="add a hat",
+        image_edit_provider_config=config,
+        image_edit_optional_request_params={},
+        custom_llm_provider="openai",
+        litellm_params=GenericLiteLLMParams(),
+        logging_obj=Mock(),
+        timeout=10.0,
+        client=client,
+    )
+
+    assert config.transform_calls == ["sync"]
+    assert captured["body"] == {"transformed_by": "sync"}
+    assert response.data[0].b64_json == "sync"

From 3920cf4dfe458b6d4982e2a64c162de21defa2a0 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 23:45:37 -0700
Subject: [PATCH 055/164] fix(image_handling): tell callers when the image host
 did not resolve instead of blaming the URL policy

validate_url raises HostResolutionError, a SSRFError subclass, for the two
DNS outcomes (lookup failed, no addresses). The image fetch helper maps
that to a "host could not be resolved" message and keeps the
user_url_allowed_hosts hint for the policy verdicts it can actually fix.
---
 .../prompt_templates/image_handling.py        | 14 +++--
 litellm/litellm_core_utils/url_utils.py       |  8 ++-
 .../litellm_core_utils/test_image_handling.py | 56 +++++++++++--------
 .../litellm_core_utils/test_url_utils.py      | 19 ++++++-
 4 files changed, 64 insertions(+), 33 deletions(-)

diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py
index 4e7dfdb4017..99efbd13320 100644
--- a/litellm/litellm_core_utils/prompt_templates/image_handling.py
+++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py
@@ -15,7 +15,7 @@ import litellm
 from litellm import verbose_logger
 from litellm.caching.caching import InMemoryCache
 from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB
-from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get
+from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError, async_safe_get, safe_get
 from litellm.types.llms.openai import AllMessageValues
 
 MAX_IMGS_IN_MEMORY: Final = 10
@@ -78,8 +78,12 @@ def _process_image_response(response: Response, url: str) -> str:
     return result
 
 
-def _url_policy_rejection(url: str, verdict: SSRFError) -> "litellm.ImageFetchError":
-    verbose_logger.warning("Image fetch of %s rejected by the URL policy: %s", url, verdict)
+def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError":
+    verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict)
+    if isinstance(verdict, HostResolutionError):
+        return litellm.ImageFetchError(
+            f"Error: Unable to fetch image from URL. The image host could not be resolved. url={url}"
+        )
     return litellm.ImageFetchError(
         "Error: Unable to fetch image from URL. The proxy's URL policy rejected this host; "
         f"an admin can allow it with `user_url_allowed_hosts` in general_settings. url={url}"
@@ -108,7 +112,7 @@ async def async_convert_url_to_base64(url: str) -> str:
         except litellm.ImageFetchError:
             raise
         except SSRFError as e:
-            raise _url_policy_rejection(url, e) from e
+            raise _rejected_image_fetch(url, e) from e
         except Exception:
             pass
     raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}")
@@ -136,7 +140,7 @@ def convert_url_to_base64(url: str) -> str:
         except litellm.ImageFetchError:
             raise
         except SSRFError as e:
-            raise _url_policy_rejection(url, e) from e
+            raise _rejected_image_fetch(url, e) from e
         except Exception as e:
             verbose_logger.exception(e)
     raise litellm.ImageFetchError(
diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py
index fa070a648f5..f4388d5fd12 100644
--- a/litellm/litellm_core_utils/url_utils.py
+++ b/litellm/litellm_core_utils/url_utils.py
@@ -93,6 +93,10 @@ class SSRFError(ValueError):
     """Raised when a URL targets a blocked network."""
 
 
+class HostResolutionError(SSRFError):
+    pass
+
+
 def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str:
     """Percent-encode one user-controlled URL path segment.
 
@@ -324,10 +328,10 @@ def validate_url(url: str) -> tuple[str, str]:
     try:
         addrinfo: Final = socket.getaddrinfo(hostname, effective_port, proto=socket.IPPROTO_TCP)
     except socket.gaierror as e:
-        raise SSRFError(f"DNS resolution failed for '{hostname}': {e}")
+        raise HostResolutionError(f"DNS resolution failed for '{hostname}': {e}")
 
     if not addrinfo:
-        raise SSRFError(f"No addresses found for '{hostname}'")
+        raise HostResolutionError(f"No addresses found for '{hostname}'")
 
     if not is_allowlisted:
         for addrinfo_entry in addrinfo:
diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py
index d8ccc9251dd..3a909298aba 100644
--- a/tests/test_litellm/litellm_core_utils/test_image_handling.py
+++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py
@@ -17,7 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import (
     async_inline_remote_media,
     convert_url_to_base64,
 )
-from litellm.litellm_core_utils.url_utils import SSRFError
+from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError
 
 
 @pytest.fixture(autouse=True)
@@ -115,9 +115,7 @@ class StreamingLargeImageClient:
             request=Request("GET", url),
         )
         # Mock the iter_bytes method to return our generator
-        response.iter_bytes = lambda chunk_size=8192: generate_chunks(
-            size_bytes, chunk_size
-        )
+        response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size)
         return response
 
 
@@ -215,9 +213,7 @@ def test_streaming_download_handles_petabyte_file(monkeypatch):
     """
     # Simulate a 1 petabyte file (1,000,000 GB)
     # Without streaming protection, this would cause OOM or hang indefinitely
-    client = StreamingLargeImageClient(
-        size_mb=1_000_000_000, include_content_length=False
-    )
+    client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False)
     monkeypatch.setattr(litellm, "module_level_client", client)
 
     with pytest.raises(litellm.ImageFetchError) as excinfo:
@@ -429,20 +425,32 @@ async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fail
 
 
 _SSRF_VERDICTS = (
-    "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, "
-    "add the host to `user_url_allowed_hosts` in general_settings.",
-    "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known",
-    "No addresses found for 'internal.example'",
+    (
+        SSRFError(
+            "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, "
+            "add the host to `user_url_allowed_hosts` in general_settings."
+        ),
+        "The proxy's URL policy rejected this host; an admin can allow it with `user_url_allowed_hosts`",
+    ),
+    (
+        HostResolutionError(
+            "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known"
+        ),
+        "The image host could not be resolved",
+    ),
+    (HostResolutionError("No addresses found for 'internal.example'"), "The image host could not be resolved"),
 )
 
 
-def _assert_one_verdict_free_message(messages, url):
-    assert len(set(messages)) == 1
-    assert "10.0.0.8" not in messages[0]
-    assert "DNS" not in messages[0]
-    assert "No addresses" not in messages[0]
-    assert "user_url_allowed_hosts" in messages[0]
-    assert url in messages[0]
+def _assert_verdict_free_messages(messages, url):
+    for message, (verdict, expected_guidance) in zip(messages, _SSRF_VERDICTS, strict=True):
+        assert expected_guidance in message
+        assert url in message
+        assert "10.0.0.8" not in message
+        assert "DNS" not in message
+        assert "No addresses" not in message
+        if isinstance(verdict, HostResolutionError):
+            assert "user_url_allowed_hosts" not in message
 
 
 async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch):
@@ -450,11 +458,11 @@ async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_r
     messages = []
     url = f"http://internal.example/{uuid.uuid4()}.png"
 
-    for verdict in _SSRF_VERDICTS:
+    for verdict, _ in _SSRF_VERDICTS:
 
         async def block(client, fetched_url, verdict=verdict, **kwargs):
             attempts.append(fetched_url)
-            raise SSRFError(verdict)
+            raise verdict
 
         monkeypatch.setattr(image_handling, "async_safe_get", block)
         with pytest.raises(litellm.ImageFetchError) as raised:
@@ -462,7 +470,7 @@ async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_r
         messages.append(raised.value.message)
 
     assert attempts == [url] * len(_SSRF_VERDICTS)
-    _assert_one_verdict_free_message(messages, url)
+    _assert_verdict_free_messages(messages, url)
 
 
 def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch):
@@ -470,11 +478,11 @@ def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeyp
     messages = []
     url = f"http://internal.example/{uuid.uuid4()}.png"
 
-    for verdict in _SSRF_VERDICTS:
+    for verdict, _ in _SSRF_VERDICTS:
 
         def block(client, fetched_url, verdict=verdict, **kwargs):
             attempts.append(fetched_url)
-            raise SSRFError(verdict)
+            raise verdict
 
         monkeypatch.setattr(image_handling, "safe_get", block)
         with pytest.raises(litellm.ImageFetchError) as raised:
@@ -482,7 +490,7 @@ def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeyp
         messages.append(raised.value.message)
 
     assert attempts == [url] * len(_SSRF_VERDICTS)
-    _assert_one_verdict_free_message(messages, url)
+    _assert_verdict_free_messages(messages, url)
 
 
 async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch):
diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py
index fccdc1a2a0a..151c7ceb5bd 100644
--- a/tests/test_litellm/litellm_core_utils/test_url_utils.py
+++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py
@@ -9,6 +9,7 @@ import pytest
 import litellm
 from litellm.litellm_core_utils import url_utils
 from litellm.litellm_core_utils.url_utils import (
+    HostResolutionError,
     SSRFError,
     _is_blocked_ip,
     assert_same_origin,
@@ -161,10 +162,24 @@ class TestValidateUrl:
         assert "/path" in rewritten
         assert "key=value" in rewritten
 
-    def test_dns_failure_raises(self, mock_dns_failure):
-        with pytest.raises(SSRFError, match="DNS resolution failed"):
+    def test_dns_failure_raises_a_host_resolution_error(self, mock_dns_failure):
+        with pytest.raises(HostResolutionError, match="DNS resolution failed"):
             validate_url("http://this-domain-does-not-exist-xyz123.invalid/test")
 
+    def test_empty_resolution_raises_a_host_resolution_error(self, monkeypatch):
+        monkeypatch.setattr(url_utils.socket, "getaddrinfo", lambda *args, **kwargs: [])
+        with pytest.raises(HostResolutionError, match="No addresses found"):
+            validate_url("http://this-domain-resolves-to-nothing.invalid/test")
+
+    def test_blocked_address_is_not_a_host_resolution_error(self, monkeypatch):
+        def fake(host, port, *a, **kw):
+            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.8", port or 80))]
+
+        monkeypatch.setattr(url_utils.socket, "getaddrinfo", fake)
+        with pytest.raises(SSRFError, match="blocked address") as raised:
+            validate_url("http://internal.example/test")
+        assert not isinstance(raised.value, HostResolutionError)
+
     def test_blocks_localhost_hostname(self, monkeypatch):
         def fake(host, port, *a, **kw):
             return [

From de2ba3fab1daab25ba2517ea1a93cf7120f996f5 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 4 Sep 2026 23:50:01 -0700
Subject: [PATCH 056/164] fix(make check): lint the test tree on tests-only
 changes like CI does

CI's required lint job runs ruff with ruff-tests.toml over tests/ and the
test-quality budget gate, but scripts/pre_commit_lint.sh only triggered make
lint on litellm/ files, so a tests-only commit passed make check with a no-op
note and then failed CI (a duplicate test name, ruff F811, did exactly that).

When tests/ Python files are in scope and no litellm/ files are, run ruff
with ruff-tests.toml over the changed test files and make lint-test-quality,
with the matching partial-staging warning, summary line, and no-op condition.
---
 scripts/pre_commit_lint.sh                 | 25 +++++-
 tests/test_litellm/test_pre_commit_lint.py | 98 ++++++++++++++++++++--
 2 files changed, 114 insertions(+), 9 deletions(-)

diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh
index d38a0eee3de..ad1b4e793ab 100755
--- a/scripts/pre_commit_lint.sh
+++ b/scripts/pre_commit_lint.sh
@@ -12,6 +12,8 @@
 #   - litellm/ Python  -> `make lint` (test-linting.yml's lint job)
 #   - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step)
 #                         + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests)
+#   - tests/ Python    -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's
+#                         test-tree ruff and test-quality budget steps)
 #   - dashboard        -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint)
 #   - proxy/types      -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift (check-ui-api-types.yml)
 #
@@ -88,15 +90,18 @@ existing_files() {
 
 litellm_py_pattern='^litellm/.*\.py$'
 e2e_py_pattern='^tests/e2e/.*\.py$'
+tests_py_pattern='^tests/.*\.py$'
 spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$'
 ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$'
 ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$'
 
-# CI's lint job (test-linting.yml) only inspects litellm/, so a tests-only or
-# scripts-only commit can't turn it red; scope the trigger there to skip the slow
-# make lint when it couldn't catch anything.
+# CI's lint job (test-linting.yml) is make lint: litellm/ checks plus two test-tree
+# steps (ruff over ruff-tests.toml and the test-quality budget). Trigger the slow
+# make lint on litellm/ files only; a tests-only commit runs just those two steps.
 litellm_py_files=$(scope_match "$litellm_py_pattern")
 e2e_py_files=$(scope_match "$e2e_py_pattern")
+tests_py_changed=$(scope_match "$tests_py_pattern")
+tests_py_files=$(printf '%s\n' "$tests_py_changed" | existing_files)
 # ruff format (and CI's format step) skip enterprise; the rest of make lint covers it.
 fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files)
 # check-ui-api-types.yml triggers on any file under litellm/proxy or litellm/types
@@ -136,6 +141,7 @@ if [ -n "$staged" ]; then
     }
     warn_skipped "Python lint (make lint)" "$litellm_py_pattern" "$litellm_py_files"
     warn_skipped "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_pattern" "$e2e_py_files"
+    warn_skipped "test-tree lint (ruff-tests.toml + test-quality budget)" "$tests_py_pattern" "$tests_py_changed"
     warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_changed"
     warn_skipped "dashboard API-type sync (npm run gen:api)" "$spec_pattern" "$spec_files"
 fi
@@ -225,6 +231,16 @@ if [ -n "$e2e_py_files" ]; then
         || { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make check." >&2; status=1; }
 fi
 
+if [ -n "$tests_py_changed" ] && [ -z "$litellm_py_files" ]; then
+    if [ -n "$tests_py_files" ]; then
+        echo "check: linting the test tree (ruff check --config ruff-tests.toml, scoped tests files)"
+        printf '%s\n' "$tests_py_files" | xargs uv run --no-sync ruff check --config ruff-tests.toml \
+            || { echo "✗ Test-tree ruff failed. Fix the errors above, then re-run make check." >&2; status=1; }
+    fi
+    echo "check: checking the test-quality budget (make lint-test-quality)"
+    make lint-test-quality || { echo "✗ Test-quality budget failed. Fix the errors above, then re-run make check." >&2; status=1; }
+fi
+
 dashboard_checks() {
     echo "check: linting dashboard (prettier + eslint + lint budgets)"
     if [ ! -d ui/litellm-dashboard/node_modules ]; then
@@ -313,10 +329,11 @@ summary_item() {
 echo "check: summary"
 summary_item "Python lint (make lint)" "$litellm_py_files" "no litellm/ Python files in scope"
 summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_files" "no tests/e2e Python files in scope"
+summary_item "test-tree lint (ruff-tests.toml + test-quality budget)" "$tests_py_changed" "no tests/ Python files in scope"
 summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope"
 summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope"
 
-if [ -z "$litellm_py_files$e2e_py_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then
+if [ -z "$litellm_py_files$e2e_py_files$tests_py_changed$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then
     echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2
     printf '%s\n' "$scope" | sed 's/^/    /' >&2
     echo "  A pass here is a no-op, not a lint verdict." >&2
diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py
index aa2260e89ea..22ae0e1bfaa 100644
--- a/tests/test_litellm/test_pre_commit_lint.py
+++ b/tests/test_litellm/test_pre_commit_lint.py
@@ -40,6 +40,9 @@ case "$*" in
             sleep 60
         fi
         ;;
+    lint-test-quality)
+        [ "${STUB_FAIL:-}" = "test-quality" ] && exit 1
+        ;;
 esac
 exit 0
 """
@@ -69,6 +72,10 @@ case "$*" in
     *orjson*)
         [ -n "${STUB_BARRIER_DIR:-}" ] && barrier_sync genapi "python dashboard"
         ;;
+    "run --no-sync ruff check --config ruff-tests.toml"*)
+        [ -n "${STUB_ARGS_DIR:-}" ] && echo "$*" >> "$STUB_ARGS_DIR/ruff_tests.args"
+        [ "${STUB_FAIL:-}" = "tests-ruff" ] && exit 1
+        ;;
 esac
 exit 0
 """
@@ -153,6 +160,13 @@ def _set_base_ref(repo: Path) -> None:
     )
 
 
+def _stage_file(repo: Path, relative: str, body: str) -> None:
+    path = repo / relative
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(body)
+    subprocess.run(["git", "add", relative], cwd=repo, check=True)
+
+
 def test_nothing_staged_scopes_to_working_tree_diff_and_runs_checks(tmp_path: Path) -> None:
     repo, bin_dir = _sandbox(tmp_path)
     _commit_all(repo, "base")
@@ -405,6 +419,7 @@ def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> No
     assert "ran:     dashboard lint (prettier + eslint + lint budgets)" in proc.stdout
     assert "ran:     dashboard API-type sync (npm run gen:api)" in proc.stdout
     assert "skipped: tests/e2e checks (basedpyright + raw HTTP client ban) (no tests/e2e Python files in scope)" in proc.stdout
+    assert "skipped: test-tree lint (ruff-tests.toml + test-quality budget) (no tests/ Python files in scope)" in proc.stdout
     assert "check: PASS" in proc.stdout
     assert "check: FAIL" not in proc.stdout
 
@@ -412,20 +427,93 @@ def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> No
 def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty_log(tmp_path: Path) -> None:
     repo, bin_dir = _sandbox(tmp_path)
     _commit_all(repo, "base")
-    tests_dir = repo / "tests" / "test_litellm"
-    tests_dir.mkdir(parents=True)
-    (tests_dir / "test_x.py").write_text("def test_x() -> None: ...\n")
-    subprocess.run(["git", "add", "tests"], cwd=repo, check=True)
+    _stage_file(repo, "scripts/tool.py", "def main() -> None: ...\n")
     proc = _run(repo, bin_dir, {})
     assert proc.returncode == 0, proc.stdout + proc.stderr
     assert "no gating lint check matches the files in scope, so nothing ran" in proc.stdout
-    assert "tests/test_litellm/test_x.py" in proc.stdout
+    assert "scripts/tool.py" in proc.stdout
     assert "a no-op, not a lint verdict" in proc.stdout
     assert "check: PASS" in proc.stdout
     assert "linting Python" not in proc.stdout
     log = (repo / ".git" / "pre_commit_lint.log").read_text()
     assert "check: summary" in log
     assert "skipped: Python lint (make lint) (no litellm/ Python files in scope)" in log
+    assert "skipped: test-tree lint (ruff-tests.toml + test-quality budget) (no tests/ Python files in scope)" in log
+
+
+def test_tests_only_change_runs_test_tree_ruff_on_staged_files_and_the_quality_gate(tmp_path: Path) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    _commit_all(repo, "base")
+    args_dir = tmp_path / "args"
+    args_dir.mkdir()
+    _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n")
+    _stage_file(repo, "tests/test_b.py", "def test_b() -> None: ...\n")
+    _stage_file(repo, "tests/fixtures/data.json", "{}\n")
+    proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)})
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+    ruff_args = (args_dir / "ruff_tests.args").read_text().splitlines()
+    assert ruff_args == ["run --no-sync ruff check --config ruff-tests.toml tests/test_a.py tests/test_b.py"]
+    assert "ran:     test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout
+    assert "no gating lint check matches" not in proc.stdout
+    assert "linting Python" not in proc.stdout
+    assert "check: PASS" in proc.stdout
+
+
+@pytest.mark.parametrize(
+    ("fail", "message"),
+    [
+        ("tests-ruff", "Test-tree ruff failed"),
+        ("test-quality", "Test-quality budget failed"),
+    ],
+)
+def test_a_failing_test_tree_check_fails_a_tests_only_run(tmp_path: Path, fail: str, message: str) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    _commit_all(repo, "base")
+    _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n")
+    proc = _run(repo, bin_dir, {"STUB_FAIL": fail})
+    assert proc.returncode == 1
+    assert message in proc.stdout + proc.stderr
+    assert "check: FAIL" in proc.stdout
+
+
+def test_tests_changed_alongside_litellm_files_defer_to_make_lint(tmp_path: Path) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    _commit_all(repo, "base")
+    args_dir = tmp_path / "args"
+    args_dir.mkdir()
+    _stage_file(repo, "litellm/foo.py", "x = 2\n")
+    _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n")
+    proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "test-quality"})
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+    assert "linting Python" in proc.stdout
+    assert not (args_dir / "ruff_tests.args").exists()
+    assert "ran:     test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout
+
+
+def test_deleted_test_file_still_runs_the_quality_gate_without_feeding_ruff_the_missing_file(tmp_path: Path) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n")
+    _commit_all(repo, "base")
+    args_dir = tmp_path / "args"
+    args_dir.mkdir()
+    subprocess.run(["git", "rm", "-q", "tests/test_a.py"], cwd=repo, check=True)
+    proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "test-quality"})
+    assert proc.returncode == 1
+    assert "Test-quality budget failed" in proc.stdout + proc.stderr
+    assert not (args_dir / "ruff_tests.args").exists()
+
+
+def test_partial_staging_warns_when_test_files_are_left_unstaged(tmp_path: Path) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n")
+    _commit_all(repo, "base")
+    _stage_file(repo, "notes.md", "hi\n")
+    (repo / "tests" / "test_a.py").write_text("def test_a() -> None:\n    assert True\n")
+    proc = _run(repo, bin_dir, {})
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+    assert "SKIPPED test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout
+    assert "tests/test_a.py" in proc.stdout
+    assert "make lint-test-quality" not in proc.stdout
 
 
 def test_run_queues_through_the_machine_wide_gate_slot_lock(tmp_path: Path) -> None:

From d22962248c1c659279e7389e6c2e1d1b640b400a Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 00:07:25 -0700
Subject: [PATCH 057/164] fix(check): run CI's whole-tree test ruff and widen
 the test-tree trigger

The scoped xargs list missed a ruff-tests.toml rule change and skipped ruff on
deletions, so the block now runs test-linting.yml's exact command over tests/.
ruff-tests.toml, test-quality-budget.json, and scripts/check_test_quality.py
trigger the block too, and it sits after the background launches so the
dashboard and gen:api jobs overlap it.
---
 scripts/pre_commit_lint.sh                 |  36 ++++----
 tests/test_litellm/test_pre_commit_lint.py | 102 ++++++++++++++++-----
 2 files changed, 96 insertions(+), 42 deletions(-)

diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh
index ad1b4e793ab..ab84d2d518a 100755
--- a/scripts/pre_commit_lint.sh
+++ b/scripts/pre_commit_lint.sh
@@ -12,7 +12,8 @@
 #   - litellm/ Python  -> `make lint` (test-linting.yml's lint job)
 #   - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step)
 #                         + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests)
-#   - tests/ Python    -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's
+#   - tests/ Python, ruff-tests.toml, test-quality-budget.json, scripts/check_test_quality.py
+#                      -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's
 #                         test-tree ruff and test-quality budget steps)
 #   - dashboard        -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint)
 #   - proxy/types      -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift (check-ui-api-types.yml)
@@ -90,18 +91,17 @@ existing_files() {
 
 litellm_py_pattern='^litellm/.*\.py$'
 e2e_py_pattern='^tests/e2e/.*\.py$'
-tests_py_pattern='^tests/.*\.py$'
+test_tree_pattern='^(tests/.*\.py|ruff-tests\.toml|test-quality-budget\.json|scripts/check_test_quality\.py)$'
 spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$'
 ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$'
 ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$'
 
 # CI's lint job (test-linting.yml) is make lint: litellm/ checks plus two test-tree
 # steps (ruff over ruff-tests.toml and the test-quality budget). Trigger the slow
-# make lint on litellm/ files only; a tests-only commit runs just those two steps.
+# make lint on litellm/ files only; without them, the test-tree steps run on their own below.
 litellm_py_files=$(scope_match "$litellm_py_pattern")
 e2e_py_files=$(scope_match "$e2e_py_pattern")
-tests_py_changed=$(scope_match "$tests_py_pattern")
-tests_py_files=$(printf '%s\n' "$tests_py_changed" | existing_files)
+test_tree_files=$(scope_match "$test_tree_pattern")
 # ruff format (and CI's format step) skip enterprise; the rest of make lint covers it.
 fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files)
 # check-ui-api-types.yml triggers on any file under litellm/proxy or litellm/types
@@ -141,7 +141,7 @@ if [ -n "$staged" ]; then
     }
     warn_skipped "Python lint (make lint)" "$litellm_py_pattern" "$litellm_py_files"
     warn_skipped "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_pattern" "$e2e_py_files"
-    warn_skipped "test-tree lint (ruff-tests.toml + test-quality budget)" "$tests_py_pattern" "$tests_py_changed"
+    warn_skipped "test-tree lint (ruff-tests.toml + test-quality budget)" "$test_tree_pattern" "$test_tree_files"
     warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_changed"
     warn_skipped "dashboard API-type sync (npm run gen:api)" "$spec_pattern" "$spec_files"
 fi
@@ -231,16 +231,6 @@ if [ -n "$e2e_py_files" ]; then
         || { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make check." >&2; status=1; }
 fi
 
-if [ -n "$tests_py_changed" ] && [ -z "$litellm_py_files" ]; then
-    if [ -n "$tests_py_files" ]; then
-        echo "check: linting the test tree (ruff check --config ruff-tests.toml, scoped tests files)"
-        printf '%s\n' "$tests_py_files" | xargs uv run --no-sync ruff check --config ruff-tests.toml \
-            || { echo "✗ Test-tree ruff failed. Fix the errors above, then re-run make check." >&2; status=1; }
-    fi
-    echo "check: checking the test-quality budget (make lint-test-quality)"
-    make lint-test-quality || { echo "✗ Test-quality budget failed. Fix the errors above, then re-run make check." >&2; status=1; }
-fi
-
 dashboard_checks() {
     echo "check: linting dashboard (prettier + eslint + lint budgets)"
     if [ ! -d ui/litellm-dashboard/node_modules ]; then
@@ -304,6 +294,15 @@ if [ -n "$spec_files" ]; then
     set +m
 fi
 
+if [ -n "$test_tree_files" ] && [ -z "$litellm_py_files" ]; then
+    echo "check: linting the test tree (ruff check --config ruff-tests.toml tests)"
+    uv run --no-sync ruff check --config ruff-tests.toml tests \
+        || { echo "✗ Test-tree ruff failed. Fix the errors above, then re-run make check." >&2; status=1; }
+    echo "check: checking the test-quality budget (make lint-test-quality)"
+    make lint-test-quality \
+        || { echo "✗ Test-quality budget failed. Fix the errors above, then re-run make check." >&2; status=1; }
+fi
+
 if [ -n "${python_pid:-}" ]; then
     wait "$python_pid" || status=1
     cat "$python_log"; rm -f "$python_log"
@@ -329,11 +328,12 @@ summary_item() {
 echo "check: summary"
 summary_item "Python lint (make lint)" "$litellm_py_files" "no litellm/ Python files in scope"
 summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_files" "no tests/e2e Python files in scope"
-summary_item "test-tree lint (ruff-tests.toml + test-quality budget)" "$tests_py_changed" "no tests/ Python files in scope"
+summary_item "test-tree lint (ruff-tests.toml + test-quality budget)" "$test_tree_files" \
+    "no tests/ Python files or test-tree lint inputs in scope"
 summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope"
 summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope"
 
-if [ -z "$litellm_py_files$e2e_py_files$tests_py_changed$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then
+if [ -z "$litellm_py_files$e2e_py_files$test_tree_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then
     echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2
     printf '%s\n' "$scope" | sed 's/^/    /' >&2
     echo "  A pass here is a no-op, not a lint verdict." >&2
diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py
index 22ae0e1bfaa..e17abf2fa9f 100644
--- a/tests/test_litellm/test_pre_commit_lint.py
+++ b/tests/test_litellm/test_pre_commit_lint.py
@@ -11,6 +11,12 @@ import pytest
 
 ROOT = Path(__file__).resolve().parents[2]
 SCRIPT = ROOT / "scripts" / "pre_commit_lint.sh"
+WHOLE_TREE_RUFF = "run --no-sync ruff check --config ruff-tests.toml tests"
+TEST_TREE_RAN = "ran:     test-tree lint (ruff-tests.toml + test-quality budget)"
+TEST_TREE_SKIPPED = (
+    "skipped: test-tree lint (ruff-tests.toml + test-quality budget) "
+    "(no tests/ Python files or test-tree lint inputs in scope)"
+)
 
 BARRIER_HELPER = """barrier_sync() {
     touch "$STUB_BARRIER_DIR/$1.started"
@@ -30,6 +36,7 @@ BARRIER_HELPER = """barrier_sync() {
 
 MAKE_STUB = """#!/bin/sh
 . "$STUB_BIN/barrier.sh"
+[ -n "${STUB_ARGS_DIR:-}" ] && echo "$*" >> "$STUB_ARGS_DIR/make.args"
 case "$*" in
     lint)
         [ "${STUB_FAIL:-}" = "make-lint" ] && exit 1
@@ -419,7 +426,7 @@ def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> No
     assert "ran:     dashboard lint (prettier + eslint + lint budgets)" in proc.stdout
     assert "ran:     dashboard API-type sync (npm run gen:api)" in proc.stdout
     assert "skipped: tests/e2e checks (basedpyright + raw HTTP client ban) (no tests/e2e Python files in scope)" in proc.stdout
-    assert "skipped: test-tree lint (ruff-tests.toml + test-quality budget) (no tests/ Python files in scope)" in proc.stdout
+    assert TEST_TREE_SKIPPED in proc.stdout
     assert "check: PASS" in proc.stdout
     assert "check: FAIL" not in proc.stdout
 
@@ -438,41 +445,83 @@ def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty
     log = (repo / ".git" / "pre_commit_lint.log").read_text()
     assert "check: summary" in log
     assert "skipped: Python lint (make lint) (no litellm/ Python files in scope)" in log
-    assert "skipped: test-tree lint (ruff-tests.toml + test-quality budget) (no tests/ Python files in scope)" in log
+    assert TEST_TREE_SKIPPED in log
 
 
-def test_tests_only_change_runs_test_tree_ruff_on_staged_files_and_the_quality_gate(tmp_path: Path) -> None:
+def _recorded(args_dir: Path, name: str) -> list[str]:
+    path = args_dir / name
+    return path.read_text().splitlines() if path.exists() else []
+
+
+def test_tests_only_change_runs_the_whole_test_tree_ruff_and_the_quality_gate(tmp_path: Path) -> None:
     repo, bin_dir = _sandbox(tmp_path)
     _commit_all(repo, "base")
     args_dir = tmp_path / "args"
     args_dir.mkdir()
     _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n")
-    _stage_file(repo, "tests/test_b.py", "def test_b() -> None: ...\n")
     _stage_file(repo, "tests/fixtures/data.json", "{}\n")
     proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)})
     assert proc.returncode == 0, proc.stdout + proc.stderr
-    ruff_args = (args_dir / "ruff_tests.args").read_text().splitlines()
-    assert ruff_args == ["run --no-sync ruff check --config ruff-tests.toml tests/test_a.py tests/test_b.py"]
-    assert "ran:     test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout
+    assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF]
+    assert _recorded(args_dir, "make.args") == ["lint-test-quality"]
+    assert TEST_TREE_RAN in proc.stdout
     assert "no gating lint check matches" not in proc.stdout
     assert "linting Python" not in proc.stdout
     assert "check: PASS" in proc.stdout
 
 
 @pytest.mark.parametrize(
-    ("fail", "message"),
-    [
-        ("tests-ruff", "Test-tree ruff failed"),
-        ("test-quality", "Test-quality budget failed"),
-    ],
+    "changed",
+    ["ruff-tests.toml", "test-quality-budget.json", "scripts/check_test_quality.py", "tests/e2e/test_x.py"],
 )
-def test_a_failing_test_tree_check_fails_a_tests_only_run(tmp_path: Path, fail: str, message: str) -> None:
+def test_test_tree_lint_inputs_trigger_the_test_tree_checks(tmp_path: Path, changed: str) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    _commit_all(repo, "base")
+    args_dir = tmp_path / "args"
+    args_dir.mkdir()
+    _stage_file(repo, changed, "x = 1\n")
+    proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)})
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+    assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF]
+    assert "lint-test-quality" in _recorded(args_dir, "make.args")
+    assert TEST_TREE_RAN in proc.stdout
+
+
+def test_nothing_staged_tests_only_working_tree_change_runs_the_test_tree_checks(tmp_path: Path) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n")
+    _commit_all(repo, "base")
+    _set_base_ref(repo)
+    args_dir = tmp_path / "args"
+    args_dir.mkdir()
+    (repo / "tests" / "test_a.py").write_text("def test_a() -> None:\n    assert True\n")
+    proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)})
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+    assert "nothing staged; scoping to the working tree's diff" in proc.stdout
+    assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF]
+    assert _recorded(args_dir, "make.args") == ["lint-test-quality"]
+
+
+def test_a_failing_test_tree_ruff_fails_the_run_and_still_runs_the_quality_gate(tmp_path: Path) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    _commit_all(repo, "base")
+    args_dir = tmp_path / "args"
+    args_dir.mkdir()
+    _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n")
+    proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "tests-ruff"})
+    assert proc.returncode == 1
+    assert "Test-tree ruff failed" in proc.stdout + proc.stderr
+    assert "check: FAIL" in proc.stdout
+    assert _recorded(args_dir, "make.args") == ["lint-test-quality"]
+
+
+def test_a_failing_quality_gate_fails_a_tests_only_run(tmp_path: Path) -> None:
     repo, bin_dir = _sandbox(tmp_path)
     _commit_all(repo, "base")
     _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n")
-    proc = _run(repo, bin_dir, {"STUB_FAIL": fail})
+    proc = _run(repo, bin_dir, {"STUB_FAIL": "test-quality"})
     assert proc.returncode == 1
-    assert message in proc.stdout + proc.stderr
+    assert "Test-quality budget failed" in proc.stdout + proc.stderr
     assert "check: FAIL" in proc.stdout
 
 
@@ -486,34 +535,39 @@ def test_tests_changed_alongside_litellm_files_defer_to_make_lint(tmp_path: Path
     proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "test-quality"})
     assert proc.returncode == 0, proc.stdout + proc.stderr
     assert "linting Python" in proc.stdout
-    assert not (args_dir / "ruff_tests.args").exists()
-    assert "ran:     test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout
+    assert _recorded(args_dir, "ruff_tests.args") == []
+    assert _recorded(args_dir, "make.args") == ["lint"]
+    assert TEST_TREE_RAN in proc.stdout
 
 
-def test_deleted_test_file_still_runs_the_quality_gate_without_feeding_ruff_the_missing_file(tmp_path: Path) -> None:
+def test_deleted_test_file_still_runs_the_test_tree_checks(tmp_path: Path) -> None:
     repo, bin_dir = _sandbox(tmp_path)
     _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n")
     _commit_all(repo, "base")
     args_dir = tmp_path / "args"
     args_dir.mkdir()
     subprocess.run(["git", "rm", "-q", "tests/test_a.py"], cwd=repo, check=True)
-    proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "test-quality"})
-    assert proc.returncode == 1
-    assert "Test-quality budget failed" in proc.stdout + proc.stderr
-    assert not (args_dir / "ruff_tests.args").exists()
+    proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)})
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+    assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF]
+    assert _recorded(args_dir, "make.args") == ["lint-test-quality"]
+    assert TEST_TREE_RAN in proc.stdout
 
 
 def test_partial_staging_warns_when_test_files_are_left_unstaged(tmp_path: Path) -> None:
     repo, bin_dir = _sandbox(tmp_path)
     _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n")
     _commit_all(repo, "base")
+    args_dir = tmp_path / "args"
+    args_dir.mkdir()
     _stage_file(repo, "notes.md", "hi\n")
     (repo / "tests" / "test_a.py").write_text("def test_a() -> None:\n    assert True\n")
-    proc = _run(repo, bin_dir, {})
+    proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)})
     assert proc.returncode == 0, proc.stdout + proc.stderr
     assert "SKIPPED test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout
     assert "tests/test_a.py" in proc.stdout
-    assert "make lint-test-quality" not in proc.stdout
+    assert _recorded(args_dir, "ruff_tests.args") == []
+    assert _recorded(args_dir, "make.args") == []
 
 
 def test_run_queues_through_the_machine_wide_gate_slot_lock(tmp_path: Path) -> None:

From 36b0d80d3a1459f116f07e45b205e953e21c8978 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 00:09:00 -0700
Subject: [PATCH 058/164] test(cost): pin the nested-reasoning helper's clamps

The strip in text_tokens_without_nested_reasoning is capped at the reasoning
share, the reported text, and the over-sum past completion_tokens. Dropping the
caps to a bare over-sum passed every existing test, so this pins each cap with
a parametrized helper test plus one billing test where text over-reports past
the reasoning share and only the nested share may be netted out
---
 .../llm_cost_calc/test_llm_cost_calc_utils.py | 23 ++++++++++++++
 tests/test_litellm/types/test_types_utils.py  | 31 ++++++++++++++++++-
 2 files changed, 53 insertions(+), 1 deletion(-)

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 e8eae73df3f..c9fe53651cd 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
@@ -4821,6 +4821,29 @@ def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tok
     assert completion_cost == pytest.approx(44 * info["output_cost_per_token"])
 
 
+def test_generic_cost_per_token_strips_only_the_reasoning_share_when_text_over_reports(
+    _local_model_cost_map: None,
+) -> None:
+    """Text over-reported past the reasoning share keeps its extra tokens billed; only the nested reasoning is netted out."""
+
+    model = "gpt-realtime-2.1-mini"
+    usage = Usage(
+        prompt_tokens=120,
+        completion_tokens=100,
+        total_tokens=220,
+        completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=100, audio_tokens=70, reasoning_tokens=10),
+    )
+
+    _, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai")
+    breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage)
+
+    info = litellm.get_model_info(model=model, custom_llm_provider="openai")
+    assert breakdown.reasoning_cost == pytest.approx(10 * info["output_cost_per_token"])
+    assert completion_cost == pytest.approx(
+        100 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"]
+    )
+
+
 def test_generic_cost_per_token_bills_nested_reasoning_once_beside_audio_output(_local_model_cost_map: None) -> None:
     """Audio-output realtime usage nests reasoning inside text_tokens next to audio_tokens; text is billed net of it."""
 
diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py
index 554604ab200..5f44ba1773e 100644
--- a/tests/test_litellm/types/test_types_utils.py
+++ b/tests/test_litellm/types/test_types_utils.py
@@ -3,7 +3,7 @@ from typing import Final
 import pytest
 
 
-from litellm.types.utils import HiddenParams, all_litellm_params
+from litellm.types.utils import HiddenParams, all_litellm_params, text_tokens_without_nested_reasoning
 
 
 def test_rust_is_a_known_litellm_param():
@@ -768,3 +768,32 @@ def test_image_response_keeps_background():
     response = ImageResponse(created=1, data=[{"b64_json": "aGk="}], background="transparent", output_format="png")
     assert response.background == "transparent"
     assert response.model_dump()["background"] == "transparent"
+
+
+@pytest.mark.parametrize(
+    ("completion_tokens", "text_tokens", "reasoning_tokens", "other_modality_tokens", "expected_text_tokens"),
+    (
+        pytest.param(50, 30, 20, 0, 30, id="details_sum_to_completion_is_a_no_op"),
+        pytest.param(34, 30, 24, 0, 10, id="strip_is_capped_at_the_over_sum"),
+        pytest.param(100, 100, 10, 70, 90, id="only_the_reasoning_share_is_stripped_when_text_over_reports_further"),
+        pytest.param(10, 5, 20, 0, 0, id="text_never_goes_negative_when_reasoning_exceeds_it"),
+    ),
+)
+def test_text_tokens_without_nested_reasoning_clamps(
+    completion_tokens: int,
+    text_tokens: int,
+    reasoning_tokens: int,
+    other_modality_tokens: int,
+    expected_text_tokens: int,
+) -> None:
+    """The strip never exceeds the reasoning share, the reported text, or the over-sum past completion_tokens."""
+
+    assert (
+        text_tokens_without_nested_reasoning(
+            completion_tokens=completion_tokens,
+            text_tokens=text_tokens,
+            reasoning_tokens=reasoning_tokens,
+            other_modality_tokens=other_modality_tokens,
+        )
+        == expected_text_tokens
+    )

From 827554954d305b7c8c24f524c961cf7a27e3a029 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 00:16:14 -0700
Subject: [PATCH 059/164] fix(image_handling): answer every SSRF rejection with
 one message so error text cannot probe internal hostnames

---
 .../prompt_templates/image_handling.py        | 10 ++---
 litellm/litellm_core_utils/url_utils.py       |  8 +---
 .../litellm_core_utils/test_image_handling.py | 41 ++++++++-----------
 .../litellm_core_utils/test_url_utils.py      | 19 +--------
 4 files changed, 24 insertions(+), 54 deletions(-)

diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py
index 99efbd13320..24f3b8bca7f 100644
--- a/litellm/litellm_core_utils/prompt_templates/image_handling.py
+++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py
@@ -15,7 +15,7 @@ import litellm
 from litellm import verbose_logger
 from litellm.caching.caching import InMemoryCache
 from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB
-from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError, async_safe_get, safe_get
+from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get
 from litellm.types.llms.openai import AllMessageValues
 
 MAX_IMGS_IN_MEMORY: Final = 10
@@ -80,13 +80,9 @@ def _process_image_response(response: Response, url: str) -> str:
 
 def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError":
     verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict)
-    if isinstance(verdict, HostResolutionError):
-        return litellm.ImageFetchError(
-            f"Error: Unable to fetch image from URL. The image host could not be resolved. url={url}"
-        )
     return litellm.ImageFetchError(
-        "Error: Unable to fetch image from URL. The proxy's URL policy rejected this host; "
-        f"an admin can allow it with `user_url_allowed_hosts` in general_settings. url={url}"
+        "Error: Unable to fetch image from URL. The proxy could not resolve this host or its URL policy rejected it; "
+        f"an admin can check the proxy log and `user_url_allowed_hosts` in general_settings. url={url}"
     )
 
 
diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py
index f4388d5fd12..fa070a648f5 100644
--- a/litellm/litellm_core_utils/url_utils.py
+++ b/litellm/litellm_core_utils/url_utils.py
@@ -93,10 +93,6 @@ class SSRFError(ValueError):
     """Raised when a URL targets a blocked network."""
 
 
-class HostResolutionError(SSRFError):
-    pass
-
-
 def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str:
     """Percent-encode one user-controlled URL path segment.
 
@@ -328,10 +324,10 @@ def validate_url(url: str) -> tuple[str, str]:
     try:
         addrinfo: Final = socket.getaddrinfo(hostname, effective_port, proto=socket.IPPROTO_TCP)
     except socket.gaierror as e:
-        raise HostResolutionError(f"DNS resolution failed for '{hostname}': {e}")
+        raise SSRFError(f"DNS resolution failed for '{hostname}': {e}")
 
     if not addrinfo:
-        raise HostResolutionError(f"No addresses found for '{hostname}'")
+        raise SSRFError(f"No addresses found for '{hostname}'")
 
     if not is_allowlisted:
         for addrinfo_entry in addrinfo:
diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py
index 3a909298aba..8fa4bd6c14d 100644
--- a/tests/test_litellm/litellm_core_utils/test_image_handling.py
+++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py
@@ -17,7 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import (
     async_inline_remote_media,
     convert_url_to_base64,
 )
-from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError
+from litellm.litellm_core_utils.url_utils import SSRFError
 
 
 @pytest.fixture(autouse=True)
@@ -425,32 +425,25 @@ async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fail
 
 
 _SSRF_VERDICTS = (
-    (
-        SSRFError(
-            "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, "
-            "add the host to `user_url_allowed_hosts` in general_settings."
-        ),
-        "The proxy's URL policy rejected this host; an admin can allow it with `user_url_allowed_hosts`",
+    SSRFError(
+        "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, "
+        "add the host to `user_url_allowed_hosts` in general_settings."
     ),
-    (
-        HostResolutionError(
-            "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known"
-        ),
-        "The image host could not be resolved",
-    ),
-    (HostResolutionError("No addresses found for 'internal.example'"), "The image host could not be resolved"),
+    SSRFError("DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known"),
+    SSRFError("No addresses found for 'internal.example'"),
 )
 
 
 def _assert_verdict_free_messages(messages, url):
-    for message, (verdict, expected_guidance) in zip(messages, _SSRF_VERDICTS, strict=True):
-        assert expected_guidance in message
-        assert url in message
-        assert "10.0.0.8" not in message
-        assert "DNS" not in message
-        assert "No addresses" not in message
-        if isinstance(verdict, HostResolutionError):
-            assert "user_url_allowed_hosts" not in message
+    assert len(messages) == len(_SSRF_VERDICTS)
+    assert len(set(messages)) == 1, "a caller must not be able to tell a blocked host from one that does not resolve"
+    message = messages[0]
+    assert "The proxy could not resolve this host or its URL policy rejected it" in message
+    assert "user_url_allowed_hosts" in message
+    assert url in message
+    assert "10.0.0.8" not in message
+    assert "DNS" not in message
+    assert "No addresses" not in message
 
 
 async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch):
@@ -458,7 +451,7 @@ async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_r
     messages = []
     url = f"http://internal.example/{uuid.uuid4()}.png"
 
-    for verdict, _ in _SSRF_VERDICTS:
+    for verdict in _SSRF_VERDICTS:
 
         async def block(client, fetched_url, verdict=verdict, **kwargs):
             attempts.append(fetched_url)
@@ -478,7 +471,7 @@ def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeyp
     messages = []
     url = f"http://internal.example/{uuid.uuid4()}.png"
 
-    for verdict, _ in _SSRF_VERDICTS:
+    for verdict in _SSRF_VERDICTS:
 
         def block(client, fetched_url, verdict=verdict, **kwargs):
             attempts.append(fetched_url)
diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py
index 151c7ceb5bd..fccdc1a2a0a 100644
--- a/tests/test_litellm/litellm_core_utils/test_url_utils.py
+++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py
@@ -9,7 +9,6 @@ import pytest
 import litellm
 from litellm.litellm_core_utils import url_utils
 from litellm.litellm_core_utils.url_utils import (
-    HostResolutionError,
     SSRFError,
     _is_blocked_ip,
     assert_same_origin,
@@ -162,24 +161,10 @@ class TestValidateUrl:
         assert "/path" in rewritten
         assert "key=value" in rewritten
 
-    def test_dns_failure_raises_a_host_resolution_error(self, mock_dns_failure):
-        with pytest.raises(HostResolutionError, match="DNS resolution failed"):
+    def test_dns_failure_raises(self, mock_dns_failure):
+        with pytest.raises(SSRFError, match="DNS resolution failed"):
             validate_url("http://this-domain-does-not-exist-xyz123.invalid/test")
 
-    def test_empty_resolution_raises_a_host_resolution_error(self, monkeypatch):
-        monkeypatch.setattr(url_utils.socket, "getaddrinfo", lambda *args, **kwargs: [])
-        with pytest.raises(HostResolutionError, match="No addresses found"):
-            validate_url("http://this-domain-resolves-to-nothing.invalid/test")
-
-    def test_blocked_address_is_not_a_host_resolution_error(self, monkeypatch):
-        def fake(host, port, *a, **kw):
-            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.8", port or 80))]
-
-        monkeypatch.setattr(url_utils.socket, "getaddrinfo", fake)
-        with pytest.raises(SSRFError, match="blocked address") as raised:
-            validate_url("http://internal.example/test")
-        assert not isinstance(raised.value, HostResolutionError)
-
     def test_blocks_localhost_hostname(self, monkeypatch):
         def fake(host, port, *a, **kw):
             return [

From e3366dddf44c7450a856ed472c2966798413ef36 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 00:26:16 -0700
Subject: [PATCH 060/164] fix(check): trigger the test-tree checks on the gate
 script and drop the scope comment

---
 scripts/pre_commit_lint.sh                 | 8 +++-----
 tests/test_litellm/test_pre_commit_lint.py | 8 +++++++-
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh
index ab84d2d518a..f245803408c 100755
--- a/scripts/pre_commit_lint.sh
+++ b/scripts/pre_commit_lint.sh
@@ -12,7 +12,8 @@
 #   - litellm/ Python  -> `make lint` (test-linting.yml's lint job)
 #   - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step)
 #                         + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests)
-#   - tests/ Python, ruff-tests.toml, test-quality-budget.json, scripts/check_test_quality.py
+#   - tests/ Python, ruff-tests.toml, test-quality-budget.json, scripts/check_test_quality.py,
+#     scripts/test_quality_gate.py
 #                      -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's
 #                         test-tree ruff and test-quality budget steps)
 #   - dashboard        -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint)
@@ -91,14 +92,11 @@ existing_files() {
 
 litellm_py_pattern='^litellm/.*\.py$'
 e2e_py_pattern='^tests/e2e/.*\.py$'
-test_tree_pattern='^(tests/.*\.py|ruff-tests\.toml|test-quality-budget\.json|scripts/check_test_quality\.py)$'
+test_tree_pattern='^(tests/.*\.py|ruff-tests\.toml|test-quality-budget\.json|scripts/(check_test_quality|test_quality_gate)\.py)$'
 spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$'
 ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$'
 ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$'
 
-# CI's lint job (test-linting.yml) is make lint: litellm/ checks plus two test-tree
-# steps (ruff over ruff-tests.toml and the test-quality budget). Trigger the slow
-# make lint on litellm/ files only; without them, the test-tree steps run on their own below.
 litellm_py_files=$(scope_match "$litellm_py_pattern")
 e2e_py_files=$(scope_match "$e2e_py_pattern")
 test_tree_files=$(scope_match "$test_tree_pattern")
diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py
index e17abf2fa9f..b84cb8aa657 100644
--- a/tests/test_litellm/test_pre_commit_lint.py
+++ b/tests/test_litellm/test_pre_commit_lint.py
@@ -472,7 +472,13 @@ def test_tests_only_change_runs_the_whole_test_tree_ruff_and_the_quality_gate(tm
 
 @pytest.mark.parametrize(
     "changed",
-    ["ruff-tests.toml", "test-quality-budget.json", "scripts/check_test_quality.py", "tests/e2e/test_x.py"],
+    [
+        "ruff-tests.toml",
+        "test-quality-budget.json",
+        "scripts/check_test_quality.py",
+        "scripts/test_quality_gate.py",
+        "tests/e2e/test_x.py",
+    ],
 )
 def test_test_tree_lint_inputs_trigger_the_test_tree_checks(tmp_path: Path, changed: str) -> None:
     repo, bin_dir = _sandbox(tmp_path)

From 2d2b5dabf27405fd413dfbfd71aef5a304775a75 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 01:44:16 -0700
Subject: [PATCH 061/164] fix(test-quality-gate): tear the base worktree down
 on SIGTERM and SIGHUP

---
 scripts/test_quality_gate.py                 | 22 ++++--
 tests/test_litellm/test_test_quality_gate.py | 77 ++++++++++++++++++++
 2 files changed, 92 insertions(+), 7 deletions(-)

diff --git a/scripts/test_quality_gate.py b/scripts/test_quality_gate.py
index 7d34b194f1c..4f79c9488b1 100644
--- a/scripts/test_quality_gate.py
+++ b/scripts/test_quality_gate.py
@@ -34,13 +34,14 @@ import argparse
 import json
 import re
 import shutil
+import signal
 import subprocess
 import sys
 import tempfile
 from collections import Counter
 from collections.abc import Mapping, Sequence
 from pathlib import Path
-from types import MappingProxyType
+from types import FrameType, MappingProxyType
 from typing import Final, NamedTuple
 
 REPO_ROOT: Final = Path(__file__).resolve().parent.parent
@@ -48,6 +49,7 @@ CHECKER: Final = REPO_ROOT / "scripts" / "check_test_quality.py"
 BUDGET_PATH: Final = REPO_ROOT / "test-quality-budget.json"
 TARGET: Final = "tests"
 DEFAULT_BASE: Final = "origin/litellm_internal_staging"
+TERMINATION_SIGNALS: Final = (signal.SIGTERM, signal.SIGHUP)
 
 _HUNK: Final = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@", re.MULTILINE)
 _FILE_HEADER: Final = re.compile(r"^\+\+\+ b/(.+)$", re.MULTILINE)
@@ -116,22 +118,28 @@ def count_by_rule(violations: Sequence[Violation]) -> Mapping[str, int]:
     return MappingProxyType(dict(Counter(v.code for v in violations)))
 
 
-def base_counts(ref: str) -> Mapping[str, int]:
+def _exit_on_termination(signum: int, _frame: FrameType | None) -> None:
+    raise SystemExit(128 + signum)
+
+
+def base_counts(ref: str, repo_root: Path = REPO_ROOT, checker: Path = CHECKER) -> Mapping[str, int]:
     """Rule counts at `ref`, measured with the *current* rule logic rather than
     whatever the checker looked like at that commit."""
+    for termination in TERMINATION_SIGNALS:
+        signal.signal(termination, _exit_on_termination)
     parent: Final = Path(tempfile.mkdtemp(prefix="tq_base_"))
     worktree: Final = parent / "wt"
     try:
-        _run(["git", "worktree", "add", "--detach", str(worktree), ref])
+        _run(["git", "worktree", "add", "--detach", str(worktree), ref], cwd=repo_root)
         (worktree / "scripts").mkdir(parents=True, exist_ok=True)
-        checker: Final = worktree / "scripts" / "check_test_quality.py"
-        shutil.copy(CHECKER, checker)
-        return count_by_rule(_check(worktree, checker))
+        base_checker: Final = worktree / "scripts" / "check_test_quality.py"
+        shutil.copy(checker, base_checker)
+        return count_by_rule(_check(worktree, base_checker))
     finally:
         # Teardown must never raise, or it masks the real error when the body failed.
         subprocess.run(
             ["git", "worktree", "remove", "--force", str(worktree)],
-            cwd=REPO_ROOT, capture_output=True, text=True,
+            cwd=repo_root, capture_output=True, text=True,
         )
         shutil.rmtree(parent, ignore_errors=True)
 
diff --git a/tests/test_litellm/test_test_quality_gate.py b/tests/test_litellm/test_test_quality_gate.py
index 8cce6bc735a..0664153f671 100644
--- a/tests/test_litellm/test_test_quality_gate.py
+++ b/tests/test_litellm/test_test_quality_gate.py
@@ -9,7 +9,13 @@ file:line.
 """
 
 import importlib.util
+import os
+import signal
+import subprocess
 import sys
+import time
+from collections.abc import Callable
+from contextlib import suppress
 from pathlib import Path
 
 _REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -23,6 +29,15 @@ _spec.loader.exec_module(gate)
 
 _BUDGET = {"TQ001": {"limit": 10}, "TQ003": {"limit": 5}}
 
+_SCAN_BASE = (
+    "import importlib.util, pathlib, sys\n"
+    "spec = importlib.util.spec_from_file_location('test_quality_gate', sys.argv[1])\n"
+    "gate = importlib.util.module_from_spec(spec)\n"
+    "sys.modules[spec.name] = gate\n"
+    "spec.loader.exec_module(gate)\n"
+    "gate.base_counts('HEAD', repo_root=pathlib.Path(sys.argv[2]), checker=pathlib.Path(sys.argv[3]))\n"
+)
+
 
 def test_a_rule_within_its_limit_is_not_a_breach():
     assert gate.evaluate({"TQ001": 10}, {"TQ001": 10}, _BUDGET) == ()
@@ -146,3 +161,65 @@ def test_the_shipped_budget_covers_every_rule_the_checker_can_emit():
     budget = json.loads((_REPO_ROOT / "test-quality-budget.json").read_text())
     assert set(budget) == {"TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ008"}
     assert all(spec["limit"] >= 0 for spec in budget.values())
+
+
+def _git(cwd: Path, *args: str) -> str:
+    proc = subprocess.run(["git", *args], cwd=cwd, capture_output=True, text=True)
+    assert proc.returncode == 0, proc.stderr
+    return proc.stdout.strip()
+
+
+def _committed_repo(tmp_path: Path) -> Path:
+    repo = tmp_path / "repo"
+    (repo / "tests").mkdir(parents=True)
+    (repo / "tests" / "test_seed.py").write_text("def test_seed():\n    assert True\n")
+    _git(repo, "init", "-q", "-b", "main")
+    _git(repo, "config", "user.email", "gate@example.com")
+    _git(repo, "config", "user.name", "gate")
+    _git(repo, "config", "commit.gpgsign", "false")
+    _git(repo, "add", "-A")
+    _git(repo, "commit", "-q", "-m", "seed")
+    return repo
+
+
+def _wait_until(predicate: Callable[[], bool], timeout_seconds: float) -> bool:
+    deadline = time.monotonic() + timeout_seconds
+    while time.monotonic() < deadline:
+        if predicate():
+            return True
+        time.sleep(0.05)
+    return predicate()
+
+
+def _reap(process: subprocess.Popen[bytes]) -> None:
+    with suppress(subprocess.TimeoutExpired):
+        process.wait(timeout=10)
+    if process.poll() is None:
+        process.kill()
+        process.wait(timeout=10)
+
+
+def _registered_worktrees(repo: Path) -> int:
+    listing = _git(repo, "worktree", "list", "--porcelain")
+    return sum(line.startswith("worktree ") for line in listing.splitlines())
+
+
+def test_a_terminated_base_scan_still_removes_its_worktree(tmp_path: Path) -> None:
+    repo = _committed_repo(tmp_path)
+    scanning = tmp_path / "scanning"
+    slow_checker = tmp_path / "slow_checker.py"
+    slow_checker.write_text(f"import pathlib, time\npathlib.Path({str(scanning)!r}).touch()\ntime.sleep(30)\n")
+    temp_dir = tmp_path / "tmp"
+    temp_dir.mkdir()
+    scan = subprocess.Popen(
+        [sys.executable, "-c", _SCAN_BASE, str(_MODULE_PATH), str(repo), str(slow_checker)],
+        env={**os.environ, "TMPDIR": str(temp_dir)},
+    )
+    try:
+        assert _wait_until(scanning.exists, 30), "the base scan never reached the checker"
+        scan.send_signal(signal.SIGTERM)
+        assert scan.wait(timeout=30) == 128 + signal.SIGTERM
+    finally:
+        _reap(scan)
+    assert _registered_worktrees(repo) == 1
+    assert list(temp_dir.iterdir()) == []

From 952f082e3ee80eb6bec0855a752ee3c101652744 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 02:07:14 -0700
Subject: [PATCH 062/164] fix(test-quality-gate): keep a termination signal the
 parent already ignores ignored

The SIGTERM/SIGHUP teardown handlers were installed unconditionally, so a base scan started
under nohup (SIGHUP inherited as SIG_IGN) would start dying on hangups it was told to ignore.
Install them only where the disposition is still the default, and cover the ignored case with a
regression test that hangs up a scan started with SIGHUP ignored and expects it to finish.
---
 scripts/test_quality_gate.py                 |  9 +++-
 tests/test_litellm/test_test_quality_gate.py | 56 ++++++++++++++++----
 2 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/scripts/test_quality_gate.py b/scripts/test_quality_gate.py
index 4f79c9488b1..c3316915b5e 100644
--- a/scripts/test_quality_gate.py
+++ b/scripts/test_quality_gate.py
@@ -122,11 +122,16 @@ def _exit_on_termination(signum: int, _frame: FrameType | None) -> None:
     raise SystemExit(128 + signum)
 
 
+def _install_termination_handlers() -> None:
+    for termination in TERMINATION_SIGNALS:
+        if signal.getsignal(termination) == signal.SIG_DFL:
+            signal.signal(termination, _exit_on_termination)
+
+
 def base_counts(ref: str, repo_root: Path = REPO_ROOT, checker: Path = CHECKER) -> Mapping[str, int]:
     """Rule counts at `ref`, measured with the *current* rule logic rather than
     whatever the checker looked like at that commit."""
-    for termination in TERMINATION_SIGNALS:
-        signal.signal(termination, _exit_on_termination)
+    _install_termination_handlers()
     parent: Final = Path(tempfile.mkdtemp(prefix="tq_base_"))
     worktree: Final = parent / "wt"
     try:
diff --git a/tests/test_litellm/test_test_quality_gate.py b/tests/test_litellm/test_test_quality_gate.py
index 0664153f671..b3511eabbfb 100644
--- a/tests/test_litellm/test_test_quality_gate.py
+++ b/tests/test_litellm/test_test_quality_gate.py
@@ -17,6 +17,7 @@ import time
 from collections.abc import Callable
 from contextlib import suppress
 from pathlib import Path
+from typing import NamedTuple
 
 _REPO_ROOT = Path(__file__).resolve().parents[2]
 _MODULE_PATH = _REPO_ROOT / "scripts" / "test_quality_gate.py"
@@ -37,6 +38,7 @@ _SCAN_BASE = (
     "spec.loader.exec_module(gate)\n"
     "gate.base_counts('HEAD', repo_root=pathlib.Path(sys.argv[2]), checker=pathlib.Path(sys.argv[3]))\n"
 )
+_SCAN_BASE_WITH_SIGHUP_IGNORED = "import signal\nsignal.signal(signal.SIGHUP, signal.SIG_IGN)\n" + _SCAN_BASE
 
 
 def test_a_rule_within_its_limit_is_not_a_breach():
@@ -204,22 +206,56 @@ def _registered_worktrees(repo: Path) -> int:
     return sum(line.startswith("worktree ") for line in listing.splitlines())
 
 
-def test_a_terminated_base_scan_still_removes_its_worktree(tmp_path: Path) -> None:
+class _StalledScan(NamedTuple):
+    process: subprocess.Popen[bytes]
+    repo: Path
+    release: Path
+    temp_dir: Path
+
+
+def _base_scan_stalled_in_its_checker(tmp_path: Path, driver: str) -> _StalledScan:
     repo = _committed_repo(tmp_path)
     scanning = tmp_path / "scanning"
+    release = tmp_path / "release"
     slow_checker = tmp_path / "slow_checker.py"
-    slow_checker.write_text(f"import pathlib, time\npathlib.Path({str(scanning)!r}).touch()\ntime.sleep(30)\n")
+    slow_checker.write_text(
+        "import pathlib, time\n"
+        f"pathlib.Path({str(scanning)!r}).touch()\n"
+        f"while not pathlib.Path({str(release)!r}).exists():\n"
+        "    time.sleep(0.05)\n"
+    )
     temp_dir = tmp_path / "tmp"
     temp_dir.mkdir()
     scan = subprocess.Popen(
-        [sys.executable, "-c", _SCAN_BASE, str(_MODULE_PATH), str(repo), str(slow_checker)],
+        [sys.executable, "-c", driver, str(_MODULE_PATH), str(repo), str(slow_checker)],
         env={**os.environ, "TMPDIR": str(temp_dir)},
     )
-    try:
-        assert _wait_until(scanning.exists, 30), "the base scan never reached the checker"
-        scan.send_signal(signal.SIGTERM)
-        assert scan.wait(timeout=30) == 128 + signal.SIGTERM
-    finally:
+    if not _wait_until(scanning.exists, 30):
         _reap(scan)
-    assert _registered_worktrees(repo) == 1
-    assert list(temp_dir.iterdir()) == []
+        raise AssertionError("the base scan never reached the checker")
+    return _StalledScan(scan, repo, release, temp_dir)
+
+
+def test_a_terminated_base_scan_still_removes_its_worktree(tmp_path: Path) -> None:
+    stalled = _base_scan_stalled_in_its_checker(tmp_path, _SCAN_BASE)
+    try:
+        stalled.process.send_signal(signal.SIGTERM)
+        assert stalled.process.wait(timeout=30) == 128 + signal.SIGTERM
+    finally:
+        _reap(stalled.process)
+    assert _registered_worktrees(stalled.repo) == 1
+    assert list(stalled.temp_dir.iterdir()) == []
+
+
+def test_a_base_scan_keeps_ignoring_the_hangup_its_parent_ignored(tmp_path: Path) -> None:
+    stalled = _base_scan_stalled_in_its_checker(tmp_path, _SCAN_BASE_WITH_SIGHUP_IGNORED)
+    try:
+        stalled.process.send_signal(signal.SIGHUP)
+        time.sleep(1)
+        assert stalled.process.poll() is None, "a hangup the parent ignored killed the scan"
+        stalled.release.touch()
+        assert stalled.process.wait(timeout=30) == 0
+    finally:
+        _reap(stalled.process)
+    assert _registered_worktrees(stalled.repo) == 1
+    assert list(stalled.temp_dir.iterdir()) == []

From d9929379003a2b5ea2d6c584fb9c1088a7e6aab7 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 03:17:11 -0700
Subject: [PATCH 063/164] fix(responses): read reasoning support from the cost
 map instead of model-name rules

---
 .../llms/openai/responses/transformation.py   |  7 -------
 ...odel_prices_and_context_window_backup.json |  4 ++++
 model_prices_and_context_window.json          |  4 ++++
 .../test_openai_responses_transformation.py   |  2 ++
 .../test_litellm/test_model_prices_schema.py  | 21 +++++++++++++++++++
 5 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py
index cdfbbb5be5c..123e9a1dda4 100644
--- a/litellm/llms/openai/responses/transformation.py
+++ b/litellm/llms/openai/responses/transformation.py
@@ -117,14 +117,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
 
         return OpenAIGPT5Config.effort_resolves_to_none(model, effort)
 
-    @staticmethod
-    def _is_o_series_name(model: str) -> bool:
-        base: Final = model.split("/")[-1]
-        return len(base) > 1 and base[0] == "o" and base[1].isdigit()
-
     def _supports_reasoning_param(self, model: str) -> bool:
-        if self._is_gpt_5_model(model=model) or self._is_o_series_name(model=model):
-            return True
         base: Final = model.split("/")[-1]
         if base not in litellm.open_ai_chat_completion_models:
             return True
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 2459ed940e0..99e728a30bc 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -37115,6 +37115,7 @@
         "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,
@@ -37155,6 +37156,7 @@
         "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,
@@ -37366,6 +37368,7 @@
         "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,
@@ -37406,6 +37409,7 @@
         "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,
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 2459ed940e0..99e728a30bc 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -37115,6 +37115,7 @@
         "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,
@@ -37155,6 +37156,7 @@
         "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,
@@ -37366,6 +37368,7 @@
         "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,
@@ -37406,6 +37409,7 @@
         "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,
diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
index bf382d4d8ce..cc884fd7dc1 100644
--- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
+++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
@@ -2034,9 +2034,11 @@ class TestReasoningFollowsModelSupport:
             ("gpt-4o", False),
             ("gpt-4.1", False),
             ("gpt-4o-mini", False),
+            ("gpt-5-search-api", False),
             ("gpt-5.6", True),
             ("o3", True),
             ("o3-deep-research", True),
+            ("o4-mini-deep-research", True),
             ("codex-mini-latest", True),
             ("computer-use-preview", True),
         ],
diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py
index 6114d1d8aba..79609032fbd 100644
--- a/tests/test_litellm/test_model_prices_schema.py
+++ b/tests/test_litellm/test_model_prices_schema.py
@@ -173,3 +173,24 @@ def test_dated_variants_carry_base_alias_service_tier_pricing(prices: dict):
         "sync the tier keys so service-tier requests against pinned snapshots are not "
         "billed at standard rates:\n" + "\n".join(drifted)
     )
+
+
+def is_openai_o_series(name: str) -> bool:
+    base = name.split("/")[-1]
+    return len(base) > 1 and base[0] == "o" and base[1].isdigit()
+
+
+def test_openai_o_series_entries_carry_supports_reasoning(prices: dict):
+    unflagged = [
+        name
+        for name, entry in prices.items()
+        if isinstance(entry, dict)
+        and entry.get("litellm_provider") == "openai"
+        and is_openai_o_series(name)
+        and entry.get("supports_reasoning") is not True
+    ]
+    assert unflagged == [], (
+        "OpenAI o-series models are reasoning models, and the Responses API drops the "
+        "`reasoning` param for any mapped OpenAI model whose entry lacks supports_reasoning; "
+        "flag these entries:\n" + "\n".join(unflagged)
+    )

From 292f926fcd5b4882e3cda1e28c6fa9211a7860d4 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 15:56:17 -0700
Subject: [PATCH 064/164] fix(oci): stream Cohere tool-calling answers once

OCI Cohere restates the whole assistant text on the chunk that carries
the tool calls and again on the terminal chunk that carries chatHistory.
Only the terminal restatement was dropped, so a tool-calling turn streamed
the text twice. Treat a toolCalls-bearing chunk as a restatement too and
drop its text once deltas were already emitted.

Resolves LIT-6819
---
 litellm/llms/oci/chat/cohere.py               | 52 ++++-------------
 .../oci/chat/test_oci_cohere_tool_calls.py    | 57 +++++++++++++++++++
 2 files changed, 69 insertions(+), 40 deletions(-)

diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py
index 6e9bb83b0a0..1b494ebad47 100644
--- a/litellm/llms/oci/chat/cohere.py
+++ b/litellm/llms/oci/chat/cohere.py
@@ -290,19 +290,14 @@ def handle_cohere_stream_chunk(
 ) -> ModelResponseStream:
     """Parse a single Cohere SSE chunk into a LiteLLM ModelResponseStream.
 
-    ``prior_tool_calls_emitted`` lets the caller signal whether tool calls
-    were already emitted in earlier chunks of the same stream. When set, the
-    terminal consolidation chunk's tool calls are suppressed (they would
-    duplicate prior deltas); otherwise they are passed through so a stream
-    that delivers tool calls only on the terminal chunk doesn't silently
-    drop them.
-
-    ``prior_text_emitted`` plays the analogous role for the ``text`` field:
-    when set, the terminal consolidation chunk's ``text`` is suppressed
-    (it would re-emit the full assembled response on top of prior deltas);
-    when unset (e.g. a degenerate stream that delivers the entire response
-    in a single SSE event carrying both ``chatHistory`` and ``finishReason``),
-    the text is passed through so the response content isn't silently lost.
+    OCI Cohere streams the answer as single-token ``text`` deltas, then restates
+    the whole assembled ``text`` on every chunk that carries ``toolCalls`` or
+    ``chatHistory`` (the tool-calls event and the terminal event). Once the
+    caller reports that earlier chunks already emitted text
+    (``prior_text_emitted``), those restatements are dropped so the client does
+    not see the answer twice; a stream whose only text lives on such a chunk
+    keeps it. ``prior_tool_calls_emitted`` plays the same role for the tool
+    calls the terminal ``chatHistory`` chunk repeats.
     """
     try:
         typed_chunk: Final = CohereStreamChunk.model_validate(dict_chunk)
@@ -315,33 +310,10 @@ def handle_cohere_stream_chunk(
     if typed_chunk.index is None:
         typed_chunk.index = 0
 
-    # OCI Cohere's terminal SSE event re-sends the full assembled response in
-    # `text` alongside a populated `chatHistory` and a non-null `finishReason`.
-    # Emitting that text would concatenate the whole response onto the
-    # already-streamed deltas. We require both signals to be present so that a
-    # future API change which adds `chatHistory` to intermediate chunks (or a
-    # rare early-populated case) doesn't silently drop legitimate token deltas.
-    is_terminal_consolidation: Final = typed_chunk.chatHistory is not None and typed_chunk.finishReason is not None
-    # On non-terminal text-free chunks (e.g. tool-call-only or keep-alive
-    # chunks) emit ``content=None`` rather than ``content=""`` so downstream
-    # stream-mergers that distinguish "no text in this delta" from "an
-    # explicitly empty text delta" behave correctly.
-    #
-    # We only suppress the terminal chunk's ``text`` when the caller has
-    # confirmed that text deltas were already emitted earlier — otherwise
-    # (e.g. a degenerate stream that delivers the whole response in a
-    # single SSE event), passing it through is the only chance to surface it.
-    text: Final[str | None] = None if (is_terminal_consolidation and prior_text_emitted) else typed_chunk.text
-
-    # Tool calls on the terminal consolidation chunk (whether from
-    # `typed_chunk.toolCalls` or from `chatHistory`) typically restate what
-    # was already streamed in intermediate chunks. Re-emitting them would
-    # mint fresh `uuid4` IDs and cause downstream consumers to execute each
-    # tool call twice. We only suppress when the caller has confirmed that
-    # tool calls were already emitted earlier — otherwise (e.g. a short
-    # response that delivers tool calls exclusively on the terminal chunk),
-    # passing them through is the only chance to surface them.
-    cohere_tool_calls = None if (is_terminal_consolidation and prior_tool_calls_emitted) else typed_chunk.toolCalls
+    restates_text: Final = typed_chunk.chatHistory is not None or typed_chunk.toolCalls is not None
+    restates_tool_calls: Final = typed_chunk.chatHistory is not None
+    text: Final[str | None] = None if (restates_text and prior_text_emitted) else typed_chunk.text
+    cohere_tool_calls: Final = None if (restates_tool_calls and prior_tool_calls_emitted) else typed_chunk.toolCalls
 
     tool_calls: list[dict[str, object]] | None = None
     if cohere_tool_calls:
diff --git a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py
index 5dd44d72d68..97b2fd88cb9 100644
--- a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py
+++ b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py
@@ -445,6 +445,63 @@ class TestOCICohereToolCalls:
         assert result.choices[0].index == 0
         assert result.choices[0].finish_reason == "stop"  # COMPLETE is mapped to stop
 
+    _TOOL_TURN_TEXT = "I will use the tool to find out the weather in Paris."
+    _TOOL_TURN_DELTAS = ["I", " will", " use", " the", " tool", " to", " find", " out", " the", " weather", " in", " Paris", "."]
+    _TOOL_TURN_CALLS = [{"name": "get_weather", "parameters": {"city": "Paris"}}]
+    _TOOL_TURN_HISTORY = [
+        {"role": "USER", "message": "Briefly say what you will do, then find out the weather in Paris using the tool."},
+        {"role": "CHATBOT", "message": _TOOL_TURN_TEXT, "toolCalls": _TOOL_TURN_CALLS},
+    ]
+    _TOOL_TURN_TERMINAL_TOGETHER = [
+        {
+            "apiFormat": "COHERE",
+            "text": _TOOL_TURN_TEXT,
+            "chatHistory": _TOOL_TURN_HISTORY,
+            "finishReason": "COMPLETE",
+            "toolCalls": _TOOL_TURN_CALLS,
+        },
+    ]
+    _TOOL_TURN_TERMINAL_SPLIT = [
+        {"apiFormat": "COHERE", "text": _TOOL_TURN_TEXT, "chatHistory": _TOOL_TURN_HISTORY, "toolCalls": _TOOL_TURN_CALLS},
+        {"apiFormat": "COHERE", "finishReason": "COMPLETE"},
+    ]
+
+    @staticmethod
+    def _drain_cohere_stream(events):
+        wrapper = OCIStreamWrapper(completion_stream=MagicMock(), model="cohere.command-a-03-2025", logging_obj=MagicMock())
+        chunks = [wrapper.chunk_creator(f"data: {json.dumps(event)}") for event in events]
+        content = "".join(chunk.choices[0].delta.content or "" for chunk in chunks)
+        tool_calls = [call for chunk in chunks for call in (chunk.choices[0].delta.tool_calls or [])]
+        finish_reasons = [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason]
+        return content, tool_calls, finish_reasons
+
+    @pytest.mark.parametrize("terminal_events", [_TOOL_TURN_TERMINAL_TOGETHER, _TOOL_TURN_TERMINAL_SPLIT])
+    def test_cohere_tool_turn_streams_the_answer_once(self, terminal_events):
+        """OCI restates the whole answer on the tool-calls chunk and again on the terminal chunk;
+        the client must read it exactly once, with one tool call and one finish reason."""
+        deltas = [{"apiFormat": "COHERE", "text": token} for token in self._TOOL_TURN_DELTAS]
+        tool_calls_event = {"apiFormat": "COHERE", "text": self._TOOL_TURN_TEXT, "toolCalls": self._TOOL_TURN_CALLS}
+
+        content, tool_calls, finish_reasons = self._drain_cohere_stream([*deltas, tool_calls_event, *terminal_events])
+
+        assert content == self._TOOL_TURN_TEXT
+        assert [(call["function"]["name"], call["function"]["arguments"]) for call in tool_calls] == [
+            ("get_weather", '{"city": "Paris"}')
+        ]
+        assert finish_reasons == ["stop"]
+
+    def test_cohere_tool_turn_without_preamble_deltas_keeps_the_only_text(self):
+        """When the tool-calls chunk carries the only copy of the text, dropping it would lose the answer."""
+        tool_calls_event = {"apiFormat": "COHERE", "text": self._TOOL_TURN_TEXT, "toolCalls": self._TOOL_TURN_CALLS}
+
+        content, tool_calls, finish_reasons = self._drain_cohere_stream(
+            [tool_calls_event, *self._TOOL_TURN_TERMINAL_TOGETHER]
+        )
+
+        assert content == self._TOOL_TURN_TEXT
+        assert len(tool_calls) == 1
+        assert finish_reasons == ["stop"]
+
     def test_cohere_parameter_mapping_excludes_tool_choice(self):
         """Test that tool_choice is excluded from Cohere parameter mapping"""
         config = OCIChatConfig()

From f62130a479bf8c95fd5b364ee60cc2be0a57cdff Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 16:24:36 -0700
Subject: [PATCH 065/164] fix(responses): floor reasoning support on the
 bundled cost map and resolve fine-tuned ids

A live cost map older than this release, or a proxy whose map fetch lags, could
strip `reasoning` from a model this release knows accepts it. The bundled map is
now the floor: any OpenAI entry it flags as reasoning keeps the param whatever
the live map says. Fine-tuned ids with an empty suffix (`ft:gpt-4o-2024-08-06:org::id`)
now resolve to their base entry instead of failing open, `chat-latest` carries
the flag, and the schema test keeps every codex, deep-research, and chat-latest
entry flagged. The none-effort check goes through a public wrapper so the
responses config stops importing a private helper.
---
 .../litellm_core_utils/get_model_cost_map.py  |  8 ++-
 .../llms/openai/responses/transformation.py   | 57 ++++++++++++++-----
 ...odel_prices_and_context_window_backup.json |  1 +
 litellm/utils.py                              | 12 +++-
 model_prices_and_context_window.json          |  1 +
 .../test_openai_responses_transformation.py   | 29 +++++++++-
 .../test_litellm/test_model_prices_schema.py  | 29 +++++++---
 tests/test_litellm/test_utils.py              |  5 ++
 8 files changed, 116 insertions(+), 26 deletions(-)

diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py
index 9cba5db8ab7..ba8738c8de0 100644
--- a/litellm/litellm_core_utils/get_model_cost_map.py
+++ b/litellm/litellm_core_utils/get_model_cost_map.py
@@ -53,12 +53,14 @@ class GetModelCostMap:
 
     _backup_model_count: int = -1  # -1 = not yet loaded
 
+    @staticmethod
+    def read_local_model_cost_map_text() -> str:
+        return files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text(encoding="utf-8")
+
     @staticmethod
     def load_local_model_cost_map() -> dict:
         """Load the local backup model cost map bundled with the package."""
-        content: Final = json.loads(
-            files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text(encoding="utf-8")
-        )
+        content: Final = json.loads(GetModelCostMap.read_local_model_cost_map_text())
         return content
 
     @classmethod
diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py
index 123e9a1dda4..189cc7fe956 100644
--- a/litellm/llms/openai/responses/transformation.py
+++ b/litellm/llms/openai/responses/transformation.py
@@ -1,15 +1,17 @@
 from collections.abc import Mapping, Sequence
+from functools import lru_cache
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints
 
 import httpx
 from openai.types.responses import ResponseReasoningItem
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel, TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, TypedDict
 
 import litellm
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.core_helpers import process_response_headers
+from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
 from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
     _safe_convert_created_field,
 )
@@ -42,6 +44,30 @@ _PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: Final = frozenset({LlmProviders.
 _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI})
 
 
+class _ReasoningSupportEntry(BaseModel):
+    litellm_provider: str | None = None
+    supports_reasoning: bool | None = None
+
+
+_BUNDLED_COST_MAP: Final = TypeAdapter(dict[str, _ReasoningSupportEntry])
+
+
+@lru_cache(maxsize=1)
+def _bundled_openai_reasoning_models() -> frozenset[str]:
+    """OpenAI models the cost map shipped with this release flags as reasoning models.
+
+    The live map can lag this release (a pinned mirror, or a proxy on newer code than the
+    map it fetches), and a lagging entry must never strip `reasoning` from a model this
+    release knows accepts it.
+    """
+    bundled: Final = _BUNDLED_COST_MAP.validate_json(GetModelCostMap.read_local_model_cost_map_text())
+    return frozenset(
+        name
+        for name, entry in bundled.items()
+        if entry.litellm_provider == LlmProviders.OPENAI.value and entry.supports_reasoning is True
+    )
+
+
 class _DeleteResponseBody(TypedDict):
     """Decoded body of the Responses API delete call."""
 
@@ -95,13 +121,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
     @staticmethod
     def _supports_reasoning_effort_none(model: str) -> bool:
         """Return True if the model supports reasoning.effort='none'."""
-        from litellm.utils import _supports_factory
+        from litellm.utils import supports_none_reasoning_effort
 
-        return _supports_factory(
-            model=model,
-            custom_llm_provider=None,
-            key="supports_none_reasoning_effort",
-        )
+        return supports_none_reasoning_effort(model=model, custom_llm_provider=None)
 
     @staticmethod
     def _effort_resolves_to_none(model: str, effort: str | None) -> bool:
@@ -117,11 +139,17 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
 
         return OpenAIGPT5Config.effort_resolves_to_none(model, effort)
 
-    def _supports_reasoning_param(self, model: str) -> bool:
-        base: Final = model.split("/")[-1]
-        if base not in litellm.open_ai_chat_completion_models:
+    @staticmethod
+    def _supports_reasoning_param(model: str) -> bool:
+        from litellm.utils import _get_model_info_helper
+
+        try:
+            info: Final = _get_model_info_helper(
+                model=model.split("/")[-1], custom_llm_provider=LlmProviders.OPENAI.value
+            )
+        except Exception:
             return True
-        return litellm.supports_reasoning(model=base, custom_llm_provider=self.custom_llm_provider.value)
+        return info["key"] in _bundled_openai_reasoning_models() or info.get("supports_reasoning") is True
 
     @staticmethod
     def _enforce_min_max_output_tokens(max_output_tokens: "int | None") -> "int | None":
@@ -182,7 +210,8 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
             else:
                 raise litellm.UnsupportedParamsError(
                     message=(
-                        f"{model} doesn't support the `reasoning` parameter. "
+                        f"{model} doesn't support the `reasoning` parameter "
+                        "(its model cost map entry lacks `supports_reasoning`). "
                         "To drop unsupported params set `litellm.drop_params = True`"
                     ),
                     status_code=400,
@@ -500,7 +529,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
         processed_headers: Final = process_response_headers(raw_response_headers)
         try:
             response = ResponsesAPIResponse.model_validate(raw_response_json)
-        except Exception:
+        except ValidationError:
             verbose_logger.debug(
                 "Error constructing ResponsesAPIResponse: %s, using model_construct", raw_response_json
             )
@@ -892,7 +921,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
 
         try:
             response = ResponsesAPIResponse.model_validate(raw_response_json)
-        except Exception:
+        except ValidationError:
             verbose_logger.debug(
                 "Error constructing ResponsesAPIResponse: %s, using model_construct", raw_response_json
             )
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 99e728a30bc..c97858b4e29 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -30698,6 +30698,7 @@
         "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,
diff --git a/litellm/utils.py b/litellm/utils.py
index 9d20d32d147..07735cc1f87 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -2805,6 +2805,13 @@ def supports_reasoning(model: str, custom_llm_provider: str | None = None) -> bo
     return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_reasoning")
 
 
+def supports_none_reasoning_effort(model: str, custom_llm_provider: str | None = None) -> bool:
+    """
+    Check if the given model accepts reasoning effort "none" and return a boolean value.
+    """
+    return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_none_reasoning_effort")
+
+
 def supports_native_structured_output(model: str, custom_llm_provider: str | None = None) -> bool:
     """
     Check if the given model supports native structured outputs and return a boolean value.
@@ -5210,13 +5217,16 @@ def _strip_openai_finetune_model_name(model_name: str) -> str:
     input: ft:gpt-3.5-turbo:my-org:custom_suffix:id
     output: ft:gpt-3.5-turbo
 
+    input: ft:gpt-4o-2024-08-06:my-org::id (OpenAI leaves the suffix empty when none was set)
+    output: ft:gpt-4o-2024-08-06
+
     Args:
     model_name (str): The full model name
 
     Returns:
     str: The stripped model name
     """
-    return re.sub(r"(:[^:]+){3}$", "", model_name)
+    return re.sub(r"(:[^:]*){3}$", "", model_name)
 
 
 def _strip_model_name(model: str, custom_llm_provider: str | None) -> str:
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 99e728a30bc..c97858b4e29 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -30698,6 +30698,7 @@
         "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,
diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
index cc884fd7dc1..66271de7d04 100644
--- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
+++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
@@ -2034,12 +2034,14 @@ class TestReasoningFollowsModelSupport:
             ("gpt-4o", False),
             ("gpt-4.1", False),
             ("gpt-4o-mini", False),
-            ("gpt-5-search-api", False),
+            ("ft:gpt-4o-2024-08-06:my-org::abc123", False),
+            ("chat-latest", True),
             ("gpt-5.6", True),
             ("o3", True),
             ("o3-deep-research", True),
             ("o4-mini-deep-research", True),
             ("codex-mini-latest", True),
+            ("ft:o4-mini-2025-04-16:my-org::abc123", True),
             ("computer-use-preview", True),
         ],
     )
@@ -2051,6 +2053,30 @@ class TestReasoningFollowsModelSupport:
         )
         assert ("reasoning" in mapped) is reasoning_survives
 
+    @pytest.mark.parametrize(
+        "model, reasoning_survives",
+        [
+            ("gpt-4o", False),
+            ("chat-latest", True),
+            ("o3", True),
+            ("o3-deep-research", True),
+        ],
+    )
+    def test_a_cost_map_older_than_this_release_never_strips_a_known_reasoning_model(
+        self, local_model_cost_map, monkeypatch, model, reasoning_survives
+    ):
+        lagging = {
+            name: {field: value for field, value in entry.items() if field != "supports_reasoning"}
+            for name, entry in litellm.model_cost.items()
+        }
+        monkeypatch.setattr(litellm, "model_cost", lagging)
+        mapped = OpenAIResponsesAPIConfig().map_openai_params(
+            response_api_optional_params={"reasoning": {"effort": "medium"}},
+            model=model,
+            drop_params=True,
+        )
+        assert ("reasoning" in mapped) is reasoning_survives
+
     def test_without_drop_params_the_error_is_litellms_400(self, local_model_cost_map, monkeypatch):
         monkeypatch.setattr(litellm, "drop_params", False)
         with pytest.raises(litellm.UnsupportedParamsError) as excinfo:
@@ -2060,6 +2086,7 @@ class TestReasoningFollowsModelSupport:
                 drop_params=False,
             )
         assert excinfo.value.status_code == 400
+        assert "cost map" in str(excinfo.value)
 
     def test_azure_deployments_keep_reasoning(self, local_model_cost_map):
         mapped = AzureOpenAIResponsesAPIConfig().map_openai_params(
diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py
index 79609032fbd..3f0992275eb 100644
--- a/tests/test_litellm/test_model_prices_schema.py
+++ b/tests/test_litellm/test_model_prices_schema.py
@@ -8,6 +8,8 @@ from pathlib import Path
 import jsonschema
 import pytest
 
+from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name
+
 REPO_ROOT = Path(__file__).parents[2]
 GENERATOR_PATH = REPO_ROOT / "ci_cd" / "generate_model_prices_schema.py"
 PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json"
@@ -175,22 +177,35 @@ def test_dated_variants_carry_base_alias_service_tier_pricing(prices: dict):
     )
 
 
+OPENAI_REASONING_FAMILY_MARKERS = ("codex", "deep-research", "chat-latest")
+
+
 def is_openai_o_series(name: str) -> bool:
-    base = name.split("/")[-1]
-    return len(base) > 1 and base[0] == "o" and base[1].isdigit()
+    return len(name) > 1 and name[0] == "o" and name[1].isdigit()
 
 
-def test_openai_o_series_entries_carry_supports_reasoning(prices: dict):
+def is_openai_reasoning_family(name: str) -> bool:
+    base = name.split("/")[-1].removeprefix("ft:")
+    if "search-api" in base:
+        return False
+    return (
+        is_openai_o_series(base)
+        or is_gpt_reasoning_series_name(base)
+        or any(marker in base for marker in OPENAI_REASONING_FAMILY_MARKERS)
+    )
+
+
+def test_openai_reasoning_family_entries_carry_supports_reasoning(prices: dict):
     unflagged = [
         name
         for name, entry in prices.items()
         if isinstance(entry, dict)
         and entry.get("litellm_provider") == "openai"
-        and is_openai_o_series(name)
+        and is_openai_reasoning_family(name)
         and entry.get("supports_reasoning") is not True
     ]
     assert unflagged == [], (
-        "OpenAI o-series models are reasoning models, and the Responses API drops the "
-        "`reasoning` param for any mapped OpenAI model whose entry lacks supports_reasoning; "
-        "flag these entries:\n" + "\n".join(unflagged)
+        "OpenAI o-series, gpt-5+, codex, deep-research, and chat-latest models are reasoning "
+        "models, and the Responses API drops the `reasoning` param for any mapped OpenAI model "
+        "whose entry lacks supports_reasoning; flag these entries:\n" + "\n".join(unflagged)
     )
diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py
index 14907e17b1b..6a90593bea9 100644
--- a/tests/test_litellm/test_utils.py
+++ b/tests/test_litellm/test_utils.py
@@ -198,6 +198,11 @@ def test_get_model_info_resolves_provider_prefixed_model_ids(local_model_cost_ma
     assert via_provider["mode"] == "responses"
 
 
+def test_get_model_info_strips_openai_finetune_ids_without_a_custom_suffix(local_model_cost_map):
+    info = litellm.get_model_info(model="ft:gpt-4o-2024-08-06:my-org::abc123", custom_llm_provider="openai")
+    assert info["key"] == "ft:gpt-4o-2024-08-06"
+
+
 def test_provider_prefixed_lookup_never_outranks_an_existing_row(local_model_cost_map):
     """The provider-prefixed candidate is tried last, after every candidate that
     already existed, so no model that resolves today can change answer. `perplexity/sonar`

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 066/164] 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 43bb55d84934762a775bef3c32f6ed5ad8bc5ce6 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 16:38:23 -0700
Subject: [PATCH 067/164] style(oci): wrap the Cohere tool-turn test fixtures
 to 120 columns

---
 .../llms/oci/chat/test_oci_cohere_tool_calls.py   | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py
index 97b2fd88cb9..729a2d25f41 100644
--- a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py
+++ b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py
@@ -446,7 +446,9 @@ class TestOCICohereToolCalls:
         assert result.choices[0].finish_reason == "stop"  # COMPLETE is mapped to stop
 
     _TOOL_TURN_TEXT = "I will use the tool to find out the weather in Paris."
-    _TOOL_TURN_DELTAS = ["I", " will", " use", " the", " tool", " to", " find", " out", " the", " weather", " in", " Paris", "."]
+    _TOOL_TURN_DELTAS = [
+        "I", " will", " use", " the", " tool", " to", " find", " out", " the", " weather", " in", " Paris", ".",
+    ]
     _TOOL_TURN_CALLS = [{"name": "get_weather", "parameters": {"city": "Paris"}}]
     _TOOL_TURN_HISTORY = [
         {"role": "USER", "message": "Briefly say what you will do, then find out the weather in Paris using the tool."},
@@ -462,13 +464,20 @@ class TestOCICohereToolCalls:
         },
     ]
     _TOOL_TURN_TERMINAL_SPLIT = [
-        {"apiFormat": "COHERE", "text": _TOOL_TURN_TEXT, "chatHistory": _TOOL_TURN_HISTORY, "toolCalls": _TOOL_TURN_CALLS},
+        {
+            "apiFormat": "COHERE",
+            "text": _TOOL_TURN_TEXT,
+            "chatHistory": _TOOL_TURN_HISTORY,
+            "toolCalls": _TOOL_TURN_CALLS,
+        },
         {"apiFormat": "COHERE", "finishReason": "COMPLETE"},
     ]
 
     @staticmethod
     def _drain_cohere_stream(events):
-        wrapper = OCIStreamWrapper(completion_stream=MagicMock(), model="cohere.command-a-03-2025", logging_obj=MagicMock())
+        wrapper = OCIStreamWrapper(
+            completion_stream=MagicMock(), model="cohere.command-a-03-2025", logging_obj=MagicMock()
+        )
         chunks = [wrapper.chunk_creator(f"data: {json.dumps(event)}") for event in events]
         content = "".join(chunk.choices[0].delta.content or "" for chunk in chunks)
         tool_calls = [call for chunk in chunks for call in (chunk.choices[0].delta.tool_calls or [])]

From 1975a54b0418d37f89aba149730e7b6b94729629 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 16:49:00 -0700
Subject: [PATCH 068/164] fix: declare medium as the only reasoning effort
 chat-latest accepts

OpenAI rejects every reasoning.effort on chat-latest except medium. With supports_reasoning set and no declared levels the entry resolved to None, so /model_group/info and the dashboard effort pickers had nothing to narrow the offered levels with
---
 litellm/model_prices_and_context_window_backup.json | 3 +++
 model_prices_and_context_window.json                | 3 +++
 tests/test_litellm/test_model_prices_schema.py      | 8 ++++++++
 3 files changed, 14 insertions(+)

diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index c97858b4e29..1ce5d08ccbc 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -30681,6 +30681,9 @@
         "max_tokens": 128000,
         "mode": "chat",
         "output_cost_per_token": 3e-05,
+        "reasoning_effort_levels": [
+            "medium"
+        ],
         "source": "https://developers.openai.com/api/docs/models/chat-latest",
         "supported_endpoints": [
             "/v1/chat/completions",
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index c97858b4e29..1ce5d08ccbc 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -30681,6 +30681,9 @@
         "max_tokens": 128000,
         "mode": "chat",
         "output_cost_per_token": 3e-05,
+        "reasoning_effort_levels": [
+            "medium"
+        ],
         "source": "https://developers.openai.com/api/docs/models/chat-latest",
         "supported_endpoints": [
             "/v1/chat/completions",
diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py
index 3f0992275eb..c2c22c25998 100644
--- a/tests/test_litellm/test_model_prices_schema.py
+++ b/tests/test_litellm/test_model_prices_schema.py
@@ -9,6 +9,7 @@ import jsonschema
 import pytest
 
 from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name
+from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts
 
 REPO_ROOT = Path(__file__).parents[2]
 GENERATOR_PATH = REPO_ROOT / "ci_cd" / "generate_model_prices_schema.py"
@@ -209,3 +210,10 @@ def test_openai_reasoning_family_entries_carry_supports_reasoning(prices: dict):
         "models, and the Responses API drops the `reasoning` param for any mapped OpenAI model "
         "whose entry lacks supports_reasoning; flag these entries:\n" + "\n".join(unflagged)
     )
+
+
+def test_chat_latest_declares_the_one_effort_openai_accepts(prices: dict):
+    """OpenAI rejects every reasoning.effort on chat-latest except medium, and a reasoning entry
+    with no declared levels resolves to None, which lets /model_group/info and the dashboard offer
+    levels the upstream will 400 on."""
+    assert resolve_supported_reasoning_efforts(prices["chat-latest"], deployment_is_mapped=True) == ("medium",)

From 347ea8f7ca08b2c5ce8fde4430f867b47ace4489 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 16:54:50 -0700
Subject: [PATCH 069/164] fix: keep provider id and metadata on Responses API
 bridged chat completions

---
 .../transformation.py                         |  23 +++-
 ...responses_transformation_transformation.py | 106 ++++++++++++++++++
 2 files changed, 126 insertions(+), 3 deletions(-)

diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py
index d70f947469a..331cd081fcd 100644
--- a/litellm/completion_extras/litellm_responses_transformation/transformation.py
+++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py
@@ -33,7 +33,7 @@ from litellm.responses.sse_output_recovery import (
     record_output_item_chunk,
     record_output_text_chunk,
 )
-from litellm.responses.utils import normalize_responses_api_stream_options
+from litellm.responses.utils import ResponsesAPIRequestUtils, normalize_responses_api_stream_options
 from litellm.types.llms.openai import (
     REASONING_EFFORT,
     ChatCompletionAnnotation,
@@ -54,7 +54,7 @@ if TYPE_CHECKING:
     )
     from pydantic import BaseModel
 
-    from litellm import LiteLLMLoggingObj, ModelResponse
+    from litellm import LiteLLMLoggingObj
     from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
     from litellm.types.llms.openai import (
         ALL_RESPONSES_API_TOOL_PARAMS,
@@ -69,6 +69,15 @@ if TYPE_CHECKING:
     from litellm.types.utils import Choices
 
 
+_CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage"))
+
+
+def _upstream_response_id(response_id: str | None) -> str | None:
+    if response_id is None:
+        return None
+    return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id)
+
+
 class _ReasoningSummaryText(TypedDict):
     type: str
     text: str
@@ -904,6 +913,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
             ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_response.usage),
         )
 
+        model_response.id = _upstream_response_id(raw_response.id) or raw_response.id
+        provider_extras: Final = raw_response.model_extra.items() if raw_response.model_extra else ()
+        for key, value in provider_extras:
+            if key not in _CHAT_COMPLETION_FIELDS and value is not None:
+                setattr(model_response, key, value)
+
         # Preserve hidden params from the ResponsesAPIResponse, especially the headers
         # which contain important provider information like x-request-id
         raw_response_hidden_params: Final = getattr(raw_response, "_hidden_params", {})
@@ -1359,14 +1374,16 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
         if event_type == "response.created":
             # Initial response creation event
             verbose_logger.debug("Chat provider: response.created -> %s", parsed_chunk)
+            created_response: Final = parsed_chunk.get("response")
             return ModelResponseStream(
+                id=_upstream_response_id(created_response.get("id")) if created_response else None,
                 choices=[
                     StreamingChoices(
                         index=0,
                         delta=Delta(content=""),
                         finish_reason=None,
                     )
-                ]
+                ],
             )
         elif event_type == "response.output_item.added":
             # New output item added
diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
index 6590718878d..97b18382a68 100644
--- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
+++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
@@ -3952,3 +3952,109 @@ def test_convert_chat_completion_messages_to_responses_api_tool_result_with_tool
 
     function_call_output = next(item for item in response if item.get("type") == "function_call_output")
     assert function_call_output["output"] == [{"type": "input_text", "text": "1 tool found"}]
+
+
+def _litellm_encoded_response_id(upstream_id: str) -> str:
+    """The id a Responses API call hands back through LiteLLM: the provider's id wrapped with the
+    deployment it came from, the way ``/v1/responses`` clients see it."""
+    import base64
+
+    tagged = f"litellm:custom_llm_provider:azure;model_id:deployment-1;response_id:{upstream_id}"
+    return "resp_" + base64.b64encode(tagged.encode()).decode()
+
+
+def test_transform_response_keeps_upstream_id_and_provider_extras():
+    """A chat completion bridged through the Responses API must answer with the provider's own
+    response id, decoded out of the deployment-tagged id LiteLLM wraps around it, and every
+    top-level field the provider adds beyond the Responses schema (Azure's ``content_filters``,
+    ``service_tier``), the way the native chat path passes unknown top-level fields through,
+    instead of a locally minted ``chatcmpl-`` id and nothing else."""
+    from unittest.mock import Mock
+
+    from litellm.completion_extras.litellm_responses_transformation.transformation import (
+        LiteLLMResponsesTransformationHandler,
+    )
+    from litellm.types.llms.openai import ResponsesAPIResponse
+    from litellm.types.utils import ModelResponse, Usage
+
+    content_filters = [
+        {"blocked": False, "source_type": "prompt", "content_filter_results": {"hate": {"filtered": False}}}
+    ]
+    raw_response = ResponsesAPIResponse.model_validate(
+        {
+            "id": _litellm_encoded_response_id("resp_azure_123"),
+            "created_at": 1734366691,
+            "object": "response",
+            "model": "gpt-5.6",
+            "status": "completed",
+            "output": [
+                {
+                    "type": "function_call",
+                    "id": "fc_1",
+                    "call_id": "call_1",
+                    "name": "lookup_weather",
+                    "arguments": '{"city": "Seattle"}',
+                    "status": "completed",
+                }
+            ],
+            "parallel_tool_calls": True,
+            "tool_choice": "auto",
+            "tools": [],
+            "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
+            "service_tier": "default",
+            "content_filters": content_filters,
+            "max_tool_calls": None,
+        }
+    )
+    model_response = ModelResponse(
+        id="chatcmpl-local",
+        created=1734366691,
+        model=None,
+        object="chat.completion",
+        choices=[],
+        usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0),
+    )
+
+    result = LiteLLMResponsesTransformationHandler().transform_response(
+        model="gpt-5.6",
+        raw_response=raw_response,
+        model_response=model_response,
+        logging_obj=Mock(),
+        request_data={"model": "gpt-5.6"},
+        messages=[{"role": "user", "content": "What is the weather in Seattle?"}],
+        optional_params={},
+        litellm_params={},
+        encoding=Mock(),
+    )
+    dumped = result.model_dump()
+
+    assert dumped["id"] == "resp_azure_123"
+    assert dumped["object"] == "chat.completion"
+    assert dumped["service_tier"] == "default"
+    assert dumped["content_filters"] == content_filters
+    assert "max_tool_calls" not in dumped, "a null provider field must not appear as a null top-level key"
+    assert "output" not in dumped and "status" not in dumped, (
+        "Responses schema fields must not leak into the chat response"
+    )
+    assert dumped["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "lookup_weather"
+
+
+def test_streaming_chunks_carry_the_upstream_response_id():
+    """Every streamed chunk of a bridged chat completion must carry the provider's response id
+    from ``response.created`` rather than a locally minted ``chatcmpl-`` id, so a client can
+    correlate the stream with the provider's request the same way the non-streaming path does."""
+    from litellm.completion_extras.litellm_responses_transformation.transformation import (
+        OpenAiResponsesToChatCompletionStreamIterator,
+    )
+
+    iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True)
+    encoded_id = _litellm_encoded_response_id("resp_azure_stream")
+    events = [
+        {"type": "response.created", "response": {"id": encoded_id, "output": []}},
+        {"type": "response.output_text.delta", "delta": "Hel"},
+        {"type": "response.completed", "response": {"id": encoded_id, "output": [{"type": "message"}]}},
+    ]
+
+    ids = [iterator.chunk_parser(event).id for event in events]
+
+    assert ids == ["resp_azure_stream"] * len(events), f"streamed chunks did not carry the upstream id: {ids}"

From 089b4b8ff458681bcbfa9741e85952e3c6aa0ece Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 16:59:24 -0700
Subject: [PATCH 070/164] test: drop docstrings from the bridge id regression
 tests

---
 ..._litellm_responses_transformation_transformation.py | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
index 97b18382a68..1599449382b 100644
--- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
+++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
@@ -3955,8 +3955,6 @@ def test_convert_chat_completion_messages_to_responses_api_tool_result_with_tool
 
 
 def _litellm_encoded_response_id(upstream_id: str) -> str:
-    """The id a Responses API call hands back through LiteLLM: the provider's id wrapped with the
-    deployment it came from, the way ``/v1/responses`` clients see it."""
     import base64
 
     tagged = f"litellm:custom_llm_provider:azure;model_id:deployment-1;response_id:{upstream_id}"
@@ -3964,11 +3962,6 @@ def _litellm_encoded_response_id(upstream_id: str) -> str:
 
 
 def test_transform_response_keeps_upstream_id_and_provider_extras():
-    """A chat completion bridged through the Responses API must answer with the provider's own
-    response id, decoded out of the deployment-tagged id LiteLLM wraps around it, and every
-    top-level field the provider adds beyond the Responses schema (Azure's ``content_filters``,
-    ``service_tier``), the way the native chat path passes unknown top-level fields through,
-    instead of a locally minted ``chatcmpl-`` id and nothing else."""
     from unittest.mock import Mock
 
     from litellm.completion_extras.litellm_responses_transformation.transformation import (
@@ -4040,9 +4033,6 @@ def test_transform_response_keeps_upstream_id_and_provider_extras():
 
 
 def test_streaming_chunks_carry_the_upstream_response_id():
-    """Every streamed chunk of a bridged chat completion must carry the provider's response id
-    from ``response.created`` rather than a locally minted ``chatcmpl-`` id, so a client can
-    correlate the stream with the provider's request the same way the non-streaming path does."""
     from litellm.completion_extras.litellm_responses_transformation.transformation import (
         OpenAiResponsesToChatCompletionStreamIterator,
     )

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 071/164] 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 d748cf40b7c989a344fe1f59335c3b93b2285a14 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 18:53:37 -0700
Subject: [PATCH 072/164] fix(responses): only drop reasoning when it carries
 an effort and let an explicit map flag win

---
 .../llms/openai/responses/transformation.py   | 16 +++++++--
 .../test_openai_responses_transformation.py   | 34 +++++++++++++++++--
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py
index 189cc7fe956..926de3e8854 100644
--- a/litellm/llms/openai/responses/transformation.py
+++ b/litellm/llms/openai/responses/transformation.py
@@ -149,7 +149,17 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
             )
         except Exception:
             return True
-        return info["key"] in _bundled_openai_reasoning_models() or info.get("supports_reasoning") is True
+        declared: Final = info.get("supports_reasoning")
+        if declared is not None:
+            return declared
+        return info["key"] in _bundled_openai_reasoning_models()
+
+    @staticmethod
+    def _requests_reasoning_effort(reasoning: object) -> bool:
+        effort: Final = (
+            reasoning.get("effort") if isinstance(reasoning, Mapping) else getattr(reasoning, "effort", None)
+        )
+        return effort is not None
 
     @staticmethod
     def _enforce_min_max_output_tokens(max_output_tokens: "int | None") -> "int | None":
@@ -202,7 +212,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
 
         if (
             self.custom_llm_provider == LlmProviders.OPENAI
-            and params.get("reasoning") is not None
+            and self._requests_reasoning_effort(params.get("reasoning"))
             and not self._supports_reasoning_param(model=model)
         ):
             if drop_params or litellm.drop_params:
@@ -210,7 +220,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
             else:
                 raise litellm.UnsupportedParamsError(
                     message=(
-                        f"{model} doesn't support the `reasoning` parameter "
+                        f"{model} doesn't support `reasoning.effort` "
                         "(its model cost map entry lacks `supports_reasoning`). "
                         "To drop unsupported params set `litellm.drop_params = True`"
                     ),
diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
index 66271de7d04..c5902b32a06 100644
--- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
+++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
@@ -2086,12 +2086,42 @@ class TestReasoningFollowsModelSupport:
                 drop_params=False,
             )
         assert excinfo.value.status_code == 400
+        assert "reasoning.effort" in str(excinfo.value)
         assert "cost map" in str(excinfo.value)
 
-    def test_azure_deployments_keep_reasoning(self, local_model_cost_map):
+    @pytest.mark.parametrize("drop_params", [True, False])
+    @pytest.mark.parametrize(
+        "reasoning",
+        [{"summary": "auto"}, {"effort": None, "summary": "auto"}, {}],
+    )
+    def test_reasoning_without_an_effort_passes_through_on_non_reasoning_models(
+        self, local_model_cost_map, monkeypatch, drop_params, reasoning
+    ):
+        monkeypatch.setattr(litellm, "drop_params", drop_params)
+        mapped = OpenAIResponsesAPIConfig().map_openai_params(
+            response_api_optional_params={"reasoning": dict(reasoning)},
+            model="gpt-4o",
+            drop_params=drop_params,
+        )
+        assert mapped["reasoning"] == reasoning
+
+    def test_an_explicit_supports_reasoning_false_beats_the_bundled_floor(self, local_model_cost_map, monkeypatch):
+        overridden = {
+            name: ({**entry, "supports_reasoning": False} if name == "o3" else entry)
+            for name, entry in litellm.model_cost.items()
+        }
+        monkeypatch.setattr(litellm, "model_cost", overridden)
+        mapped = OpenAIResponsesAPIConfig().map_openai_params(
+            response_api_optional_params={"reasoning": {"effort": "medium"}},
+            model="o3",
+            drop_params=True,
+        )
+        assert "reasoning" not in mapped
+
+    def test_azure_deployments_keep_reasoning_even_on_a_non_reasoning_model_name(self, local_model_cost_map):
         mapped = AzureOpenAIResponsesAPIConfig().map_openai_params(
             response_api_optional_params={"reasoning": {"effort": "medium"}},
-            model="my-o3-deployment",
+            model="gpt-4o",
             drop_params=True,
         )
         assert mapped["reasoning"] == {"effort": "medium"}

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 073/164] fix(batches): claim the batch cost spend row in the
 database before charging

The cost callback used to look for an existing `<batch id>_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 074/164] 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 075/164] 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 64cbe6d0aa332bb6295dd37d6eb2066629c3af07 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 19:21:40 -0700
Subject: [PATCH 076/164] fix(bridge): carry provider metadata on streamed
 chats and keep served ids in spend logs

The terminal chunk of a bridged streaming chat now carries the same provider
fields the non-streaming response does (service_tier, content_filters), so
streaming clients see them on the final chunk

The passthrough drops Responses API bookkeeping (background, top_logprobs,
store, ...) by subtracting the OpenAI SDK's Response schema from the fields it
copies instead of a hand-kept denylist

Spend log rows for /v1/messages calls served through the Responses adapter keep
the id the client was handed instead of the decoded upstream id
---
 .../transformation.py                         |  25 ++++-
 litellm/litellm_core_utils/litellm_logging.py |   9 +-
 ...responses_transformation_transformation.py | 105 +++++++++++++++++-
 .../test_litellm_logging.py                   |  33 ++++++
 4 files changed, 162 insertions(+), 10 deletions(-)

diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py
index 331cd081fcd..87350b5479c 100644
--- a/litellm/completion_extras/litellm_responses_transformation/transformation.py
+++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py
@@ -5,8 +5,11 @@ Handler for transforming /chat/completions api requests to litellm.responses req
 import json
 import os
 from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args
 
+from openai.types.chat import ChatCompletion
+from openai.types.responses import Response
 from openai.types.responses.custom_tool_param import CustomToolParam
 from openai.types.responses.response_input_param import (
     FunctionCallOutput,
@@ -43,6 +46,7 @@ from litellm.types.llms.openai import (
     ChatCompletionToolParamFunctionChunk,
     Reasoning,
     ResponsesAPIOptionalRequestParams,
+    ResponsesAPIResponse,
     ResponsesAPIStreamEvents,
 )
 from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
@@ -70,6 +74,19 @@ if TYPE_CHECKING:
 
 
 _CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage"))
+_RESPONSES_API_ONLY_FIELDS: Final = frozenset((*Response.model_fields, *ResponsesAPIResponse.model_fields)) - frozenset(
+    ChatCompletion.model_fields
+)
+
+
+def _provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]:
+    return MappingProxyType(
+        {
+            key: value
+            for key, value in (response_fields.items() if response_fields else ())
+            if value is not None and key not in _CHAT_COMPLETION_FIELDS and key not in _RESPONSES_API_ONLY_FIELDS
+        }
+    )
 
 
 def _upstream_response_id(response_id: str | None) -> str | None:
@@ -914,10 +931,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         )
 
         model_response.id = _upstream_response_id(raw_response.id) or raw_response.id
-        provider_extras: Final = raw_response.model_extra.items() if raw_response.model_extra else ()
-        for key, value in provider_extras:
-            if key not in _CHAT_COMPLETION_FIELDS and value is not None:
-                setattr(model_response, key, value)
+        for key, value in _provider_metadata(raw_response.model_extra).items():
+            setattr(model_response, key, value)
 
         # Preserve hidden params from the ResponsesAPIResponse, especially the headers
         # which contain important provider information like x-request-id
@@ -1551,6 +1566,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                 from litellm.responses.utils import ResponseAPILoggingUtils
 
                 usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response_data.get("usage"))
+            provider_metadata: Final = _provider_metadata(response_data)
             return ModelResponseStream(
                 choices=[
                     StreamingChoices(
@@ -1563,6 +1579,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                     )
                 ],
                 usage=usage,
+                provider_specific_fields=dict(provider_metadata) or None,  # mutable-ok: field is typed dict
             )
         else:
             pass
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index c31c4323157..dbd58766488 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -3918,11 +3918,12 @@ class Logging(LiteLLMLoggingBaseClass):
             LiteLLMResponsesTransformationHandler,
         )
 
+        served_id: Final = _provider_response_id(result)
         try:
-            return LiteLLMResponsesTransformationHandler().transform_response(
+            translated: Final = LiteLLMResponsesTransformationHandler().transform_response(
                 model=self.model,
                 raw_response=result,
-                model_response=litellm.ModelResponse(id=_provider_response_id(result)),
+                model_response=litellm.ModelResponse(id=served_id),
                 logging_obj=self,
                 request_data={},
                 messages=[],
@@ -3930,6 +3931,8 @@ class Logging(LiteLLMLoggingBaseClass):
                 litellm_params={},
                 encoding=litellm.encoding,
             )
+            translated.id = served_id or translated.id
+            return translated
         except Exception as e:
             verbose_logger.debug(
                 "Responses API -> ModelResponse translation failed for "
@@ -3937,7 +3940,7 @@ class Logging(LiteLLMLoggingBaseClass):
                 "usage-only ModelResponse to keep the spend_logs row.",
                 str(e),
             )
-            model_response: Final = litellm.ModelResponse(id=_provider_response_id(result))
+            model_response: Final = litellm.ModelResponse(id=served_id)
             model_response.model = self.model
             usage: Final = getattr(result, "usage", None)
             if usage is not None and ResponseAPILoggingUtils._is_response_api_usage(usage):
diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
index 1599449382b..d4d47b145d1 100644
--- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
+++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
@@ -3955,10 +3955,11 @@ def test_convert_chat_completion_messages_to_responses_api_tool_result_with_tool
 
 
 def _litellm_encoded_response_id(upstream_id: str) -> str:
-    import base64
+    from litellm.responses.utils import ResponsesAPIRequestUtils
 
-    tagged = f"litellm:custom_llm_provider:azure;model_id:deployment-1;response_id:{upstream_id}"
-    return "resp_" + base64.b64encode(tagged.encode()).decode()
+    return ResponsesAPIRequestUtils._build_responses_api_response_id(
+        custom_llm_provider="azure", model_id="deployment-1", response_id=upstream_id
+    )
 
 
 def test_transform_response_keeps_upstream_id_and_provider_extras():
@@ -3997,6 +3998,9 @@ def test_transform_response_keeps_upstream_id_and_provider_extras():
             "service_tier": "default",
             "content_filters": content_filters,
             "max_tool_calls": None,
+            "background": False,
+            "top_logprobs": 0,
+            "store": True,
         }
     )
     model_response = ModelResponse(
@@ -4029,9 +4033,63 @@ def test_transform_response_keeps_upstream_id_and_provider_extras():
     assert "output" not in dumped and "status" not in dumped, (
         "Responses schema fields must not leak into the chat response"
     )
+    assert not {"background", "top_logprobs", "store"} & dumped.keys(), (
+        "Responses API bookkeeping must not ride along as chat metadata"
+    )
     assert dumped["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "lookup_weather"
 
 
+def test_bridged_response_is_priced_by_the_reported_service_tier():
+    from unittest.mock import Mock
+
+    from litellm.completion_extras.litellm_responses_transformation.transformation import (
+        LiteLLMResponsesTransformationHandler,
+    )
+    from litellm.types.llms.openai import ResponsesAPIResponse
+    from litellm.types.utils import ModelResponse
+
+    raw_response = ResponsesAPIResponse.model_validate(
+        {
+            "id": "resp_flex",
+            "created_at": 1734366691,
+            "object": "response",
+            "model": "gpt-5.4",
+            "status": "completed",
+            "output": [
+                {
+                    "type": "message",
+                    "id": "msg_1",
+                    "role": "assistant",
+                    "status": "completed",
+                    "content": [{"type": "output_text", "text": "hi", "annotations": []}],
+                }
+            ],
+            "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100},
+            "service_tier": "flex",
+        }
+    )
+
+    result = LiteLLMResponsesTransformationHandler().transform_response(
+        model="gpt-5.4",
+        raw_response=raw_response,
+        model_response=ModelResponse(),
+        logging_obj=Mock(),
+        request_data={"model": "gpt-5.4"},
+        messages=[{"role": "user", "content": "hi"}],
+        optional_params={},
+        litellm_params={},
+        encoding=Mock(),
+    )
+    pricing = litellm.model_cost["gpt-5.4"]
+    flex_cost = 1000 * pricing["input_cost_per_token_flex"] + 100 * pricing["output_cost_per_token_flex"]
+    standard_cost = 1000 * pricing["input_cost_per_token"] + 100 * pricing["output_cost_per_token"]
+
+    cost = litellm.completion_cost(completion_response=result, custom_llm_provider="openai")
+
+    assert cost == pytest.approx(flex_cost)
+    assert cost < standard_cost
+
+
 def test_streaming_chunks_carry_the_upstream_response_id():
     from litellm.completion_extras.litellm_responses_transformation.transformation import (
         OpenAiResponsesToChatCompletionStreamIterator,
@@ -4048,3 +4106,44 @@ def test_streaming_chunks_carry_the_upstream_response_id():
     ids = [iterator.chunk_parser(event).id for event in events]
 
     assert ids == ["resp_azure_stream"] * len(events), f"streamed chunks did not carry the upstream id: {ids}"
+
+
+def test_streaming_final_chunk_carries_provider_metadata():
+    from unittest.mock import MagicMock
+
+    from litellm.completion_extras.litellm_responses_transformation.transformation import (
+        OpenAiResponsesToChatCompletionStreamIterator,
+    )
+    from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
+
+    iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True)
+    content_filters = [{"blocked": False, "source_type": "completion", "content_filter_results": {}}]
+    events = [
+        {"type": "response.created", "response": {"id": "resp_azure_stream", "output": []}},
+        {"type": "response.output_text.delta", "delta": "Hello"},
+        {
+            "type": "response.completed",
+            "response": {
+                "id": "resp_azure_stream",
+                "output": [{"type": "message"}],
+                "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4},
+                "service_tier": "default",
+                "content_filters": content_filters,
+                "background": False,
+            },
+        },
+    ]
+    stream = CustomStreamWrapper(
+        completion_stream=iter([iterator.chunk_parser(event) for event in events]),
+        model="gpt-5.6",
+        custom_llm_provider="azure",
+        logging_obj=MagicMock(),
+    )
+
+    chunks = [chunk.model_dump() for chunk in stream]
+
+    assert chunks[-1]["choices"][0]["finish_reason"] == "stop"
+    assert chunks[-1]["service_tier"] == "default"
+    assert chunks[-1]["content_filters"] == content_filters
+    assert "background" not in chunks[-1]
+    assert all("service_tier" not in chunk for chunk in chunks[:-1])
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..b5d56f390a9 100644
--- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
@@ -4391,6 +4391,39 @@ def test_handle_anthropic_messages_response_logging_translates_bare_responses_ap
     assert result.usage.total_tokens == 18  # type: ignore[attr-defined]
 
 
+def test_handle_anthropic_messages_response_logging_keeps_the_served_response_id():
+    from openai.types.responses import ResponseOutputMessage, ResponseOutputText
+
+    from litellm.responses.utils import ResponsesAPIRequestUtils
+    from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
+
+    served_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
+        custom_llm_provider="openai", model_id="deployment-1", response_id="resp_upstream"
+    )
+    logging_obj = _anthropic_messages_logging_obj()
+    result = logging_obj._handle_anthropic_messages_response_logging(
+        result=ResponsesAPIResponse(
+            id=served_id,
+            created_at=1700000000,
+            output=[
+                ResponseOutputMessage(
+                    id="msg-1",
+                    type="message",
+                    role="assistant",
+                    status="completed",
+                    content=[ResponseOutputText(annotations=[], text="hi", type="output_text")],
+                )
+            ],
+            usage=ResponseAPIUsage(input_tokens=2, output_tokens=1, total_tokens=3),
+            service_tier="flex",
+        )
+    )
+
+    assert isinstance(result, ModelResponse)
+    assert result.id == served_id, "the spend log row must keep the id the caller was served"
+    assert result.service_tier == "flex"
+
+
 def test_handle_anthropic_messages_response_logging_passes_model_response_through():
     """Anthropic-native path already yields a ModelResponse; it must be returned unchanged."""
     logging_obj = _anthropic_messages_logging_obj()

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 077/164] 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 078/164] fix(spend): let a batch's charge survive an older
 proxy's $0 poll row

A proxy running the old code wrote <batch id>_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 079/164] 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 8da43835a6be47f58a45ce2967c43f50867401c9 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 22:26:09 -0700
Subject: [PATCH 080/164] test(ci): guard against two workflow jobs publishing
 one check-run name

A ruleset's required status check names a check run and GitHub matches it by
that name alone, so two jobs publishing the same name leave the gate unable to
say which job proved it. The new code-quality check reads every workflow,
expands matrix values and local reusable-workflow calls the way Actions does,
and fails when one name has more than one job behind it.
---
 .github/workflows/test-code-quality.yml       |   6 +
 .../check_workflow_job_name_collisions.py     | 179 +++++++++++++++
 .../test_workflow_job_name_collisions.py      | 211 ++++++++++++++++++
 3 files changed, 396 insertions(+)
 create mode 100644 tests/code_coverage_tests/check_workflow_job_name_collisions.py
 create mode 100644 tests/code_coverage_tests/test_workflow_job_name_collisions.py

diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml
index 36fd656c231..e6d2264fbf0 100644
--- a/.github/workflows/test-code-quality.yml
+++ b/.github/workflows/test-code-quality.yml
@@ -74,6 +74,12 @@ jobs:
       - name: check_workflow_startup_safety
         run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_startup_safety.py
 
+      - name: check_workflow_job_name_collisions
+        run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_job_name_collisions.py
+
+      - name: test_workflow_job_name_collisions
+        run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_workflow_job_name_collisions.py
+
       - name: test_e2e_changed_gate
         run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py
 
diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
new file mode 100644
index 00000000000..42e11a3e81d
--- /dev/null
+++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
@@ -0,0 +1,179 @@
+"""Catch two workflow jobs that publish check runs under the same name.
+
+A ruleset's required status check names a check run and GitHub matches it by that
+name alone. When two jobs publish the same name the required context stops
+mapping to the job that proves it: the commit carries two check runs under one
+name and nothing says which one the ruleset required. Both being green hides the
+clash completely, so the context quietly stops meaning what the ruleset intended.
+
+`.github/workflows/auto-close-duplicates.yml` shipped a job id `test` while
+`.github/workflows/test-mcp.yml` already published the required `test` context,
+and commit ed5761daef4ae17152446d182c860630c38b7268 carried both check runs.
+This invariant has to be enforced here because CI cannot enforce it on itself.
+
+A job publishes its `name:` when it sets one and its job id otherwise. Two shapes
+expand that further. A `${{ matrix.key }}` reference publishes one check run per
+value the matrix supplies, so two shard lists that overlap collide even though
+their templates read differently. A job calling a local reusable workflow
+publishes one check run per job of the callee, named `<caller> / <callee>`, which
+is why a caller's name never collides with a plain job that happens to match it.
+A reference nothing resolves stays in the string, so two jobs carrying the same
+unresolved template still compare equal and their collision is still caught.
+"""
+
+import itertools
+import operator
+import re
+import sys
+from collections.abc import Iterator, Mapping, Sequence
+from pathlib import Path
+from typing import Final
+
+import yaml
+from pydantic import BaseModel, Field, ValidationError
+
+REPO_ROOT: Final = Path(__file__).resolve().parent.parent.parent
+WORKFLOWS_DIR: Final = REPO_ROOT / ".github" / "workflows"
+MATRIX_REF: Final = re.compile(r"\$\{\{\s*matrix\.(?P<key>[\w-]+)\s*\}\}")
+LOCAL_CALL_PREFIX: Final = "./"
+
+
+class CheckRunNameCollision(Exception):
+    pass
+
+
+class Job(BaseModel):
+    name: str | None = None
+    uses: str | None = None
+    strategy: Mapping[str, object] = Field(default_factory=dict)
+
+
+class Workflow(BaseModel):
+    jobs: Mapping[str, Job] = Field(default_factory=dict)
+
+
+def parse(text: str) -> tuple[Workflow, object] | None:
+    """The workflow plus its raw `on:` value, or None when the file is not a workflow."""
+    parsed: Final = yaml.safe_load(text)
+    if not isinstance(parsed, dict):
+        return None
+    try:
+        return Workflow.model_validate(parsed), parsed.get(True, parsed.get("on"))
+    except ValidationError:
+        return None
+
+
+def events(raw_on: object) -> frozenset[str]:
+    if isinstance(raw_on, Mapping):
+        return frozenset(str(key) for key in raw_on)
+    if isinstance(raw_on, str):
+        return frozenset({raw_on})
+    if isinstance(raw_on, Sequence):
+        return frozenset(str(event) for event in raw_on)
+    return frozenset()
+
+
+def publishes_check_runs(raw_on: object) -> bool:
+    """A `workflow_call`-only workflow posts its check runs through callers, never itself."""
+    return events(raw_on) != frozenset({"workflow_call"})
+
+
+def matrix_values(job: Job, key: str) -> tuple[str, ...]:
+    matrix: Final = job.strategy.get("matrix")
+    if not isinstance(matrix, Mapping):
+        return ()
+    listed: Final = matrix.get(key)
+    rows: Final = matrix.get("include")
+    from_list: Final = (
+        tuple(str(value) for value in listed if isinstance(value, (str, int, float)))
+        if isinstance(listed, Sequence) and not isinstance(listed, str)
+        else ()
+    )
+    from_rows: Final = (
+        tuple(str(row[key]) for row in rows if isinstance(row, Mapping) and isinstance(row.get(key), (str, int, float)))
+        if isinstance(rows, Sequence) and not isinstance(rows, str)
+        else ()
+    )
+    return tuple(dict.fromkeys(from_list + from_rows))
+
+
+def substituted(template: str, values: Mapping[str, str]) -> str:
+    return MATRIX_REF.sub(lambda ref: values.get(ref.group("key"), ref.group(0)), template)
+
+
+def expand(template: str, job: Job) -> tuple[str, ...]:
+    keys: Final = tuple(dict.fromkeys(ref.group("key") for ref in MATRIX_REF.finditer(template)))
+    candidates: Final = tuple((key, matrix_values(job, key)) for key in keys)
+    resolvable: Final = tuple((key, values) for key, values in candidates if values)
+    if not resolvable:
+        return (template,)
+    return tuple(
+        substituted(template, dict(zip((key for key, _ in resolvable), combination)))
+        for combination in itertools.product(*(values for _, values in resolvable))
+    )
+
+
+def callee_path(job: Job) -> str | None:
+    if job.uses is None or not job.uses.startswith(LOCAL_CALL_PREFIX):
+        return None
+    return job.uses[len(LOCAL_CALL_PREFIX) :].split("@")[0]
+
+
+def job_names(job_id: str, job: Job, workflows: Mapping[str, Workflow]) -> tuple[str, ...]:
+    prefixes: Final = expand(job.name or job_id, job)
+    callee: Final = workflows.get(callee_path(job) or "")
+    if callee is None:
+        return prefixes
+    suffixes: Final = tuple(
+        name
+        for callee_id, callee_job in callee.jobs.items()
+        for name in expand(callee_job.name or callee_id, callee_job)
+    )
+    return tuple(f"{prefix} / {suffix}" for prefix in prefixes for suffix in suffixes)
+
+
+def published(sources: Mapping[str, str]) -> Iterator[tuple[str, str]]:
+    parsed: Final = {rel: entry for rel, text in sources.items() if (entry := parse(text)) is not None}
+    workflows: Final = {rel: workflow for rel, (workflow, _) in parsed.items()}
+    for rel, (workflow, raw_on) in parsed.items():
+        if not publishes_check_runs(raw_on):
+            continue
+        for job_id, job in workflow.jobs.items():
+            for name in job_names(job_id, job, workflows):
+                yield name, f"{rel} job `{job_id}`"
+
+
+def owners_by_name(sources: Mapping[str, str]) -> Iterator[tuple[str, tuple[str, ...]]]:
+    for name, pairs in itertools.groupby(sorted(published(sources)), key=operator.itemgetter(0)):
+        yield name, tuple(owner for _, owner in pairs)
+
+
+def collisions(sources: Mapping[str, str]) -> tuple[str, ...]:
+    return tuple(
+        f"`{name}` is published by {len(owners)} jobs: {', '.join(owners)}. A required status check matching "
+        f"that name cannot say which job proves it; give one of them a distinct `name:` or job id."
+        for name, owners in owners_by_name(sources)
+        if len(owners) > 1
+    )
+
+
+def workflow_sources() -> Mapping[str, str]:
+    """Repo-relative posix paths to text, the keys `uses: ./...` resolves against."""
+    return {path.relative_to(REPO_ROOT).as_posix(): path.read_text() for path in sorted(WORKFLOWS_DIR.glob("*.y*ml"))}
+
+
+def main() -> None:
+    sources: Final = workflow_sources()
+    found: Final = collisions(sources)
+    if found:
+        raise CheckRunNameCollision("Check-run names are not unique:\n  - " + "\n  - ".join(found))
+
+    print(f"Check-run names are unique across {len(sources)} workflows")
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except CheckRunNameCollision as exc:
+        print(f"ERROR: {exc}", file=sys.stderr)
+        sys.exit(1)
diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
new file mode 100644
index 00000000000..35c1480ecd8
--- /dev/null
+++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
@@ -0,0 +1,211 @@
+import sys
+from pathlib import Path
+from typing import Final
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from check_workflow_job_name_collisions import callee_path, collisions, parse, workflow_sources  # noqa: E402
+
+
+def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None:
+    assert collisions(workflow_sources()) == ()
+
+
+def test_every_workflow_in_the_repo_parses_into_jobs() -> None:
+    unparsed: Final = tuple(
+        rel for rel, source in workflow_sources().items() if (entry := parse(source)) is None or not entry[0].jobs
+    )
+
+    assert unparsed == ()
+
+
+def test_every_local_reusable_call_in_the_repo_resolves_to_a_workflow() -> None:
+    sources: Final = workflow_sources()
+    parsed: Final = tuple(entry for text in sources.values() if (entry := parse(text)) is not None)
+    unresolved: Final = tuple(
+        job.uses
+        for workflow, _ in parsed
+        for job in workflow.jobs.values()
+        if (callee := callee_path(job)) is not None and callee not in sources
+    )
+
+    assert unresolved == ()
+
+
+def test_two_jobs_falling_back_to_the_same_job_id_collide() -> None:
+    sources: Final = {
+        "a.yml": "on: pull_request\njobs:\n  test:\n    runs-on: ubuntu-latest\n",
+        "b.yml": "on: pull_request\njobs:\n  test:\n    runs-on: ubuntu-latest\n",
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "`test` is published by 2 jobs" in found[0]
+    assert "a.yml job `test`" in found[0] and "b.yml job `test`" in found[0]
+
+
+def test_an_explicit_name_overrides_the_job_id_and_clears_the_collision() -> None:
+    sources: Final = {
+        "a.yml": "on: pull_request\njobs:\n  test:\n    name: Sweep tests\n    runs-on: ubuntu-latest\n",
+        "b.yml": "on: pull_request\njobs:\n  test:\n    runs-on: ubuntu-latest\n",
+    }
+
+    assert collisions(sources) == ()
+
+
+def test_an_explicit_name_matching_another_job_id_collides() -> None:
+    sources: Final = {
+        "a.yml": "on: pull_request\njobs:\n  sweep:\n    name: test\n    runs-on: ubuntu-latest\n",
+        "b.yml": "on: pull_request\njobs:\n  test:\n    runs-on: ubuntu-latest\n",
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "`test` is published by 2 jobs" in found[0]
+
+
+def test_two_callers_of_one_reusable_workflow_collide_on_a_shared_matrix_value() -> None:
+    base: Final = "on:\n  workflow_call:\njobs:\n  run:\n    name: Run tests\n    runs-on: ubuntu-latest\n"
+    caller: Final = (
+        "on: pull_request\n"
+        "jobs:\n"
+        "  {job}:\n"
+        "    name: ${{{{ matrix.shard }}}}\n"
+        "    uses: ./.github/workflows/base.yml\n"
+        "    strategy:\n"
+        "      matrix:\n"
+        "        include:\n"
+        "          - shard: {shard}\n"
+    )
+    sources: Final = {
+        ".github/workflows/base.yml": base,
+        "unit.yml": caller.format(job="unit", shard="proxy-auth"),
+        "proxy-db.yml": caller.format(job="proxy-db", shard="proxy-auth"),
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "`proxy-auth / Run tests` is published by 2 jobs" in found[0]
+
+
+def test_distinct_matrix_values_through_one_reusable_workflow_do_not_collide() -> None:
+    base: Final = "on:\n  workflow_call:\njobs:\n  run:\n    name: Run tests\n    runs-on: ubuntu-latest\n"
+    caller: Final = (
+        "on: pull_request\n"
+        "jobs:\n"
+        "  {job}:\n"
+        "    name: ${{{{ matrix.shard }}}}\n"
+        "    uses: ./.github/workflows/base.yml\n"
+        "    strategy:\n"
+        "      matrix:\n"
+        "        include:\n"
+        "          - shard: {shard}\n"
+    )
+    sources: Final = {
+        ".github/workflows/base.yml": base,
+        "unit.yml": caller.format(job="unit", shard="proxy-auth"),
+        "proxy-db.yml": caller.format(job="proxy-db", shard="budgets"),
+    }
+
+    assert collisions(sources) == ()
+
+
+def test_a_reusable_caller_does_not_collide_with_a_plain_job_of_the_same_name() -> None:
+    sources: Final = {
+        ".github/workflows/base.yml": (
+            "on:\n  workflow_call:\njobs:\n  run:\n    name: Run tests\n    runs-on: ubuntu-latest\n"
+        ),
+        "unit.yml": (
+            "on: pull_request\n"
+            "jobs:\n"
+            "  unit:\n"
+            "    name: ${{ matrix.shard }}\n"
+            "    uses: ./.github/workflows/base.yml\n"
+            "    strategy:\n"
+            "      matrix:\n"
+            "        include:\n"
+            "          - shard: proxy-behavior\n"
+        ),
+        "postgres.yml": (
+            "on: pull_request\n"
+            "jobs:\n"
+            "  postgres:\n"
+            "    name: ${{ matrix.shard }}\n"
+            "    runs-on: ubuntu-latest\n"
+            "    strategy:\n"
+            "      matrix:\n"
+            "        include:\n"
+            "          - shard: proxy-behavior\n"
+        ),
+    }
+
+    assert collisions(sources) == ()
+
+
+def test_a_workflow_call_only_workflow_publishes_nothing_of_its_own() -> None:
+    sources: Final = {
+        "base.yml": "on:\n  workflow_call:\njobs:\n  run:\n    runs-on: ubuntu-latest\n",
+        "other.yml": "on:\n  workflow_call:\njobs:\n  run:\n    runs-on: ubuntu-latest\n",
+    }
+
+    assert collisions(sources) == ()
+
+
+def test_a_workflow_call_workflow_that_also_runs_on_pull_request_still_publishes() -> None:
+    sources: Final = {
+        "base.yml": "on:\n  workflow_call:\n  pull_request:\njobs:\n  run:\n    runs-on: ubuntu-latest\n",
+        "other.yml": "on: pull_request\njobs:\n  run:\n    runs-on: ubuntu-latest\n",
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "`run` is published by 2 jobs" in found[0]
+
+
+def test_a_matrix_list_supplies_values_the_same_way_include_rows_do() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\n"
+            "jobs:\n"
+            "  build:\n"
+            "    name: Analyze (${{ matrix.language }})\n"
+            "    runs-on: ubuntu-latest\n"
+            "    strategy:\n"
+            "      matrix:\n"
+            "        language: [python, go]\n"
+        ),
+        "b.yml": "on: pull_request\njobs:\n  go:\n    name: Analyze (go)\n    runs-on: ubuntu-latest\n",
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "`Analyze (go)` is published by 2 jobs" in found[0]
+
+
+def test_two_jobs_sharing_one_unresolvable_template_still_collide() -> None:
+    template: Final = (
+        "on: pull_request\njobs:\n  {job}:\n    name: ${{{{ matrix.shard }}}}\n    runs-on: ubuntu-latest\n"
+    )
+    sources: Final = {
+        "a.yml": template.format(job="one"),
+        "b.yml": template.format(job="two"),
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "is published by 2 jobs" in found[0]
+
+
+def test_a_file_that_is_not_a_workflow_is_ignored() -> None:
+    sources: Final = {
+        "notes.yml": "just a string\n",
+        "a.yml": "on: pull_request\njobs:\n  test:\n    runs-on: ubuntu-latest\n",
+    }
+
+    assert collisions(sources) == ()

From 44b7a76439f1896736e61679b4cd0b5e66dd297a Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 22:26:38 -0700
Subject: [PATCH 081/164] fix(ci): stop the auto-close duplicates job colliding
 with the required test check

Its job id was `test`, the same check-run name test-mcp.yml publishes and
guard-internal-staging requires, so a commit touching the sweep carried two
check runs called `test`. Renaming it to `sweep-tests` leaves the required
context with exactly one job behind it.
---
 .github/workflows/auto-close-duplicates.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/auto-close-duplicates.yml b/.github/workflows/auto-close-duplicates.yml
index d8256917805..9a362b4c85c 100644
--- a/.github/workflows/auto-close-duplicates.yml
+++ b/.github/workflows/auto-close-duplicates.yml
@@ -22,7 +22,7 @@ on:
 permissions: {}
 
 jobs:
-  test:
+  sweep-tests:
     if: github.event_name == 'pull_request'
     runs-on: ubuntu-latest
     timeout-minutes: 5

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 082/164] 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 083/164] 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/<model> 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/<model>`` 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/<model>`` 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 084/164] 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 085/164] 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 baee7d8175a1e46085a477f45257dc8b9006b4c1 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 22:47:23 -0700
Subject: [PATCH 086/164] fix: evaluate job name expressions in the check-run
 collision guard

The guard only substituted a bare `${{ matrix.key }}`, so any name built from a
larger expression stayed in the string as its own template. `_test-unit-base.yml`
names its job with a ternary over `format()`, which meant every shard published
an opaque name and 23 of the 33 required contexts, all of them `<shard> / Run
tests`, were invisible to the very check meant to protect them.

Job names are now evaluated per matrix combination over the pieces a name can
hold: string literals, `matrix.<key>`, `format()`, `==` and `!=`, and the
`<cond> && <a> || <b>` idiom. All 33 required contexts now resolve, and nothing
in the repo leaves an expression unresolved. An expression the evaluator does not
understand still falls back to its verbatim template, so two jobs sharing one
stays a collision.

The test also drops its `sys.path.insert`, which the test-quality budget counts
under TQ003; pytest already puts the file's own directory on the path.
---
 .../check_workflow_job_name_collisions.py     | 109 +++++++++++++++---
 .../test_workflow_job_name_collisions.py      |  63 +++++++++-
 2 files changed, 152 insertions(+), 20 deletions(-)

diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
index 42e11a3e81d..849a906c3b8 100644
--- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
@@ -12,13 +12,17 @@ and commit ed5761daef4ae17152446d182c860630c38b7268 carried both check runs.
 This invariant has to be enforced here because CI cannot enforce it on itself.
 
 A job publishes its `name:` when it sets one and its job id otherwise. Two shapes
-expand that further. A `${{ matrix.key }}` reference publishes one check run per
-value the matrix supplies, so two shard lists that overlap collide even though
-their templates read differently. A job calling a local reusable workflow
-publishes one check run per job of the callee, named `<caller> / <callee>`, which
-is why a caller's name never collides with a plain job that happens to match it.
-A reference nothing resolves stays in the string, so two jobs carrying the same
-unresolved template still compare equal and their collision is still caught.
+expand that further. A name carrying `${{ ... }}` publishes one check run per
+matrix combination it reads, so two shard lists that overlap collide even though
+their templates read differently. Each expression is evaluated per combination
+over the pieces a job name can hold: string literals, `matrix.<key>`, `format()`,
+`==` and `!=`, and the `<cond> && <a> || <b>` idiom, which is how the shards
+reach their real `<shard> / Run tests` names rather than staying opaque. A job
+calling a local reusable workflow publishes one check run per job of the callee,
+named `<caller> / <callee>`, which is why a caller's name never collides with a
+plain job that happens to match it. An expression nothing resolves stays in the
+string, so two jobs carrying the same unresolved template still compare equal and
+their collision is still caught.
 """
 
 import itertools
@@ -27,6 +31,7 @@ import re
 import sys
 from collections.abc import Iterator, Mapping, Sequence
 from pathlib import Path
+from types import MappingProxyType
 from typing import Final
 
 import yaml
@@ -34,7 +39,13 @@ from pydantic import BaseModel, Field, ValidationError
 
 REPO_ROOT: Final = Path(__file__).resolve().parent.parent.parent
 WORKFLOWS_DIR: Final = REPO_ROOT / ".github" / "workflows"
-MATRIX_REF: Final = re.compile(r"\$\{\{\s*matrix\.(?P<key>[\w-]+)\s*\}\}")
+EXPRESSION: Final = re.compile(r"\$\{\{(?P<body>.*?)\}\}", re.DOTALL)
+MATRIX_KEY: Final = re.compile(r"matrix\.(?P<key>[\w-]+)")
+MATRIX_REF: Final = re.compile(r"^matrix\.(?P<key>[\w-]+)$")
+LITERAL: Final = re.compile(r"^'(?P<text>[^']*)'$")
+FORMAT_CALL: Final = re.compile(r"^format\((?P<args>.*)\)$", re.DOTALL)
+COMPARISON: Final = re.compile(r"^(?P<left>.+?)\s*(?P<operator>==|!=)\s*(?P<right>.+)$", re.DOTALL)
+NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({})
 LOCAL_CALL_PREFIX: Final = "./"
 
 
@@ -97,19 +108,85 @@ def matrix_values(job: Job, key: str) -> tuple[str, ...]:
     return tuple(dict.fromkeys(from_list + from_rows))
 
 
-def substituted(template: str, values: Mapping[str, str]) -> str:
-    return MATRIX_REF.sub(lambda ref: values.get(ref.group("key"), ref.group(0)), template)
+def scanned(state: tuple[int, bool], char: str) -> tuple[int, bool]:
+    depth, quoted = state
+    if char == "'":
+        return depth, not quoted
+    if quoted:
+        return depth, quoted
+    return depth + int(char == "(") - int(char == ")"), quoted
+
+
+def split_outside(text: str, token: str) -> tuple[str, ...]:
+    """`text` cut on every `token` that sits outside quotes and parentheses."""
+    states: Final = tuple(itertools.accumulate(text, scanned, initial=(0, False)))
+    cuts: Final = tuple(
+        index
+        for index in range(len(text) - len(token) + 1)
+        if text.startswith(token, index) and states[index] == (0, False)
+    )
+    starts: Final = (0, *(cut + len(token) for cut in cuts))
+    return tuple(text[start:end] for start, end in zip(starts, (*cuts, len(text))))
+
+
+def value_of(text: str, values: Mapping[str, str]) -> str | None:
+    expression: Final = text.strip()
+    literal: Final = LITERAL.match(expression)
+    if literal is not None:
+        return literal.group("text")
+    reference: Final = MATRIX_REF.match(expression)
+    if reference is not None:
+        return values.get(reference.group("key"))
+    call: Final = FORMAT_CALL.match(expression)
+    if call is None:
+        return None
+    arguments: Final = tuple(value_of(part, values) for part in split_outside(call.group("args"), ","))
+    resolved: Final = tuple(argument for argument in arguments if argument is not None)
+    if not resolved or len(resolved) != len(arguments):
+        return None
+    return resolved[0].format(*resolved[1:])
+
+
+def holds(condition: str, values: Mapping[str, str]) -> bool | None:
+    comparison: Final = COMPARISON.match(condition.strip())
+    if comparison is None:
+        return None
+    left: Final = value_of(comparison.group("left"), values)
+    right: Final = value_of(comparison.group("right"), values)
+    if left is None or right is None:
+        return None
+    return (left == right) == (comparison.group("operator") == "==")
+
+
+def evaluate(body: str, values: Mapping[str, str]) -> str | None:
+    """The single string this expression yields, or None when its shape is not understood."""
+    branches: Final = tuple(split_outside(alternative, "&&") for alternative in split_outside(body, "||"))
+    outcomes: Final = tuple(tuple(holds(part, values) for part in branch[:-1]) for branch in branches)
+    if any(outcome is None for branch in outcomes for outcome in branch):
+        return None
+    taken: Final = next((branch[-1] for branch, outcome in zip(branches, outcomes) if all(outcome)), None)
+    return None if taken is None else value_of(taken, values)
+
+
+def resolved_span(span: re.Match[str], values: Mapping[str, str]) -> str:
+    substitution: Final = evaluate(span.group("body"), values)
+    return span.group(0) if substitution is None else substitution
+
+
+def rendered(template: str, values: Mapping[str, str]) -> str:
+    return EXPRESSION.sub(lambda span: resolved_span(span, values), template)
 
 
 def expand(template: str, job: Job) -> tuple[str, ...]:
-    keys: Final = tuple(dict.fromkeys(ref.group("key") for ref in MATRIX_REF.finditer(template)))
-    candidates: Final = tuple((key, matrix_values(job, key)) for key in keys)
-    resolvable: Final = tuple((key, values) for key, values in candidates if values)
+    keys: Final = tuple(dict.fromkeys(ref.group("key") for ref in MATRIX_KEY.finditer(template)))
+    resolvable: Final = tuple((key, values) for key in keys if (values := matrix_values(job, key)))
     if not resolvable:
-        return (template,)
+        return (rendered(template, NO_MATRIX),)
     return tuple(
-        substituted(template, dict(zip((key for key, _ in resolvable), combination)))
-        for combination in itertools.product(*(values for _, values in resolvable))
+        dict.fromkeys(
+            rendered(template, dict(zip((key for key, _ in resolvable), combination)))
+            for combination in itertools.product(*(values for _, values in resolvable))
+        )
     )
 
 
diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
index 35c1480ecd8..96820d50040 100644
--- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
@@ -1,10 +1,30 @@
-import sys
-from pathlib import Path
 from typing import Final
 
-sys.path.insert(0, str(Path(__file__).resolve().parent))
+from check_workflow_job_name_collisions import callee_path, collisions, parse, published, workflow_sources
 
-from check_workflow_job_name_collisions import callee_path, collisions, parse, workflow_sources  # noqa: E402
+REUSABLE_BASE: Final = """on:
+  workflow_call:
+jobs:
+  run:
+    name: >-
+      ${{ matrix.python-version == '3.12' && 'Run tests'
+      || format('Run tests (Python {0})', matrix.python-version) }}
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.12", "3.13"]
+"""
+
+SHARD_CALLER: Final = """on: pull_request
+jobs:
+  unit:
+    name: ${{ matrix.shard }}
+    uses: ./.github/workflows/base.yml
+    strategy:
+      matrix:
+        include:
+          - shard: core-utils
+"""
 
 
 def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None:
@@ -209,3 +229,38 @@ def test_a_file_that_is_not_a_workflow_is_ignored() -> None:
     }
 
     assert collisions(sources) == ()
+
+
+def test_a_conditional_name_expands_to_the_branch_each_matrix_value_takes() -> None:
+    names: Final = frozenset(
+        name for name, _ in published({".github/workflows/base.yml": REUSABLE_BASE, "unit.yml": SHARD_CALLER})
+    )
+
+    assert names == frozenset({"core-utils / Run tests", "core-utils / Run tests (Python 3.13)"})
+
+
+def test_a_conditional_name_never_publishes_the_branch_its_condition_rules_out() -> None:
+    names: Final = frozenset(
+        name for name, _ in published({".github/workflows/base.yml": REUSABLE_BASE, "unit.yml": SHARD_CALLER})
+    )
+
+    assert "core-utils / Run tests (Python 3.12)" not in names
+
+
+def test_a_conditional_reusable_name_collides_with_a_plain_job_publishing_the_same_name() -> None:
+    sources: Final = {
+        ".github/workflows/base.yml": REUSABLE_BASE,
+        "unit.yml": SHARD_CALLER,
+        "postgres.yml": ("on: pull_request\njobs:\n  legacy:\n    name: core-utils / Run tests\n"),
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "`core-utils / Run tests` is published by 2 jobs" in found[0]
+
+
+def test_every_workflow_in_the_repo_resolves_every_expression_in_its_job_names() -> None:
+    unresolved: Final = tuple(f"{owner}: {name}" for name, owner in published(workflow_sources()) if "${{" in name)
+
+    assert unresolved == ()

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 087/164] 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 d6a727fe0f52e8d280235d1cc98c43462643d150 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 23:02:32 -0700
Subject: [PATCH 088/164] fix: keep matrix include rows whole when expanding
 job names

The guard read each matrix key's values independently and crossed them, so
a job name reading two keys off one include row published pairs no job ever
runs, which could fail a valid workflow on a required check

It now builds the combinations GitHub builds: the listed keys crossed, each
include row folded into the combinations it overwrites nothing in, and a row
that fits nowhere standing on its own
---
 .../check_workflow_job_name_collisions.py     | 73 +++++++++++++------
 .../test_workflow_job_name_collisions.py      | 46 ++++++++++++
 2 files changed, 95 insertions(+), 24 deletions(-)

diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
index 849a906c3b8..c683de7f08b 100644
--- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
@@ -13,8 +13,9 @@ This invariant has to be enforced here because CI cannot enforce it on itself.
 
 A job publishes its `name:` when it sets one and its job id otherwise. Two shapes
 expand that further. A name carrying `${{ ... }}` publishes one check run per
-matrix combination it reads, so two shard lists that overlap collide even though
-their templates read differently. Each expression is evaluated per combination
+combination the matrix produces, with each `include` row's values held together
+rather than crossed with the other rows', so two shard lists that overlap collide
+even though their templates read differently. Each expression is evaluated per combination
 over the pieces a job name can hold: string literals, `matrix.<key>`, `format()`,
 `==` and `!=`, and the `<cond> && <a> || <b>` idiom, which is how the shards
 reach their real `<shard> / Run tests` names rather than staying opaque. A job
@@ -40,12 +41,13 @@ from pydantic import BaseModel, Field, ValidationError
 REPO_ROOT: Final = Path(__file__).resolve().parent.parent.parent
 WORKFLOWS_DIR: Final = REPO_ROOT / ".github" / "workflows"
 EXPRESSION: Final = re.compile(r"\$\{\{(?P<body>.*?)\}\}", re.DOTALL)
-MATRIX_KEY: Final = re.compile(r"matrix\.(?P<key>[\w-]+)")
 MATRIX_REF: Final = re.compile(r"^matrix\.(?P<key>[\w-]+)$")
 LITERAL: Final = re.compile(r"^'(?P<text>[^']*)'$")
 FORMAT_CALL: Final = re.compile(r"^format\((?P<args>.*)\)$", re.DOTALL)
 COMPARISON: Final = re.compile(r"^(?P<left>.+?)\s*(?P<operator>==|!=)\s*(?P<right>.+)$", re.DOTALL)
 NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({})
+SCALAR: Final = (str, int, float)
+MATRIX_DIRECTIVES: Final = frozenset({"include", "exclude"})
 LOCAL_CALL_PREFIX: Final = "./"
 
 
@@ -89,23 +91,52 @@ def publishes_check_runs(raw_on: object) -> bool:
     return events(raw_on) != frozenset({"workflow_call"})
 
 
-def matrix_values(job: Job, key: str) -> tuple[str, ...]:
+def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, ...]], ...]:
+    return tuple(
+        (str(key), tuple(str(value) for value in values if isinstance(value, SCALAR)))
+        for key, values in matrix.items()
+        if str(key) not in MATRIX_DIRECTIVES and isinstance(values, Sequence) and not isinstance(values, str)
+    )
+
+
+def include_rows(matrix: Mapping[str, object]) -> tuple[Mapping[str, str], ...]:
+    rows: Final = matrix.get("include")
+    if not isinstance(rows, Sequence) or isinstance(rows, str):
+        return ()
+    return tuple(
+        MappingProxyType({str(key): str(value) for key, value in row.items() if isinstance(value, SCALAR)})
+        for row in rows
+        if isinstance(row, Mapping)
+    )
+
+
+def extends(row: Mapping[str, str], combination: Mapping[str, str]) -> bool:
+    """GitHub folds an `include` row into a combination only where it overwrites no listed value."""
+    return all(combination[key] == value for key, value in row.items() if key in combination)
+
+
+def extended(combination: Mapping[str, str], rows: Sequence[Mapping[str, str]]) -> Mapping[str, str]:
+    additions: Final = {key: value for row in rows if extends(row, combination) for key, value in row.items()}
+    return MappingProxyType({**combination, **additions})
+
+
+def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...]:
+    """One mapping per job the matrix produces, each `include` row's values staying together."""
     matrix: Final = job.strategy.get("matrix")
     if not isinstance(matrix, Mapping):
         return ()
-    listed: Final = matrix.get(key)
-    rows: Final = matrix.get("include")
-    from_list: Final = (
-        tuple(str(value) for value in listed if isinstance(value, (str, int, float)))
-        if isinstance(listed, Sequence) and not isinstance(listed, str)
+    listed: Final = listed_values(matrix)
+    rows: Final = include_rows(matrix)
+    crossed: Final = (
+        tuple(
+            MappingProxyType(dict(zip((key for key, _ in listed), values)))
+            for values in itertools.product(*(values for _, values in listed))
+        )
+        if listed
         else ()
     )
-    from_rows: Final = (
-        tuple(str(row[key]) for row in rows if isinstance(row, Mapping) and isinstance(row.get(key), (str, int, float)))
-        if isinstance(rows, Sequence) and not isinstance(rows, str)
-        else ()
-    )
-    return tuple(dict.fromkeys(from_list + from_rows))
+    standalone: Final = tuple(row for row in rows if not any(extends(row, combination) for combination in crossed))
+    return (*(extended(combination, rows) for combination in crossed), *standalone)
 
 
 def scanned(state: tuple[int, bool], char: str) -> tuple[int, bool]:
@@ -178,16 +209,10 @@ def rendered(template: str, values: Mapping[str, str]) -> str:
 
 
 def expand(template: str, job: Job) -> tuple[str, ...]:
-    keys: Final = tuple(dict.fromkeys(ref.group("key") for ref in MATRIX_KEY.finditer(template)))
-    resolvable: Final = tuple((key, values) for key in keys if (values := matrix_values(job, key)))
-    if not resolvable:
+    combinations: Final = matrix_combinations(job)
+    if not combinations:
         return (rendered(template, NO_MATRIX),)
-    return tuple(
-        dict.fromkeys(
-            rendered(template, dict(zip((key for key, _ in resolvable), combination)))
-            for combination in itertools.product(*(values for _, values in resolvable))
-        )
-    )
+    return tuple(dict.fromkeys(rendered(template, values) for values in combinations))
 
 
 def callee_path(job: Job) -> str | None:
diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
index 96820d50040..359072326bd 100644
--- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
@@ -27,6 +27,33 @@ jobs:
 """
 
 
+CORRELATED_ROWS: Final = """on: pull_request
+jobs:
+  unit:
+    name: ${{ matrix.shard }} on ${{ matrix.test-path }}
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        include:
+          - shard: core-utils
+            test-path: tests/core
+          - shard: proxy
+            test-path: tests/proxy
+"""
+
+LISTED_PLUS_ROW: Final = """on: pull_request
+jobs:
+  unit:
+    name: ${{ matrix.python-version }} ${{ matrix.label }}
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.12", "3.13"]
+        include:
+          - label: fast
+"""
+
+
 def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None:
     assert collisions(workflow_sources()) == ()
 
@@ -260,6 +287,25 @@ def test_a_conditional_reusable_name_collides_with_a_plain_job_publishing_the_sa
     assert "`core-utils / Run tests` is published by 2 jobs" in found[0]
 
 
+def test_a_name_reading_two_matrix_keys_publishes_only_the_pairs_each_include_row_holds() -> None:
+    names: Final = frozenset(name for name, _ in published({"unit.yml": CORRELATED_ROWS}))
+
+    assert names == frozenset({"core-utils on tests/core", "proxy on tests/proxy"})
+
+
+def test_a_name_reading_two_matrix_keys_never_publishes_a_pair_no_include_row_holds() -> None:
+    names: Final = frozenset(name for name, _ in published({"unit.yml": CORRELATED_ROWS}))
+
+    assert "core-utils on tests/proxy" not in names
+    assert "proxy on tests/core" not in names
+
+
+def test_an_include_row_carrying_no_listed_key_extends_every_listed_combination() -> None:
+    names: Final = frozenset(name for name, _ in published({"unit.yml": LISTED_PLUS_ROW}))
+
+    assert names == frozenset({"3.12 fast", "3.13 fast"})
+
+
 def test_every_workflow_in_the_repo_resolves_every_expression_in_its_job_names() -> None:
     unresolved: Final = tuple(f"{owner}: {name}" for name, owner in published(workflow_sources()) if "${{" in name)
 

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 089/164] 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 9f379b36b9c484ff045fadfcec213f47857670db Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 23:18:18 -0700
Subject: [PATCH 090/164] refactor: return the collision check's failure
 instead of raising it

The checker raised a custom exception and caught it two lines down in the
same module, which is the throw-then-catch the repo's coding guide rules
out. `main` now prints the same message and returns the exit code, so the
collision list stays a value the whole way out
---
 .../check_workflow_job_name_collisions.py        | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
index c683de7f08b..9dc13c2ea0c 100644
--- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
@@ -51,10 +51,6 @@ MATRIX_DIRECTIVES: Final = frozenset({"include", "exclude"})
 LOCAL_CALL_PREFIX: Final = "./"
 
 
-class CheckRunNameCollision(Exception):
-    pass
-
-
 class Job(BaseModel):
     name: str | None = None
     uses: str | None = None
@@ -264,18 +260,16 @@ def workflow_sources() -> Mapping[str, str]:
     return {path.relative_to(REPO_ROOT).as_posix(): path.read_text() for path in sorted(WORKFLOWS_DIR.glob("*.y*ml"))}
 
 
-def main() -> None:
+def main() -> int:
     sources: Final = workflow_sources()
     found: Final = collisions(sources)
     if found:
-        raise CheckRunNameCollision("Check-run names are not unique:\n  - " + "\n  - ".join(found))
+        print("ERROR: Check-run names are not unique:\n  - " + "\n  - ".join(found), file=sys.stderr)
+        return 1
 
     print(f"Check-run names are unique across {len(sources)} workflows")
+    return 0
 
 
 if __name__ == "__main__":
-    try:
-        main()
-    except CheckRunNameCollision as exc:
-        print(f"ERROR: {exc}", file=sys.stderr)
-        sys.exit(1)
+    sys.exit(main())

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 091/164] 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 092/164] 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 093/164] 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 094/164] 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 095/164] 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 096/164] 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 5bbc83e3defa48db50546aa81ba60a9cdb45ed58 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 6 Sep 2026 00:04:31 -0700
Subject: [PATCH 097/164] fix: model the check-run names GitHub really
 publishes

The collision sweep read a job's name as its `name:` or bare job id, which is
wrong for a matrix job that sets no name: GitHub publishes `build (3.12)`, one
per combination. That missed real duplicates and invented ones that don't exist.
It also crossed every matrix value while ignoring `exclude`, so it checked
combinations no job ever runs.

Four smaller gaps went with it. Boolean matrix values reached a name as `True`
rather than `true`. A `format()` whose arguments cannot fill its placeholders
raised straight out of the script instead of leaving the name unresolved. A job
calling a reusable workflow only ever chained one level, and a call outside the
repo fell back to the caller's own name, which GitHub never posts. A job whose
`name:` was not a string failed validation and silently dropped every job in
that file, so the sweep now renders any scalar and reports a file it cannot read
instead of skipping it.
---
 .../check_workflow_job_name_collisions.py     | 174 +++++++++++-----
 .../test_workflow_job_name_collisions.py      | 196 +++++++++++++++++-
 2 files changed, 319 insertions(+), 51 deletions(-)

diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
index 9dc13c2ea0c..4eef3df1c11 100644
--- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
@@ -11,19 +11,27 @@ clash completely, so the context quietly stops meaning what the ruleset intended
 and commit ed5761daef4ae17152446d182c860630c38b7268 carried both check runs.
 This invariant has to be enforced here because CI cannot enforce it on itself.
 
-A job publishes its `name:` when it sets one and its job id otherwise. Two shapes
-expand that further. A name carrying `${{ ... }}` publishes one check run per
-combination the matrix produces, with each `include` row's values held together
-rather than crossed with the other rows', so two shard lists that overlap collide
-even though their templates read differently. Each expression is evaluated per combination
-over the pieces a job name can hold: string literals, `matrix.<key>`, `format()`,
-`==` and `!=`, and the `<cond> && <a> || <b>` idiom, which is how the shards
-reach their real `<shard> / Run tests` names rather than staying opaque. A job
-calling a local reusable workflow publishes one check run per job of the callee,
-named `<caller> / <callee>`, which is why a caller's name never collides with a
-plain job that happens to match it. An expression nothing resolves stays in the
-string, so two jobs carrying the same unresolved template still compare equal and
-their collision is still caught.
+A job publishes its `name:` when it sets one, and otherwise its job id plus the
+values of the combination it runs, the way GitHub writes `build (3.12)`. A name
+carrying `${{ ... }}` publishes one check run per combination the matrix
+produces: `exclude` rows drop combinations before `include` rows fold into the
+survivors, and each `include` row's values stay together rather than crossing
+with the other rows', so two shard lists that overlap collide even though their
+templates read differently. Each expression is evaluated per combination over the
+pieces a job name can hold: string literals, `matrix.<key>`, `format()`, `==` and
+`!=`, and the `<cond> && <a> || <b>` idiom, which is how the shards reach their
+real `<shard> / Run tests` names rather than staying opaque. An expression
+nothing resolves stays in the string, so two jobs carrying the same unresolved
+template still compare equal and their collision is still caught.
+
+A job calling a local reusable workflow publishes one check run per job of the
+callee, named `<caller> / <callee>` and chained through however many levels of
+local calls it takes, which is why a caller's name never collides with a plain
+job that happens to match it. A call this sweep cannot follow, to a workflow
+outside the repo or to one that is not there, publishes nothing rather than a
+name GitHub never posts. A file under `.github/workflows/` that does not read as
+a workflow at all is reported rather than skipped, since skipping it silently
+would hide every job it holds.
 """
 
 import itertools
@@ -31,6 +39,7 @@ import operator
 import re
 import sys
 from collections.abc import Iterator, Mapping, Sequence
+from dataclasses import dataclass
 from pathlib import Path
 from types import MappingProxyType
 from typing import Final
@@ -46,13 +55,19 @@ LITERAL: Final = re.compile(r"^'(?P<text>[^']*)'$")
 FORMAT_CALL: Final = re.compile(r"^format\((?P<args>.*)\)$", re.DOTALL)
 COMPARISON: Final = re.compile(r"^(?P<left>.+?)\s*(?P<operator>==|!=)\s*(?P<right>.+)$", re.DOTALL)
 NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({})
+NO_CALLERS: Final[frozenset[str]] = frozenset()
 SCALAR: Final = (str, int, float)
 MATRIX_DIRECTIVES: Final = frozenset({"include", "exclude"})
 LOCAL_CALL_PREFIX: Final = "./"
 
 
+@dataclass(frozen=True, slots=True)
+class Unreadable:
+    reason: str
+
+
 class Job(BaseModel):
-    name: str | None = None
+    name: object = None
     uses: str | None = None
     strategy: Mapping[str, object] = Field(default_factory=dict)
 
@@ -61,15 +76,20 @@ class Workflow(BaseModel):
     jobs: Mapping[str, Job] = Field(default_factory=dict)
 
 
-def parse(text: str) -> tuple[Workflow, object] | None:
-    """The workflow plus its raw `on:` value, or None when the file is not a workflow."""
-    parsed: Final = yaml.safe_load(text)
+def scalar_text(value: object) -> str:
+    """A YAML scalar the way GitHub renders it, so `true` never reaches a name as `True`."""
+    return str(value).lower() if isinstance(value, bool) else str(value)
+
+
+def parse(source: str) -> tuple[Workflow, object] | Unreadable:
+    """The workflow plus its raw `on:` value, or why the file does not read as one."""
+    parsed: Final = yaml.safe_load(source)
     if not isinstance(parsed, dict):
-        return None
+        return Unreadable("its top level is not a mapping of workflow keys")
     try:
         return Workflow.model_validate(parsed), parsed.get(True, parsed.get("on"))
-    except ValidationError:
-        return None
+    except ValidationError as error:
+        return Unreadable(f"{error.error_count()} of its job definitions have a shape GitHub would reject")
 
 
 def events(raw_on: object) -> frozenset[str]:
@@ -89,23 +109,28 @@ def publishes_check_runs(raw_on: object) -> bool:
 
 def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, ...]], ...]:
     return tuple(
-        (str(key), tuple(str(value) for value in values if isinstance(value, SCALAR)))
+        (str(key), tuple(scalar_text(value) for value in values if isinstance(value, SCALAR)))
         for key, values in matrix.items()
         if str(key) not in MATRIX_DIRECTIVES and isinstance(values, Sequence) and not isinstance(values, str)
     )
 
 
-def include_rows(matrix: Mapping[str, object]) -> tuple[Mapping[str, str], ...]:
-    rows: Final = matrix.get("include")
+def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...]:
+    rows: Final = matrix.get(directive)
     if not isinstance(rows, Sequence) or isinstance(rows, str):
         return ()
     return tuple(
-        MappingProxyType({str(key): str(value) for key, value in row.items() if isinstance(value, SCALAR)})
+        MappingProxyType({str(key): scalar_text(value) for key, value in row.items() if isinstance(value, SCALAR)})
         for row in rows
         if isinstance(row, Mapping)
     )
 
 
+def drops(row: Mapping[str, str], combination: Mapping[str, str]) -> bool:
+    """GitHub removes a combination that carries every value one `exclude` row names."""
+    return all(combination.get(key) == value for key, value in row.items())
+
+
 def extends(row: Mapping[str, str], combination: Mapping[str, str]) -> bool:
     """GitHub folds an `include` row into a combination only where it overwrites no listed value."""
     return all(combination[key] == value for key, value in row.items() if key in combination)
@@ -116,23 +141,29 @@ def extended(combination: Mapping[str, str], rows: Sequence[Mapping[str, str]])
     return MappingProxyType({**combination, **additions})
 
 
+def crossed_values(listed: Sequence[tuple[str, tuple[str, ...]]]) -> tuple[Mapping[str, str], ...]:
+    if not listed:
+        return ()
+    return tuple(
+        MappingProxyType(dict(zip((key for key, _ in listed), values)))
+        for values in itertools.product(*(values for _, values in listed))
+    )
+
+
 def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...]:
-    """One mapping per job the matrix produces, each `include` row's values staying together."""
+    """One mapping per job the matrix produces, `exclude` applied before `include` as GitHub does."""
     matrix: Final = job.strategy.get("matrix")
     if not isinstance(matrix, Mapping):
         return ()
-    listed: Final = listed_values(matrix)
-    rows: Final = include_rows(matrix)
-    crossed: Final = (
-        tuple(
-            MappingProxyType(dict(zip((key for key, _ in listed), values)))
-            for values in itertools.product(*(values for _, values in listed))
-        )
-        if listed
-        else ()
+    rows: Final = directive_rows(matrix, "include")
+    dropped: Final = directive_rows(matrix, "exclude")
+    kept: Final = tuple(
+        combination
+        for combination in crossed_values(listed_values(matrix))
+        if not any(drops(row, combination) for row in dropped)
     )
-    standalone: Final = tuple(row for row in rows if not any(extends(row, combination) for combination in crossed))
-    return (*(extended(combination, rows) for combination in crossed), *standalone)
+    standalone: Final = tuple(row for row in rows if not any(extends(row, combination) for combination in kept))
+    return (*(extended(combination, rows) for combination in kept), *standalone)
 
 
 def scanned(state: tuple[int, bool], char: str) -> tuple[int, bool]:
@@ -156,6 +187,14 @@ def split_outside(text: str, token: str) -> tuple[str, ...]:
     return tuple(text[start:end] for start, end in zip(starts, (*cuts, len(text))))
 
 
+def formatted(template: str, arguments: Sequence[str]) -> str | None:
+    """A `format()` whose placeholders the arguments cannot fill resolves to nothing, never a crash."""
+    try:
+        return template.format(*arguments)
+    except (IndexError, KeyError, ValueError):
+        return None
+
+
 def value_of(text: str, values: Mapping[str, str]) -> str | None:
     expression: Final = text.strip()
     literal: Final = LITERAL.match(expression)
@@ -171,7 +210,7 @@ def value_of(text: str, values: Mapping[str, str]) -> str | None:
     resolved: Final = tuple(argument for argument in arguments if argument is not None)
     if not resolved or len(resolved) != len(arguments):
         return None
-    return resolved[0].format(*resolved[1:])
+    return formatted(resolved[0], resolved[1:])
 
 
 def holds(condition: str, values: Mapping[str, str]) -> bool | None:
@@ -211,27 +250,58 @@ def expand(template: str, job: Job) -> tuple[str, ...]:
     return tuple(dict.fromkeys(rendered(template, values) for values in combinations))
 
 
+def suffixed(job_id: str, combination: Mapping[str, str]) -> str:
+    """The name GitHub gives a job with no `name:`, its id plus the combination it runs."""
+    return f"{job_id} ({', '.join(combination.values())})" if combination else job_id
+
+
+def published_names(job_id: str, job: Job) -> tuple[str, ...]:
+    if job.name is not None:
+        return expand(scalar_text(job.name), job)
+    suffixes: Final = tuple(dict.fromkeys(suffixed(job_id, values) for values in matrix_combinations(job)))
+    return suffixes or (job_id,)
+
+
 def callee_path(job: Job) -> str | None:
     if job.uses is None or not job.uses.startswith(LOCAL_CALL_PREFIX):
         return None
     return job.uses[len(LOCAL_CALL_PREFIX) :].split("@")[0]
 
 
-def job_names(job_id: str, job: Job, workflows: Mapping[str, Workflow]) -> tuple[str, ...]:
-    prefixes: Final = expand(job.name or job_id, job)
-    callee: Final = workflows.get(callee_path(job) or "")
-    if callee is None:
+def job_names(
+    job_id: str, job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str] = NO_CALLERS
+) -> tuple[str, ...]:
+    prefixes: Final = published_names(job_id, job)
+    if job.uses is None:
         return prefixes
+    path: Final = callee_path(job)
+    callee: Final = workflows.get(path or "")
+    if path is None or callee is None or path in callers:
+        return ()
     suffixes: Final = tuple(
         name
         for callee_id, callee_job in callee.jobs.items()
-        for name in expand(callee_job.name or callee_id, callee_job)
+        for name in job_names(callee_id, callee_job, workflows, callers | {path})
     )
     return tuple(f"{prefix} / {suffix}" for prefix in prefixes for suffix in suffixes)
 
 
+def readable(sources: Mapping[str, str]) -> Mapping[str, tuple[Workflow, object]]:
+    parsed: Final = {rel: parse(source) for rel, source in sources.items()}
+    return MappingProxyType({rel: entry for rel, entry in parsed.items() if not isinstance(entry, Unreadable)})
+
+
+def unreadable(sources: Mapping[str, str]) -> tuple[str, ...]:
+    parsed: Final = {rel: parse(source) for rel, source in sources.items()}
+    return tuple(
+        f"{rel} sits in the workflows directory but {entry.reason}, so none of its jobs were checked."
+        for rel, entry in sorted(parsed.items())
+        if isinstance(entry, Unreadable)
+    )
+
+
 def published(sources: Mapping[str, str]) -> Iterator[tuple[str, str]]:
-    parsed: Final = {rel: entry for rel, text in sources.items() if (entry := parse(text)) is not None}
+    parsed: Final = readable(sources)
     workflows: Final = {rel: workflow for rel, (workflow, _) in parsed.items()}
     for rel, (workflow, raw_on) in parsed.items():
         if not publishes_check_runs(raw_on):
@@ -260,16 +330,26 @@ def workflow_sources() -> Mapping[str, str]:
     return {path.relative_to(REPO_ROOT).as_posix(): path.read_text() for path in sorted(WORKFLOWS_DIR.glob("*.y*ml"))}
 
 
-def main() -> int:
-    sources: Final = workflow_sources()
+def report(header: str, problems: Sequence[str]) -> None:
+    if problems:
+        print(f"ERROR: {header}:\n  - " + "\n  - ".join(problems), file=sys.stderr)
+
+
+def exit_code(sources: Mapping[str, str]) -> int:
+    unread: Final = unreadable(sources)
     found: Final = collisions(sources)
-    if found:
-        print("ERROR: Check-run names are not unique:\n  - " + "\n  - ".join(found), file=sys.stderr)
+    report("Some workflows could not be read", unread)
+    report("Check-run names are not unique", found)
+    if unread or found:
         return 1
 
     print(f"Check-run names are unique across {len(sources)} workflows")
     return 0
 
 
+def main() -> int:
+    return exit_code(workflow_sources())
+
+
 if __name__ == "__main__":
     sys.exit(main())
diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
index 359072326bd..3dd8a3d4088 100644
--- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
@@ -1,6 +1,15 @@
 from typing import Final
 
-from check_workflow_job_name_collisions import callee_path, collisions, parse, published, workflow_sources
+from check_workflow_job_name_collisions import (
+    Unreadable,
+    callee_path,
+    collisions,
+    exit_code,
+    parse,
+    published,
+    unreadable,
+    workflow_sources,
+)
 
 REUSABLE_BASE: Final = """on:
   workflow_call:
@@ -53,6 +62,62 @@ jobs:
           - label: fast
 """
 
+NAMELESS_MATRIX: Final = """on: pull_request
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.12", "3.13"]
+"""
+
+EXCLUDED_PAIR: Final = """on: pull_request
+jobs:
+  unit:
+    name: ${{ matrix.os }}-${{ matrix.python-version }}
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        os: [ubuntu, macos]
+        python-version: ["3.12", "3.13"]
+        exclude:
+          - os: macos
+            python-version: "3.13"
+"""
+
+EXCLUDED_KEY: Final = """on: pull_request
+jobs:
+  unit:
+    name: ${{ matrix.os }}-${{ matrix.python-version }}
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        os: [ubuntu, macos]
+        python-version: ["3.12", "3.13"]
+        exclude:
+          - os: macos
+"""
+
+BOOLEAN_MATRIX: Final = """on: pull_request
+jobs:
+  unit:
+    name: cache ${{ matrix.cached }}
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        cached: [true, false]
+"""
+
+UNFILLABLE_FORMAT: Final = """on: pull_request
+jobs:
+  unit:
+    name: ${{ format('{0} {1}', matrix.shard) }}
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        shard: [core]
+"""
+
 
 def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None:
     assert collisions(workflow_sources()) == ()
@@ -60,7 +125,9 @@ def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None:
 
 def test_every_workflow_in_the_repo_parses_into_jobs() -> None:
     unparsed: Final = tuple(
-        rel for rel, source in workflow_sources().items() if (entry := parse(source)) is None or not entry[0].jobs
+        rel
+        for rel, source in workflow_sources().items()
+        if isinstance(entry := parse(source), Unreadable) or not entry[0].jobs
     )
 
     assert unparsed == ()
@@ -68,7 +135,7 @@ def test_every_workflow_in_the_repo_parses_into_jobs() -> None:
 
 def test_every_local_reusable_call_in_the_repo_resolves_to_a_workflow() -> None:
     sources: Final = workflow_sources()
-    parsed: Final = tuple(entry for text in sources.values() if (entry := parse(text)) is not None)
+    parsed: Final = tuple(entry for text in sources.values() if not isinstance(entry := parse(text), Unreadable))
     unresolved: Final = tuple(
         job.uses
         for workflow, _ in parsed
@@ -249,15 +316,28 @@ def test_two_jobs_sharing_one_unresolvable_template_still_collide() -> None:
     assert "is published by 2 jobs" in found[0]
 
 
-def test_a_file_that_is_not_a_workflow_is_ignored() -> None:
+def test_a_file_that_is_not_a_workflow_is_reported_rather_than_skipped() -> None:
     sources: Final = {
         "notes.yml": "just a string\n",
         "a.yml": "on: pull_request\njobs:\n  test:\n    runs-on: ubuntu-latest\n",
     }
 
+    found: Final = unreadable(sources)
+
+    assert len(found) == 1
+    assert "notes.yml" in found[0]
     assert collisions(sources) == ()
 
 
+def test_a_workflow_holding_a_job_shape_github_would_reject_is_reported() -> None:
+    sources: Final = {"a.yml": "on: pull_request\njobs:\n  test:\n    uses: [not, a, string]\n"}
+
+    found: Final = unreadable(sources)
+
+    assert len(found) == 1
+    assert "a.yml" in found[0]
+
+
 def test_a_conditional_name_expands_to_the_branch_each_matrix_value_takes() -> None:
     names: Final = frozenset(
         name for name, _ in published({".github/workflows/base.yml": REUSABLE_BASE, "unit.yml": SHARD_CALLER})
@@ -310,3 +390,111 @@ def test_every_workflow_in_the_repo_resolves_every_expression_in_its_job_names()
     unresolved: Final = tuple(f"{owner}: {name}" for name, owner in published(workflow_sources()) if "${{" in name)
 
     assert unresolved == ()
+
+
+def test_a_matrix_job_with_no_name_publishes_the_id_and_values_github_appends() -> None:
+    names: Final = frozenset(name for name, _ in published({"a.yml": NAMELESS_MATRIX}))
+
+    assert names == frozenset({"build (3.12)", "build (3.13)"})
+
+
+def test_a_matrix_job_with_no_name_does_not_collide_with_a_plain_job_carrying_its_id() -> None:
+    sources: Final = {
+        "a.yml": NAMELESS_MATRIX,
+        "b.yml": "on: pull_request\njobs:\n  build:\n    runs-on: ubuntu-latest\n",
+    }
+
+    assert collisions(sources) == ()
+
+
+def test_a_matrix_job_with_no_name_collides_with_the_suffixed_name_github_writes() -> None:
+    sources: Final = {
+        "a.yml": NAMELESS_MATRIX,
+        "b.yml": "on: pull_request\njobs:\n  legacy:\n    name: build (3.13)\n    runs-on: ubuntu-latest\n",
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "`build (3.13)` is published by 2 jobs" in found[0]
+
+
+def test_an_excluded_combination_publishes_no_check_run() -> None:
+    names: Final = frozenset(name for name, _ in published({"unit.yml": EXCLUDED_PAIR}))
+
+    assert names == frozenset({"ubuntu-3.12", "ubuntu-3.13", "macos-3.12"})
+
+
+def test_an_exclude_row_naming_one_key_drops_every_combination_carrying_it() -> None:
+    names: Final = frozenset(name for name, _ in published({"unit.yml": EXCLUDED_KEY}))
+
+    assert names == frozenset({"ubuntu-3.12", "ubuntu-3.13"})
+
+
+def test_a_boolean_matrix_value_renders_the_way_github_writes_it() -> None:
+    names: Final = frozenset(name for name, _ in published({"unit.yml": BOOLEAN_MATRIX}))
+
+    assert names == frozenset({"cache true", "cache false"})
+
+
+def test_a_format_call_its_arguments_cannot_fill_leaves_the_name_unresolved() -> None:
+    names: Final = frozenset(name for name, _ in published({"unit.yml": UNFILLABLE_FORMAT}))
+
+    assert names == frozenset({"${{ format('{0} {1}', matrix.shard) }}"})
+
+
+def test_a_call_to_a_workflow_outside_the_repo_publishes_no_name() -> None:
+    sources: Final = {
+        "a.yml": "on: pull_request\njobs:\n  unit:\n    uses: BerriAI/other/.github/workflows/base.yml@main\n",
+        "b.yml": "on: pull_request\njobs:\n  unit:\n    runs-on: ubuntu-latest\n",
+    }
+
+    assert frozenset(name for name, _ in published(sources)) == frozenset({"unit"})
+    assert collisions(sources) == ()
+
+
+def test_a_chain_of_local_reusable_calls_publishes_every_level_of_the_chain() -> None:
+    sources: Final = {
+        ".github/workflows/leaf.yml": (
+            "on:\n  workflow_call:\njobs:\n  run:\n    name: Leaf\n    runs-on: ubuntu-latest\n"
+        ),
+        ".github/workflows/mid.yml": (
+            "on:\n  workflow_call:\njobs:\n  call:\n    name: Mid\n    uses: ./.github/workflows/leaf.yml\n"
+        ),
+        "top.yml": "on: pull_request\njobs:\n  top:\n    name: Top\n    uses: ./.github/workflows/mid.yml\n",
+    }
+
+    names: Final = frozenset(name for name, _ in published(sources))
+
+    assert names == frozenset({"Top / Mid / Leaf"})
+
+
+def test_a_job_name_that_is_not_a_string_still_publishes_the_value_github_renders() -> None:
+    sources: Final = {
+        "a.yml": "on: pull_request\njobs:\n  sweep:\n    name: 2024\n    runs-on: ubuntu-latest\n",
+        "b.yml": 'on: pull_request\njobs:\n  other:\n    name: "2024"\n    runs-on: ubuntu-latest\n',
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "`2024` is published by 2 jobs" in found[0]
+
+
+def test_the_check_fails_when_a_file_in_the_workflows_directory_cannot_be_read() -> None:
+    assert exit_code({"notes.yml": "just a string\n"}) == 1
+
+
+def test_the_check_fails_when_two_jobs_publish_one_check_run_name() -> None:
+    plain: Final = "on: pull_request\njobs:\n  test:\n    runs-on: ubuntu-latest\n"
+
+    assert exit_code({"a.yml": plain, "b.yml": plain}) == 1
+
+
+def test_the_check_passes_when_every_file_reads_and_every_name_is_unique() -> None:
+    sources: Final = {
+        "a.yml": "on: pull_request\njobs:\n  test:\n    runs-on: ubuntu-latest\n",
+        "b.yml": "on: pull_request\njobs:\n  sweep:\n    runs-on: ubuntu-latest\n",
+    }
+
+    assert exit_code(sources) == 0

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 098/164] 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 a08fb489cdb2375b25b19fdca0a0f17a90942574 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Sun, 6 Sep 2026 00:38:07 -0700
Subject: [PATCH 099/164] test(ui): accept trailing slashes in migration
 sidebar links

---
 tests/e2e/ui/tests/migration/migratedPages.spec.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/e2e/ui/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts
index 547330190bd..bc96d9a6c66 100644
--- a/tests/e2e/ui/tests/migration/migratedPages.spec.ts
+++ b/tests/e2e/ui/tests/migration/migratedPages.spec.ts
@@ -29,12 +29,12 @@ async function expectRendered(page: Page) {
 }
 
 /**
- * Click a migrated page's sidebar link. Migrated items render as <a href=".../ui/<segment>">;
+ * Click a migrated page's sidebar link. Migrated items render as <a href=".../ui/<segment>/">;
  * nested ones live under collapsible groups whose children only render while the
  * group is open, so expand collapsed groups until the link is clickable.
  */
 async function clickSidebar(page: Page, segment: string) {
-  const link = sidebar(page).locator(`a[href$="/ui/${segment}"]`).first();
+  const link = sidebar(page).locator(`a[href$="/ui/${segment}"], a[href$="/ui/${segment}/"]`).first();
   const collapsedGroups = sidebar(page).getByRole("button", { expanded: false });
   for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) {
     const stillCollapsed = await collapsedGroups.count();

From 6b493cb61ebe17928ee48ba90686d442687e3674 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Sun, 6 Sep 2026 00:43:13 -0700
Subject: [PATCH 100/164] test(guardrails): pin the presidio spend-log record
 and the UI's masked-entity persistence

The presidio suite proved masking happened by reading the served answer, and the
UI suite proved the wizard's "Select All & Mask" produced a row. Neither checked
the thing an operator actually looks at afterwards: the audit trail.

Adds an e2e test asserting the spend log carries the pre_call guardrail record
for a masked request: status, provider, per-entity masked counts, and the
detected-entity list the dashboard's guardrail panel renders its scores from.
It keys off the x-litellm-applied-guardrails response header rather than masked
text in the answer, because whether the model echoes the prompt back is a model
decision, not a guardrail one. Both a mutation that stops writing the record and
one that empties the entity list turn it red.

Two records land on one log, pre_call and post_call, so the assertion selects on
mode as well as name; picking by name alone could hand it the empty post_call
record depending on write order.

On the UI side, the Presidio wizard test now reads the stored guardrail back and
asserts every persisted entity carries the MASK action. A row appearing in the
table did not prove the entity selection survived the save, so a wizard that
persisted an empty pii_entities_config would have passed.

SpendLogRow gains a typed metadata field. guardrail_response is left as object
because each provider writes its own shape there (presidio a list of entities,
bedrock an assessment object, a failed run the exception string); a union narrow
enough to be useful would fail to parse the others and break every suite that
reads a spend log. The caller validates the shape it expects with a TypeAdapter.

run_e2e.sh now pins PROXY_BASE_URL to the stack's own origin. The proxy builds
its post-login redirect from that variable when it is set, so a value inherited
from a developer's .env sent the browser off the relocated stack and the suite's
login step timed out on every port but 4000.
---
 tests/e2e/coverage_registry/guardrail.yaml    |   1 +
 .../guardrails/test_presidio_masking_e2e.py   | 124 +++++++++++++++++-
 tests/e2e/models.py                           |  22 ++++
 tests/e2e/ui/run_e2e.sh                       |   1 +
 .../ui/tests/guardrails/guardrails.spec.ts    |  12 ++
 5 files changed, 156 insertions(+), 4 deletions(-)

diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml
index c64fd6150af..81832bebf49 100644
--- a/tests/e2e/coverage_registry/guardrail.yaml
+++ b/tests/e2e/coverage_registry/guardrail.yaml
@@ -3,6 +3,7 @@
 - {id: guardrail.presidio.pre_call.masks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "PII masking pre-call; data-leak blast radius"}
 - {id: guardrail.presidio.post_call.masks, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Mask PII in model output"}
 - {id: guardrail.presidio.logging_only.masks, module: guardrail, tier: P0, hook_point: logging_only, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Redact in logs without blocking"}
+- {id: guardrail.presidio.pre_call.logs_masked_entities, module: guardrail, tier: P0, hook_point: pre_call, assertions: [logs_masked_entities], exercised_on: [chat_completions], source: "guardrail_hooks/presidio.py", rationale: "A masking run must record itself on the spend log: the dashboard's guardrail panel renders the masked-entity counts and per-entity scores straight off metadata.guardrail_information, so a run that masks but records nothing leaves an operator unable to audit it"}
 - {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"}
 - {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"}
 - {id: guardrail.litellm_content_filter.pre_call.allows, module: guardrail, tier: P0, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Team disable_global_guardrails bypasses default-on content filter"}
diff --git a/tests/e2e/guardrails/test_presidio_masking_e2e.py b/tests/e2e/guardrails/test_presidio_masking_e2e.py
index c6d87473c21..4edf6c4ceb6 100644
--- a/tests/e2e/guardrails/test_presidio_masking_e2e.py
+++ b/tests/e2e/guardrails/test_presidio_masking_e2e.py
@@ -26,16 +26,22 @@ from __future__ import annotations
 import os
 import time
 from collections.abc import Callable
-from typing import Literal
+from typing import Final, Literal
 
 import pytest
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 
 from e2e_config import unique_marker
-from e2e_http import Result, Success
+from e2e_http import Result, StreamingResponse, Success
 from guardrails_client import GuardrailMode, GuardrailsClient, PiiAction, PiiEntity, PresidioParamsBody
 from lifecycle import ResourceManager
-from models import AnthropicMessagesResponse, ChatResponse
+from models import (
+    AnthropicMessagesResponse,
+    ChatResponse,
+    GuardrailEntityMatch,
+    GuardrailRunRecord,
+    SpendLogRow,
+)
 
 pytestmark = pytest.mark.e2e
 
@@ -276,3 +282,113 @@ class TestPresidioPostCallMasking:
                     f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
                 )
             time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
+
+
+_LOGGED_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"}
+
+_ENTITY_LIST_ADAPTER: Final = TypeAdapter(list[GuardrailEntityMatch])
+
+
+def _applied_guardrails(outcome: StreamingResponse) -> str:
+    return outcome.headers.get("x-litellm-applied-guardrails", "")
+
+
+def _guardrail_records(row: SpendLogRow) -> tuple[GuardrailRunRecord, ...]:
+    metadata = row.metadata
+    if metadata is None:
+        return ()
+    return tuple(metadata.guardrail_information or ())
+
+
+def _poll_until_guardrail_applied(
+    client: GuardrailsClient, key: str, guardrail_name: str, prompt: str
+) -> StreamingResponse:
+    deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
+    last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128)
+    while time.monotonic() < deadline:
+        if last.ok and guardrail_name in _applied_guardrails(last):
+            return last
+        time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
+        last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128)
+    return last
+
+
+class TestPresidioSpendLogRecord:
+    @pytest.mark.covers(
+        "guardrail.presidio.pre_call.logs_masked_entities",
+        exercised_on=["chat_completions"],
+    )
+    def test_masking_run_is_recorded_on_the_spend_log(
+        self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
+    ) -> None:
+        name = f"e2e-presidio-log-{unique_marker()}"
+        _register_presidio(client, resources, name=name, entities=_LOGGED_ENTITIES)
+
+        email = _fake_email()
+        prompt = _pii_prompt(unique_marker(), email)
+
+        outcome = _poll_until_guardrail_applied(client, scoped_key, name, prompt)
+        assert outcome.ok, f"the guarded call must be served, got {outcome.status_code}: {outcome.body[:400]}"
+        assert name in _applied_guardrails(outcome), (
+            "the response must carry x-litellm-applied-guardrails naming the guardrail; without it "
+            f"the 200 only proves the guardrail never attached. Got {_applied_guardrails(outcome)!r}"
+        )
+
+        request_id = ChatResponse.model_validate_json(outcome.body).id
+        assert request_id, f"the served response must carry an id to look the spend log up by: {outcome.body[:400]}"
+
+        rows = client.proxy.poll_logs_for_request_id(
+            request_id,
+            predicate=lambda logged: bool(_guardrail_records(logged[0])),
+        )
+        assert rows, f"no spend log row ever appeared for request {request_id}"
+
+        records = _guardrail_records(rows[0])
+        assert records, (
+            f"the spend log for {request_id} carries no guardrail_information, so the dashboard's "
+            "guardrail panel would render nothing for a request the guardrail demonstrably ran on"
+        )
+
+        record = next(
+            (entry for entry in records if entry.guardrail_name == name and entry.guardrail_mode == "pre_call"),
+            None,
+        )
+        assert record is not None, (
+            "guardrail_information carries no pre_call record for this guardrail, only "
+            f"{[(entry.guardrail_name, entry.guardrail_mode) for entry in records]}"
+        )
+        assert record.guardrail_status == "success", (
+            f"the recorded status must be success for a run that masked and served, got {record.guardrail_status!r}"
+        )
+        assert record.guardrail_provider == "presidio", (
+            f"the record must attribute the run to presidio so the dashboard picks the right "
+            f"renderer, got {record.guardrail_provider!r}"
+        )
+
+        counts = record.masked_entity_count or {}
+        assert _LOGGED_ENTITIES.keys() <= counts.keys(), (
+            f"every entity the guardrail was configured to mask must appear in masked_entity_count, got {counts}"
+        )
+        assert all(counts[entity] >= 1 for entity in _LOGGED_ENTITIES), (
+            f"each masked entity must be counted at least once, got {counts}"
+        )
+
+        assert not isinstance(record.guardrail_response, str), (
+            "guardrail_response must be the detected-entity list, but it is the string "
+            f"{record.guardrail_response!r}. A redaction marker here means the proxy is not storing "
+            "prompts in spend logs, which also strips the per-entity scores from the guardrail panel"
+        )
+        entities = _ENTITY_LIST_ADAPTER.validate_python(record.guardrail_response)
+        assert entities, (
+            "guardrail_response must carry the detected entities; the dashboard's Detected Entities "
+            "list and its per-entity scores are rendered from exactly this array"
+        )
+        assert {entity.entity_type for entity in entities} >= _LOGGED_ENTITIES.keys(), (
+            f"the detected entities must cover what was masked, got "
+            f"{sorted(entity.entity_type for entity in entities)}"
+        )
+        for entity in entities:
+            assert 0.0 < entity.score <= 1.0, f"{entity.entity_type} carries an out-of-range score: {entity.score}"
+            assert 0 <= entity.start < entity.end, (
+                f"{entity.entity_type} carries a degenerate span: {entity.start}-{entity.end}"
+            )
diff --git a/tests/e2e/models.py b/tests/e2e/models.py
index 2c6c0e9bbd4..d5be98045a8 100644
--- a/tests/e2e/models.py
+++ b/tests/e2e/models.py
@@ -573,6 +573,27 @@ class OcrResponse(BaseModel):
 # ---------- spend logs ----------
 
 
+class GuardrailEntityMatch(BaseModel):
+    entity_type: str
+    score: float
+    start: int
+    end: int
+
+
+class GuardrailRunRecord(BaseModel):
+    guardrail_name: str | None = None
+    guardrail_mode: str | None = None
+    guardrail_status: str | None = None
+    guardrail_provider: str | None = None
+    masked_entity_count: dict[str, int] | None = None
+    guardrail_response: object | None = None
+
+
+class SpendLogMetadata(BaseModel):
+    applied_guardrails: list[str] | None = None
+    guardrail_information: list[GuardrailRunRecord] | None = None
+
+
 class SpendLogRow(BaseModel):
     request_id: str | None = None
     api_key: str | None = None
@@ -589,6 +610,7 @@ class SpendLogRow(BaseModel):
     completion_tokens: int | None = None
     total_tokens: int | None = None
     request_tags: list[str] | None = None
+    metadata: SpendLogMetadata | None = None
 
 
 class SpendLogs(RootModel[list[SpendLogRow]]):
diff --git a/tests/e2e/ui/run_e2e.sh b/tests/e2e/ui/run_e2e.sh
index 67e3225f668..a369d8c8f7a 100755
--- a/tests/e2e/ui/run_e2e.sh
+++ b/tests/e2e/ui/run_e2e.sh
@@ -151,6 +151,7 @@ export DISABLE_SCHEMA_UPDATE="true"
 export E2E_UI_BASE_URL="${E2E_UI_BASE_URL:-http://127.0.0.1:${PROXY_PORT}}"
 # Ensure the proxy serves UI at /ui (not behind a subpath)
 export SERVER_ROOT_PATH=""
+export PROXY_BASE_URL="$E2E_UI_BASE_URL"
 # Boot with an external logout URL so proxyLogoutUrl.spec.ts can assert the
 # redirect. This same value is exported to the Playwright process below (the
 # spec's skip guard reads it). Safe for the rest of the suite — nothing else
diff --git a/tests/e2e/ui/tests/guardrails/guardrails.spec.ts b/tests/e2e/ui/tests/guardrails/guardrails.spec.ts
index 1e43c7a2b22..0267eeae909 100644
--- a/tests/e2e/ui/tests/guardrails/guardrails.spec.ts
+++ b/tests/e2e/ui/tests/guardrails/guardrails.spec.ts
@@ -7,6 +7,7 @@ import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traff
 interface StoredGuardrail {
   guardrail_id: string;
   guardrail_name: string | null;
+  litellm_params?: { pii_entities_config?: Record<string, string> | null } | null;
 }
 
 async function listGuardrails(page: PlaywrightPage): Promise<StoredGuardrail[]> {
@@ -249,6 +250,17 @@ test.describe("Guardrails", () => {
     const row = page.getByRole("row").filter({ hasText: guardrailName });
     await expect(row).toHaveCount(1, { timeout: 15_000 });
 
+    const stored = await findGuardrail(page, guardrailName);
+    const piiConfig = stored?.litellm_params?.pii_entities_config ?? {};
+    expect(
+      Object.keys(piiConfig).length,
+      "Select All & Mask persisted no PII entities, so the guardrail would mask nothing",
+    ).toBeGreaterThan(0);
+    expect(
+      Object.entries(piiConfig).filter(([, action]) => action !== "MASK"),
+      "Select All & Mask must persist every entity with the MASK action",
+    ).toEqual([]);
+
     await navigateToPage(page, Page.Teams);
     await dismissFeedbackPopup(page);
     await clickTeamId(page, E2E_TEAM_NO_ADMIN_ID);

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 101/164] 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 b1cc1e8dae77a6ff1d0517bcedde40cddd55a9d4 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Sun, 6 Sep 2026 00:59:18 -0700
Subject: [PATCH 102/164] test(ui): assert navigation through accessible links
 and page content

---
 tests/e2e/ui/fixtures/migratedPages.ts        | 190 +++++++++++++-----
 tests/e2e/ui/helpers/navigation.ts            |  24 +++
 tests/e2e/ui/tests/migration/README.md        |  31 ++-
 .../ui/tests/migration/migratedPages.spec.ts  | 126 +++++-------
 tests/e2e/ui/tests/navigation/sidebar.spec.ts |  59 +++++-
 5 files changed, 282 insertions(+), 148 deletions(-)

diff --git a/tests/e2e/ui/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts
index 58939ca2b9a..bce09b49e10 100644
--- a/tests/e2e/ui/fixtures/migratedPages.ts
+++ b/tests/e2e/ui/fixtures/migratedPages.ts
@@ -1,50 +1,142 @@
-/**
- * Source of truth for the App Router migration E2E suites.
- *
- * Add an entry (legacy sidebar page id -> route segment) once a page's migration
- * has MERGED to the branch under test. Consumers pick it up automatically:
- *   - migration smoke (tests/migration/migratedPages.spec.ts), via MIGRATED_E2E_SEGMENTS:
- *       default mount:           npm run e2e:migration
- *       server-root-path mount:  SERVER_ROOT_PATH=/<root> npm run e2e:migration:root
- *   - navigation specs that assert per-page URLs (tests/navigation/sidebar.spec.ts)
- *
- * Keep this in lockstep with MIGRATED_PAGES in src/utils/migratedPages.ts.
- */
-export const MIGRATED_E2E_PAGES: Record<string, string> = {
-  "api-keys": "api-keys",
-  models: "models-and-endpoints",
-  api_ref: "api-reference",
-  "llm-playground": "playground",
-  projects: "projects",
-  "access-groups": "access-groups",
-  budgets: "budgets",
-  workflows: "workflows",
-  "guardrails-monitor": "guardrails-monitor",
-  "mcp-servers": "mcp-servers",
-  "search-tools": "search-tools",
-  "tag-management": "tag-management",
-  "vector-stores": "vector-stores",
-  memory: "memory",
-  policies: "policies",
-  guardrails: "guardrails",
-  prompts: "prompts",
-  "tool-policies": "tool-policies",
-  skills: "skills",
-  caching: "caching",
-  "cost-tracking": "cost-tracking",
-  "transform-request": "transform-request",
-  "ui-theme": "ui-theme",
-  logs: "logs",
-  "admin-panel": "admin-panel",
-  "logging-and-alerts": "logging-and-alerts",
-  "model-hub-table": "model-hub-table",
-  new_usage: "usage",
-  usage: "old-usage",
-  agents: "agents",
-  "router-settings": "router-settings",
-  users: "users",
-  teams: "teams",
-  organizations: "organizations",
-};
+export type MigratedPage = Readonly<{
+  segment: string;
+  linkName: string | RegExp;
+  group?: string;
+  content: Readonly<{ role: "heading" | "tab" | "button"; name: string }> | Readonly<{ text: string }>;
+  unlicensedText?: string;
+}>;
 
-export const MIGRATED_E2E_SEGMENTS: string[] = [...new Set(Object.values(MIGRATED_E2E_PAGES))];
+export const MIGRATED_E2E_PAGES: Readonly<Record<string, MigratedPage>> = {
+  "api-keys": { segment: "api-keys", linkName: "Virtual Keys", content: { role: "heading", name: "Virtual Keys" } },
+  models: {
+    segment: "models-and-endpoints",
+    linkName: "Models + Endpoints",
+    content: { role: "heading", name: "Model Management" },
+  },
+  api_ref: {
+    segment: "api-reference",
+    linkName: "API Reference",
+    content: { role: "heading", name: "OpenAI Compatible Proxy: API Reference" },
+  },
+  "llm-playground": { segment: "playground", linkName: "Playground", content: { role: "tab", name: "Chat" } },
+  projects: {
+    segment: "projects",
+    linkName: /^Projects(?: Beta)?$/,
+    content: { role: "heading", name: "Projects" },
+  },
+  "access-groups": {
+    segment: "access-groups",
+    linkName: "Access Groups",
+    content: { role: "heading", name: "Access Groups" },
+  },
+  budgets: { segment: "budgets", linkName: "Budgets", content: { role: "heading", name: "Budgets" } },
+  workflows: {
+    segment: "workflows",
+    linkName: "Workflow Runs",
+    group: "Agentic",
+    content: { text: "Workflow Runs" },
+  },
+  "guardrails-monitor": {
+    segment: "guardrails-monitor",
+    linkName: "Guardrails Monitor",
+    content: { role: "heading", name: "Guardrails Monitor" },
+  },
+  "mcp-servers": {
+    segment: "mcp-servers",
+    linkName: "MCP Servers",
+    content: { role: "heading", name: "MCP Servers" },
+  },
+  "search-tools": {
+    segment: "search-tools",
+    linkName: "Search Tools",
+    group: "Tools",
+    content: { role: "heading", name: "Search Tools" },
+  },
+  "tag-management": {
+    segment: "tag-management",
+    linkName: "Tag Management",
+    group: "Experimental",
+    content: { role: "heading", name: "Tag Management" },
+  },
+  "vector-stores": {
+    segment: "vector-stores",
+    linkName: "Vector Stores",
+    group: "Tools",
+    content: { role: "heading", name: "Vector Store Management" },
+  },
+  memory: { segment: "memory", linkName: "Memory", group: "Agentic", content: { role: "heading", name: "Memory" } },
+  policies: { segment: "policies", linkName: "Policies", content: { role: "tab", name: "Policy Simulator" } },
+  guardrails: { segment: "guardrails", linkName: "Guardrails", content: { role: "tab", name: "Guardrails" } },
+  prompts: {
+    segment: "prompts",
+    linkName: "Prompts",
+    group: "Experimental",
+    content: { role: "button", name: "Add New Prompt" },
+  },
+  "tool-policies": {
+    segment: "tool-policies",
+    linkName: "Tool Policies",
+    group: "Tools",
+    content: { role: "heading", name: "Tool Policies" },
+  },
+  skills: { segment: "skills", linkName: "Skills", content: { role: "heading", name: "Skills" } },
+  caching: { segment: "caching", linkName: "Response Cache", content: { role: "tab", name: "Cache Settings" } },
+  "cost-tracking": {
+    segment: "cost-tracking",
+    linkName: "Cost Tracking",
+    group: "Settings",
+    content: { text: "Cost Tracking Settings" },
+  },
+  "transform-request": {
+    segment: "transform-request",
+    linkName: "API Playground",
+    group: "Experimental",
+    content: { role: "heading", name: "Playground" },
+  },
+  "ui-theme": {
+    segment: "ui-theme",
+    linkName: "UI Theme",
+    group: "Settings",
+    content: { role: "heading", name: "UI Theme Customization" },
+  },
+  logs: { segment: "logs", linkName: "Logs", content: { role: "heading", name: "Request Logs" } },
+  "admin-panel": {
+    segment: "admin-panel",
+    linkName: "Admin Settings",
+    group: "Settings",
+    content: { role: "heading", name: "Admin Access" },
+  },
+  "logging-and-alerts": {
+    segment: "logging-and-alerts",
+    linkName: "Logging & Alerts",
+    group: "Settings",
+    content: { role: "tab", name: "Logging Callbacks" },
+  },
+  "model-hub-table": {
+    segment: "model-hub-table",
+    linkName: "AI Hub",
+    content: { role: "heading", name: "AI Hub" },
+  },
+  new_usage: { segment: "usage", linkName: "Usage", content: { role: "heading", name: "Usage View" } },
+  usage: {
+    segment: "old-usage",
+    linkName: "Old Usage",
+    group: "Experimental",
+    content: { role: "tab", name: "All Up" },
+  },
+  agents: { segment: "agents", linkName: "Agents", group: "Agentic", content: { role: "heading", name: "Agents" } },
+  "router-settings": {
+    segment: "router-settings",
+    linkName: "Router Settings",
+    group: "Settings",
+    content: { role: "heading", name: "Routing Settings" },
+  },
+  users: { segment: "users", linkName: "Internal Users", content: { role: "tab", name: "Users" } },
+  teams: { segment: "teams", linkName: "Teams", content: { role: "heading", name: "Teams" } },
+  organizations: {
+    segment: "organizations",
+    linkName: "Organizations",
+    content: { text: "Click on an organization ID to view its details." },
+    unlicensedText: "This is a LiteLLM Enterprise feature, and requires a valid key to use. Get a trial key here.",
+  },
+};
diff --git a/tests/e2e/ui/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts
index a58ece16f9c..4a7c4e7baa9 100644
--- a/tests/e2e/ui/helpers/navigation.ts
+++ b/tests/e2e/ui/helpers/navigation.ts
@@ -1,5 +1,29 @@
 import { Page } from "../fixtures/pages";
 import { Page as PlaywrightPage, expect } from "@playwright/test";
+import { UI_BASE_URL } from "../constants";
+
+export const sidebarLink = (page: PlaywrightPage, name: string | RegExp) =>
+  page.getByRole("complementary").getByRole("link", { name, exact: true });
+
+export async function clickSidebarLink(page: PlaywrightPage, name: string | RegExp, groupName?: string): Promise<void> {
+  const link = sidebarLink(page, name);
+  if (groupName && !(await link.isVisible())) {
+    const group = page.getByRole("complementary").getByRole("button", { name: groupName, exact: true });
+    await expect(group).toBeVisible();
+    if ((await group.getAttribute("aria-expanded")) === "false") {
+      await group.click();
+    }
+  }
+  await link.click();
+}
+
+export async function expectUiRoute(page: PlaywrightPage, segment: string): Promise<void> {
+  const root = (process.env.SERVER_ROOT_PATH ?? "").replace(/\/+$/, "");
+  const expected = new URL(`${root}/ui/${segment}`, UI_BASE_URL);
+  await expect(page, `navigate to ${expected.pathname}`).toHaveURL(
+    (url) => url.origin === expected.origin && url.pathname.replace(/\/+$/, "") === expected.pathname,
+  );
+}
 
 /**
  * Navigates to a specific page using the page query parameter.
diff --git a/tests/e2e/ui/tests/migration/README.md b/tests/e2e/ui/tests/migration/README.md
index d6b33598ec4..59933502463 100644
--- a/tests/e2e/ui/tests/migration/README.md
+++ b/tests/e2e/ui/tests/migration/README.md
@@ -1,17 +1,25 @@
 # App Router migration smoke
 
 A growing E2E smoke for pages migrated from the legacy `?page=` switch to App
-Router path routes. For each migrated page it clicks the page's sidebar link, checks
-the URL is the path route and the page renders, reloads it, then clicks off to a
-legacy page and back to confirm navigation still works. It runs in two situations:
-the default mount and a non-root `SERVER_ROOT_PATH` mount.
+Router path routes. For each page it clicks the sidebar link by its accessible
+name, verifies the destination's content, reloads it, then visits Virtual Keys
+and returns. It runs at the default mount and a non-root `SERVER_ROOT_PATH` mount
+
+Link selection does not depend on `href` formatting. URL assertions compare the
+origin and pathname, allowing a trailing slash, query string, and fragment while
+rejecting another route or mount. Reloads must return a successful document,
+and each journey must finish without uncaught browser errors
 
 ## Adding a page
 
-When a page's migration merges, add its route segment to
-`tests/e2e/ui/fixtures/migratedPages.ts` (keep it in lockstep with `MIGRATED_PAGES`
-in `ui/litellm-dashboard/src/utils/migratedPages.ts`). Both suites pick it up
-automatically.
+Add an entry to `tests/e2e/ui/fixtures/migratedPages.ts`, keyed by the legacy page
+ID. Specify its route segment, accessible link name, sidebar group if collapsed,
+and distinctive visible content such as a heading or tab. Keep expectations
+independent of the application's route table so an incorrect destination fails
+the test. Both navigation suites use this fixture
+
+For a licensed-only page, `unlicensedText` describes the expected upgrade notice.
+The authenticated session's license claim determines which content must render
 
 ## Running
 
@@ -31,4 +39,9 @@ SERVER_ROOT_PATH=/litellm npm run e2e:migration:root
 ```
 
 `globalSetup` logs in once per role; the admin storage state is reused for these
-tests. Under a non-root mount it logs in at `${SERVER_ROOT_PATH}/ui/login`.
+tests. Under a non-root mount it logs in at `${SERVER_ROOT_PATH}/ui/login`
+
+`tests/navigation/sidebar.spec.ts` also checks the navigation helpers against
+equivalent link formats on the live dashboard and a deep link containing a query
+string and fragment. The link-format cases change only the rendered `href`
+attribute to exercise the locator contract; destination pages and APIs remain live
diff --git a/tests/e2e/ui/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts
index bc96d9a6c66..f2dbd2dbc77 100644
--- a/tests/e2e/ui/tests/migration/migratedPages.spec.ts
+++ b/tests/e2e/ui/tests/migration/migratedPages.spec.ts
@@ -1,105 +1,73 @@
 import { test, expect, type Page } from "@playwright/test";
-import { MIGRATED_E2E_SEGMENTS } from "../../fixtures/migratedPages";
+import { MIGRATED_E2E_PAGES, type MigratedPage } from "../../fixtures/migratedPages";
 import { ADMIN_STORAGE_PATH } from "../../constants";
-import { dismissFeedbackPopup } from "../../helpers/navigation";
+import { clickSidebarLink, dismissFeedbackPopup, expectUiRoute, sidebarLink } from "../../helpers/navigation";
+import { proxyIsPremium } from "../../helpers/premium";
 
-/**
- * App Router migration smoke as a user journey: start where the proxy lands you,
- * click a migrated page in the sidebar, confirm it routed and rendered, reload it
- * (the check a wrong server_root_path breaks), bounce to a legacy page and back,
- * and, once two pages are migrated, navigate directly between two migrated pages.
- *
- * Driven by MIGRATED_E2E_SEGMENTS, so it grows as pages are migrated. Set
- * SERVER_ROOT_PATH (e.g. "/litellm") to exercise the non-root mount; leave it
- * unset for the default mount. Boot the proxy with the matching value first.
- */
-const ROOT = process.env.SERVER_ROOT_PATH ?? "";
+const ROOT = (process.env.SERVER_ROOT_PATH ?? "").replace(/\/+$/, "");
+const apiKeys = MIGRATED_E2E_PAGES["api-keys"];
 
-const esc = (s: string) => s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
-const pathRe = (segment: string) => new RegExp(`${esc(ROOT)}/ui/${esc(segment)}/?($|\\?)`);
-// Scope nav lookups to the sidebar (a `complementary` landmark). The top bar
-// now renders a breadcrumb whose current-page item is also a "Virtual Keys"
-// link, so an unscoped locator would match two elements.
-const sidebar = (page: Page) => page.getByRole("complementary");
-const virtualKeysLink = (page: Page) => sidebar(page).getByRole("link", { name: "Virtual Keys", exact: true });
-
-/** The dashboard shell is present (sidebar rendered); page didn't 404 / crash. */
-async function expectRendered(page: Page) {
-  await expect(virtualKeysLink(page)).toBeVisible({ timeout: 20_000 });
+async function expectContent(page: Page, destination: MigratedPage): Promise<void> {
+  await expect(sidebarLink(page, apiKeys.linkName)).toBeVisible({ timeout: 20_000 });
+  const main = page.getByRole("main");
+  if (destination.unlicensedText && !proxyIsPremium()) {
+    await expect(main.getByText(destination.unlicensedText, { exact: true })).toBeVisible();
+    return;
+  }
+  const content = destination.content;
+  const landmark =
+    "role" in content
+      ? main.getByRole(content.role, { name: content.name, exact: true })
+      : main.getByText(content.text, { exact: true });
+  await expect(landmark).toBeVisible();
 }
 
-/**
- * Click a migrated page's sidebar link. Migrated items render as <a href=".../ui/<segment>/">;
- * nested ones live under collapsible groups whose children only render while the
- * group is open, so expand collapsed groups until the link is clickable.
- */
-async function clickSidebar(page: Page, segment: string) {
-  const link = sidebar(page).locator(`a[href$="/ui/${segment}"], a[href$="/ui/${segment}/"]`).first();
-  const collapsedGroups = sidebar(page).getByRole("button", { expanded: false });
-  for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) {
-    const stillCollapsed = await collapsedGroups.count();
-    if (stillCollapsed === 0) break;
-    await collapsedGroups.first().click();
-    await expect(collapsedGroups).toHaveCount(stillCollapsed - 1);
-  }
-  await link.click();
+async function navigateToDestination(page: Page, destination: MigratedPage): Promise<void> {
+  await clickSidebarLink(page, destination.linkName, destination.group);
+  await expectUiRoute(page, destination.segment);
+  await dismissFeedbackPopup(page);
+  await expectContent(page, destination);
 }
 
 test.use({ storageState: ADMIN_STORAGE_PATH });
 
 test.describe("App Router migrated pages", () => {
-  for (const segment of MIGRATED_E2E_SEGMENTS) {
-    test(`${segment}: sidebar nav, reload, and round-trip via the api-keys landing`, async ({ page }) => {
+  for (const destination of Object.values(MIGRATED_E2E_PAGES)) {
+    test(`${destination.segment}: sidebar nav, reload, and round-trip via the api-keys landing`, async ({ page }) => {
       const pageErrors: string[] = [];
-      page.on("pageerror", (e) => pageErrors.push(String(e)));
+      page.on("pageerror", (error) => pageErrors.push(String(error)));
 
-      // 1. Start where the proxy lands us.
-      await page.goto(`${ROOT}/ui/`);
+      const landing = await page.goto(`${ROOT}/ui/`);
+      expect(landing?.ok(), "dashboard document loads successfully").toBe(true);
       await dismissFeedbackPopup(page);
-      await expectRendered(page);
+      await expectContent(page, apiKeys);
 
-      // 2. Click the migrated page in the sidebar -> path route + rendered.
-      await clickSidebar(page, segment);
-      await expect(page).toHaveURL(pathRe(segment));
-      await expectRendered(page);
-      // 3. Reload the path route directly; a wrong server_root_path 404s here.
-      await page.reload();
+      await navigateToDestination(page, destination);
+
+      const reloaded = await page.reload();
+      expect(reloaded?.ok(), `${destination.segment} document loads on reload`).toBe(true);
+      await expectUiRoute(page, destination.segment);
       await dismissFeedbackPopup(page);
-      await expect(page).toHaveURL(pathRe(segment));
-      await expectRendered(page);
-      // 4. Click the Virtual Keys sidebar link to the api-keys landing (now a path route), then back.
-      await virtualKeysLink(page).click();
-      await expect(page).toHaveURL(pathRe("api-keys"));
-      await dismissFeedbackPopup(page);
-      await expectRendered(page);
-      // 5. Click back to the migrated page.
-      await clickSidebar(page, segment);
-      await expect(page).toHaveURL(pathRe(segment));
-      await expectRendered(page);
-      expect(pageErrors, `page errors during ${segment} journey`).toEqual([]);
+      await expectContent(page, destination);
+
+      await navigateToDestination(page, apiKeys);
+      await navigateToDestination(page, destination);
+      expect(pageErrors, `page errors during ${destination.segment} journey`).toEqual([]);
     });
   }
 
   test("navigates directly between two migrated pages", async ({ page }) => {
-    test.skip(MIGRATED_E2E_SEGMENTS.length < 2, "needs >= 2 migrated pages");
-    const [first, second] = MIGRATED_E2E_SEGMENTS;
     const pageErrors: string[] = [];
-    page.on("pageerror", (e) => pageErrors.push(String(e)));
+    page.on("pageerror", (error) => pageErrors.push(String(error)));
 
-    await page.goto(`${ROOT}/ui/`);
+    const landing = await page.goto(`${ROOT}/ui/`);
+    expect(landing?.ok(), "dashboard document loads successfully").toBe(true);
     await dismissFeedbackPopup(page);
+    await expectContent(page, apiKeys);
 
-    await clickSidebar(page, first);
-    await expect(page).toHaveURL(pathRe(first));
-    await expectRendered(page);
-    await clickSidebar(page, second);
-    await expect(page).toHaveURL(pathRe(second));
-    await expectRendered(page);
-    // Back to the first migrated page.
-    await clickSidebar(page, first);
-    await expect(page).toHaveURL(pathRe(first));
-    await expectRendered(page);
-
-    expect(pageErrors, "page errors during migrated -> migrated nav").toEqual([]);
+    for (const destination of [apiKeys, MIGRATED_E2E_PAGES.models, apiKeys]) {
+      await navigateToDestination(page, destination);
+    }
+    expect(pageErrors, "page errors during migrated page navigation").toEqual([]);
   });
 });
diff --git a/tests/e2e/ui/tests/navigation/sidebar.spec.ts b/tests/e2e/ui/tests/navigation/sidebar.spec.ts
index b220dc09ae2..eaa4985c659 100644
--- a/tests/e2e/ui/tests/navigation/sidebar.spec.ts
+++ b/tests/e2e/ui/tests/navigation/sidebar.spec.ts
@@ -3,7 +3,13 @@ import { Role } from "../../fixtures/roles";
 import { ADMIN_STORAGE_PATH } from "../../constants";
 import { Page } from "../../fixtures/pages";
 import { menuLabelToPage } from "../../fixtures/menuMappings";
-import { navigateToPage } from "../../helpers/navigation";
+import {
+  clickSidebarLink,
+  dismissFeedbackPopup,
+  expectUiRoute,
+  navigateToPage,
+  sidebarLink,
+} from "../../helpers/navigation";
 import { MIGRATED_E2E_PAGES } from "../../fixtures/migratedPages";
 import type { Page as PlaywrightPage } from "@playwright/test";
 
@@ -11,7 +17,7 @@ const sidebarButtons = {
   [Role.ProxyAdmin]: [
     "Virtual Keys",
     "Playground",
-    "Models",
+    "Models + Endpoints",
     "Usage",
     "Teams",
     "Internal Users",
@@ -22,9 +28,9 @@ const sidebarButtons = {
 
 /** Migrated pages live at a path route; legacy pages keep the ?page= query param. */
 async function expectPageUrl(page: PlaywrightPage, pageKey: string): Promise<void> {
-  const migratedSegment = MIGRATED_E2E_PAGES[pageKey];
-  if (migratedSegment) {
-    await expect(page).toHaveURL(new RegExp(`/ui/${migratedSegment}/?($|\\?)`));
+  const migratedPage = MIGRATED_E2E_PAGES[pageKey];
+  if (migratedPage) {
+    await expectUiRoute(page, migratedPage.segment);
   } else {
     await expect(page).toHaveURL(new RegExp(`[?&]page=${pageKey}(&|$)`));
   }
@@ -51,12 +57,7 @@ for (const { role, storage } of roles) {
           throw new Error(`No page mapping found for menu label: ${buttonLabel}`);
         }
 
-        // Sidebar items are links inside the `complementary` landmark; scoping
-        // there avoids the top-bar breadcrumb, which also links the page name.
-        const tab = page.getByRole("complementary").getByRole("link", { name: buttonLabel });
-        await expect(tab).toBeVisible();
-
-        await tab.click();
+        await clickSidebarLink(page, buttonLabel);
 
         await expectPageUrl(page, expectedPage);
       }
@@ -81,5 +82,41 @@ for (const { role, storage } of roles) {
       await navigateToPage(page, Page.LlmPlayground);
       await expectPageUrl(page, Page.LlmPlayground);
     });
+
+    for (const format of ["without trailing slash", "absolute with query and fragment", "relative"] as const) {
+      test(`sidebar locator tolerates hrefs ${format}`, async ({ page }) => {
+        await page.goto("/ui/");
+        await dismissFeedbackPopup(page);
+        const link = sidebarLink(page, "Models + Endpoints");
+        await expect(link).toBeVisible();
+        const destination = new URL("/ui/models-and-endpoints/", page.url());
+        const href =
+          format === "without trailing slash"
+            ? destination.pathname.replace(/\/$/, "")
+            : format === "relative"
+              ? "./models-and-endpoints/"
+              : `${destination.href}?source=navigation-smoke#overview`;
+
+        await link.evaluate((element, value) => element.setAttribute("href", value), href);
+        await expect(link).toHaveAttribute("href", href);
+        await clickSidebarLink(page, "Models + Endpoints");
+
+        await expectUiRoute(page, "models-and-endpoints");
+        await expect(
+          page.getByRole("main").getByRole("heading", { name: "Model Management", exact: true }),
+        ).toBeVisible();
+      });
+    }
+
+    test("route assertion tolerates a query string and fragment on a deep link", async ({ page }) => {
+      const response = await page.goto("/ui/models-and-endpoints/?source=navigation-smoke#overview");
+      expect(response?.ok()).toBe(true);
+      await expectUiRoute(page, "models-and-endpoints");
+      await expect(
+        page.getByRole("main").getByRole("heading", { name: "Model Management", exact: true }),
+      ).toBeVisible();
+      expect(new URL(page.url()).search).toBe("?source=navigation-smoke");
+      expect(new URL(page.url()).hash).toBe("#overview");
+    });
   });
 }

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 103/164] 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 5ca9e2605081d26503786cdd116e7fc8245c8903 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 6 Sep 2026 01:09:36 -0700
Subject: [PATCH 104/164] fix(ci): leave check-run names the sweep cannot
 resolve out of the comparison

A job name holding an expression the sweep could not resolve was compared as
if it were the published name. Two jobs whose names differ per matrix value or
per caller input were reported as a collision, and a matrix that was itself an
expression collapsed onto the bare job id and did the same.

Model what a job publishes as known names beside the reasons the rest stay
unknown. Anything the sweep cannot work out contributes no name and is
reported as a note instead of guessed at. An expression over contexts that are
fixed for the whole run still compares, so two jobs sharing one of those are
still caught.
---
 .../check_workflow_job_name_collisions.py     | 173 ++++++++++++++----
 .../test_workflow_job_name_collisions.py      | 151 ++++++++++++++-
 2 files changed, 280 insertions(+), 44 deletions(-)

diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
index 4eef3df1c11..1e81acc4f90 100644
--- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
@@ -20,16 +20,24 @@ with the other rows', so two shard lists that overlap collide even though their
 templates read differently. Each expression is evaluated per combination over the
 pieces a job name can hold: string literals, `matrix.<key>`, `format()`, `==` and
 `!=`, and the `<cond> && <a> || <b>` idiom, which is how the shards reach their
-real `<shard> / Run tests` names rather than staying opaque. An expression
-nothing resolves stays in the string, so two jobs carrying the same unresolved
-template still compare equal and their collision is still caught.
+real `<shard> / Run tests` names rather than staying opaque.
+
+Whatever the sweep cannot work out is left out of the comparison and reported
+instead of guessed, because a guess that lands wrong fails a workflow GitHub
+would have published perfectly well. A name left holding an expression over
+`inputs`, `needs` or `matrix` reads differently from every job that runs it, so
+it is one of those; an expression over `github` and the other contexts fixed for
+the whole run reads the same everywhere, so two jobs carrying it still clash and
+the check still says so. A matrix that is itself an expression or that lists
+values which are not scalars, and a call this sweep cannot follow, go in the same
+bucket. The cost is that a real clash hiding behind one of them goes unseen,
+which leaves a merge no worse off than before this check existed, where the
+opposite direction would block work that was fine.
 
 A job calling a local reusable workflow publishes one check run per job of the
 callee, named `<caller> / <callee>` and chained through however many levels of
 local calls it takes, which is why a caller's name never collides with a plain
-job that happens to match it. A call this sweep cannot follow, to a workflow
-outside the repo or to one that is not there, publishes nothing rather than a
-name GitHub never posts. A file under `.github/workflows/` that does not read as
+job that happens to match it. A file under `.github/workflows/` that does not read as
 a workflow at all is reported rather than skipped, since skipping it silently
 would hide every job it holds.
 """
@@ -54,6 +62,8 @@ MATRIX_REF: Final = re.compile(r"^matrix\.(?P<key>[\w-]+)$")
 LITERAL: Final = re.compile(r"^'(?P<text>[^']*)'$")
 FORMAT_CALL: Final = re.compile(r"^format\((?P<args>.*)\)$", re.DOTALL)
 COMPARISON: Final = re.compile(r"^(?P<left>.+?)\s*(?P<operator>==|!=)\s*(?P<right>.+)$", re.DOTALL)
+CONTEXT_REF: Final = re.compile(r"\b(?P<root>[a-z][\w-]*)\s*\.")
+RUN_FIXED_CONTEXTS: Final = frozenset({"github", "vars", "env", "runner", "secrets"})
 NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({})
 NO_CALLERS: Final[frozenset[str]] = frozenset()
 SCALAR: Final = (str, int, float)
@@ -66,6 +76,19 @@ class Unreadable:
     reason: str
 
 
+@dataclass(frozen=True, slots=True)
+class Opaque:
+    reason: str
+
+
+@dataclass(frozen=True, slots=True)
+class Names:
+    """The check-run names a job publishes, beside the reasons the rest of them stay unknown."""
+
+    known: tuple[str, ...] = ()
+    unknown: tuple[str, ...] = ()
+
+
 class Job(BaseModel):
     name: object = None
     uses: str | None = None
@@ -107,12 +130,23 @@ def publishes_check_runs(raw_on: object) -> bool:
     return events(raw_on) != frozenset({"workflow_call"})
 
 
-def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, ...]], ...]:
-    return tuple(
-        (str(key), tuple(scalar_text(value) for value in values if isinstance(value, SCALAR)))
-        for key, values in matrix.items()
-        if str(key) not in MATRIX_DIRECTIVES and isinstance(values, Sequence) and not isinstance(values, str)
+def scalar_list(value: object) -> tuple[str, ...] | Opaque:
+    """One matrix key's values, or why the combinations it produces cannot be worked out."""
+    if not isinstance(value, Sequence) or isinstance(value, str):
+        return Opaque("a matrix key holds something other than a list of values")
+    if any(not isinstance(item, SCALAR) for item in value):
+        return Opaque("a matrix key lists values that are not plain scalars")
+    return tuple(scalar_text(item) for item in value)
+
+
+def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, ...]], ...] | Opaque:
+    listed: Final = tuple(
+        (str(key), scalar_list(values)) for key, values in matrix.items() if str(key) not in MATRIX_DIRECTIVES
     )
+    opaque: Final = next((values for _, values in listed if isinstance(values, Opaque)), None)
+    if opaque is not None:
+        return opaque
+    return tuple((key, values) for key, values in listed if not isinstance(values, Opaque))
 
 
 def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...]:
@@ -150,17 +184,20 @@ def crossed_values(listed: Sequence[tuple[str, tuple[str, ...]]]) -> tuple[Mappi
     )
 
 
-def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...]:
+def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...] | Opaque:
     """One mapping per job the matrix produces, `exclude` applied before `include` as GitHub does."""
     matrix: Final = job.strategy.get("matrix")
-    if not isinstance(matrix, Mapping):
+    if matrix is None:
         return ()
+    if not isinstance(matrix, Mapping):
+        return Opaque("the matrix itself comes from an expression")
+    listed: Final = listed_values(matrix)
+    if isinstance(listed, Opaque):
+        return listed
     rows: Final = directive_rows(matrix, "include")
     dropped: Final = directive_rows(matrix, "exclude")
     kept: Final = tuple(
-        combination
-        for combination in crossed_values(listed_values(matrix))
-        if not any(drops(row, combination) for row in dropped)
+        combination for combination in crossed_values(listed) if not any(drops(row, combination) for row in dropped)
     )
     standalone: Final = tuple(row for row in rows if not any(extends(row, combination) for combination in kept))
     return (*(extended(combination, rows) for combination in kept), *standalone)
@@ -243,11 +280,33 @@ def rendered(template: str, values: Mapping[str, str]) -> str:
     return EXPRESSION.sub(lambda span: resolved_span(span, values), template)
 
 
-def expand(template: str, job: Job) -> tuple[str, ...]:
+def stable(name: str) -> bool:
+    """An expression over the run's own contexts reads the same from every job, so two of them still clash."""
+    return all(
+        reference.group("root") in RUN_FIXED_CONTEXTS
+        for span in EXPRESSION.finditer(name)
+        for reference in CONTEXT_REF.finditer(span.group("body"))
+    )
+
+
+def comparable(name: str) -> bool:
+    """A leftover expression over `inputs`, `needs` or `matrix` names a different check run per job."""
+    return EXPRESSION.search(name) is None or stable(name)
+
+
+def settled(names: Sequence[str]) -> Names:
+    return Names(
+        tuple(name for name in names if comparable(name)),
+        tuple(f"its name stays `{name}`" for name in names if not comparable(name)),
+    )
+
+
+def expand(template: str, job: Job) -> Names:
     combinations: Final = matrix_combinations(job)
-    if not combinations:
-        return (rendered(template, NO_MATRIX),)
-    return tuple(dict.fromkeys(rendered(template, values) for values in combinations))
+    if isinstance(combinations, Opaque):
+        return Names((), (combinations.reason,))
+    over: Final = combinations or (NO_MATRIX,)
+    return settled(tuple(dict.fromkeys(rendered(template, values) for values in over)))
 
 
 def suffixed(job_id: str, combination: Mapping[str, str]) -> str:
@@ -255,11 +314,14 @@ def suffixed(job_id: str, combination: Mapping[str, str]) -> str:
     return f"{job_id} ({', '.join(combination.values())})" if combination else job_id
 
 
-def published_names(job_id: str, job: Job) -> tuple[str, ...]:
+def published_names(job_id: str, job: Job) -> Names:
     if job.name is not None:
         return expand(scalar_text(job.name), job)
-    suffixes: Final = tuple(dict.fromkeys(suffixed(job_id, values) for values in matrix_combinations(job)))
-    return suffixes or (job_id,)
+    combinations: Final = matrix_combinations(job)
+    if isinstance(combinations, Opaque):
+        return Names((), (combinations.reason,))
+    suffixes: Final = tuple(dict.fromkeys(suffixed(job_id, values) for values in combinations))
+    return Names(suffixes or (job_id,))
 
 
 def callee_path(job: Job) -> str | None:
@@ -268,22 +330,40 @@ def callee_path(job: Job) -> str | None:
     return job.uses[len(LOCAL_CALL_PREFIX) :].split("@")[0]
 
 
-def job_names(
-    job_id: str, job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str] = NO_CALLERS
-) -> tuple[str, ...]:
+def joined(groups: Sequence[Names]) -> Names:
+    return Names(
+        tuple(name for group in groups for name in group.known),
+        tuple(reason for group in groups for reason in group.unknown),
+    )
+
+
+def call_blocker(job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str]) -> str | None:
+    path: Final = callee_path(job)
+    if path is None:
+        return "it calls a reusable workflow outside this repository"
+    if path in callers:
+        return f"its call to {path} loops back on itself"
+    return None if path in workflows else f"it calls {path}, which this checkout does not hold"
+
+
+def job_names(job_id: str, job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str] = NO_CALLERS) -> Names:
     prefixes: Final = published_names(job_id, job)
     if job.uses is None:
         return prefixes
-    path: Final = callee_path(job)
-    callee: Final = workflows.get(path or "")
-    if path is None or callee is None or path in callers:
-        return ()
-    suffixes: Final = tuple(
-        name
-        for callee_id, callee_job in callee.jobs.items()
-        for name in job_names(callee_id, callee_job, workflows, callers | {path})
+    blocker: Final = call_blocker(job, workflows, callers)
+    if blocker is not None:
+        return Names((), (*prefixes.unknown, blocker))
+    path: Final = callee_path(job) or ""
+    suffixes: Final = joined(
+        tuple(
+            job_names(callee_id, callee_job, workflows, callers | {path})
+            for callee_id, callee_job in workflows[path].jobs.items()
+        )
+    )
+    return Names(
+        tuple(f"{prefix} / {suffix}" for prefix in prefixes.known for suffix in suffixes.known),
+        (*prefixes.unknown, *suffixes.unknown),
     )
-    return tuple(f"{prefix} / {suffix}" for prefix in prefixes for suffix in suffixes)
 
 
 def readable(sources: Mapping[str, str]) -> Mapping[str, tuple[Workflow, object]]:
@@ -300,15 +380,29 @@ def unreadable(sources: Mapping[str, str]) -> tuple[str, ...]:
     )
 
 
-def published(sources: Mapping[str, str]) -> Iterator[tuple[str, str]]:
+def scanned_jobs(sources: Mapping[str, str]) -> Iterator[tuple[str, str, Names]]:
     parsed: Final = readable(sources)
     workflows: Final = {rel: workflow for rel, (workflow, _) in parsed.items()}
     for rel, (workflow, raw_on) in parsed.items():
         if not publishes_check_runs(raw_on):
             continue
         for job_id, job in workflow.jobs.items():
-            for name in job_names(job_id, job, workflows):
-                yield name, f"{rel} job `{job_id}`"
+            yield rel, job_id, job_names(job_id, job, workflows)
+
+
+def published(sources: Mapping[str, str]) -> Iterator[tuple[str, str]]:
+    for rel, job_id, names in scanned_jobs(sources):
+        for name in names.known:
+            yield name, f"{rel} job `{job_id}`"
+
+
+def blind_spots(sources: Mapping[str, str]) -> tuple[str, ...]:
+    """Jobs whose published names GitHub decides at run time, which no offline sweep can compare."""
+    return tuple(
+        f"{rel} job `{job_id}` publishes a name this check cannot work out because {reason}."
+        for rel, job_id, names in scanned_jobs(sources)
+        for reason in sorted(names.unknown)
+    )
 
 
 def owners_by_name(sources: Mapping[str, str]) -> Iterator[tuple[str, tuple[str, ...]]]:
@@ -338,6 +432,9 @@ def report(header: str, problems: Sequence[str]) -> None:
 def exit_code(sources: Mapping[str, str]) -> int:
     unread: Final = unreadable(sources)
     found: Final = collisions(sources)
+    blind: Final = blind_spots(sources)
+    if blind:
+        print("NOTE: names left out of the comparison:\n  - " + "\n  - ".join(blind))
     report("Some workflows could not be read", unread)
     report("Check-run names are not unique", found)
     if unread or found:
diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
index 3dd8a3d4088..39e61ff48d1 100644
--- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
@@ -2,6 +2,7 @@ from typing import Final
 
 from check_workflow_job_name_collisions import (
     Unreadable,
+    blind_spots,
     callee_path,
     collisions,
     exit_code,
@@ -301,9 +302,9 @@ def test_a_matrix_list_supplies_values_the_same_way_include_rows_do() -> None:
     assert "`Analyze (go)` is published by 2 jobs" in found[0]
 
 
-def test_two_jobs_sharing_one_unresolvable_template_still_collide() -> None:
+def test_two_jobs_sharing_a_template_over_the_run_itself_still_collide() -> None:
     template: Final = (
-        "on: pull_request\njobs:\n  {job}:\n    name: ${{{{ matrix.shard }}}}\n    runs-on: ubuntu-latest\n"
+        "on: pull_request\njobs:\n  {job}:\n    name: ${{{{ github.event_name }}}}-build\n    runs-on: ubuntu-latest\n"
     )
     sources: Final = {
         "a.yml": template.format(job="one"),
@@ -316,6 +317,19 @@ def test_two_jobs_sharing_one_unresolvable_template_still_collide() -> None:
     assert "is published by 2 jobs" in found[0]
 
 
+def test_two_jobs_sharing_a_template_that_reads_per_job_are_not_called_a_collision() -> None:
+    template: Final = (
+        "on: pull_request\njobs:\n  {job}:\n    name: ${{{{ matrix.shard }}}}\n    runs-on: ubuntu-latest\n"
+    )
+    sources: Final = {
+        "a.yml": template.format(job="one"),
+        "b.yml": template.format(job="two"),
+    }
+
+    assert collisions(sources) == ()
+    assert len(blind_spots(sources)) == 2
+
+
 def test_a_file_that_is_not_a_workflow_is_reported_rather_than_skipped() -> None:
     sources: Final = {
         "notes.yml": "just a string\n",
@@ -437,13 +451,14 @@ def test_a_boolean_matrix_value_renders_the_way_github_writes_it() -> None:
     assert names == frozenset({"cache true", "cache false"})
 
 
-def test_a_format_call_its_arguments_cannot_fill_leaves_the_name_unresolved() -> None:
-    names: Final = frozenset(name for name, _ in published({"unit.yml": UNFILLABLE_FORMAT}))
+def test_a_format_call_its_arguments_cannot_fill_publishes_nothing_to_compare() -> None:
+    sources: Final = {"unit.yml": UNFILLABLE_FORMAT}
 
-    assert names == frozenset({"${{ format('{0} {1}', matrix.shard) }}"})
+    assert frozenset(name for name, _ in published(sources)) == frozenset()
+    assert "its name stays" in blind_spots(sources)[0]
 
 
-def test_a_call_to_a_workflow_outside_the_repo_publishes_no_name() -> None:
+def test_a_call_to_a_workflow_outside_the_repo_is_reported_rather_than_guessed() -> None:
     sources: Final = {
         "a.yml": "on: pull_request\njobs:\n  unit:\n    uses: BerriAI/other/.github/workflows/base.yml@main\n",
         "b.yml": "on: pull_request\njobs:\n  unit:\n    runs-on: ubuntu-latest\n",
@@ -451,6 +466,7 @@ def test_a_call_to_a_workflow_outside_the_repo_publishes_no_name() -> None:
 
     assert frozenset(name for name, _ in published(sources)) == frozenset({"unit"})
     assert collisions(sources) == ()
+    assert "outside this repository" in blind_spots(sources)[0]
 
 
 def test_a_chain_of_local_reusable_calls_publishes_every_level_of_the_chain() -> None:
@@ -498,3 +514,126 @@ def test_the_check_passes_when_every_file_reads_and_every_name_is_unique() -> No
     }
 
     assert exit_code(sources) == 0
+
+
+def test_two_callers_of_one_reusable_workflow_named_from_its_inputs_do_not_collide() -> None:
+    sources: Final = {
+        ".github/workflows/callee.yml": (
+            "on:\n  workflow_call:\njobs:\n  run:\n    name: ${{ inputs.suite }}\n    runs-on: ubuntu-latest\n"
+        ),
+        "caller.yml": (
+            "on: pull_request\njobs:\n"
+            "  alpha:\n    name: A\n    uses: ./.github/workflows/callee.yml\n    with:\n      suite: alpha\n"
+            "  beta:\n    name: A\n    uses: ./.github/workflows/callee.yml\n    with:\n      suite: beta\n"
+        ),
+    }
+
+    assert collisions(sources) == ()
+    assert len(blind_spots(sources)) == 2
+
+
+def test_a_matrix_that_is_itself_an_expression_never_collapses_onto_the_bare_job_id() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\njobs:\n  build:\n    strategy:\n"
+            "      matrix: ${{ fromJson(needs.plan.outputs.matrix) }}\n    runs-on: ubuntu-latest\n"
+        ),
+        "b.yml": "on: pull_request\njobs:\n  build:\n    runs-on: ubuntu-latest\n",
+    }
+
+    assert frozenset(name for name, _ in published(sources)) == frozenset({"build"})
+    assert collisions(sources) == ()
+    assert "the matrix itself comes from an expression" in blind_spots(sources)[0]
+
+
+def test_a_matrix_listing_objects_never_collapses_onto_the_bare_job_id() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\njobs:\n  build:\n    strategy:\n      matrix:\n        target:\n"
+            "          - os: ubuntu\n          - os: windows\n    runs-on: ubuntu-latest\n"
+        ),
+        "b.yml": "on: pull_request\njobs:\n  build:\n    runs-on: ubuntu-latest\n",
+    }
+
+    assert frozenset(name for name, _ in published(sources)) == frozenset({"build"})
+    assert collisions(sources) == ()
+    assert "not plain scalars" in blind_spots(sources)[0]
+
+
+def test_a_call_to_a_workflow_file_the_checkout_does_not_hold_is_reported() -> None:
+    sources: Final = {"a.yml": "on: pull_request\njobs:\n  unit:\n    uses: ./.github/workflows/gone.yml\n"}
+
+    assert collisions(sources) == ()
+    assert "which this checkout does not hold" in blind_spots(sources)[0]
+
+
+def test_reusable_workflows_calling_each_other_in_a_loop_are_reported_not_followed() -> None:
+    sources: Final = {
+        ".github/workflows/a.yml": (
+            "on:\n  workflow_call:\njobs:\n  call:\n    name: A\n    uses: ./.github/workflows/b.yml\n"
+        ),
+        ".github/workflows/b.yml": (
+            "on:\n  workflow_call:\njobs:\n  call:\n    name: B\n    uses: ./.github/workflows/a.yml\n"
+        ),
+        "top.yml": "on: pull_request\njobs:\n  top:\n    name: Top\n    uses: ./.github/workflows/a.yml\n",
+    }
+
+    assert collisions(sources) == ()
+    assert any("loops back on itself" in spot for spot in blind_spots(sources))
+
+
+def test_a_caller_still_publishes_the_callee_jobs_it_can_read() -> None:
+    sources: Final = {
+        ".github/workflows/callee.yml": (
+            "on:\n  workflow_call:\njobs:\n"
+            "  lint:\n    name: Lint\n    runs-on: ubuntu-latest\n"
+            "  suite:\n    name: ${{ inputs.suite }}\n    runs-on: ubuntu-latest\n"
+        ),
+        "caller.yml": "on: pull_request\njobs:\n  call:\n    name: A\n    uses: ./.github/workflows/callee.yml\n",
+    }
+
+    assert frozenset(name for name, _ in published(sources)) == frozenset({"A / Lint"})
+    assert len(blind_spots(sources)) == 1
+
+
+def test_a_name_the_check_cannot_work_out_is_reported_without_failing_the_check() -> None:
+    sources: Final = {
+        "a.yml": "on: pull_request\njobs:\n  unit:\n    uses: BerriAI/other/.github/workflows/base.yml@main\n",
+    }
+
+    assert blind_spots(sources) != ()
+    assert exit_code(sources) == 0
+
+
+def test_a_caller_whose_own_name_is_unreadable_publishes_none_of_its_callee_names() -> None:
+    sources: Final = {
+        ".github/workflows/callee.yml": (
+            "on:\n  workflow_call:\njobs:\n  lint:\n    name: Lint\n    runs-on: ubuntu-latest\n"
+        ),
+        "caller.yml": (
+            "on: pull_request\njobs:\n  call:\n    name: ${{ matrix.suite }}\n"
+            "    uses: ./.github/workflows/callee.yml\n"
+        ),
+        "other.yml": "on: pull_request\njobs:\n  plain:\n    name: Lint\n    runs-on: ubuntu-latest\n",
+    }
+
+    assert frozenset(name for name, _ in published(sources)) == frozenset({"Lint"})
+    assert collisions(sources) == ()
+    assert "its name stays" in blind_spots(sources)[0]
+
+
+def test_an_include_row_naming_a_listed_key_extends_only_the_combinations_it_matches() -> None:
+    sources: Final = {
+        "unit.yml": (
+            "on: pull_request\njobs:\n  unit:\n"
+            "    name: ${{ matrix.python-version }} ${{ matrix.label }}\n"
+            "    runs-on: ubuntu-latest\n    strategy:\n      matrix:\n"
+            '        python-version: ["3.12", "3.13"]\n'
+            "        include:\n"
+            '          - python-version: "3.12"\n'
+            "            label: fast\n"
+        )
+    }
+
+    assert frozenset(name for name, _ in published(sources)) == frozenset({"3.12 fast"})
+    assert len(blind_spots(sources)) == 1

From 904542a559c1f297c63a5b299557d340c6f3f350 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 6 Sep 2026 01:27:03 -0700
Subject: [PATCH 105/164] fix(ci): stop guessing at names built from contexts
 the sweep cannot read

Three ways the sweep could fail a workflow GitHub would publish fine.

`github.workflow` and `github.job` were counted as fixed for the whole run, so
two jobs naming themselves after the workflow they sit in were reported as a
collision. `runner` and `vars` were wrong the same way. Drop the exception
entirely: a name still holding an expression is one GitHub resolves per job, so
it is nothing to compare, which is what the rest of the module already does.

`format()` was resolved with Python's semantics, so an attribute lookup crashed
the script and a width specifier padded a name GitHub never pads. Fill `{0}`
holes and escaped braces, and treat anything richer as unresolved.

A matrix `include` or `exclude` row holding a value that is not a scalar lost
that key and became an empty row, which excludes every combination. Report the
row instead of quietly reshaping the matrix around it.
---
 .../check_workflow_job_name_collisions.py     | 52 +++++++-------
 .../test_workflow_job_name_collisions.py      | 70 +++++++++++++++++--
 2 files changed, 91 insertions(+), 31 deletions(-)

diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
index 1e81acc4f90..9976bd9aa41 100644
--- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
@@ -24,13 +24,12 @@ real `<shard> / Run tests` names rather than staying opaque.
 
 Whatever the sweep cannot work out is left out of the comparison and reported
 instead of guessed, because a guess that lands wrong fails a workflow GitHub
-would have published perfectly well. A name left holding an expression over
-`inputs`, `needs` or `matrix` reads differently from every job that runs it, so
-it is one of those; an expression over `github` and the other contexts fixed for
-the whole run reads the same everywhere, so two jobs carrying it still clash and
-the check still says so. A matrix that is itself an expression or that lists
-values which are not scalars, and a call this sweep cannot follow, go in the same
-bucket. The cost is that a real clash hiding behind one of them goes unseen,
+would have published perfectly well. A name still holding an expression once the
+combination is filled in is one GitHub resolves per job, so it is one of those:
+guessing that two jobs sharing such a template clash would fail workflows over a
+context this sweep cannot read. A matrix that is itself an expression or that
+lists values which are not scalars, an `include` or `exclude` row shaped the same
+way, and a call this sweep cannot follow, go in the same bucket. The cost is that a real clash hiding behind one of them goes unseen,
 which leaves a merge no worse off than before this check existed, where the
 opposite direction would block work that was fine.
 
@@ -62,8 +61,7 @@ MATRIX_REF: Final = re.compile(r"^matrix\.(?P<key>[\w-]+)$")
 LITERAL: Final = re.compile(r"^'(?P<text>[^']*)'$")
 FORMAT_CALL: Final = re.compile(r"^format\((?P<args>.*)\)$", re.DOTALL)
 COMPARISON: Final = re.compile(r"^(?P<left>.+?)\s*(?P<operator>==|!=)\s*(?P<right>.+)$", re.DOTALL)
-CONTEXT_REF: Final = re.compile(r"\b(?P<root>[a-z][\w-]*)\s*\.")
-RUN_FIXED_CONTEXTS: Final = frozenset({"github", "vars", "env", "runner", "secrets"})
+GITHUB_PLACEHOLDER: Final = re.compile(r"\{\{|\}\}|\{\d+\}")
 NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({})
 NO_CALLERS: Final[frozenset[str]] = frozenset()
 SCALAR: Final = (str, int, float)
@@ -149,15 +147,17 @@ def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, .
     return tuple((key, values) for key, values in listed if not isinstance(values, Opaque))
 
 
-def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...]:
+def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...] | Opaque:
+    """One `include` or `exclude` row, or why the combinations they shape cannot be worked out."""
     rows: Final = matrix.get(directive)
     if not isinstance(rows, Sequence) or isinstance(rows, str):
         return ()
-    return tuple(
-        MappingProxyType({str(key): scalar_text(value) for key, value in row.items() if isinstance(value, SCALAR)})
-        for row in rows
-        if isinstance(row, Mapping)
-    )
+    mappings: Final = tuple(row for row in rows if isinstance(row, Mapping))
+    if len(mappings) != len(rows):
+        return Opaque(f"a matrix `{directive}` row is not a mapping of values")
+    if any(not isinstance(value, SCALAR) for row in mappings for value in row.values()):
+        return Opaque(f"a matrix `{directive}` row holds a value that is not a plain scalar")
+    return tuple(MappingProxyType({str(key): scalar_text(value) for key, value in row.items()}) for row in mappings)
 
 
 def drops(row: Mapping[str, str], combination: Mapping[str, str]) -> bool:
@@ -195,7 +195,11 @@ def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...] | Opaque:
     if isinstance(listed, Opaque):
         return listed
     rows: Final = directive_rows(matrix, "include")
+    if isinstance(rows, Opaque):
+        return rows
     dropped: Final = directive_rows(matrix, "exclude")
+    if isinstance(dropped, Opaque):
+        return dropped
     kept: Final = tuple(
         combination for combination in crossed_values(listed) if not any(drops(row, combination) for row in dropped)
     )
@@ -225,7 +229,10 @@ def split_outside(text: str, token: str) -> tuple[str, ...]:
 
 
 def formatted(template: str, arguments: Sequence[str]) -> str | None:
-    """A `format()` whose placeholders the arguments cannot fill resolves to nothing, never a crash."""
+    """GitHub's `format()` fills `{0}`-style holes and escapes braces, so anything richer resolves to nothing."""
+    residue: Final = GITHUB_PLACEHOLDER.sub("", template)
+    if "{" in residue or "}" in residue:
+        return None
     try:
         return template.format(*arguments)
     except (IndexError, KeyError, ValueError):
@@ -280,18 +287,9 @@ def rendered(template: str, values: Mapping[str, str]) -> str:
     return EXPRESSION.sub(lambda span: resolved_span(span, values), template)
 
 
-def stable(name: str) -> bool:
-    """An expression over the run's own contexts reads the same from every job, so two of them still clash."""
-    return all(
-        reference.group("root") in RUN_FIXED_CONTEXTS
-        for span in EXPRESSION.finditer(name)
-        for reference in CONTEXT_REF.finditer(span.group("body"))
-    )
-
-
 def comparable(name: str) -> bool:
-    """A leftover expression over `inputs`, `needs` or `matrix` names a different check run per job."""
-    return EXPRESSION.search(name) is None or stable(name)
+    """A name still holding an expression is one GitHub resolves per job, so it is nothing to compare."""
+    return EXPRESSION.search(name) is None
 
 
 def settled(names: Sequence[str]) -> Names:
diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
index 39e61ff48d1..44a5adbbb33 100644
--- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
@@ -302,7 +302,7 @@ def test_a_matrix_list_supplies_values_the_same_way_include_rows_do() -> None:
     assert "`Analyze (go)` is published by 2 jobs" in found[0]
 
 
-def test_two_jobs_sharing_a_template_over_the_run_itself_still_collide() -> None:
+def test_two_jobs_sharing_a_template_over_the_run_itself_are_reported_rather_than_guessed() -> None:
     template: Final = (
         "on: pull_request\njobs:\n  {job}:\n    name: ${{{{ github.event_name }}}}-build\n    runs-on: ubuntu-latest\n"
     )
@@ -311,10 +311,72 @@ def test_two_jobs_sharing_a_template_over_the_run_itself_still_collide() -> None
         "b.yml": template.format(job="two"),
     }
 
-    found: Final = collisions(sources)
+    assert collisions(sources) == ()
+    assert len(blind_spots(sources)) == 2
 
-    assert len(found) == 1
-    assert "is published by 2 jobs" in found[0]
+
+def test_a_name_reading_the_workflow_it_sits_in_is_not_called_a_collision() -> None:
+    template: Final = (
+        "on: pull_request\njobs:\n  {job}:\n    name: ${{{{ github.workflow }}}} / build\n    runs-on: ubuntu-latest\n"
+    )
+    sources: Final = {
+        "a.yml": template.format(job="one"),
+        "b.yml": template.format(job="two"),
+    }
+
+    assert collisions(sources) == ()
+    assert exit_code(sources) == 0
+
+
+def test_a_format_call_python_accepts_but_github_does_not_publishes_nothing_to_compare() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\njobs:\n  one:\n    name: ${{ format('{0.real}', matrix.shard) }}\n"
+            "    runs-on: ubuntu-latest\n    strategy:\n      matrix:\n        shard: [core]\n"
+        ),
+    }
+
+    assert collisions(sources) == ()
+    assert blind_spots(sources) != ()
+
+
+def test_a_format_call_padding_its_argument_publishes_nothing_to_compare() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\njobs:\n  one:\n    name: ${{ format('{0:>8}', matrix.shard) }}\n"
+            "    runs-on: ubuntu-latest\n    strategy:\n      matrix:\n        shard: [core]\n"
+        ),
+        "b.yml": "on: pull_request\njobs:\n  two:\n    name: '    core'\n    runs-on: ubuntu-latest\n",
+    }
+
+    assert collisions(sources) == ()
+    assert blind_spots(sources) != ()
+
+
+def test_an_exclude_row_that_is_not_a_mapping_is_reported_rather_than_skipped() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\njobs:\n  build:\n    runs-on: ubuntu-latest\n"
+            "    strategy:\n      matrix:\n        v: [1, 2]\n        exclude:\n          - oops\n"
+        ),
+        "b.yml": "on: pull_request\njobs:\n  other:\n    name: build (1)\n    runs-on: ubuntu-latest\n",
+    }
+
+    assert collisions(sources) == ()
+    assert blind_spots(sources) != ()
+
+
+def test_an_exclude_row_holding_a_non_scalar_never_drops_every_combination() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\njobs:\n  build:\n    runs-on: ubuntu-latest\n"
+            "    strategy:\n      matrix:\n        v: [1, 2]\n        exclude:\n          - cfg: {k: 1}\n"
+        ),
+        "b.yml": "on: pull_request\njobs:\n  build:\n    runs-on: ubuntu-latest\n",
+    }
+
+    assert collisions(sources) == ()
+    assert blind_spots(sources) != ()
 
 
 def test_two_jobs_sharing_a_template_that_reads_per_job_are_not_called_a_collision() -> None:

From aa2c41f489eb3eba1de9983d641239b131057daa Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 6 Sep 2026 01:59:03 -0700
Subject: [PATCH 106/164] fix(ci): stop the name-collision check failing legal
 workflows

An expression at `jobs.<id>.strategy` is legal on GitHub, but the model
required a mapping there, so a workflow using one made the whole file
unreadable and turned code-quality red. That job's names are now a blind
spot like any other name the sweep cannot work out offline.

A matrix whose `name:` holds no matrix value publishes that one name once
per combination, which leaves a required context just as ambiguous as two
jobs sharing a name, so it now reports instead of deduping.

A file that does not parse as one YAML document is reported the way the
module already promised, rather than escaping as a traceback.
---
 .../check_workflow_job_name_collisions.py     | 52 ++++++++++++-----
 .../test_workflow_job_name_collisions.py      | 58 +++++++++++++++++++
 2 files changed, 94 insertions(+), 16 deletions(-)

diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
index 9976bd9aa41..941039e8340 100644
--- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
@@ -1,10 +1,12 @@
-"""Catch two workflow jobs that publish check runs under the same name.
+"""Catch workflow jobs that publish check runs under the same name.
 
 A ruleset's required status check names a check run and GitHub matches it by that
 name alone. When two jobs publish the same name the required context stops
 mapping to the job that proves it: the commit carries two check runs under one
 name and nothing says which one the ruleset required. Both being green hides the
 clash completely, so the context quietly stops meaning what the ruleset intended.
+One job lands in the same place when its `name:` holds no matrix value, since
+every combination it runs then reports under that one name.
 
 `.github/workflows/auto-close-duplicates.yml` shipped a job id `test` while
 `.github/workflows/test-mcp.yml` already published the required `test` context,
@@ -29,16 +31,17 @@ combination is filled in is one GitHub resolves per job, so it is one of those:
 guessing that two jobs sharing such a template clash would fail workflows over a
 context this sweep cannot read. A matrix that is itself an expression or that
 lists values which are not scalars, an `include` or `exclude` row shaped the same
-way, and a call this sweep cannot follow, go in the same bucket. The cost is that a real clash hiding behind one of them goes unseen,
-which leaves a merge no worse off than before this check existed, where the
-opposite direction would block work that was fine.
+way, a whole `strategy:` that comes from an expression, and a call this sweep
+cannot follow, go in the same bucket. The cost is that a real clash hiding behind
+one of them goes unseen, which leaves a merge no worse off than before this check
+existed, where the opposite direction would block work that was fine.
 
 A job calling a local reusable workflow publishes one check run per job of the
 callee, named `<caller> / <callee>` and chained through however many levels of
 local calls it takes, which is why a caller's name never collides with a plain
-job that happens to match it. A file under `.github/workflows/` that does not read as
-a workflow at all is reported rather than skipped, since skipping it silently
-would hide every job it holds.
+job that happens to match it. A file under `.github/workflows/` that does not
+read as one workflow at all is reported rather than skipped, since skipping it
+silently would hide every job it holds.
 """
 
 import itertools
@@ -90,7 +93,7 @@ class Names:
 class Job(BaseModel):
     name: object = None
     uses: str | None = None
-    strategy: Mapping[str, object] = Field(default_factory=dict)
+    strategy: object = Field(default_factory=dict)
 
 
 class Workflow(BaseModel):
@@ -104,7 +107,10 @@ def scalar_text(value: object) -> str:
 
 def parse(source: str) -> tuple[Workflow, object] | Unreadable:
     """The workflow plus its raw `on:` value, or why the file does not read as one."""
-    parsed: Final = yaml.safe_load(source)
+    try:
+        parsed: Final = yaml.safe_load(source)
+    except yaml.YAMLError:
+        return Unreadable("it does not read as one YAML document")
     if not isinstance(parsed, dict):
         return Unreadable("its top level is not a mapping of workflow keys")
     try:
@@ -186,6 +192,8 @@ def crossed_values(listed: Sequence[tuple[str, tuple[str, ...]]]) -> tuple[Mappi
 
 def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...] | Opaque:
     """One mapping per job the matrix produces, `exclude` applied before `include` as GitHub does."""
+    if not isinstance(job.strategy, Mapping):
+        return Opaque("its whole `strategy` comes from an expression")
     matrix: Final = job.strategy.get("matrix")
     if matrix is None:
         return ()
@@ -304,7 +312,7 @@ def expand(template: str, job: Job) -> Names:
     if isinstance(combinations, Opaque):
         return Names((), (combinations.reason,))
     over: Final = combinations or (NO_MATRIX,)
-    return settled(tuple(dict.fromkeys(rendered(template, values) for values in over)))
+    return settled(tuple(rendered(template, values) for values in over))
 
 
 def suffixed(job_id: str, combination: Mapping[str, str]) -> str:
@@ -408,13 +416,25 @@ def owners_by_name(sources: Mapping[str, str]) -> Iterator[tuple[str, tuple[str,
         yield name, tuple(owner for _, owner in pairs)
 
 
+def clash(name: str, owners: Sequence[str]) -> str | None:
+    """Why one name is ambiguous, whether two jobs carry it or one job repeats it over its matrix."""
+    jobs: Final = tuple(dict.fromkeys(owners))
+    if len(jobs) > 1:
+        return (
+            f"`{name}` is published by {len(jobs)} jobs: {', '.join(jobs)}. A required status check matching "
+            f"that name cannot say which job proves it; give one of them a distinct `name:` or job id."
+        )
+    if len(owners) > 1:
+        return (
+            f"`{name}` is published {len(owners)} times by {jobs[0]}, once per matrix combination. A required "
+            f"status check matching that name cannot say which run proves it; put a matrix value in its `name:`."
+        )
+    return None
+
+
 def collisions(sources: Mapping[str, str]) -> tuple[str, ...]:
-    return tuple(
-        f"`{name}` is published by {len(owners)} jobs: {', '.join(owners)}. A required status check matching "
-        f"that name cannot say which job proves it; give one of them a distinct `name:` or job id."
-        for name, owners in owners_by_name(sources)
-        if len(owners) > 1
-    )
+    found: Final = tuple(clash(name, owners) for name, owners in owners_by_name(sources))
+    return tuple(message for message in found if message is not None)
 
 
 def workflow_sources() -> Mapping[str, str]:
diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
index 44a5adbbb33..ab6b1d2965e 100644
--- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
@@ -699,3 +699,61 @@ def test_an_include_row_naming_a_listed_key_extends_only_the_combinations_it_mat
 
     assert frozenset(name for name, _ in published(sources)) == frozenset({"3.12 fast"})
     assert len(blind_spots(sources)) == 1
+
+
+def test_a_job_whose_whole_strategy_is_an_expression_is_reported_rather_than_rejecting_the_file() -> None:
+    sources: Final = {
+        "plan.yml": (
+            "on: pull_request\njobs:\n  plan:\n    name: Plan\n    runs-on: ubuntu-latest\n"
+            "  fan:\n    strategy: ${{ fromJSON(needs.plan.outputs.strategy) }}\n    runs-on: ubuntu-latest\n"
+        )
+    }
+
+    assert unreadable(sources) == ()
+    assert frozenset(name for name, _ in published(sources)) == frozenset({"Plan"})
+    assert "`strategy` comes from an expression" in blind_spots(sources)[0]
+    assert exit_code(sources) == 0
+
+
+def test_one_job_publishing_one_name_for_every_matrix_combination_is_a_collision() -> None:
+    sources: Final = {
+        "unit.yml": (
+            "on: pull_request\njobs:\n  build:\n    name: Run tests\n    runs-on: ubuntu-latest\n"
+            '    strategy:\n      matrix:\n        python-version: ["3.12", "3.13"]\n'
+        )
+    }
+
+    found: Final = collisions(sources)
+    assert len(found) == 1
+    assert "`Run tests` is published 2 times by unit.yml job `build`" in found[0]
+    assert exit_code(sources) == 1
+
+
+def test_a_name_carrying_a_matrix_value_publishes_one_name_per_combination_without_colliding() -> None:
+    sources: Final = {
+        "unit.yml": (
+            "on: pull_request\njobs:\n  build:\n    name: Run tests ${{ matrix.python-version }}\n"
+            '    runs-on: ubuntu-latest\n    strategy:\n      matrix:\n        python-version: ["3.12", "3.13"]\n'
+        )
+    }
+
+    assert frozenset(name for name, _ in published(sources)) == frozenset({"Run tests 3.12", "Run tests 3.13"})
+    assert collisions(sources) == ()
+    assert exit_code(sources) == 0
+
+
+def test_a_file_that_is_not_valid_yaml_is_reported_rather_than_raising() -> None:
+    sources: Final = {"broken.yml": "jobs:\n  build: [\n"}
+
+    assert unreadable(sources) == (
+        "broken.yml sits in the workflows directory but it does not read as one YAML "
+        "document, so none of its jobs were checked.",
+    )
+    assert exit_code(sources) == 1
+
+
+def test_a_file_holding_two_yaml_documents_is_reported_rather_than_raising() -> None:
+    sources: Final = {"two.yml": "on: pull_request\n---\non: push\n"}
+
+    assert len(unreadable(sources)) == 1
+    assert exit_code(sources) == 1

From 5a27e11263b2e267475ab39e287ead25920d43d9 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 6 Sep 2026 02:32:50 -0700
Subject: [PATCH 107/164] fix(ci): compare names one workflow run settles the
 same way

A `name:` whose only leftover expressions read a `github.` property other
than `github.job` is filled in identically for every job of the run that
publishes it, so two jobs of one workflow carrying it land on the same
check run. Those names now compare against the other jobs of their own
file instead of sitting in the blind-spot bucket. They stay out of the
comparison across files, where two workflows can run on different events
---
 .../check_workflow_job_name_collisions.py     | 67 ++++++++++++--
 .../test_workflow_job_name_collisions.py      | 91 ++++++++++++++++++-
 2 files changed, 148 insertions(+), 10 deletions(-)

diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
index 941039e8340..482f8386536 100644
--- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
@@ -27,12 +27,16 @@ real `<shard> / Run tests` names rather than staying opaque.
 Whatever the sweep cannot work out is left out of the comparison and reported
 instead of guessed, because a guess that lands wrong fails a workflow GitHub
 would have published perfectly well. A name still holding an expression once the
-combination is filled in is one GitHub resolves per job, so it is one of those:
-guessing that two jobs sharing such a template clash would fail workflows over a
-context this sweep cannot read. A matrix that is itself an expression or that
-lists values which are not scalars, an `include` or `exclude` row shaped the same
-way, a whole `strategy:` that comes from an expression, and a call this sweep
-cannot follow, go in the same bucket. The cost is that a real clash hiding behind
+combination is filled in is usually one GitHub resolves per job, so it is one of
+those: guessing that two jobs sharing such a template clash would fail workflows
+over a context this sweep cannot read. The exception is a name whose leftover
+expressions all read a `github.` property other than `github.job`, which one run
+fills in the same way for every job in it, so those are compared against the
+other jobs of their own workflow and stay out of the comparison across files,
+where two workflows can run on different events. A matrix that is itself an
+expression or that lists values which are not scalars, an `include` or `exclude`
+row shaped the same way, a whole `strategy:` that comes from an expression, and a
+call this sweep cannot follow, go in the same bucket. The cost is that a real clash hiding behind
 one of them goes unseen, which leaves a merge no worse off than before this check
 existed, where the opposite direction would block work that was fine.
 
@@ -64,6 +68,7 @@ MATRIX_REF: Final = re.compile(r"^matrix\.(?P<key>[\w-]+)$")
 LITERAL: Final = re.compile(r"^'(?P<text>[^']*)'$")
 FORMAT_CALL: Final = re.compile(r"^format\((?P<args>.*)\)$", re.DOTALL)
 COMPARISON: Final = re.compile(r"^(?P<left>.+?)\s*(?P<operator>==|!=)\s*(?P<right>.+)$", re.DOTALL)
+RUN_WIDE: Final = re.compile(r"^github\.(?!job\b)[\w.]+$")
 GITHUB_PLACEHOLDER: Final = re.compile(r"\{\{|\}\}|\{\d+\}")
 NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({})
 NO_CALLERS: Final[frozenset[str]] = frozenset()
@@ -88,6 +93,7 @@ class Names:
 
     known: tuple[str, ...] = ()
     unknown: tuple[str, ...] = ()
+    local: tuple[str, ...] = ()
 
 
 class Job(BaseModel):
@@ -300,10 +306,17 @@ def comparable(name: str) -> bool:
     return EXPRESSION.search(name) is None
 
 
+def run_wide(name: str) -> bool:
+    """A name whose leftover expressions one workflow run fills in the same way for every job in it."""
+    return all(RUN_WIDE.match(span.group("body").strip()) is not None for span in EXPRESSION.finditer(name))
+
+
 def settled(names: Sequence[str]) -> Names:
+    unresolved: Final = tuple(name for name in names if not comparable(name))
     return Names(
         tuple(name for name in names if comparable(name)),
-        tuple(f"its name stays `{name}`" for name in names if not comparable(name)),
+        tuple(f"its name stays `{name}`" for name in unresolved if not run_wide(name)),
+        tuple(name for name in unresolved if run_wide(name)),
     )
 
 
@@ -340,9 +353,15 @@ def joined(groups: Sequence[Names]) -> Names:
     return Names(
         tuple(name for group in groups for name in group.known),
         tuple(reason for group in groups for reason in group.unknown),
+        tuple(name for group in groups for name in group.local),
     )
 
 
+def tagged(names: Names) -> tuple[tuple[str, bool], ...]:
+    """Each name a job publishes beside whether only its own workflow's run settles it."""
+    return (*((name, False) for name in names.known), *((name, True) for name in names.local))
+
+
 def call_blocker(job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str]) -> str | None:
     path: Final = callee_path(job)
     if path is None:
@@ -366,9 +385,15 @@ def job_names(job_id: str, job: Job, workflows: Mapping[str, Workflow], callers:
             for callee_id, callee_job in workflows[path].jobs.items()
         )
     )
+    composed: Final = tuple(
+        (f"{prefix} / {suffix}", prefix_local or suffix_local)
+        for prefix, prefix_local in tagged(prefixes)
+        for suffix, suffix_local in tagged(suffixes)
+    )
     return Names(
-        tuple(f"{prefix} / {suffix}" for prefix in prefixes.known for suffix in suffixes.known),
+        tuple(name for name, is_local in composed if not is_local),
         (*prefixes.unknown, *suffixes.unknown),
+        tuple(name for name, is_local in composed if is_local),
     )
 
 
@@ -432,9 +457,33 @@ def clash(name: str, owners: Sequence[str]) -> str | None:
     return None
 
 
+def local_published(sources: Mapping[str, str]) -> Iterator[tuple[tuple[str, str], str]]:
+    """Names their own workflow's run settles, keyed by the file whose run settles them."""
+    for rel, job_id, names in scanned_jobs(sources):
+        for name in names.local:
+            yield (rel, name), f"job `{job_id}`"
+
+
+def local_clash(rel: str, name: str, owners: Sequence[str]) -> str | None:
+    """Why one workflow's own run lands several of its jobs on one check run."""
+    if len(owners) < 2:
+        return None
+    jobs: Final = tuple(dict.fromkeys(owners))
+    return (
+        f"`{name}` is published {len(owners)} times inside {rel}, by {', '.join(jobs)}. One run fills that "
+        f"expression in the same way throughout, so they all land on one check run; make the names differ."
+    )
+
+
+def local_clashes(sources: Mapping[str, str]) -> tuple[str, ...]:
+    grouped: Final = itertools.groupby(sorted(local_published(sources)), key=operator.itemgetter(0))
+    found: Final = tuple(local_clash(rel, name, tuple(owner for _, owner in pairs)) for (rel, name), pairs in grouped)
+    return tuple(message for message in found if message is not None)
+
+
 def collisions(sources: Mapping[str, str]) -> tuple[str, ...]:
     found: Final = tuple(clash(name, owners) for name, owners in owners_by_name(sources))
-    return tuple(message for message in found if message is not None)
+    return (*(message for message in found if message is not None), *local_clashes(sources))
 
 
 def workflow_sources() -> Mapping[str, str]:
diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
index ab6b1d2965e..aeb10985e1b 100644
--- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
@@ -302,7 +302,7 @@ def test_a_matrix_list_supplies_values_the_same_way_include_rows_do() -> None:
     assert "`Analyze (go)` is published by 2 jobs" in found[0]
 
 
-def test_two_jobs_sharing_a_template_over_the_run_itself_are_reported_rather_than_guessed() -> None:
+def test_two_workflows_sharing_a_run_wide_template_are_not_called_a_collision() -> None:
     template: Final = (
         "on: pull_request\njobs:\n  {job}:\n    name: ${{{{ github.event_name }}}}-build\n    runs-on: ubuntu-latest\n"
     )
@@ -311,10 +311,99 @@ def test_two_jobs_sharing_a_template_over_the_run_itself_are_reported_rather_tha
         "b.yml": template.format(job="two"),
     }
 
+    assert collisions(sources) == ()
+    assert blind_spots(sources) == ()
+
+
+def test_two_jobs_of_one_workflow_sharing_a_run_wide_template_are_a_collision() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\njobs:\n"
+            "  one:\n    name: ${{ github.event_name }}-build\n    runs-on: ubuntu-latest\n"
+            "  two:\n    name: ${{ github.event_name }}-build\n    runs-on: ubuntu-latest\n"
+        ),
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "is published 2 times inside a.yml, by job `one`, job `two`" in found[0]
+    assert exit_code(sources) == 1
+
+
+def test_a_run_wide_template_carrying_a_matrix_value_does_not_collide_inside_one_workflow() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\njobs:\n"
+            "  one:\n    name: ${{ github.event_name }}-${{ matrix.shard }}\n    runs-on: ubuntu-latest\n"
+            "    strategy:\n      matrix:\n        shard: [core, extras]\n"
+        ),
+    }
+
+    assert collisions(sources) == ()
+    assert blind_spots(sources) == ()
+
+
+def test_a_run_wide_name_repeated_over_a_matrix_by_one_job_is_a_collision() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\njobs:\n"
+            "  one:\n    name: ${{ github.event_name }}-build\n    runs-on: ubuntu-latest\n"
+            "    strategy:\n      matrix:\n        shard: [core, extras]\n"
+        ),
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "is published 2 times inside a.yml, by job `one`" in found[0]
+
+
+def test_a_name_reading_the_job_it_sits_in_stays_out_of_the_comparison() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\njobs:\n"
+            "  one:\n    name: ${{ github.job }}-build\n    runs-on: ubuntu-latest\n"
+            "  two:\n    name: ${{ github.job }}-build\n    runs-on: ubuntu-latest\n"
+        ),
+    }
+
     assert collisions(sources) == ()
     assert len(blind_spots(sources)) == 2
 
 
+def test_a_run_wide_caller_name_collides_through_the_workflow_it_calls() -> None:
+    sources: Final = {
+        ".github/workflows/a.yml": (
+            "on: pull_request\njobs:\n"
+            "  one:\n    name: ${{ github.event_name }}\n    uses: ./.github/workflows/c.yml\n"
+            "  two:\n    name: ${{ github.event_name }}\n    uses: ./.github/workflows/c.yml\n"
+        ),
+        ".github/workflows/c.yml": "on:\n  workflow_call:\njobs:\n  build:\n    runs-on: ubuntu-latest\n",
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "github.event_name }} / build` is published 2 times inside .github/workflows/a.yml" in found[0]
+
+
+def test_a_run_wide_name_inside_a_called_workflow_collides_under_the_caller() -> None:
+    sources: Final = {
+        ".github/workflows/a.yml": ("on: pull_request\njobs:\n  one:\n    uses: ./.github/workflows/c.yml\n"),
+        ".github/workflows/c.yml": (
+            "on:\n  workflow_call:\njobs:\n"
+            "  build:\n    name: ${{ github.event_name }}\n    runs-on: ubuntu-latest\n"
+            "  lint:\n    name: ${{ github.event_name }}\n    runs-on: ubuntu-latest\n"
+        ),
+    }
+
+    found: Final = collisions(sources)
+
+    assert len(found) == 1
+    assert "`one / ${{ github.event_name }}` is published 2 times inside .github/workflows/a.yml" in found[0]
+
+
 def test_a_name_reading_the_workflow_it_sits_in_is_not_called_a_collision() -> None:
     template: Final = (
         "on: pull_request\njobs:\n  {job}:\n    name: ${{{{ github.workflow }}}} / build\n    runs-on: ubuntu-latest\n"

From 274a489c463452c78e3e69db2017ec2e3ad59315 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Sun, 6 Sep 2026 02:36:46 -0700
Subject: [PATCH 108/164] test(e2e/ui): automate the RC checklist's Presidio
 guardrail walk

The guardrail section of the release checklist is done by hand every cut:
create a Presidio guardrail through the wizard, send a sentence with PII from
the playground, then open Logs and check the guardrail caught it. Nothing
covered that path, so a break anywhere along it surfaced only when someone
happened to repeat the steps.

Adds a spec that walks it once and turns the eyeball checks into assertions.
The strongest of them is the leak check: it reads the request back and fails if
the stored prompt still carries the raw address or number, which is what the
manual step is really looking for.

The stack gains a Presidio stand-in that answers the two routes the guardrail
calls, detecting a fixed regex set with a Luhn check on card numbers. Real
Presidio's detection quality is Presidio's business, and pinning the UI lane to
it would mean two heavy containers with spaCy models on every CI run for a test
that is about LiteLLM's integration. The real analyzer stays covered in the
Python lane. The stand-in returns the same entities at the same spans as the
real one for the checklist's sentence, and driving the real guardrail against it
produces the same record shape, so a test written against it is written against
the product's real behavior.

Making the stand-in return the text unmasked turns the spec red on the raw
address reaching the spend log, so the leak assertion reads live data.

run_e2e.sh and both CircleCI UI jobs start the stand-in alongside the mock LLM.
Its port is overridable like the others so two checkouts can run at once.
---
 .circleci/config.yml                          |   8 +
 tests/e2e/ui/constants.ts                     |   2 +
 .../fixtures/mock_presidio_server/server.py   | 107 ++++++++++++
 tests/e2e/ui/run_e2e.sh                       |  21 ++-
 .../guardrails/presidioUserStory.spec.ts      | 158 ++++++++++++++++++
 5 files changed, 292 insertions(+), 4 deletions(-)
 create mode 100644 tests/e2e/ui/fixtures/mock_presidio_server/server.py
 create mode 100644 tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 6e368a3debe..276555c0cf7 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -2648,6 +2648,10 @@ jobs:
           name: Start mock LLM server
           command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py
           background: true
+      - run:
+          name: Start mock Presidio server
+          command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py
+          background: true
       - run:
           name: Start LiteLLM proxy
           environment:
@@ -2778,6 +2782,10 @@ jobs:
           name: Start mock LLM server
           command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py
           background: true
+      - run:
+          name: Start mock Presidio server
+          command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py
+          background: true
       - run:
           name: Start LiteLLM proxy under a server root path
           environment:
diff --git a/tests/e2e/ui/constants.ts b/tests/e2e/ui/constants.ts
index 3a62252915c..71774c95d24 100644
--- a/tests/e2e/ui/constants.ts
+++ b/tests/e2e/ui/constants.ts
@@ -15,6 +15,8 @@ export const UI_BASE_URL = (
 // writable path (the image runner already exports TMPDIR) to relocate them.
 export const ARTIFACT_DIR = process.env.E2E_UI_ARTIFACT_DIR || ".";
 
+export const MOCK_PRESIDIO_URL = (process.env.E2E_MOCK_PRESIDIO_URL || "http://127.0.0.1:8091").replace(/\/+$/, "");
+
 const storagePath = (name: string): string => path.join(ARTIFACT_DIR, name);
 
 // Storage state paths for each role
diff --git a/tests/e2e/ui/fixtures/mock_presidio_server/server.py b/tests/e2e/ui/fixtures/mock_presidio_server/server.py
new file mode 100644
index 00000000000..0fc97c4a065
--- /dev/null
+++ b/tests/e2e/ui/fixtures/mock_presidio_server/server.py
@@ -0,0 +1,107 @@
+"""
+Mock Presidio analyzer + anonymizer for UI e2e tests.
+Serves POST /analyze and POST /anonymize over a fixed set of regex recognizers.
+"""
+
+import os
+import re
+
+import uvicorn
+from fastapi import FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+
+
+def _luhn_ok(candidate):
+    digits = [int(c) for c in candidate if c.isdigit()]
+    doubled = [d * 2 - 9 if d * 2 > 9 else d * 2 for d in digits[-2::-2]]
+    return len(digits) >= 13 and (sum(digits[::-1][::2]) + sum(doubled)) % 10 == 0
+
+
+RECOGNIZERS = (
+    ("EMAIL_ADDRESS", re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.-]+\b"), 1.0, None),
+    ("CREDIT_CARD", re.compile(r"\b(?:\d[ -]?){13,19}\b"), 1.0, _luhn_ok),
+    ("US_SSN", re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), 0.85, None),
+    ("PHONE_NUMBER", re.compile(r"\b(?:\+?\d{1,2}[ .-]?)?(?:\(\d{3}\)|\d{3})[ .-]?\d{3}[ .-]?\d{4}\b"), 0.75, None),
+    ("IP_ADDRESS", re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"), 0.6, None),
+    ("URL", re.compile(r"\bhttps?://[^\s]+"), 0.5, None),
+)
+
+app = FastAPI(title="Mock Presidio Server")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+def _detect(text, wanted):
+    found = [
+        {"entity_type": name, "start": m.start(), "end": m.end(), "score": score}
+        for name, pattern, score, validate in RECOGNIZERS
+        if wanted is None or name in wanted
+        for m in pattern.finditer(text)
+        if validate is None or validate(m.group())
+    ]
+    ranked = sorted(found, key=lambda r: (-r["score"], r["start"]))
+    kept = []
+    for candidate in ranked:
+        overlaps = any(candidate["start"] < k["end"] and k["start"] < candidate["end"] for k in kept)
+        if not overlaps:
+            kept.append(candidate)
+    return sorted(kept, key=lambda r: r["start"])
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok"}
+
+
+@app.post("/analyze")
+async def analyze(request: Request):
+    body = await request.json()
+    text = body.get("text", "") or ""
+    entities = body.get("entities")
+    wanted = set(entities) if entities else None
+    threshold = body.get("score_threshold") or 0.0
+    return [
+        {**result, "analysis_explanation": None, "recognition_metadata": {"recognizer_name": "MockRecognizer"}}
+        for result in _detect(text, wanted)
+        if result["score"] >= threshold
+    ]
+
+
+@app.post("/anonymize")
+async def anonymize(request: Request):
+    body = await request.json()
+    text = body.get("text", "") or ""
+    results = sorted(body.get("analyzer_results") or [], key=lambda r: r["start"])
+
+    pieces = []
+    items = []
+    cursor = 0
+    for result in results:
+        start, end = result["start"], result["end"]
+        if start < cursor:
+            continue
+        placeholder = f"<{result['entity_type']}>"
+        pieces.append(text[cursor:start])
+        masked_start = sum(len(p) for p in pieces)
+        pieces.append(placeholder)
+        items.append(
+            {
+                "operator": "replace",
+                "entity_type": result["entity_type"],
+                "start": masked_start,
+                "end": masked_start + len(placeholder),
+                "text": placeholder,
+            }
+        )
+        cursor = end
+    pieces.append(text[cursor:])
+
+    return {"text": "".join(pieces), "items": list(reversed(items))}
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="127.0.0.1", port=int(os.environ.get("MOCK_PRESIDIO_PORT", "8091")))
diff --git a/tests/e2e/ui/run_e2e.sh b/tests/e2e/ui/run_e2e.sh
index a369d8c8f7a..ad0454b6347 100755
--- a/tests/e2e/ui/run_e2e.sh
+++ b/tests/e2e/ui/run_e2e.sh
@@ -14,7 +14,7 @@ set -euo pipefail
 #
 # Ports default to 4000 / 5432 / 8090 and can be moved when another checkout
 # already holds them:
-#   PROXY_PORT=4100 POSTGRES_PORT=5532 MOCK_LLM_PORT=8190 ./run_e2e.sh
+#   PROXY_PORT=4100 POSTGRES_PORT=5532 MOCK_LLM_PORT=8190 MOCK_PRESIDIO_PORT=8191 ./run_e2e.sh
 #
 # In CI (CI=true), expects:
 #   - PostgreSQL already running on 127.0.0.1:5432
@@ -29,6 +29,7 @@ DASHBOARD_DIR="$REPO_ROOT/ui/litellm-dashboard"
 IS_CI="${CI:-false}"
 CONTAINER_NAME="litellm-e2e-postgres-$$"
 MOCK_PID=""
+MOCK_PRESIDIO_PID=""
 PROXY_PID=""
 PROXY_LOG=""
 
@@ -40,7 +41,8 @@ PROXY_LOG=""
 PROXY_PORT="${PROXY_PORT:-4000}"
 POSTGRES_PORT="${POSTGRES_PORT:-5432}"
 MOCK_LLM_PORT="${MOCK_LLM_PORT:-8090}"
-export MOCK_LLM_PORT
+MOCK_PRESIDIO_PORT="${MOCK_PRESIDIO_PORT:-8091}"
+export MOCK_LLM_PORT MOCK_PRESIDIO_PORT
 
 # --- Ensure common tool paths are available (local dev only) ---
 if [ "$IS_CI" = "false" ]; then
@@ -82,6 +84,7 @@ fi
 cleanup() {
   echo "Cleaning up..."
   [ -n "$MOCK_PID" ] && kill "$MOCK_PID" 2>/dev/null || true
+  [ -n "$MOCK_PRESIDIO_PID" ] && kill "$MOCK_PRESIDIO_PID" 2>/dev/null || true
   [ -n "$PROXY_PID" ] && kill "$PROXY_PID" 2>/dev/null || true
   [ -n "$PROXY_LOG" ] && rm -f "$PROXY_LOG" || true
   if [ "$IS_CI" = "false" ]; then
@@ -110,9 +113,9 @@ if [ "$IS_CI" = "false" ]; then
   # to someone else's :5432 (a psql session, a running app, a Prisma engine
   # talking to a remote database) aborts the run with "port 5432 is in use"
   # while nothing is actually bound locally.
-  for port in "$PROXY_PORT" "$POSTGRES_PORT" "$MOCK_LLM_PORT"; do
+  for port in "$PROXY_PORT" "$POSTGRES_PORT" "$MOCK_LLM_PORT" "$MOCK_PRESIDIO_PORT"; do
     if lsof -nP -iTCP:"$port" -sTCP:LISTEN >/dev/null 2>&1; then
-      echo "Error: port $port is in use (override with PROXY_PORT / POSTGRES_PORT / MOCK_LLM_PORT)"
+      echo "Error: port $port is in use (override with PROXY_PORT / POSTGRES_PORT / MOCK_LLM_PORT / MOCK_PRESIDIO_PORT)"
       exit 1
     fi
   done
@@ -143,6 +146,7 @@ fi
 # --- Credentials ---
 export LITELLM_MASTER_KEY="sk-1234"
 export MOCK_LLM_URL="http://127.0.0.1:${MOCK_LLM_PORT}/v1"
+export E2E_MOCK_PRESIDIO_URL="http://127.0.0.1:${MOCK_PRESIDIO_PORT}"
 export DISABLE_SCHEMA_UPDATE="true"
 # The suite resolves its target from E2E_UI_BASE_URL (constants.ts), which
 # otherwise defaults to :4000 -- so without this a relocated stack would be
@@ -204,6 +208,15 @@ for i in $(seq 1 15); do
   sleep 1
 done
 
+echo "=== Starting mock Presidio server ==="
+uv run --no-sync python "$SCRIPT_DIR/fixtures/mock_presidio_server/server.py" &
+MOCK_PRESIDIO_PID=$!
+
+for i in $(seq 1 15); do
+  if curl -sf http://127.0.0.1:${MOCK_PRESIDIO_PORT}/health >/dev/null 2>&1; then break; fi
+  sleep 1
+done
+
 # --- LiteLLM proxy ---
 echo "=== Starting LiteLLM proxy ==="
 cd "$REPO_ROOT"
diff --git a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts
new file mode 100644
index 00000000000..975b343287a
--- /dev/null
+++ b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts
@@ -0,0 +1,158 @@
+import { test, expect, type Page as PlaywrightPage } from "@playwright/test";
+import { ADMIN_STORAGE_PATH, MOCK_PRESIDIO_URL } from "../../constants";
+import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation";
+import { Page } from "../../fixtures/pages";
+import { CHAT_MODEL_A, masterKey, rootPath, waitForSpendLogByPrompt } from "../../helpers/traffic";
+import { openPlayground, selectModel, sendButton, onlyVisible } from "../../helpers/playground";
+
+const RAW_EMAIL = "jane.doe@example.com";
+const RAW_PHONE = "555-867-5309";
+
+const visibleTestId = (page: PlaywrightPage, id: string) => page.getByTestId(id).filter({ visible: true });
+
+const requestLogsRows = (page: PlaywrightPage) =>
+  page.locator("table").filter({ visible: true }).first().locator("tbody tr");
+
+async function createPresidioGuardrail(page: PlaywrightPage, guardrailName: string): Promise<void> {
+  await navigateToPage(page, Page.Guardrails);
+  await dismissFeedbackPopup(page);
+
+  await page.getByRole("button", { name: /Add New Guardrail/i }).click();
+  await page.getByRole("menuitem", { name: "Add Provider Guardrail" }).click();
+
+  const dialog = page.getByRole("dialog", { name: "Create guardrail" });
+  await expect(dialog).toBeVisible({ timeout: 10_000 });
+
+  await dialog.getByLabel("Guardrail Name").fill(guardrailName);
+
+  const providerSelect = dialog.getByRole("combobox", { name: "Guardrail Provider" });
+  await providerSelect.click();
+  await providerSelect.fill("Presidio");
+  await page.getByRole("option", { name: "Presidio PII" }).click();
+
+  await dialog.getByLabel("Mode", { exact: true }).click();
+  await page.keyboard.type("pre_call");
+  await expect(page.getByRole("option", { name: "pre_call" })).toBeAttached({ timeout: 5_000 });
+  await page.keyboard.press("Enter");
+  await expect(dialog.getByText("pre_call", { exact: true })).toBeVisible({ timeout: 5_000 });
+  await dialog.getByText("Create guardrail", { exact: true }).click();
+
+  await dialog.getByLabel("presidio_analyzer_api_base").fill(MOCK_PRESIDIO_URL);
+  await dialog.getByLabel("presidio_anonymizer_api_base").fill(MOCK_PRESIDIO_URL);
+
+  await dialog.getByRole("button", { name: "Next" }).click();
+  await expect(dialog.getByText("Configure PII Protection")).toBeVisible({ timeout: 10_000 });
+  await dialog.getByRole("button", { name: "Select All & Mask" }).click();
+
+  await dialog.getByRole("button", { name: "Create Guardrail" }).click();
+  await expect(page.getByText("Guardrail created successfully").first()).toBeVisible({ timeout: 15_000 });
+  await expect(page.getByRole("row").filter({ hasText: guardrailName })).toHaveCount(1, { timeout: 15_000 });
+}
+
+async function deleteGuardrail(page: PlaywrightPage, guardrailName: string): Promise<void> {
+  await navigateToPage(page, Page.Guardrails);
+  await dismissFeedbackPopup(page);
+  const row = page.getByRole("row").filter({ hasText: guardrailName });
+  await expect(row).toHaveCount(1, { timeout: 15_000 });
+  await row.getByRole("button", { name: "Open guardrail actions" }).click();
+  await page.getByRole("menuitem", { name: "Delete" }).click();
+  const deleteModal = page.getByRole("dialog", { name: "Delete Guardrail" });
+  await expect(deleteModal).toBeVisible({ timeout: 5_000 });
+  await deleteModal.getByRole("button", { name: "Delete", exact: true }).click();
+  await expect(page.getByText(`Guardrail "${guardrailName}" deleted successfully`)).toBeVisible({ timeout: 10_000 });
+}
+
+test.describe("Presidio PII guardrail, end to end from the dashboard", () => {
+  test.use({ storageState: ADMIN_STORAGE_PATH });
+
+  test("masks PII sent from the Playground and shows the run in Logs", async ({ page, request }) => {
+    const guardrailName = `e2e-presidio-story-${Date.now()}`;
+    const marker = `case-ref-${Math.random().toString(36).slice(2, 10)}`;
+    const prompt = `${marker}. Email me at ${RAW_EMAIL} or call ${RAW_PHONE}.`;
+
+    await createPresidioGuardrail(page, guardrailName);
+
+    await openPlayground(page);
+    await selectModel(page, CHAT_MODEL_A);
+
+    const guardrailSelect = onlyVisible(page.getByPlaceholder("Select guardrails"));
+    await expect(guardrailSelect).toBeVisible({ timeout: 20_000 });
+    await guardrailSelect.click();
+    await guardrailSelect.fill(guardrailName);
+    await onlyVisible(page.getByRole("option", { name: guardrailName, exact: true })).click({ timeout: 20_000 });
+    await page.keyboard.press("Escape");
+
+    const input = onlyVisible(page.getByPlaceholder("Type your message", { exact: false }));
+    await expect(input).toBeVisible({ timeout: 15_000 });
+
+    await expect
+      .poll(
+        async () => {
+          await input.fill(prompt);
+          await sendButton(page).click();
+          const res = await request.get(`${rootPath()}/spend/logs`, {
+            headers: { Authorization: `Bearer ${masterKey()}` },
+          });
+          if (!res.ok()) return false;
+          const rows: { metadata?: { applied_guardrails?: string[] } }[] = await res.json();
+          return (Array.isArray(rows) ? rows : []).some((row) =>
+            (row.metadata?.applied_guardrails ?? []).includes(guardrailName),
+          );
+        },
+        {
+          message: `the playground never produced a request that ran ${guardrailName}`,
+          timeout: 90_000,
+          intervals: [5_000],
+        },
+      )
+      .toBe(true);
+
+    const requestId = await waitForSpendLogByPrompt(request, marker);
+
+    const stored = await request.get(`${rootPath()}/spend/logs?request_id=${requestId}`, {
+      headers: { Authorization: `Bearer ${masterKey()}` },
+    });
+    expect(stored.ok(), `spend log read failed: ${stored.status()}`).toBe(true);
+    const storedBody = JSON.stringify(await stored.json());
+    expect(storedBody, "the raw email reached the spend log, so the prompt was stored unmasked").not.toContain(
+      RAW_EMAIL,
+    );
+    expect(storedBody, "the raw phone number reached the spend log, so the prompt was stored unmasked").not.toContain(
+      RAW_PHONE,
+    );
+    expect(storedBody, "the stored prompt carries no EMAIL_ADDRESS placeholder, so nothing was masked").toContain(
+      "<EMAIL_ADDRESS>",
+    );
+    expect(storedBody, "the stored prompt carries no PHONE_NUMBER placeholder, so nothing was masked").toContain(
+      "<PHONE_NUMBER>",
+    );
+
+    await navigateToPage(page, Page.Logs);
+    await dismissFeedbackPopup(page);
+    const search = visibleTestId(page, "datatable-search");
+    await expect(search).toBeVisible({ timeout: 20_000 });
+    await search.fill(requestId);
+    const row = requestLogsRows(page).filter({ hasText: requestId });
+    await expect(row, `no logs row for request ${requestId}`).toHaveCount(1, { timeout: 30_000 });
+    await row.click();
+
+    const drawer = page.getByRole("dialog").first();
+    await expect(drawer.getByText("Guardrails & Policy Compliance")).toBeVisible({ timeout: 20_000 });
+    await expect(drawer.getByText(`Pre-call guardrail: ${guardrailName}`).first()).toBeVisible({ timeout: 20_000 });
+    await expect(drawer.getByText(`${marker}. Email me at <EMAIL_ADDRESS> or call <PHONE_NUMBER>.`)).toBeVisible({
+      timeout: 20_000,
+    });
+
+    await drawer.getByText("2 matched").first().click();
+    await expect(drawer.getByText("Detected Entities (2)").first()).toBeVisible({ timeout: 10_000 });
+    await expect(drawer.getByText("EMAIL_ADDRESS").first()).toBeVisible({ timeout: 10_000 });
+    await expect(drawer.getByText("PHONE_NUMBER").first()).toBeVisible({ timeout: 10_000 });
+    await expect(drawer.getByText("Score: 1.00").first()).toBeVisible({ timeout: 10_000 });
+    await expect(drawer.getByText("Score: 0.75").first()).toBeVisible({ timeout: 10_000 });
+
+    await expect(drawer.getByText(RAW_EMAIL)).toHaveCount(0);
+    await expect(drawer.getByText(RAW_PHONE)).toHaveCount(0);
+
+    await deleteGuardrail(page, guardrailName);
+  });
+});

From 4edf6f4dc530f88afe09d2bb986a4bdc5a962cc0 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 6 Sep 2026 02:59:16 -0700
Subject: [PATCH 109/164] fix(ci): keep an expression matrix directive out of
 the comparison

`include:` or `exclude:` written as `${{ ... }}` read back as a string, and
the sweep treated that as the directive being absent, so it expanded every
combination GitHub would have dropped. A job whose `name:` holds no matrix
value then looked like it repeated one name across combinations that never
run. An absent directive still means no rows; anything that is not a list
of rows now joins the names left out of the comparison
---
 .../check_workflow_job_name_collisions.py     |  4 ++-
 .../test_workflow_job_name_collisions.py      | 32 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
index 482f8386536..ae2c1d80c8f 100644
--- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py
@@ -162,8 +162,10 @@ def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, .
 def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...] | Opaque:
     """One `include` or `exclude` row, or why the combinations they shape cannot be worked out."""
     rows: Final = matrix.get(directive)
-    if not isinstance(rows, Sequence) or isinstance(rows, str):
+    if rows is None:
         return ()
+    if not isinstance(rows, Sequence) or isinstance(rows, str):
+        return Opaque(f"a matrix `{directive}` is itself an expression rather than a list of rows")
     mappings: Final = tuple(row for row in rows if isinstance(row, Mapping))
     if len(mappings) != len(rows):
         return Opaque(f"a matrix `{directive}` row is not a mapping of values")
diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
index aeb10985e1b..0f5ba43bd7a 100644
--- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py
+++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py
@@ -846,3 +846,35 @@ def test_a_file_holding_two_yaml_documents_is_reported_rather_than_raising() ->
 
     assert len(unreadable(sources)) == 1
     assert exit_code(sources) == 1
+
+
+def test_an_exclude_that_is_itself_an_expression_is_reported_rather_than_ignored() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\njobs:\n  one:\n    name: build\n    runs-on: ubuntu-latest\n"
+            "    strategy:\n      matrix:\n        python: ['3.11', '3.12']\n"
+            "        exclude: ${{ fromJson(vars.SKIP) }}\n"
+        ),
+    }
+
+    found: Final = blind_spots(sources)
+
+    assert collisions(sources) == ()
+    assert len(found) == 1
+    assert "a matrix `exclude` is itself an expression" in found[0]
+
+
+def test_an_include_that_is_itself_an_expression_is_reported_rather_than_ignored() -> None:
+    sources: Final = {
+        "a.yml": (
+            "on: pull_request\njobs:\n  one:\n    name: build-${{ matrix.python }}\n    runs-on: ubuntu-latest\n"
+            "    strategy:\n      matrix:\n        python: ['3.11']\n"
+            "        include: ${{ fromJson(vars.EXTRA) }}\n"
+        ),
+    }
+
+    found: Final = blind_spots(sources)
+
+    assert collisions(sources) == ()
+    assert len(found) == 1
+    assert "a matrix `include` is itself an expression" in found[0]

From b0e53bfbe4e8bbeb1e3941d23728e1b5f8beaeaa Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Sun, 6 Sep 2026 03:13:09 -0700
Subject: [PATCH 110/164] fix(e2e/ui): fail the run when the Presidio fixture
 never comes up

Both readiness loops broke out on success and fell through on timeout, so a
mock Presidio server that failed to bind left the run going with nothing
serving /analyze. The guardrail then errored at request time and the failure
surfaced as an unrelated Playwright assertion in presidioUserStory.spec.ts
rather than as the missing fixture it actually was.

Fail the local runner with the port in the message, and add the matching wait
step to both CircleCI UI jobs, which had no readiness check at all.
---
 .circleci/config.yml    | 18 ++++++++++++++++++
 tests/e2e/ui/run_e2e.sh |  7 ++++++-
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 276555c0cf7..aa851f829e4 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -2652,6 +2652,15 @@ jobs:
           name: Start mock Presidio server
           command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py
           background: true
+      - run:
+          name: Wait for mock Presidio server
+          command: |
+            for i in $(seq 1 30); do
+              if curl -sf http://127.0.0.1:8091/health >/dev/null 2>&1; then exit 0; fi
+              sleep 1
+            done
+            echo "Mock Presidio server never answered /health on port 8091" >&2
+            exit 1
       - run:
           name: Start LiteLLM proxy
           environment:
@@ -2786,6 +2795,15 @@ jobs:
           name: Start mock Presidio server
           command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py
           background: true
+      - run:
+          name: Wait for mock Presidio server
+          command: |
+            for i in $(seq 1 30); do
+              if curl -sf http://127.0.0.1:8091/health >/dev/null 2>&1; then exit 0; fi
+              sleep 1
+            done
+            echo "Mock Presidio server never answered /health on port 8091" >&2
+            exit 1
       - run:
           name: Start LiteLLM proxy under a server root path
           environment:
diff --git a/tests/e2e/ui/run_e2e.sh b/tests/e2e/ui/run_e2e.sh
index ad0454b6347..beb1bc8bf3b 100755
--- a/tests/e2e/ui/run_e2e.sh
+++ b/tests/e2e/ui/run_e2e.sh
@@ -212,10 +212,15 @@ echo "=== Starting mock Presidio server ==="
 uv run --no-sync python "$SCRIPT_DIR/fixtures/mock_presidio_server/server.py" &
 MOCK_PRESIDIO_PID=$!
 
+PRESIDIO_READY=0
 for i in $(seq 1 15); do
-  if curl -sf http://127.0.0.1:${MOCK_PRESIDIO_PORT}/health >/dev/null 2>&1; then break; fi
+  if curl -sf http://127.0.0.1:${MOCK_PRESIDIO_PORT}/health >/dev/null 2>&1; then PRESIDIO_READY=1; break; fi
   sleep 1
 done
+if [ "$PRESIDIO_READY" -ne 1 ]; then
+  echo "Mock Presidio server never answered /health on port ${MOCK_PRESIDIO_PORT}" >&2
+  exit 1
+fi
 
 # --- LiteLLM proxy ---
 echo "=== Starting LiteLLM proxy ==="

From bbc6ca6205c00dcb47a56789b4a639a29ed0e468 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Sun, 6 Sep 2026 07:04:21 -0700
Subject: [PATCH 111/164] fix(e2e/ui): stop the masked-prompt assertion
 tripping strict mode

The Logs drawer renders the masked prompt in three places, so matching it
without narrowing raised a strict mode violation instead of asserting
visibility. CI caught it as a flake: the spec failed its first attempt on
b0e53bfbe4 and passed on retry, which is a locator defect rather than a
timing one and would have gone red on any run that saw all three nodes.

Narrow to the first match, matching the guardrail-name assertion above it.
The leak checks below stay on toHaveCount(0), which is unaffected by
multiple matches and is what actually proves nothing raw reached the drawer.
---
 tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts
index 975b343287a..3ab5b820896 100644
--- a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts
+++ b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts
@@ -139,9 +139,8 @@ test.describe("Presidio PII guardrail, end to end from the dashboard", () => {
     const drawer = page.getByRole("dialog").first();
     await expect(drawer.getByText("Guardrails & Policy Compliance")).toBeVisible({ timeout: 20_000 });
     await expect(drawer.getByText(`Pre-call guardrail: ${guardrailName}`).first()).toBeVisible({ timeout: 20_000 });
-    await expect(drawer.getByText(`${marker}. Email me at <EMAIL_ADDRESS> or call <PHONE_NUMBER>.`)).toBeVisible({
-      timeout: 20_000,
-    });
+    const maskedPrompt = drawer.getByText(`${marker}. Email me at <EMAIL_ADDRESS> or call <PHONE_NUMBER>.`);
+    await expect(maskedPrompt.first()).toBeVisible({ timeout: 20_000 });
 
     await drawer.getByText("2 matched").first().click();
     await expect(drawer.getByText("Detected Entities (2)").first()).toBeVisible({ timeout: 10_000 });

From 70f1894c28398d8c7ceaace23b876a831af59a25 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Sun, 6 Sep 2026 07:21:36 -0700
Subject: [PATCH 112/164] fix(e2e/ui): wait on the visible masked prompt, not
 the first match

The Logs drawer renders the masked prompt in three nodes and the first one
in DOM order is hidden, so the previous commit's .first() traded a strict
mode violation for a locator that waits 20s on an invisible element. Local
runs against a warm stack failed on it every time, resolving the node 22
times and reporting "unexpected value hidden" each time.

Use onlyVisible, the helper this spec already uses for the playground
selectors, which filters to the visible node before taking the first. Solo
runs go 3 for 3 and the guardrails folder passes 5 of 5.

Also drop the 30s wait on the PII step added while chasing this: cold start
was never the cause, and it left the spec sitting on a dead locator longer.
---
 tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts
index 3ab5b820896..ff4d6260312 100644
--- a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts
+++ b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts
@@ -139,8 +139,8 @@ test.describe("Presidio PII guardrail, end to end from the dashboard", () => {
     const drawer = page.getByRole("dialog").first();
     await expect(drawer.getByText("Guardrails & Policy Compliance")).toBeVisible({ timeout: 20_000 });
     await expect(drawer.getByText(`Pre-call guardrail: ${guardrailName}`).first()).toBeVisible({ timeout: 20_000 });
-    const maskedPrompt = drawer.getByText(`${marker}. Email me at <EMAIL_ADDRESS> or call <PHONE_NUMBER>.`);
-    await expect(maskedPrompt.first()).toBeVisible({ timeout: 20_000 });
+    const maskedPrompt = onlyVisible(drawer.getByText(`${marker}. Email me at <EMAIL_ADDRESS> or call <PHONE_NUMBER>.`));
+    await expect(maskedPrompt).toBeVisible({ timeout: 20_000 });
 
     await drawer.getByText("2 matched").first().click();
     await expect(drawer.getByText("Detected Entities (2)").first()).toBeVisible({ timeout: 10_000 });

From fdd1c910963a9161709b1b0eb5faa5e339e7182a Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Sun, 6 Sep 2026 08:34:19 -0700
Subject: [PATCH 113/164] fix(e2e/ui): scope the drawer assertions to the
 visible, exact match

CI flagged the spec flaky twice more. Both were the same defect in different
places: the Logs drawer renders several nodes per string and the first in DOM
order is often hidden, so .first() waited 20s on an invisible element. The
entity assertions had a second problem on top, since getByText("EMAIL_ADDRESS")
substring-matched the masked prompt div, whose text contains <EMAIL_ADDRESS>,
rather than the entity chip.

Route every drawer assertion through onlyVisible, and match the entity type
and score exactly, which is what the panel renders them as: entity_type and
"Score: N.NN" each get their own span.

Verified on a live stack: 6 of 6 solo runs and the guardrails folder 5 of 5.
Mutating the analyzer to detect nothing turns the spec red on the raw address
reaching the spend log, so the assertions still carry their weight.
---
 .../guardrails/presidioUserStory.spec.ts      | 20 +++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts
index ff4d6260312..d4ed5308342 100644
--- a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts
+++ b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts
@@ -137,17 +137,17 @@ test.describe("Presidio PII guardrail, end to end from the dashboard", () => {
     await row.click();
 
     const drawer = page.getByRole("dialog").first();
-    await expect(drawer.getByText("Guardrails & Policy Compliance")).toBeVisible({ timeout: 20_000 });
-    await expect(drawer.getByText(`Pre-call guardrail: ${guardrailName}`).first()).toBeVisible({ timeout: 20_000 });
-    const maskedPrompt = onlyVisible(drawer.getByText(`${marker}. Email me at <EMAIL_ADDRESS> or call <PHONE_NUMBER>.`));
-    await expect(maskedPrompt).toBeVisible({ timeout: 20_000 });
+    await expect(onlyVisible(drawer.getByText("Guardrails & Policy Compliance"))).toBeVisible({ timeout: 20_000 });
+    await expect(onlyVisible(drawer.getByText(`Pre-call guardrail: ${guardrailName}`))).toBeVisible({ timeout: 20_000 });
+    const maskedPrompt = drawer.getByText(`${marker}. Email me at <EMAIL_ADDRESS> or call <PHONE_NUMBER>.`);
+    await expect(onlyVisible(maskedPrompt)).toBeVisible({ timeout: 20_000 });
 
-    await drawer.getByText("2 matched").first().click();
-    await expect(drawer.getByText("Detected Entities (2)").first()).toBeVisible({ timeout: 10_000 });
-    await expect(drawer.getByText("EMAIL_ADDRESS").first()).toBeVisible({ timeout: 10_000 });
-    await expect(drawer.getByText("PHONE_NUMBER").first()).toBeVisible({ timeout: 10_000 });
-    await expect(drawer.getByText("Score: 1.00").first()).toBeVisible({ timeout: 10_000 });
-    await expect(drawer.getByText("Score: 0.75").first()).toBeVisible({ timeout: 10_000 });
+    await onlyVisible(drawer.getByText("2 matched")).click();
+    await expect(onlyVisible(drawer.getByText("Detected Entities (2)"))).toBeVisible({ timeout: 10_000 });
+    await expect(onlyVisible(drawer.getByText("EMAIL_ADDRESS", { exact: true }))).toBeVisible({ timeout: 10_000 });
+    await expect(onlyVisible(drawer.getByText("PHONE_NUMBER", { exact: true }))).toBeVisible({ timeout: 10_000 });
+    await expect(onlyVisible(drawer.getByText("Score: 1.00", { exact: true }))).toBeVisible({ timeout: 10_000 });
+    await expect(onlyVisible(drawer.getByText("Score: 0.75", { exact: true }))).toBeVisible({ timeout: 10_000 });
 
     await expect(drawer.getByText(RAW_EMAIL)).toHaveCount(0);
     await expect(drawer.getByText(RAW_PHONE)).toHaveCount(0);

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 114/164] 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 <yujong@berri.ai>
Date: Mon, 7 Sep 2026 09:00:58 -0700
Subject: [PATCH 115/164] 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: "<!-- litellm-release-wheel-size -->"
-        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 <mubashir.osmani777@gmail.com>
Date: Mon, 7 Sep 2026 12:45:16 -0400
Subject: [PATCH 116/164] 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 <yuneng@berri.ai>
Date: Mon, 7 Sep 2026 09:53:18 -0700
Subject: [PATCH 117/164] 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", () => {

From f40c3842dc4cebbf1eb02379a492acb1d0c87605 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Mon, 7 Sep 2026 10:08:19 -0700
Subject: [PATCH 118/164] fix(router): serialize heuristic tuning quota filters
 for Prisma

---
 litellm/repositories/model_repository.py      |  3 +--
 .../repositories/test_repositories.py         | 22 ++++++++++++++++++-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py
index 540f197044f..d24eb8ffc62 100644
--- a/litellm/repositories/model_repository.py
+++ b/litellm/repositories/model_repository.py
@@ -4,7 +4,6 @@ Model repository for database operations on LiteLLM_ProxyModelTable.
 
 import json
 from collections.abc import Mapping, Sequence
-from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Protocol
 
 from litellm.models.model import LiteLLM_ProxyModelTable
@@ -109,7 +108,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
     async def find_all_except(self, model_id: str) -> Sequence[LiteLLM_ProxyModelTable]:
         """Find every model except the row currently being updated."""
         records: Final = await self.table.find_many(
-            where=MappingProxyType({"model_id": MappingProxyType({"not": model_id})})
+            where={"model_id": {"not": model_id}}  # mutable-ok: Prisma requires plain dicts for query serialization
         )
         return tuple(self._to_model_list(records))
 
diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py
index 758d379f22c..a890d7ceed0 100644
--- a/tests/test_litellm/repositories/test_repositories.py
+++ b/tests/test_litellm/repositories/test_repositories.py
@@ -4,10 +4,13 @@ Tests for gateway repository layer.
 
 import json
 from datetime import datetime
-from typing import Any, Dict, List, Optional
+from types import SimpleNamespace
+from typing import Any, Dict, Final, List, Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from prisma import models as prisma_models
+from prisma.builder import QueryBuilder
 
 from litellm.models.base import DomainModel
 from litellm.models.budget import LiteLLM_BudgetTable
@@ -307,6 +310,23 @@ class TestModelRepository:
         client = MockPrismaClient()
         return ModelRepository(client)
 
+    @pytest.mark.asyncio
+    async def test_find_all_except_serializes_exclusion_for_prisma(self) -> None:
+        find_many: Final = AsyncMock(return_value=[])
+        client: Final = SimpleNamespace(
+            db=SimpleNamespace(litellm_proxymodeltable=SimpleNamespace(find_many=find_many))
+        )
+
+        await ModelRepository(client).find_all_except("current-model")
+
+        find_many.assert_awaited_once()
+        query: Final = QueryBuilder(
+            method="find_many",
+            model=prisma_models.LiteLLM_ProxyModelTable,
+            arguments=find_many.call_args.kwargs,
+        ).build_query()
+        assert 'where: { model_id: { not: "current-model" } }' in " ".join(query.split())
+
     def test_table_is_wrapped_for_config_sync(self, repo):
         from litellm.proxy.common_utils.config_sync_pubsub import (
             _PublishOnWriteActions,

From 217cb126236cb2a73245f7c3ac5525f88dfc0f7c Mon Sep 17 00:00:00 2001
From: yujonglee <yujong@berri.ai>
Date: Mon, 7 Sep 2026 10:43:45 -0700
Subject: [PATCH 119/164] refactor(rust): remove per-request enablement
 arguments (#39928)

* refactor(rust): remove per-request enablement arguments

* fix(rust): remove ignored transcription enablement

* refactor(rust): remove OCR-specific bridge controls
---
 .../litellm_core_utils/get_litellm_params.py  |  9 +--
 litellm/llms/custom_httpx/llm_http_handler.py | 11 +--
 litellm/ocr/main.py                           | 11 +--
 litellm/rust_bridge/chat_completions.py       |  3 +-
 litellm/rust_bridge/configuration.py          | 33 +-------
 litellm/rust_bridge/ocr.py                    | 46 +++--------
 litellm/rust_bridge/transcription.py          |  1 -
 litellm/types/router.py                       |  1 -
 .../test_rust_bridge_messages.py              | 26 +-----
 .../test_get_litellm_params.py                | 29 -------
 .../chat/test_anthropic_chat_handler.py       |  8 +-
 .../chat/test_bedrock_converse_handler.py     | 12 ++-
 .../custom_httpx/test_llm_http_handler.py     | 22 +++--
 tests/test_litellm/ocr/test_rust_bridge.py    | 80 ++++++++-----------
 .../responses/test_rust_bridge_websocket.py   | 21 ++---
 .../rust_bridge/test_chat_completions.py      | 38 ++++-----
 .../rust_bridge/test_configuration.py         | 61 +++-----------
 .../test_audio_transcription_rust_bridge.py   |  4 +-
 tests/test_litellm/test_utils.py              | 46 -----------
 ui/litellm-dashboard/src/lib/http/schema.d.ts |  4 -
 20 files changed, 109 insertions(+), 357 deletions(-)

diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py
index 389e6f7f501..1fd79db15a6 100644
--- a/litellm/litellm_core_utils/get_litellm_params.py
+++ b/litellm/litellm_core_utils/get_litellm_params.py
@@ -21,13 +21,10 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset(
     }
 )
 
-# The per-deployment Rust opt-in.
-RUST_KWARG_KEY: Final = "rust"
-
 # Keys `completion()` forwards from its own kwargs into `get_litellm_params`,
 # which are otherwise invisible to it because that call site passes explicit
 # named arguments rather than `**kwargs`.
-FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({RUST_KWARG_KEY})
+FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS
 
 # Pre-define optional kwargs keys as frozenset for O(1) lookups
 # These are extracted from kwargs only if present, avoiding unnecessary .get() calls
@@ -58,10 +55,6 @@ OPTIONAL_KWARGS_KEYS: Final = (
             "itpm",
             "otpm",
             "use_xai_oauth",
-            # The per-deployment Rust opt-in. `all_litellm_params` keeps it out
-            # of the provider body; this keeps it *in* litellm_params, which is
-            # where the chat completions handlers read it from.
-            RUST_KWARG_KEY,
         }
     )
     | AWS_CREDENTIAL_KWARGS_KEYS
diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index f281c249c72..552b1549db8 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -163,13 +163,10 @@ from litellm.utils import (
 
 def _rust_responses_websocket_enabled(
     custom_llm_provider: str | None,
-    litellm_params: GenericLiteLLMParams,
 ) -> bool:
     from litellm.rust_bridge.configuration import rust_enabled
 
-    raw_request_override: Final = litellm_params.get("rust")
-    request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
-    return custom_llm_provider == "openai" and rust_enabled(request_override=request_override)
+    return custom_llm_provider == "openai" and rust_enabled()
 
 
 from .http_handler import get_shared_realtime_ssl_context
@@ -2403,9 +2400,7 @@ class BaseLLMHTTPHandler:
             return None
         from litellm.rust_bridge.configuration import rust_enabled
 
-        raw_request_override: Final = litellm_params.get("rust")
-        request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
-        if not rust_enabled(request_override=request_override):
+        if not rust_enabled():
             return None
         if has_agentic_hook:
             return None
@@ -6514,7 +6509,7 @@ class BaseLLMHTTPHandler:
 
             @asynccontextmanager
             async def _backend_connection():
-                if _rust_responses_websocket_enabled(custom_llm_provider, litellm_params):
+                if _rust_responses_websocket_enabled(custom_llm_provider):
                     from litellm.rust_bridge import responses_websocket as rust_responses_websocket
 
                     rust_backend: Final = await rust_responses_websocket.connect(
diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py
index 6c68971f8d5..56c2292d00c 100644
--- a/litellm/ocr/main.py
+++ b/litellm/ocr/main.py
@@ -29,6 +29,7 @@ from litellm.llms.base_llm.ocr.transformation import (
 )
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
 from litellm.rust_bridge import ocr as rust_ocr_bridge
+from litellm.rust_bridge.configuration import rust_enabled
 from litellm.types.router import GenericLiteLLMParams
 from litellm.utils import ProviderConfigManager, client
 
@@ -196,12 +197,6 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool:
     return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS
 
 
-def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool:
-    raw_request_override: Final = prepared_request.litellm_params.get("rust")
-    request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
-    return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override)
-
-
 def _rust_bridge_optional_params(
     prepared_request: _PreparedOCRRequest,
     resolve_secret: Callable[[str], str | None],
@@ -430,7 +425,7 @@ async def aocr(
         custom_llm_provider = prepared.custom_llm_provider
         completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
 
-        if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared):
+        if _rust_ocr_supported(prepared) and rust_enabled():
             from litellm.secret_managers.main import get_secret_str
 
             rust_response: Final = await _run_rust_aocr(
@@ -702,7 +697,7 @@ def ocr(
         custom_llm_provider = prepared.custom_llm_provider
         completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
 
-        if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared):
+        if _rust_ocr_supported(prepared) and rust_enabled():
             from litellm.secret_managers.main import get_secret_str
 
             rust_response: Final = _run_rust_ocr(
diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py
index c599667ab17..674bd8847f7 100644
--- a/litellm/rust_bridge/chat_completions.py
+++ b/litellm/rust_bridge/chat_completions.py
@@ -247,8 +247,7 @@ def rust_chat_completions_accepts(
         return False
     if stream:
         return False
-    request_override: Final = litellm_params.get("rust") if litellm_params is not None else None
-    if not rust_enabled(request_override=request_override if isinstance(request_override, bool) else None):
+    if not rust_enabled():
         return False
     if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params):
         verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path")
diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py
index 515ab6edef1..5582027bb5d 100644
--- a/litellm/rust_bridge/configuration.py
+++ b/litellm/rust_bridge/configuration.py
@@ -1,13 +1,11 @@
 from __future__ import annotations
 
 import os
-import warnings
 from typing import Final
 
 DEFAULT_RUST_ENABLED: Final = False
 _TRUE_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"})
 _GLOBAL_ENV_NAME: Final = "LITELLM_RUST"
-_LEGACY_OCR_ENV_NAME: Final = "LITELLM_USE_RUST_OCR"
 
 
 class _RustConfiguration:
@@ -26,49 +24,24 @@ def _parse_env_bool(value: str | None) -> bool | None:
 
 def resolve_rust_enabled(
     *,
-    request_override: bool | None,
     process_override: bool | None,
     environment_override: bool | None,
-    legacy_environment_override: bool | None = None,
     release_default: bool = DEFAULT_RUST_ENABLED,
 ) -> bool:
-    if request_override is not None:
-        return request_override
     if process_override is not None:
         return process_override
     if environment_override is not None:
         return environment_override
-    if legacy_environment_override is not None:
-        return legacy_environment_override
     return release_default
 
 
-def rust_enabled(*, request_override: bool | None = None) -> bool:
-    if request_override is not None:
-        return request_override
-    process_override: Final = _CONFIGURATION.override
-    if process_override is not None:
-        return process_override
-    global_override: Final = _parse_env_bool(os.getenv(_GLOBAL_ENV_NAME))
-    legacy_override: Final = None if global_override is not None else _parse_env_bool(os.getenv(_LEGACY_OCR_ENV_NAME))
-    if legacy_override is not None:
-        warnings.warn(
-            f"{_LEGACY_OCR_ENV_NAME} is deprecated; use {_GLOBAL_ENV_NAME} instead",
-            DeprecationWarning,
-            stacklevel=2,
-        )
+def rust_enabled() -> bool:
     return resolve_rust_enabled(
-        request_override=None,
-        process_override=None,
-        environment_override=global_override,
-        legacy_environment_override=legacy_override,
+        process_override=_CONFIGURATION.override,
+        environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)),
     )
 
 
-def rust_ocr_enabled(*, request_override: bool | None = None) -> bool:
-    return rust_enabled(request_override=request_override)
-
-
 def reset_rust_configuration() -> None:
     _CONFIGURATION.override = None
 
diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py
index 86038438f57..b7fdb5a98ef 100644
--- a/litellm/rust_bridge/ocr.py
+++ b/litellm/rust_bridge/ocr.py
@@ -7,12 +7,9 @@ from typing import Final, Protocol, cast  # noqa: TID251  # native extension exp
 
 import httpx
 
-from litellm.rust_bridge import configuration as _configuration
+from litellm.rust_bridge.bindings import NativeBinding
 from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds
 
-rust_ocr_enabled = _configuration.rust_ocr_enabled
-rust = _configuration.rust
-
 
 class RustOcr(Protocol):
     def __call__(
@@ -44,49 +41,24 @@ class RustAocr(Protocol):
         raise NotImplementedError
 
 
-class _Unset:
-    pass
+def _as_ocr(value: object) -> RustOcr | None:
+    return cast(RustOcr, value) if callable(value) else None
 
 
-_UNSET: Final[_Unset] = _Unset()
+def _as_aocr(value: object) -> RustAocr | None:
+    return cast(RustAocr, value) if callable(value) else None
 
 
-_rust_ocr_impl: RustOcr | None = None
-_rust_aocr_impl: RustAocr | None = None
-
-
-def set_rust_ocr(
-    *,
-    ocr: RustOcr | None | _Unset = _UNSET,
-    aocr: RustAocr | None | _Unset = _UNSET,
-) -> None:
-    global _rust_ocr_impl, _rust_aocr_impl
-    if not isinstance(ocr, _Unset):
-        _rust_ocr_impl = ocr
-    if not isinstance(aocr, _Unset):
-        _rust_aocr_impl = aocr
+_OCR: Final = NativeBinding("ocr", validate=_as_ocr)
+_AOCR: Final = NativeBinding("aocr", validate=_as_aocr)
 
 
 def load_rust_ocr() -> RustOcr | None:
-    if _rust_ocr_impl is not None:
-        return _rust_ocr_impl
-    from litellm.rust_bridge import get_native_bridge
-
-    native_bridge: Final = get_native_bridge()
-    if native_bridge is None:
-        return None
-    return cast(RustOcr, native_bridge.ocr)
+    return _OCR.load()
 
 
 def load_rust_aocr() -> RustAocr | None:
-    if _rust_aocr_impl is not None:
-        return _rust_aocr_impl
-    from litellm.rust_bridge import get_native_bridge
-
-    native_bridge: Final = get_native_bridge()
-    if native_bridge is None:
-        return None
-    return cast(RustAocr, getattr(native_bridge, "aocr", None))
+    return _AOCR.load()
 
 
 def ocr(
diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py
index 3d71f6f8a50..6c81786accd 100644
--- a/litellm/rust_bridge/transcription.py
+++ b/litellm/rust_bridge/transcription.py
@@ -56,7 +56,6 @@ _STATE: Final = _RustTranscriptionState()
 
 
 def configure_rust_transcription(
-    enabled: bool = True,
     *,
     transcription: RustTranscription | None | _Unset = _UNSET,
     atranscription: RustAtranscription | None | _Unset = _UNSET,
diff --git a/litellm/types/router.py b/litellm/types/router.py
index 2dad22751de..0db482d8a58 100644
--- a/litellm/types/router.py
+++ b/litellm/types/router.py
@@ -307,7 +307,6 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
     """
 
     custom_llm_provider: str | None = None
-    rust: bool | None = None
     tpm: int | None = None
     rpm: int | None = None
     itpm: int | None = None
diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py
index b2cf253d164..a30474245c6 100644
--- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py
+++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py
@@ -126,16 +126,6 @@ def test_load_rust_messages_returns_injected_impl():
     assert rust_messages.load_rust_messages() is bridge
 
 
-def test_bare_rust_still_toggles_ocr():
-    from litellm.rust_bridge.ocr import rust_ocr_enabled
-
-    litellm.rust(True)
-    assert rust_ocr_enabled() is True
-
-    litellm.rust(False)
-    assert rust_ocr_enabled() is False
-
-
 def test_load_rust_amessages_returns_injected_impl():
     bridge = RecordingAsyncMessages()
     litellm.rust(True)
@@ -214,7 +204,7 @@ async def test_amessages_wrapper_forwards_args():
 def _gate(**overrides):
     kwargs = {
         "custom_llm_provider": "azure_ai",
-        "litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True),
+        "litellm_params": GenericLiteLLMParams(api_key="sk-azure"),
         "has_agentic_hook": False,
         "model": "claude-sonnet-4-5",
         "api_key": "sk-azure",
@@ -282,18 +272,6 @@ async def test_gate_uses_process_enable_without_request_override():
     assert bridge.calls[0]["custom_llm_provider"] == "azure_ai"
 
 
-@pytest.mark.asyncio
-async def test_gate_skips_rust_when_flag_false():
-    bridge = ExplodingAsyncMessages()
-    litellm.rust(True)
-    rust_messages.set_rust_messages(amessages=bridge)
-
-    response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False))
-
-    assert response is None
-    assert bridge.calls == 0
-
-
 @pytest.mark.asyncio
 async def test_gate_invokes_rust_for_native_anthropic_provider():
     bridge = RecordingAsyncMessages()
@@ -302,7 +280,7 @@ async def test_gate_invokes_rust_for_native_anthropic_provider():
 
     response = await _gate(
         custom_llm_provider="anthropic",
-        litellm_params=GenericLiteLLMParams(api_key="sk-ant", rust=True),
+        litellm_params=GenericLiteLLMParams(api_key="sk-ant"),
         api_key="sk-ant",
         api_base="https://api.anthropic.com",
         headers={"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"},
diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py
index 956da571d43..fb4cb494bee 100644
--- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py
+++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py
@@ -215,32 +215,3 @@ class TestMetadataFallsBackToLitellmMetadata:
         assert result["metadata"] is not litellm_metadata
         result["metadata"].pop("trace_id")
         assert litellm_metadata == {"trace_id": "trace-1"}
-
-
-class TestRustOptIn:
-    """`rust: true` is a litellm param, so it has to reach `litellm_params`.
-
-    `all_litellm_params` keeps it out of the provider body; without it also
-    being carried into `litellm_params` the chat completions handlers cannot
-    see the opt-in and the Rust path is silently never taken.
-    """
-
-    def test_rust_is_an_optional_kwargs_key(self):
-        assert "rust" in _OPTIONAL_KWARGS_KEYS
-
-    def test_rust_is_forwarded_from_completion_kwargs(self):
-        from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS
-
-        assert "rust" in FORWARDED_KWARGS_KEYS
-
-    def test_rust_survives_into_litellm_params(self):
-        params = get_litellm_params(rust=True)
-        assert params["rust"] is True
-
-    def test_rust_is_absent_when_the_deployment_did_not_set_it(self):
-        assert "rust" not in get_litellm_params()
-
-    def test_rust_stays_out_of_the_provider_body(self):
-        from litellm.types.utils import all_litellm_params
-
-        assert "rust" in all_litellm_params
diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py
index 043537f8c1f..b4b173b20c3 100644
--- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py
+++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py
@@ -2256,7 +2256,7 @@ class TestRustChatCompletionsHook:
     def _reset_bridge(self, monkeypatch):
         from litellm.rust_bridge import chat_completions as bridge
 
-        monkeypatch.delenv("LITELLM_RUST", raising=False)
+        monkeypatch.setenv("LITELLM_RUST", "1")
         bridge.set_rust_chat_completions(
             chat_completions=None, achat_completions=None, decline=None
         )
@@ -2282,7 +2282,7 @@ class TestRustChatCompletionsHook:
             "logging_obj": MagicMock(),
             "optional_params": {"max_tokens": 16},
             "timeout": 30.0,
-            "litellm_params": {"rust": True},
+            "litellm_params": {},
             "acompletion": False,
             "headers": {},
             "client": None,
@@ -2366,7 +2366,8 @@ class TestRustChatCompletionsHook:
         )
         assert seen["call"][0]["optional_params"]["max_tokens"] == 7
 
-    def test_without_the_opt_in_the_core_is_never_consulted(self):
+    def test_without_the_opt_in_the_core_is_never_consulted(self, monkeypatch):
+        monkeypatch.setenv("LITELLM_RUST", "0")
         from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
         from litellm.llms.anthropic.chat.transformation import AnthropicConfig
 
@@ -2587,6 +2588,7 @@ class TestRustChatCompletionsHook:
 
     def test_pre_call_logging_still_fires_when_rust_is_not_involved(self, monkeypatch):
         """The suppression must not swallow the log on the ordinary path."""
+        monkeypatch.setenv("LITELLM_RUST", "0")
         from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
         from litellm.llms.anthropic.chat.transformation import AnthropicConfig
 
diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py
index c4d6896b17b..f34b8eb1fb9 100644
--- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py
+++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py
@@ -48,7 +48,7 @@ RESOLVED_CREDENTIALS = Credentials(
 
 @pytest.fixture(autouse=True)
 def reset_bridge(monkeypatch):
-    monkeypatch.delenv("LITELLM_RUST", raising=False)
+    monkeypatch.setenv("LITELLM_RUST", "1")
     bridge.set_rust_chat_completions(
         chat_completions=None, achat_completions=None, decline=None
     )
@@ -87,7 +87,7 @@ def _completion_kwargs(**overrides):
         "optional_params": {"maxTokens": 16},
         "acompletion": False,
         "timeout": 30.0,
-        "litellm_params": {"rust": True},
+        "litellm_params": {},
         "extra_headers": None,
         "client": None,
         "api_key": None,
@@ -157,7 +157,8 @@ def test_the_core_receives_the_untranslated_openai_messages():
     ]
 
 
-def test_without_the_opt_in_the_core_is_never_consulted():
+def test_without_the_opt_in_the_core_is_never_consulted(monkeypatch):
+    monkeypatch.setenv("LITELLM_RUST", "0")
     seen = _inject()
     try:
         _run(litellm_params={})
@@ -401,9 +402,10 @@ def test_pre_call_logging_fires_once_when_the_sync_rust_path_declines():
     assert logging_obj.pre_call.call_count == 1
 
 
-def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in():
+def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(monkeypatch):
     """The suppression must not swallow the log on a request the gate declined,
     so a deployment with no `rust` flag keeps exactly the log it always had."""
+    monkeypatch.setenv("LITELLM_RUST", "0")
     logging_obj = MagicMock()
     response = _run(
         logging_obj=logging_obj,
@@ -491,6 +493,7 @@ def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monke
     """With only `AWS_BEARER_TOKEN_BEDROCK` configured boto3 resolves no
     credentials at all. Preparing the Rust handoff must not dereference that
     None: the bearer token signs the request on its own."""
+    monkeypatch.setenv("LITELLM_RUST", "0")
     monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
     client = _sync_client_returning_converse_response()
 
@@ -520,6 +523,7 @@ def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, co
     """The deployment's AWS profile does not exist, so resolving SigV4 credentials
     raises; a bearer-token deployment must still serve the request, since the
     bearer token alone signs it."""
+    monkeypatch.setenv("LITELLM_RUST", "0")
     if configured_through == "env_var":
         monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
     else:
diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
index 023e2d8843f..5c6cc4a52da 100644
--- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
+++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
@@ -2689,20 +2689,18 @@ async def test_generic_http_handler_async_streaming_forwards_provider_response_h
 
 
 @pytest.mark.parametrize(
-    "custom_llm_provider, litellm_params, expected",
-    [
-        ("openai", GenericLiteLLMParams(rust=True), True),
-        ("openai", GenericLiteLLMParams(), False),
-        ("openai", GenericLiteLLMParams(rust=False), False),
-        ("azure", GenericLiteLLMParams(rust=True), False),
-        ("hosted_vllm", GenericLiteLLMParams(rust=True), False),
-        (None, GenericLiteLLMParams(rust=True), False),
-    ],
+    "custom_llm_provider, enabled, expected",
+    [("openai", True, True), ("openai", False, False), ("azure", True, False),
+     ("hosted_vllm", True, False), (None, True, False)],
 )
-def test_the_rust_responses_websocket_needs_both_openai_and_the_rust_flag(
-    custom_llm_provider, litellm_params, expected
+def test_the_rust_responses_websocket_needs_openai_and_process_enablement(
+    custom_llm_provider, enabled, expected, monkeypatch
 ):
-    assert _rust_responses_websocket_enabled(custom_llm_provider, litellm_params) is expected
+    from litellm.rust_bridge import configuration
+
+    configuration.reset_rust_configuration()
+    monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0")
+    assert _rust_responses_websocket_enabled(custom_llm_provider) is expected
 
 
 def test_a_plain_callback_does_not_advertise_a_pre_call_deployment_hook(monkeypatch):
diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py
index 1c2e07e0d24..d28184b96a6 100644
--- a/tests/test_litellm/ocr/test_rust_bridge.py
+++ b/tests/test_litellm/ocr/test_rust_bridge.py
@@ -17,6 +17,7 @@ from litellm.rust_bridge import configuration
 # explicitly via importlib rather than attribute traversal.
 ocr_main = importlib.import_module("litellm.ocr.main")
 rust_bridge = importlib.import_module("litellm.rust_bridge.ocr")
+rust_bridge_bindings = importlib.import_module("litellm.rust_bridge.bindings")
 rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader")
 
 MODEL = "mistral/mistral-ocr-latest"
@@ -215,11 +216,13 @@ def build_prepared_request(
 @pytest.fixture(autouse=True)
 def _reset_rust_flag():
     """Keep the global toggle isolated between tests."""
-    rust_bridge.set_rust_ocr(ocr=None, aocr=None)
+    rust_bridge._OCR.reset()
+    rust_bridge._AOCR.reset()
     configuration.reset_rust_configuration()
     rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
     yield
-    rust_bridge.set_rust_ocr(ocr=None, aocr=None)
+    rust_bridge._OCR.reset()
+    rust_bridge._AOCR.reset()
     configuration.reset_rust_configuration()
     rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
 
@@ -229,7 +232,7 @@ def fake_bridge():
     """Enable the Rust path with an injected recording bridge (no native wheel)."""
     bridge = RecordingBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge)
+    rust_bridge._OCR.override(bridge)
     return bridge
 
 
@@ -238,34 +241,14 @@ def fake_async_bridge():
     """Enable the async Rust path with an injected recording bridge."""
     bridge = RecordingAsyncBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(aocr=bridge)
+    rust_bridge._AOCR.override(bridge)
     return bridge
 
 
-def test_rust_toggles_flag():
-    assert rust_bridge.rust_ocr_enabled() is False
-    litellm.rust(True)
-    assert rust_bridge.rust_ocr_enabled() is True
-    litellm.rust(False)
-    assert rust_bridge.rust_ocr_enabled() is False
-
-
-def test_env_var_enables_rust_ocr(monkeypatch):
-    monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1")
-    with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"):
-        assert rust_bridge.rust_ocr_enabled() is True
-
-
-def test_explicit_false_overrides_process_enable():
-    litellm.rust(True)
-
-    assert ocr_main._rust_ocr_enabled(build_prepared_request(litellm_params={"rust": False})) is False
-
-
 def test_load_rust_ocr_returns_injected_impl():
     bridge = RecordingBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge)
+    rust_bridge._OCR.override(bridge)
     assert rust_bridge.load_rust_ocr() is bridge
 
 
@@ -329,7 +312,7 @@ def test_native_bridge_available_reflects_loader(monkeypatch):
 def test_load_rust_aocr_returns_injected_impl():
     bridge = RecordingAsyncBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(aocr=bridge)
+    rust_bridge._AOCR.override(bridge)
     assert rust_bridge.load_rust_aocr() is bridge
 
 
@@ -338,7 +321,8 @@ def test_toggle_without_ocr_arg_preserves_injected_impl():
     bridge = RecordingBridge()
     async_bridge = RecordingAsyncBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge)
+    rust_bridge._OCR.override(bridge)
+    rust_bridge._AOCR.override(async_bridge)
 
     litellm.rust(False)
     assert rust_bridge.load_rust_ocr() is bridge
@@ -350,16 +334,18 @@ def test_toggle_without_ocr_arg_preserves_injected_impl():
 
 def test_explicit_ocr_none_clears_injected_impl(monkeypatch):
     monkeypatch.setattr(
-        importlib.import_module("litellm.rust_bridge"),
+        rust_bridge_bindings,
         "get_native_bridge",
         lambda: None,
     )
     bridge = RecordingBridge()
     async_bridge = RecordingAsyncBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge)
+    rust_bridge._OCR.override(bridge)
+    rust_bridge._AOCR.override(async_bridge)
 
-    rust_bridge.set_rust_ocr(ocr=None, aocr=None)
+    rust_bridge._OCR.override(None)
+    rust_bridge._AOCR.override(None)
     assert rust_bridge.load_rust_ocr() is None
     assert rust_bridge.load_rust_aocr() is None
 
@@ -368,7 +354,7 @@ def test_load_rust_ocr_none_when_extension_absent(monkeypatch):
     """With no injected impl and no compiled wheel, the loader returns None so the
     caller degrades to the Python path instead of raising ImportError."""
     monkeypatch.setattr(
-        importlib.import_module("litellm.rust_bridge"),
+        rust_bridge_bindings,
         "get_native_bridge",
         lambda: None,
     )
@@ -385,7 +371,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch):
     fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE)  # type: ignore[attr-defined]
     fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE)  # type: ignore[attr-defined]
     monkeypatch.setattr(
-        importlib.import_module("litellm.rust_bridge"),
+        rust_bridge_bindings,
         "get_native_bridge",
         lambda: fake_module,
     )
@@ -406,7 +392,7 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response():
 
     litellm.rust(True)
 
-    rust_bridge.set_rust_ocr(ocr=bridge)
+    rust_bridge._OCR.override(bridge)
     response = rust_bridge.ocr(
         model="mistral-ocr-latest",
         document=DOCUMENT,
@@ -441,7 +427,7 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response():
 
     litellm.rust(True)
 
-    rust_bridge.set_rust_ocr(aocr=bridge)
+    rust_bridge._AOCR.override(bridge)
     response = await rust_bridge.aocr(
         model="mistral-ocr-maas",
         document=DOCUMENT,
@@ -470,7 +456,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response():
     bridge = RecordingBridge()
     logging_obj = RecordingLogging()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge)
+    rust_bridge._OCR.override(bridge)
 
     response = ocr_main._run_rust_ocr(
         prepared_request=build_prepared_request(
@@ -503,7 +489,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response():
 def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing():
     bridge = RecordingBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge)
+    rust_bridge._OCR.override(bridge)
 
     ocr_main._run_rust_ocr(
         prepared_request=build_prepared_request(api_key=None, timeout=None),
@@ -516,7 +502,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing():
 def test_run_rust_ocr_prefers_explicit_key_over_resolver():
     bridge = RecordingBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge)
+    rust_bridge._OCR.override(bridge)
 
     def _resolver(name: str) -> str | None:
         raise AssertionError(f"resolver should not be called for {name}")
@@ -536,7 +522,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var():
     bridge = RecordingBridge()
     resolver_calls = []
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge)
+    rust_bridge._OCR.override(bridge)
 
     def _resolver(name):
         resolver_calls.append(name)
@@ -559,7 +545,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var():
 def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
     bridge = RecordingBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge)
+    rust_bridge._OCR.override(bridge)
 
     ocr_main._run_rust_ocr(
         prepared_request=build_prepared_request(
@@ -586,7 +572,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
 def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager():
     bridge = RecordingBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge)
+    rust_bridge._OCR.override(bridge)
 
     def _resolver(name: str) -> str | None:
         return {
@@ -610,7 +596,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana
 def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager():
     bridge = RecordingBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge)
+    rust_bridge._OCR.override(bridge)
 
     ocr_main._run_rust_ocr(
         prepared_request=build_prepared_request(
@@ -628,7 +614,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager():
 def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint():
     bridge = RecordingBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge)
+    rust_bridge._OCR.override(bridge)
 
     ocr_main._run_rust_ocr(
         prepared_request=build_prepared_request(
@@ -649,7 +635,7 @@ def test_run_rust_ocr_runs_pre_call_logging():
     logging_obj = RecordingLogging()
     bridge = RecordingBridge()
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=bridge)
+    rust_bridge._OCR.override(bridge)
 
     ocr_main._run_rust_ocr(
         prepared_request=build_prepared_request(
@@ -737,7 +723,7 @@ def test_ocr_exception_type_uses_resolved_provider_context(
 
     monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(ocr=RaisingBridge())
+    rust_bridge._OCR.override(RaisingBridge())
 
     with pytest.raises(CapturedException):
         litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
@@ -783,7 +769,7 @@ async def test_aocr_exception_type_uses_resolved_provider_context(
 
     monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
     litellm.rust(True)
-    rust_bridge.set_rust_ocr(aocr=RaisingAsyncBridge())
+    rust_bridge._AOCR.override(RaisingAsyncBridge())
 
     with pytest.raises(CapturedException):
         await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
@@ -812,9 +798,7 @@ def test_ocr_does_not_route_to_rust_when_disabled():
     """With the flag off, the bridge must not be consulted even if an impl exists."""
     bridge = RecordingBridge()
     litellm.rust(False)
-    rust_bridge.set_rust_ocr(ocr=bridge)
-
-    assert rust_bridge.rust_ocr_enabled() is False
+    rust_bridge._OCR.override(bridge)
     # The impl stays available for injection, but the disabled flag gates usage,
     # so ocr() never reaches the Rust path (asserted via the enabled-path test).
     assert bridge.calls == []
diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py
index 4b446368dbe..74d96bda336 100644
--- a/tests/test_litellm/responses/test_rust_bridge_websocket.py
+++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py
@@ -4,7 +4,6 @@ import pytest
 
 from litellm.llms.custom_httpx.llm_http_handler import _rust_responses_websocket_enabled
 from litellm.rust_bridge import configuration, responses_websocket
-from litellm.types.router import GenericLiteLLMParams
 
 
 class _FakeNativeConnection:
@@ -48,22 +47,12 @@ def reset_responses_websocket():
     configuration.reset_rust_configuration()
 
 
-def test_rust_websocket_bridge_is_disabled_without_flag() -> None:
-    assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams())
-    assert not _rust_responses_websocket_enabled("anthropic", GenericLiteLLMParams(rust=True))
-    assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=True))
-
-
-def test_explicit_false_overrides_process_enable() -> None:
+def test_rust_websocket_bridge_uses_process_enablement() -> None:
+    configuration.rust(False)
+    assert not _rust_responses_websocket_enabled("openai")
     configuration.rust(True)
-
-    assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=False))
-
-
-def test_process_enable_applies_without_request_override() -> None:
-    configuration.rust(True)
-
-    assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams())
+    assert _rust_responses_websocket_enabled("openai")
+    assert not _rust_responses_websocket_enabled("anthropic")
 
 
 @pytest.mark.asyncio
diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py
index 0489f4ff017..b2fd2e6dcc0 100644
--- a/tests/test_litellm/rust_bridge/test_chat_completions.py
+++ b/tests/test_litellm/rust_bridge/test_chat_completions.py
@@ -10,8 +10,8 @@ from __future__ import annotations
 import pytest
 
 import litellm
-from litellm.rust_bridge import chat_completions as bridge
 from litellm.rust_bridge import configuration
+from litellm.rust_bridge import chat_completions as bridge
 from litellm.types.utils import ModelResponse
 
 RUST_RESPONSE = {
@@ -67,10 +67,11 @@ def _hide_native_bridge(monkeypatch):
 
 
 @pytest.fixture(autouse=True)
-def reset_bridge():
+def reset_bridge(monkeypatch):
     """Every test starts with no injected callables, and leaves none behind."""
     bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
     configuration.reset_rust_configuration()
+    monkeypatch.setenv("LITELLM_RUST", "1")
     yield
     bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
     configuration.reset_rust_configuration()
@@ -112,7 +113,7 @@ def _accepts(**overrides) -> bool:
         "messages": MESSAGES,
         "optional_params": {"max_tokens": 16},
         "custom_llm_provider": "anthropic",
-        "litellm_params": {"rust": True},
+        "litellm_params": {},
         "stream": None,
     }
     kwargs.update(overrides)
@@ -126,23 +127,16 @@ class TestGate:
         bridge.set_rust_chat_completions(decline=gate)
         assert _accepts(litellm_params={}) is False
         assert _accepts(litellm_params=None) is False
-        assert _accepts(litellm_params={"rust": False}) is False
         assert gate.calls == [], "the gate must not be consulted before opt-in"
 
     def test_accepts_when_the_deployment_opted_in_and_the_core_agrees(self, monkeypatch):
-        monkeypatch.delenv("LITELLM_RUST", raising=False)
+        monkeypatch.setenv("LITELLM_RUST", "1")
         gate = _RecordingDecline()
         bridge.set_rust_chat_completions(decline=gate)
         assert _accepts() is True
         assert gate.calls[0]["model"] == "claude-sonnet-4-5"
         assert gate.calls[0]["custom_llm_provider"] == "anthropic"
 
-    def test_explicit_false_overrides_process_enable(self):
-        bridge.set_rust_chat_completions(decline=_RecordingDecline())
-        configuration.rust(True)
-
-        assert _accepts(litellm_params={"rust": False}) is False
-
     def test_process_enable_applies_without_request_override(self):
         bridge.set_rust_chat_completions(decline=_RecordingDecline())
         configuration.rust(True)
@@ -155,7 +149,7 @@ class TestGate:
         assert _accepts(litellm_params={}) is True
 
     def test_declines_streaming_and_providers_off_the_path(self, monkeypatch):
-        monkeypatch.delenv("LITELLM_RUST", raising=False)
+        monkeypatch.setenv("LITELLM_RUST", "1")
         gate = _RecordingDecline()
         bridge.set_rust_chat_completions(decline=gate)
         assert _accepts(stream=True) is False
@@ -170,10 +164,10 @@ class TestGate:
         handed `optional_params` only, so accepting here would send the request
         to Anthropic with the abuse-detection attribution silently missing.
         """
-        monkeypatch.delenv("LITELLM_RUST", raising=False)
+        monkeypatch.setenv("LITELLM_RUST", "1")
         gate = _RecordingDecline()
         bridge.set_rust_chat_completions(decline=gate)
-        assert _accepts(litellm_params={"rust": True, "metadata": {"user_id": "u-123"}}) is False
+        assert _accepts(litellm_params={"metadata": {"user_id": "u-123"}}) is False
         assert gate.calls == [], "the core must not be consulted for a request it cannot see the key of"
 
         # Bedrock's Converse transform reads no `user_id`, and an Anthropic request
@@ -182,13 +176,13 @@ class TestGate:
             _accepts(
                 custom_llm_provider="bedrock",
                 model="bedrock/us-east-1/anthropic.claude-v2",
-                litellm_params={"rust": True, "metadata": {"user_id": "u-123"}},
+                litellm_params={"metadata": {"user_id": "u-123"}},
             )
             is True
         )
-        assert _accepts(litellm_params={"rust": True, "metadata": {"trace_id": "t-1"}}) is True
-        assert _accepts(litellm_params={"rust": True, "metadata": {"user_id": None}}) is True
-        assert _accepts(litellm_params={"rust": True, "metadata": None}) is True
+        assert _accepts(litellm_params={"metadata": {"trace_id": "t-1"}}) is True
+        assert _accepts(litellm_params={"metadata": {"user_id": None}}) is True
+        assert _accepts(litellm_params={"metadata": None}) is True
 
     def test_declines_a_bedrock_request_while_the_proxy_owns_request_metadata(self, monkeypatch):
         """`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the
@@ -196,7 +190,7 @@ class TestGate:
         evicting a caller-supplied one. The core can do neither, so an operator
         who armed `bedrock_request_metadata_fields` keeps the Python path.
         """
-        monkeypatch.delenv("LITELLM_RUST", raising=False)
+        monkeypatch.setenv("LITELLM_RUST", "1")
         gate = _RecordingDecline()
         bridge.set_rust_chat_completions(decline=gate)
         bedrock = {
@@ -213,17 +207,17 @@ class TestGate:
         assert _accepts(**bedrock) is True, "the decline follows the operator's opt-in alone"
 
     def test_declines_when_the_core_declines(self, monkeypatch):
-        monkeypatch.delenv("LITELLM_RUST", raising=False)
+        monkeypatch.setenv("LITELLM_RUST", "1")
         bridge.set_rust_chat_completions(decline=_RecordingDecline("streaming"))
         assert _accepts() is False
 
     def test_declines_when_the_bridge_is_unavailable(self, monkeypatch):
-        monkeypatch.delenv("LITELLM_RUST", raising=False)
+        monkeypatch.setenv("LITELLM_RUST", "1")
         _hide_native_bridge(monkeypatch)
         assert _accepts() is False
 
     def test_declines_when_the_gate_itself_raises(self, monkeypatch):
-        monkeypatch.delenv("LITELLM_RUST", raising=False)
+        monkeypatch.setenv("LITELLM_RUST", "1")
 
         def exploding(**_kwargs):
             raise RuntimeError("boom")
diff --git a/tests/test_litellm/rust_bridge/test_configuration.py b/tests/test_litellm/rust_bridge/test_configuration.py
index 15f69f95335..aff9d5acac1 100644
--- a/tests/test_litellm/rust_bridge/test_configuration.py
+++ b/tests/test_litellm/rust_bridge/test_configuration.py
@@ -10,7 +10,6 @@ from typing import Final
 import pytest
 
 from litellm.rust_bridge import configuration
-from litellm.rust_bridge import ocr as rust_ocr
 
 
 @pytest.fixture(autouse=True)
@@ -19,42 +18,31 @@ def _isolated_configuration(  # pyright: ignore[reportUnusedFunction]  # pytest
 ) -> Generator[None]:
     configuration.reset_rust_configuration()
     monkeypatch.delenv("LITELLM_RUST", raising=False)
-    monkeypatch.delenv("LITELLM_USE_RUST_OCR", raising=False)
-    rust_ocr.set_rust_ocr(ocr=None, aocr=None)
     yield
     configuration.reset_rust_configuration()
-    rust_ocr.set_rust_ocr(ocr=None, aocr=None)
 
 
 @pytest.mark.parametrize(
-    ("request_override", "process", "environment", "legacy_environment", "release_default", "expected"),
+    ("process", "environment", "release_default", "expected"),
     (
-        (False, True, True, True, True, False),
-        (True, False, False, False, False, True),
-        (None, False, True, True, True, False),
-        (None, True, False, False, False, True),
-        (None, None, False, True, True, False),
-        (None, None, True, False, False, True),
-        (None, None, None, False, True, False),
-        (None, None, None, True, False, True),
-        (None, None, None, None, False, False),
-        (None, None, None, None, True, True),
+        (False, True, True, False),
+        (True, False, False, True),
+        (None, False, True, False),
+        (None, True, False, True),
+        (None, None, False, False),
+        (None, None, True, True),
     ),
 )
 def test_resolution_precedence(
-    request_override: bool | None,
     process: bool | None,
     environment: bool | None,
-    legacy_environment: bool | None,
     release_default: bool,
     expected: bool,
 ) -> None:
     assert (
         configuration.resolve_rust_enabled(
-            request_override=request_override,
             process_override=process,
             environment_override=environment,
-            legacy_environment_override=legacy_environment,
             release_default=release_default,
         )
         is expected
@@ -71,7 +59,6 @@ def test_process_override_wins_over_environment(monkeypatch: pytest.MonkeyPatch)
     configuration.rust(True)
 
     assert configuration.rust_enabled() is True
-    assert configuration.rust_enabled(request_override=False) is False
 
 
 def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -83,18 +70,8 @@ def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPat
 @pytest.mark.parametrize("value", ("", " ", "sometimes", "2"))
 def test_invalid_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None:
     monkeypatch.setenv("LITELLM_RUST", value)
-    monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1")
 
     assert configuration.rust_enabled() is False
-    assert configuration.rust_ocr_enabled() is False
-
-
-@pytest.mark.parametrize("value", ("", " ", "sometimes", "2"))
-def test_invalid_legacy_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None:
-    monkeypatch.setenv("LITELLM_USE_RUST_OCR", value)
-
-    with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"):
-        assert configuration.rust_enabled() is False
 
 
 def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -104,40 +81,20 @@ def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytes
         assert executor.submit(configuration.rust_enabled).result() is True
         configuration.rust(False)
         assert executor.submit(configuration.rust_enabled).result() is False
-        assert executor.submit(configuration.rust_ocr_enabled).result() is False
         configuration.reset_rust_configuration()
         assert executor.submit(configuration.rust_enabled).result() is True
-        assert executor.submit(configuration.rust_ocr_enabled).result() is True
 
 
 def test_explicit_override_precedes_invalid_environment(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setenv("LITELLM_RUST", "sometimes")
 
-    assert configuration.rust_enabled(request_override=False) is False
     configuration.rust(True)
     assert configuration.rust_enabled() is True
 
 
-def test_legacy_ocr_environment_is_deprecated_and_global(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1")
-
-    with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"):
-        assert configuration.rust_enabled() is True
-    with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"):
-        assert configuration.rust_ocr_enabled() is True
-
-
-def test_global_environment_precedes_legacy_ocr_environment(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setenv("LITELLM_RUST", "0")
-    monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1")
-
-    assert configuration.rust_enabled() is False
-
-
-@pytest.mark.parametrize("environment_name", ("LITELLM_RUST", "LITELLM_USE_RUST_OCR"))
 @pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False")))
-def test_environment_controls_startup(environment_name: str, value: str, expected: str) -> None:
-    environment: Final = {**os.environ, environment_name: value}
+def test_environment_controls_startup(value: str, expected: str) -> None:
+    environment: Final = {**os.environ, "LITELLM_RUST": value}
     result: Final = subprocess.run(
         (
             sys.executable,
diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py
index bbeb6c38f78..112464bda22 100644
--- a/tests/test_litellm/test_audio_transcription_rust_bridge.py
+++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py
@@ -44,7 +44,7 @@ class AsyncBridge:
 
 def test_enabled_sync_bridge_receives_audio() -> None:
     bridge = SyncBridge()
-    rust_bridge.configure_rust_transcription(True, transcription=bridge)
+    rust_bridge.configure_rust_transcription(transcription=bridge)
     result = rust_bridge.transcription(
         model="mistral.voxtral-mini-3b-2507",
         audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"},
@@ -61,7 +61,7 @@ def test_enabled_sync_bridge_receives_audio() -> None:
 
 @pytest.mark.asyncio
 async def test_enabled_async_bridge() -> None:
-    rust_bridge.configure_rust_transcription(True, atranscription=AsyncBridge())
+    rust_bridge.configure_rust_transcription(atranscription=AsyncBridge())
     result = await rust_bridge.atranscription(
         model="mistral.voxtral-mini-3b-2507",
         audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"},
diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py
index d686b032ee0..271c84384b4 100644
--- a/tests/test_litellm/test_utils.py
+++ b/tests/test_litellm/test_utils.py
@@ -5239,52 +5239,6 @@ def test_client_side_timeout_marker_never_reaches_the_provider():
     )
 
 
-def test_rust_flag_not_forwarded_as_provider_param():
-    forwarded = get_non_default_completion_params({"rust": True, "temperature": 0.5})
-    assert "rust" not in forwarded
-
-
-def test_completion_does_not_leak_rust_flag_into_provider_request_body():
-    mock_response = MagicMock()
-    mock_response.model_dump.return_value = {
-        "id": "chatcmpl-1",
-        "object": "chat.completion",
-        "created": 1234567890,
-        "model": "gpt-4o-mini",
-        "choices": [
-            {
-                "index": 0,
-                "message": {"role": "assistant", "content": "hi"},
-                "finish_reason": "stop",
-            }
-        ],
-        "usage": {
-            "prompt_tokens": 1,
-            "completion_tokens": 1,
-            "total_tokens": 2,
-        },
-    }
-
-    mock_raw_response = MagicMock()
-    mock_raw_response.headers = {}
-    mock_raw_response.parse.return_value = mock_response
-
-    mock_client = MagicMock()
-    mock_client.chat.completions.with_raw_response.create.return_value = mock_raw_response
-
-    litellm.completion(
-        model="openai/gpt-4o-mini",
-        messages=[{"role": "user", "content": "hi"}],
-        rust=True,
-        api_key="sk-test",
-        client=mock_client,
-    )
-
-    create_kwargs = mock_client.chat.completions.with_raw_response.create.call_args.kwargs
-    assert "rust" not in create_kwargs
-    assert "rust" not in (create_kwargs.get("extra_body") or {})
-
-
 class _RecordingDeploymentFailureLogger(CustomLogger):
     def __init__(self) -> None:
         super().__init__()
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index b1534c19670..108e50b52a1 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -29562,8 +29562,6 @@ export interface components {
             regional_processing_uplift_multiplier_us?: number | null;
             /** Rpm */
             rpm?: number | null;
-            /** Rust */
-            rust?: boolean | null;
             /** S3 Bucket Name */
             s3_bucket_name?: string | null;
             /** S3 Encryption Key Id */
@@ -39729,8 +39727,6 @@ export interface components {
             regional_processing_uplift_multiplier_us?: number | null;
             /** Rpm */
             rpm?: number | null;
-            /** Rust */
-            rust?: boolean | null;
             /** S3 Bucket Name */
             s3_bucket_name?: string | null;
             /** S3 Encryption Key Id */

From 8da735410c7903b0e8c92033e444a76d32a03c6f Mon Sep 17 00:00:00 2001
From: yujonglee <yujong@berri.ai>
Date: Mon, 7 Sep 2026 10:44:22 -0700
Subject: [PATCH 120/164] test(ocr): trace callback lifecycle parity (#40063)

---
 .../crates/ai-gateway/src/ocr/hooks.rs        | 12 +++
 .../crates/ai-gateway/tests/ocr_lifecycle.rs  | 52 +++++++++--
 .../python-bridge/src/function_trace.rs       | 28 ++++--
 .../python-bridge/src/routes/definition.rs    |  9 +-
 .../shared/tracing/native.py                  |  3 +-
 .../strategies/trace_parity/models.py         |  1 +
 .../strategies/trace_parity/sdk/execution.py  | 74 +++++++++++----
 .../strategies/trace_parity/sdk/ocr/case.py   | 91 +++++++++++++++++--
 8 files changed, 220 insertions(+), 50 deletions(-)

diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
index 446b323db3a..ed41f1ff9e7 100644
--- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
+++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
@@ -265,6 +265,12 @@ impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLi
         Box::pin(async move { Ok(request) })
     }
 
+    #[tracing::instrument(
+        name = "success_callback",
+        target = "litellm::function_trace",
+        level = "trace",
+        skip_all
+    )]
     fn async_log_success_event<'a>(
         &'a self,
         context: &'a CallLifecycleContext,
@@ -288,6 +294,12 @@ impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLi
         })
     }
 
+    #[tracing::instrument(
+        name = "failure_callback",
+        target = "litellm::function_trace",
+        level = "trace",
+        skip_all
+    )]
     fn async_log_failure_event<'a>(
         &'a self,
         context: &'a CallLifecycleContext,
diff --git a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs
index c3a89f4394d..60e90ed2a7c 100644
--- a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs
+++ b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs
@@ -11,9 +11,13 @@ use litellm_ai_gateway::integrations::custom_logger::{
 use litellm_ai_gateway::integrations::types::RequestMetadata;
 use litellm_ai_gateway::ocr::{OcrRequest, ocr};
 use litellm_core::error::Error;
+#[cfg(feature = "trace-parity")]
+use litellm_core::observability::FunctionTrace;
 use serde_json::{Map, Value, json};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpListener, TcpStream};
+#[cfg(feature = "trace-parity")]
+use tracing::instrument::WithSubscriber;
 
 async fn read_http_headers(socket: &mut TcpStream) -> String {
     let mut request = Vec::new();
@@ -320,14 +324,17 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
         GuardrailEventHook::PreCall,
         GuardrailEventHook::DuringCall,
     ]));
-    let response = ocr(OcrRequest {
+    #[cfg(feature = "trace-parity")]
+    let trace = FunctionTrace::default();
+    let api_base = format!("http://{addr}");
+    let call = ocr(OcrRequest {
         model: "mistral-ocr-latest",
         document: json!({
             "type": "document_url",
             "document_url": "https://example.com/doc.pdf"
         }),
         api_key: Some("sk-test"),
-        api_base: Some(&format!("http://{addr}")),
+        api_base: Some(&api_base),
         custom_llm_provider: Some("mistral"),
         extra_headers: None,
         optional_params: Map::new(),
@@ -339,9 +346,10 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
             ..Default::default()
         },
         litellm_call_id: Some("ocr-call-1"),
-    })
-    .await
-    .expect("ocr request succeeds");
+    });
+    #[cfg(feature = "trace-parity")]
+    let call = call.with_subscriber(trace.dispatcher());
+    let response = call.await.expect("ocr request succeeds");
 
     assert_eq!(response["pages"][0]["markdown"], "ok");
     assert_eq!(
@@ -359,6 +367,16 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
             error_kind: None,
         }]
     );
+    #[cfg(feature = "trace-parity")]
+    assert_eq!(
+        trace
+            .events()
+            .iter()
+            .filter(|event| event.function.ends_with("_callback"))
+            .map(|event| event.function)
+            .collect::<Vec<_>>(),
+        vec!["success_callback"]
+    );
 
     let request = server.await.expect("server task completes");
     assert!(request.contains(r#""guarded_pre":true"#), "{request}");
@@ -388,14 +406,17 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
     });
 
     let logger = Arc::new(RecordingOcrLogger::default());
-    let err = ocr(OcrRequest {
+    #[cfg(feature = "trace-parity")]
+    let trace = FunctionTrace::default();
+    let api_base = format!("http://{addr}");
+    let call = ocr(OcrRequest {
         model: "mistral-ocr-latest",
         document: json!({
             "type": "document_url",
             "document_url": "https://example.com/doc.pdf"
         }),
         api_key: Some("sk-test"),
-        api_base: Some(&format!("http://{addr}")),
+        api_base: Some(&api_base),
         custom_llm_provider: Some("mistral"),
         extra_headers: None,
         optional_params: Map::new(),
@@ -404,9 +425,10 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
         guardrails: Vec::new(),
         request_metadata: RequestMetadata::default(),
         litellm_call_id: Some("ocr-call-2"),
-    })
-    .await
-    .expect_err("provider error propagates");
+    });
+    #[cfg(feature = "trace-parity")]
+    let call = call.with_subscriber(trace.dispatcher());
+    let err = call.await.expect_err("provider error propagates");
 
     assert!(matches!(err, Error::Http { status: 500, .. }));
     server.await.expect("server task completes");
@@ -421,6 +443,16 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
             error_kind: Some("HttpError".to_string()),
         }]
     );
+    #[cfg(feature = "trace-parity")]
+    assert_eq!(
+        trace
+            .events()
+            .iter()
+            .filter(|event| event.function.ends_with("_callback"))
+            .map(|event| event.function)
+            .collect::<Vec<_>>(),
+        vec!["failure_callback"]
+    );
 }
 
 #[tokio::test]
diff --git a/litellm-rust/crates/python-bridge/src/function_trace.rs b/litellm-rust/crates/python-bridge/src/function_trace.rs
index ea7d9f4993e..bc3c962f7a3 100644
--- a/litellm-rust/crates/python-bridge/src/function_trace.rs
+++ b/litellm-rust/crates/python-bridge/src/function_trace.rs
@@ -1,3 +1,4 @@
+use std::fmt::Display;
 use std::future::Future;
 
 use litellm_core::observability::{FunctionTrace, FunctionTraceEvent};
@@ -6,17 +7,32 @@ use tracing::instrument::WithSubscriber;
 
 #[derive(Serialize)]
 pub(crate) struct TracedResponse<T> {
-    response: T,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    response: Option<T>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    error: Option<String>,
     trace: Vec<FunctionTraceEvent>,
 }
 
 pub(crate) async fn capture<T, E>(
     future: impl Future<Output = Result<T, E>>,
-) -> Result<TracedResponse<T>, E> {
+) -> Result<TracedResponse<T>, E>
+where
+    E: Display,
+{
     let trace = FunctionTrace::default();
-    let response = future.with_subscriber(trace.dispatcher()).await?;
-    Ok(TracedResponse {
-        response,
-        trace: trace.events(),
+    let result = future.with_subscriber(trace.dispatcher()).await;
+    let events = trace.events();
+    Ok(match result {
+        Ok(response) => TracedResponse {
+            response: Some(response),
+            error: None,
+            trace: events,
+        },
+        Err(error) => TracedResponse {
+            response: None,
+            error: Some(error.to_string()),
+            trace: events,
+        },
     })
 }
diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs
index 3285da14d5f..bc51647cbad 100644
--- a/litellm-rust/crates/python-bridge/src/routes/definition.rs
+++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs
@@ -486,10 +486,11 @@ asyncio.run(exercise())
             let code = CString::new(
                 r#"
 result = routes.echo("traced")
-assert result == {
-    "response": "traced",
-    "trace": [{"function": "execute_echo", "depth": 0}],
-}
+assert result["response"] == "traced", result
+assert [event["function"] for event in result["trace"]] == ["execute_echo"], result
+failure = routes.echo("error")
+assert failure["error"] == "invalid request: synthetic error", failure
+assert [event["function"] for event in failure["trace"]] == ["execute_echo"], failure
 "#,
             )
             .expect("Python source should not contain null bytes");
diff --git a/tests/rust-python-harness/shared/tracing/native.py b/tests/rust-python-harness/shared/tracing/native.py
index 4f988f65294..688995cbc4b 100644
--- a/tests/rust-python-harness/shared/tracing/native.py
+++ b/tests/rust-python-harness/shared/tracing/native.py
@@ -19,7 +19,8 @@ class _TraceEventPayload(BaseModel):
 
 class TraceResponsePayload(BaseModel):
     model_config = ConfigDict(strict=True, extra="forbid")
-    response: object
+    response: object = None
+    error: str | None = None
     trace: tuple[_TraceEventPayload, ...] | list[_TraceEventPayload]
 
 
diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py
index da560f99730..04659b25382 100644
--- a/tests/rust-python-harness/strategies/trace_parity/models.py
+++ b/tests/rust-python-harness/strategies/trace_parity/models.py
@@ -16,6 +16,7 @@ TraceFailureSource = Literal["python", "rust", "harness"]
 class RouteFixture:
     kwargs: dict[str, object]
     provider_responses: tuple[RecordedHttpResponse, ...]
+    expected_failure: bool = False
 
 
 @dataclass(frozen=True, slots=True)
diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py
index f8d7c55d4e2..eb6c9233565 100644
--- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py
+++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py
@@ -2,15 +2,16 @@ from __future__ import annotations
 
 import asyncio
 from collections.abc import Awaitable
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Final, Protocol, cast
 
 from ....shared.parity.replay import replay_server
 from ....shared.reporting.models import Surface
-from ....shared.tracing.native import native_trace_events
+from ....shared.tracing.native import TraceResponsePayload, native_trace_events
 from ....shared.tracing.profiler import FunctionTraceEvent, profile_python
 from ....shared.tracing.steps import Engine, pipeline_projection
-from ..models import RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario
+from ..models import RouteFixture, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario
 from ..reporting import TraceComparisonArtifact
 
 
@@ -18,9 +19,22 @@ class SdkCall(Protocol):
     def __call__(self, **kwargs: object) -> object: ...
 
 
+@dataclass(frozen=True, slots=True)
+class _CollectedTrace:
+    events: tuple[FunctionTraceEvent, ...]
+    error: str | None = None
+
+
 def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> object:
     async def invoke_async() -> object:
-        return await cast(Awaitable[object], function(**kwargs))
+        try:
+            return await cast(Awaitable[object], function(**kwargs))
+        finally:
+            await asyncio.sleep(0)
+            from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+
+            await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
+            await GLOBAL_LOGGING_WORKER.stop()
 
     if asynchronous:
         return asyncio.run(invoke_async())
@@ -48,16 +62,30 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa
     return cast(SdkCall, getattr(owner, spec.python_entrypoints[int(asynchronous)]))
 
 
+def _python_invocation_error(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> str | None:
+    try:
+        _invoke(function, kwargs, asynchronous=asynchronous)
+    except Exception as error:
+        return f"{type(error).__name__}: {error}"
+    return None
+
+
 def _collect(
-    function: SdkCall, kwargs: dict[str, object], engine: Engine, *, asynchronous: bool
-) -> tuple[FunctionTraceEvent, ...]:
+    function: SdkCall,
+    fixture: RouteFixture,
+    engine: Engine,
+    *,
+    asynchronous: bool,
+) -> _CollectedTrace:
+    kwargs: Final = fixture.kwargs
     if engine == "rust":
-        return native_trace_events(_invoke(function, kwargs, asynchronous=asynchronous))
+        payload: Final = TraceResponsePayload.model_validate(_invoke(function, kwargs, asynchronous=asynchronous))
+        return _CollectedTrace(native_trace_events(payload), payload.error)
     import litellm
 
     with profile_python(Path(litellm.__file__).parent, threads=True) as profiler:
-        _invoke(function, kwargs, asynchronous=asynchronous)
-    return tuple(profiler.events)
+        error: Final = _python_invocation_error(function, kwargs, asynchronous=asynchronous)
+    return _CollectedTrace(tuple(profiler.events), error)
 
 
 def collect_trace(
@@ -68,22 +96,30 @@ def collect_trace(
         return function
     try:
         with replay_server() as provider:
-            fixture: Final = spec.fixture(engine, provider.url)
-            for response in fixture.provider_responses:
+            base_fixture: Final = spec.fixture(engine, provider.url)
+            for response in base_fixture.provider_responses:
                 provider.enqueue_response(response)
-            kwargs: Final = {
-                **fixture.kwargs,
-                "api_key": "test-key",
-                "api_base": provider.url,
-                **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}),
-            }
-            events: Final = _collect(function, kwargs, engine, asynchronous=asynchronous)
+            fixture: Final = RouteFixture(
+                kwargs={
+                    **base_fixture.kwargs,
+                    "api_key": "test-key",
+                    "api_base": provider.url,
+                    **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}),
+                },
+                provider_responses=base_fixture.provider_responses,
+                expected_failure=base_fixture.expected_failure,
+            )
+            collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous)
             provider.take_requests(len(fixture.provider_responses))
     except Exception as error:
         return TraceExecutionFailure(engine, f"{type(error).__name__}: {error}")
-    if not events:
+    if fixture.expected_failure and collected.error is None:
+        return TraceExecutionFailure(engine, "call succeeded but the scenario expects failure")
+    if not fixture.expected_failure and collected.error is not None:
+        return TraceExecutionFailure(engine, collected.error)
+    if not collected.events:
         return TraceExecutionFailure(engine, "trace is empty")
-    return events
+    return collected.events
 
 
 def _failure_message(result: tuple[FunctionTraceEvent, ...] | TraceExecutionFailure) -> str | None:
diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py
index fe214f45339..2a4a1b3a152 100644
--- a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py
+++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py
@@ -19,6 +19,20 @@ COMMON_MAPPINGS: Final = (
     mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"),
 )
 
+SUCCESS_CALLBACK_SYNC_MAPPING: Final = mapping(
+    rust_span="success_callback",
+    python_frame=r"BoundedLoggingThreadPoolExecutor\.submit$",
+)
+SUCCESS_CALLBACK_ASYNC_MAPPING: Final = mapping(
+    rust_span="success_callback",
+    python_frame=r"Logging\.async_success_handler$",
+)
+FAILURE_CALLBACK_MAPPING: Final = mapping(
+    rust_span="failure_callback",
+    python_frame=r"Logging\.(?:async_)?failure_handler$",
+)
+IGNORED_SUCCESS_CALLBACK_MAPPING: Final = mapping(rust_span="success_callback")
+
 SYNC_MAPPINGS: Final = (
     *COMMON_MAPPINGS,
     mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
@@ -39,6 +53,21 @@ ASYNC_MAPPINGS: Final = (
     ),
 )
 
+CALLBACK_SUCCESS_SYNC_MAPPINGS: Final = (*SYNC_MAPPINGS, SUCCESS_CALLBACK_SYNC_MAPPING)
+CALLBACK_SUCCESS_ASYNC_MAPPINGS: Final = (*ASYNC_MAPPINGS, SUCCESS_CALLBACK_ASYNC_MAPPING)
+CALLBACK_FAILURE_SYNC_MAPPINGS: Final = (
+    *COMMON_MAPPINGS,
+    mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
+    FAILURE_CALLBACK_MAPPING,
+)
+CALLBACK_FAILURE_ASYNC_MAPPINGS: Final = (
+    *COMMON_MAPPINGS,
+    mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
+    mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"),
+    FAILURE_CALLBACK_MAPPING,
+)
+
+
 AZURE_COMMON_MAPPINGS: Final = (
     *COMMON_MAPPINGS[:7],
     mapping(
@@ -99,6 +128,34 @@ def _mistral_fixture(engine: Engine, _base_url: str) -> RouteFixture:
     return _fixture(engine, "mistral/mistral-ocr-latest")
 
 
+def _callback_fixture(engine: Engine, *, failure: bool) -> RouteFixture:
+    fixture: Final = _fixture(engine, "mistral/mistral-ocr-latest")
+    provider_responses: Final = (
+        (
+            RecordedHttpResponse.from_bytes(
+                400,
+                (HttpHeader(name="content-type", value="application/json"),),
+                b'{"message":"trace callback provider failure"}',
+            ),
+        )
+        if failure
+        else fixture.provider_responses
+    )
+    return RouteFixture(
+        kwargs=fixture.kwargs,
+        provider_responses=provider_responses,
+        expected_failure=failure,
+    )
+
+
+def _mistral_callback_success_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    return _callback_fixture(engine, failure=False)
+
+
+def _mistral_callback_failure_fixture(engine: Engine, _base_url: str) -> RouteFixture:
+    return _callback_fixture(engine, failure=True)
+
+
 def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture:
     return _fixture(
         engine,
@@ -304,36 +361,50 @@ TRACE_SUITE: Final = TraceSuite(
             name="mistral",
             fixture=_mistral_fixture,
             mappings=COMMON_MAPPINGS,
-            sync_mappings=SYNC_MAPPINGS,
-            async_mappings=ASYNC_MAPPINGS,
+            sync_mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            async_mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+        ),
+        TraceScenario(
+            name="mistral-callback-success",
+            fixture=_mistral_callback_success_fixture,
+            mappings=COMMON_MAPPINGS,
+            sync_mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS,
+            async_mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS,
+        ),
+        TraceScenario(
+            name="mistral-callback-failure",
+            fixture=_mistral_callback_failure_fixture,
+            mappings=(*COMMON_MAPPINGS, FAILURE_CALLBACK_MAPPING),
+            sync_mappings=CALLBACK_FAILURE_SYNC_MAPPINGS,
+            async_mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS,
         ),
         TraceScenario(
             name="azure-ai",
             fixture=_azure_fixture,
             mappings=AZURE_COMMON_MAPPINGS,
-            sync_mappings=AZURE_SYNC_MAPPINGS,
-            async_mappings=AZURE_ASYNC_MAPPINGS,
+            sync_mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            async_mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
         ),
         TraceScenario(
             name="azure-document-intelligence",
             fixture=_azure_document_intelligence_fixture,
             mappings=DOCUMENT_INTELLIGENCE_COMMON_MAPPINGS,
-            sync_mappings=DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS,
-            async_mappings=DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS,
+            sync_mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            async_mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
         ),
         TraceScenario(
             name="vertex-ai",
             fixture=_vertex_fixture,
             mappings=VERTEX_COMMON_MAPPINGS,
-            sync_mappings=VERTEX_SYNC_MAPPINGS,
-            async_mappings=VERTEX_ASYNC_MAPPINGS,
+            sync_mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            async_mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
         ),
         TraceScenario(
             name="vertex-deepseek",
             fixture=_vertex_deepseek_fixture,
             mappings=DEEPSEEK_COMMON_MAPPINGS,
-            sync_mappings=DEEPSEEK_SYNC_MAPPINGS,
-            async_mappings=DEEPSEEK_ASYNC_MAPPINGS,
+            sync_mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
+            async_mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
         ),
     ),
 )

From 2700ffe8c978229628767a45e2c83b7c1bb0297a Mon Sep 17 00:00:00 2001
From: yucheng-berri <yucheng@berri.ai>
Date: Mon, 7 Sep 2026 11:12:34 -0700
Subject: [PATCH 121/164] fix: capture provider request id in failure logging
 payloads (#40045)

* fix: capture provider request id in failure logging payloads

* fix: include common provider request id headers

* test: remove accidental formatter churn

* test: type provider request id cases

* test: cover azure and google request id headers

* refactor: share provider request id headers from constants

* chore: remove redundant request header comment

* test: cover provider header lookup failures
---
 litellm/constants.py                          |  10 ++
 litellm/litellm_core_utils/litellm_logging.py |  27 +++++
 litellm/types/utils.py                        |   1 +
 .../test_litellm_logging.py                   | 111 +++++++++++++++++-
 4 files changed, 148 insertions(+), 1 deletion(-)

diff --git a/litellm/constants.py b/litellm/constants.py
index ce744e9c58a..d53686e5e5b 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -1901,6 +1901,16 @@ HTTP_FRAMING_HEADERS: Final[frozenset[str]] = frozenset(
     }
 )
 
+PROVIDER_REQUEST_ID_HEADERS: Final[tuple[str, ...]] = (
+    "x-amzn-requestid",
+    "x-request-id",
+    "request-id",
+    "x-ms-request-id",
+    "apim-request-id",
+    "x-goog-request-id",
+    "cf-ray",
+)
+
 # Browser-facing security headers that a malicious or misconfigured upstream
 # provider must not be able to set on the proxy's own response.
 BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 60abad17c03..ca2cca5360f 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -42,6 +42,7 @@ from litellm.caching.caching_handler import LLMCachingHandler
 from litellm.constants import (
     DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
     DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
+    PROVIDER_REQUEST_ID_HEADERS,
     SENTRY_DENYLIST,
     SENTRY_PII_DENYLIST,
 )
@@ -255,6 +256,30 @@ _in_memory_loggers: Final[list[CustomLogger]] = []
 
 _STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggingMetadata.__annotations__.keys())
 
+
+def _get_provider_request_id(original_exception: Exception) -> str | None:
+    try:
+        error_response: Final = getattr(original_exception, "response", None)
+        header_sources: Final = (
+            _get_response_headers(original_exception),
+            getattr(error_response, "headers", None),
+            getattr(original_exception, "litellm_response_headers", None),
+        )
+        return next(
+            (
+                str(value)
+                for expected_header_name in PROVIDER_REQUEST_ID_HEADERS
+                for headers in header_sources
+                if isinstance(headers, Mapping)
+                for header_name, value in headers.items()
+                if isinstance(header_name, str) and header_name.lower() == expected_header_name and value
+            ),
+            None,
+        )
+    except Exception:
+        return None
+
+
 ### GLOBAL VARIABLES ###
 
 # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys
@@ -5664,6 +5689,7 @@ class StandardLoggingPayloadSetup:
         rate_limit_category: Final = validate_rate_limit_category(getattr(original_exception, "category", None))
         rate_limit_type: Final = validate_rate_limit_type(getattr(original_exception, "rate_limit_type", None))
         budget_error: Final = original_exception if isinstance(original_exception, BudgetExceededError) else None
+        provider_request_id: Final = _get_provider_request_id(original_exception) if original_exception else None
 
         return StandardLoggingPayloadErrorInformation(
             error_code=error_status,
@@ -5671,6 +5697,7 @@ class StandardLoggingPayloadSetup:
             llm_provider=_llm_provider_in_exception,
             traceback=_redact_string(traceback_info),
             error_message=_redact_string(error_message),
+            error_provider_request_id=provider_request_id,
             error_rate_limit_category=rate_limit_category,
             error_rate_limit_type=rate_limit_type,
             error_budget_entity_type=budget_error.entity_type if budget_error else None,
diff --git a/litellm/types/utils.py b/litellm/types/utils.py
index 3a1422bb190..010ff18d166 100644
--- a/litellm/types/utils.py
+++ b/litellm/types/utils.py
@@ -3064,6 +3064,7 @@ class StandardLoggingPayloadErrorInformation(TypedDict, total=False):
     llm_provider: str | None
     traceback: str | None
     error_message: str | None
+    error_provider_request_id: ReadOnly[str | None]
     # error_rate_limit_category:
     #   For 429 / rate-limit errors, the source of the rate limit. One of the
     #   string values defined by `litellm.exceptions.RateLimitErrorCategory`
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 3748b5a9d64..0fdca755685 100644
--- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
@@ -1,7 +1,9 @@
+import asyncio
 import contextlib
+import datetime
 import os
 import sys
-import asyncio
+from typing import Literal
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -6439,6 +6441,113 @@ def test_get_standard_logging_object_payload_survives_logging_obj_without_timing
     assert payload["hidden_params"]["litellm_overhead_time_ms"] is None
 
 
+@pytest.mark.parametrize(
+    ("header_name", "header_source"),
+    (
+        ("x-amzn-RequestId", "response"),
+        ("x-request-id", "response"),
+        ("request-id", "response"),
+        ("x-ms-request-id", "response"),
+        ("apim-request-id", "response"),
+        ("x-goog-request-id", "response"),
+        ("cf-ray", "response"),
+        ("X-Request-Id", "litellm_response_headers"),
+        ("X-MS-Request-ID", "headers"),
+    ),
+)
+def test_failure_standard_logging_payload_captures_provider_request_id(
+    logging_obj: LitellmLogging,
+    header_name: str,
+    header_source: Literal["response", "litellm_response_headers", "headers"],
+):
+    from litellm.litellm_core_utils.litellm_logging import (
+        get_standard_logging_object_payload,
+    )
+
+    request_id = "provider-request-123"
+    request = httpx.Request("POST", "https://provider.example/v1/chat/completions")
+    response = httpx.Response(429, headers={header_name: request_id}, request=request)
+    provider_error = httpx.HTTPStatusError("provider error", request=request, response=response)
+    if header_source == "litellm_response_headers":
+        response.headers.clear()
+        provider_error.litellm_response_headers = {header_name: request_id}
+    elif header_source == "headers":
+        response.headers.clear()
+        provider_error.headers = {header_name: request_id}
+    now = datetime.datetime.now()
+
+    payload = get_standard_logging_object_payload(
+        kwargs={"litellm_call_id": "call-1", "model": "test-model", "messages": []},
+        init_response_obj={},
+        start_time=now,
+        end_time=now,
+        logging_obj=logging_obj,
+        status="failure",
+        original_exception=provider_error,
+    )
+
+    assert payload is not None
+    assert payload["error_information"] is not None
+    assert payload["error_information"]["error_provider_request_id"] == request_id
+
+
+def test_get_error_information_ignores_unsupported_headers() -> None:
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    request = httpx.Request("POST", "https://provider.example/v1/chat/completions")
+    response = httpx.Response(429, headers={"retry-after": "3"}, request=request)
+    provider_error = httpx.HTTPStatusError("provider error", request=request, response=response)
+
+    error_information = StandardLoggingPayloadSetup.get_error_information(provider_error)
+
+    assert error_information["error_provider_request_id"] is None
+
+
+def test_get_error_information_uses_header_precedence_and_fallback() -> None:
+    from litellm.exceptions import RateLimitError
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    request = httpx.Request("POST", "https://provider.example/v1/chat/completions")
+    response = httpx.Response(
+        429,
+        headers={"x-request-id": "response-id", "x-amzn-requestid": "amazon-id"},
+        request=request,
+    )
+    provider_error = RateLimitError(
+        message="provider error",
+        llm_provider="test-provider",
+        model="test-model",
+        response=response,
+        headers={"retry-after": "3"},
+    )
+
+    error_information = StandardLoggingPayloadSetup.get_error_information(provider_error)
+
+    assert error_information["error_provider_request_id"] == "amazon-id"
+
+
+def test_get_error_information_ignores_malformed_headers() -> None:
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    provider_error = Exception("provider error")
+    provider_error.headers = [("x-request-id", "provider-request-123")]
+
+    error_information = StandardLoggingPayloadSetup.get_error_information(provider_error)
+
+    assert error_information["error_provider_request_id"] is None
+
+
+def test_get_provider_request_id_ignores_header_lookup_errors() -> None:
+    from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id
+
+    class HeaderLookupError(Exception):
+        @property
+        def response(self) -> object:
+            raise RuntimeError("headers unavailable")
+
+    assert _get_provider_request_id(HeaderLookupError("provider error")) is None
+
+
 def test_get_standard_logging_object_payload_failure_status_keeps_overhead_none(logging_obj):
     """A post_call guardrail can fail the request after the upstream call succeeded; the failure
     payload keeps litellm_overhead_time_ms None, matching responses that carry their own _hidden_params."""

From 9acc01efce25bc83769a5f45f0c1de4991edcb3e Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng@berri.ai>
Date: Mon, 7 Sep 2026 11:18:16 -0700
Subject: [PATCH 122/164] test(e2e/ui): cover member role and budget edits,
 member permission delegation, and team guardrail removal (#40042)

* test(e2e/ui): cover member role and budget edits, member permission delegation, and team guardrail removal

Three Playwright specs for the Teams flows enterprise customers hit most, each
owning its fixtures and proving the mutation through a read-back rather than a
toast.

- teamMemberEdit: an admin edits a member's team role and per-member budget,
  and both survive a reload of the Members table
- memberPermissions: a plain member is refused /key/generate for their team,
  a team admin grants it on the Member Permissions tab, and the member then
  creates a team key that serves a real completion
- teamGuardrailRemoval: clearing a team's only guardrail on the Settings tab
  really clears it, and traffic the guardrail refused starts serving again

* test(e2e/ui): make the new team specs safe to run in parallel

Fixture ids came from Date.now(), so two repeats starting in the same
millisecond minted the same user id: one got a 409 and the loser's teardown
deleted the user the other was still signed in as. Ids now carry a random
suffix.

Also move the member-permissions setup inside the cleanup-protected block so a
half-finished setup cannot leak a team, and close both browser contexts the
test opens.
---
 .../proxy-admin/teamGuardrailRemoval.spec.ts  | 145 ++++++++++++++
 .../tests/proxy-admin/teamMemberEdit.spec.ts  | 129 +++++++++++++
 .../team-admin/memberPermissions.spec.ts      | 177 ++++++++++++++++++
 3 files changed, 451 insertions(+)
 create mode 100644 tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts
 create mode 100644 tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts
 create mode 100644 tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts

diff --git a/tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts b/tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts
new file mode 100644
index 00000000000..f7930378b1e
--- /dev/null
+++ b/tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts
@@ -0,0 +1,145 @@
+import { test, expect, type APIRequestContext, type Page as PlaywrightPage } from "@playwright/test";
+import { ADMIN_STORAGE_PATH } from "../../constants";
+import { Page } from "../../fixtures/pages";
+import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation";
+import { proxyIsPremium } from "../../helpers/premium";
+import { readBack } from "../../helpers/roundTrip";
+import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic";
+
+const auth = () => ({ Authorization: `Bearer ${masterKey()}` });
+
+async function guardrailId(request: APIRequestContext, name: string): Promise<string | undefined> {
+  const res = await request.get("/v2/guardrails/list", { headers: auth() });
+  expect(res.ok(), `GET /v2/guardrails/list (${res.status()})`).toBe(true);
+  const rows = (await res.json()).guardrails as { guardrail_id: string; guardrail_name: string | null }[];
+  return rows.find((row) => row.guardrail_name === name)?.guardrail_id;
+}
+
+async function teamGuardrails(page: PlaywrightPage, teamId: string): Promise<string[]> {
+  const body = await readBack<{ team_info: { metadata: { guardrails?: string[] } | null } }>(
+    page,
+    `/team/info?team_id=${encodeURIComponent(teamId)}`,
+  );
+  return body.team_info.metadata?.guardrails ?? [];
+}
+
+async function keywordPromptStatus(request: APIRequestContext, apiKey: string, keyword: string): Promise<number> {
+  const res = await request.post("/v1/chat/completions", {
+    headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" },
+    data: { model: CHAT_MODEL_A, messages: [{ role: "user", content: `please tell me about ${keyword}` }] },
+  });
+  return res.status();
+}
+
+test.describe("Proxy Admin - Team guardrail removal", () => {
+  test.use({ storageState: ADMIN_STORAGE_PATH });
+
+  test("Clearing a team's only guardrail on the Settings tab lets blocked traffic through again", async ({
+    page,
+    request,
+  }) => {
+    test.skip(!proxyIsPremium(), "proxy under test is unlicensed, so team guardrails are premium-gated");
+
+    const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`;
+    const guardrailName = `e2e-team-guardrail-${stamp}`;
+    const bannedKeyword = `e2eteamban${stamp}`;
+    const teamAlias = `e2e-guardrail-team-${stamp}`;
+
+    let teamId = "";
+    let teamKey = "";
+    try {
+      const guardrailRes = await request.post("/guardrails", {
+        headers: auth(),
+        data: {
+          guardrail: {
+            guardrail_name: guardrailName,
+            litellm_params: {
+              guardrail: "litellm_content_filter",
+              mode: "pre_call",
+              default_on: false,
+              blocked_words: [{ keyword: bannedKeyword, action: "BLOCK" }],
+            },
+          },
+        },
+      });
+      expect(
+        guardrailRes.ok(),
+        `POST /guardrails failed (${guardrailRes.status()}): ${await guardrailRes.text()}`,
+      ).toBe(true);
+
+      const teamRes = await request.post("/team/new", {
+        headers: auth(),
+        data: { team_alias: teamAlias, models: [CHAT_MODEL_A], metadata: { guardrails: [guardrailName] } },
+      });
+      expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true);
+      teamId = (await teamRes.json()).team_id as string;
+
+      const keyRes = await request.post("/key/generate", { headers: auth(), data: { team_id: teamId } });
+      expect(keyRes.ok(), `POST /key/generate failed (${keyRes.status()}): ${await keyRes.text()}`).toBe(true);
+      teamKey = (await keyRes.json()).key as string;
+
+      await expect
+        .poll(async () => keywordPromptStatus(request, teamKey, bannedKeyword), {
+          message: "the team's guardrail never started refusing the banned keyword",
+          timeout: 60_000,
+        })
+        .toBe(400);
+
+      await navigateToPage(page, Page.Teams);
+      await dismissFeedbackPopup(page);
+      await clickTeamId(page, teamId);
+      await page.getByRole("tab", { name: "Settings" }).click();
+      await page.getByRole("button", { name: "Edit Settings" }).click();
+
+      const chip = page.locator('[data-slot="combobox-chip"]').filter({ hasText: guardrailName });
+      await expect(chip).toBeVisible({ timeout: 10_000 });
+      await chip.locator('[data-slot="combobox-chip-remove"]').click();
+      await expect(chip).toHaveCount(0, { timeout: 10_000 });
+
+      await page.getByRole("button", { name: "Save Changes" }).click();
+
+      await expect
+        .poll(async () => teamGuardrails(page, teamId), {
+          message: "the team still carries a guardrail in /team/info after the save",
+          timeout: 20_000,
+        })
+        .toEqual([]);
+
+      await page.reload();
+      await page.getByRole("tab", { name: "Settings" }).click();
+      await page.getByRole("button", { name: "Edit Settings" }).click();
+      await expect(page.getByRole("combobox", { name: "Select guardrails" })).toBeVisible({ timeout: 15_000 });
+      await expect(
+        page.locator('[data-slot="combobox-chip"]').filter({ hasText: guardrailName }),
+        "the removed guardrail is gone from the Settings tab after a reload",
+      ).toHaveCount(0);
+
+      await expect
+        .poll(async () => keywordPromptStatus(request, teamKey, bannedKeyword), {
+          message: "the team key is still refused for a keyword whose guardrail was removed",
+          timeout: 60_000,
+        })
+        .toBe(200);
+
+      const served = await request.post("/v1/chat/completions", {
+        headers: { Authorization: `Bearer ${teamKey}`, "Content-Type": "application/json" },
+        data: {
+          model: CHAT_MODEL_A,
+          messages: [{ role: "user", content: `please tell me about ${bannedKeyword}` }],
+        },
+      });
+      expect((await served.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT);
+    } finally {
+      if (teamKey) {
+        await request.post("/key/delete", { headers: auth(), data: { keys: [teamKey] } });
+      }
+      if (teamId) {
+        await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } });
+      }
+      const id = await guardrailId(request, guardrailName);
+      if (id) {
+        await request.delete(`/guardrails/${id}`, { headers: auth() });
+      }
+    }
+  });
+});
diff --git a/tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts b/tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts
new file mode 100644
index 00000000000..c7325e78f75
--- /dev/null
+++ b/tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts
@@ -0,0 +1,129 @@
+import { test, expect, type Page as PlaywrightPage } from "@playwright/test";
+import { ADMIN_STORAGE_PATH } from "../../constants";
+import { Page } from "../../fixtures/pages";
+import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation";
+import { readBack } from "../../helpers/roundTrip";
+import { CHAT_MODEL_A, masterKey } from "../../helpers/traffic";
+
+interface TeamInfoResponse {
+  team_info: {
+    models: string[];
+    members_with_roles: { user_id?: string; role?: string }[];
+  };
+  team_memberships: {
+    user_id: string;
+    litellm_budget_table: { max_budget: number | null } | null;
+  }[];
+}
+
+const auth = () => ({ Authorization: `Bearer ${masterKey()}` });
+
+async function teamInfo(page: PlaywrightPage, teamId: string): Promise<TeamInfoResponse> {
+  return readBack<TeamInfoResponse>(page, `/team/info?team_id=${encodeURIComponent(teamId)}`);
+}
+
+function roleOf(info: TeamInfoResponse, userId: string): string | undefined {
+  return info.team_info.members_with_roles.find((member) => member.user_id === userId)?.role;
+}
+
+function budgetOf(info: TeamInfoResponse, userId: string): number | null | undefined {
+  return info.team_memberships.find((membership) => membership.user_id === userId)?.litellm_budget_table?.max_budget;
+}
+
+function otherMembers(info: TeamInfoResponse, userId: string): string[] {
+  return info.team_info.members_with_roles
+    .filter((member) => member.user_id !== userId)
+    .map((member) => `${member.user_id}:${member.role}`)
+    .sort();
+}
+
+test.describe("Proxy Admin - Team member edit", () => {
+  test.use({ storageState: ADMIN_STORAGE_PATH });
+
+  const createdTeams: string[] = [];
+  const createdUsers: string[] = [];
+
+  test.afterEach(async ({ request }) => {
+    for (const teamId of createdTeams.splice(0)) {
+      await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } });
+    }
+    for (const userId of createdUsers.splice(0)) {
+      await request.post("/user/delete", { headers: auth(), data: { user_ids: [userId] } });
+    }
+  });
+
+  test("Editing a member's role and per-member budget persists and survives a reload", async ({ page, request }) => {
+    const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`;
+    const memberId = `e2e-member-edit-${stamp}`;
+    const teamAlias = `e2e-member-edit-team-${stamp}`;
+
+    createdUsers.push(memberId);
+    const created = await request.post("/user/new", {
+      headers: auth(),
+      data: { user_id: memberId, user_role: "internal_user", auto_create_key: false },
+    });
+    expect(created.ok(), `POST /user/new failed (${created.status()}): ${await created.text()}`).toBe(true);
+
+    const teamRes = await request.post("/team/new", {
+      headers: auth(),
+      data: {
+        team_alias: teamAlias,
+        models: [CHAT_MODEL_A],
+        members_with_roles: [{ user_id: memberId, role: "admin" }],
+      },
+    });
+    expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true);
+    const teamId = (await teamRes.json()).team_id as string;
+    createdTeams.push(teamId);
+
+    const before = await teamInfo(page, teamId);
+    expect(roleOf(before, memberId), "the member starts out as a team admin").toBe("admin");
+    expect(budgetOf(before, memberId) ?? null, "the member starts out with no per-member budget").toBeNull();
+    expect(
+      otherMembers(before, memberId).length,
+      "the team has another member for the edit to leave alone",
+    ).toBeGreaterThan(0);
+
+    await navigateToPage(page, Page.Teams);
+    await dismissFeedbackPopup(page);
+    await clickTeamId(page, teamId);
+    await page.getByRole("tab", { name: "Members" }).click();
+
+    const memberRow = page.locator("tr", { hasText: memberId }).first();
+    await expect(memberRow).toBeVisible({ timeout: 10_000 });
+    await memberRow.getByTestId("edit-member").click();
+
+    const modal = page.getByRole("dialog", { name: "Edit Member" });
+    await expect(modal).toBeVisible({ timeout: 10_000 });
+
+    await modal.getByLabel(/^Role/).click();
+    await page.getByRole("option", { name: "User", exact: true }).click();
+    await modal.getByLabel(/Team Member Budget \(USD\)/).fill("5");
+    await modal.getByRole("button", { name: "Save Changes" }).click();
+
+    await expect(page.getByText("Team member updated successfully").first()).toBeVisible({ timeout: 10_000 });
+
+    await expect
+      .poll(
+        async () => {
+          const info = await teamInfo(page, teamId);
+          return [roleOf(info, memberId), budgetOf(info, memberId)];
+        },
+        { message: "the member's role and budget never landed in /team/info", timeout: 20_000 },
+      )
+      .toEqual(["user", 5]);
+
+    await page.reload();
+    await page.getByRole("tab", { name: "Members" }).click();
+    const reloadedRow = page.locator("tr", { hasText: memberId }).first();
+    await expect(reloadedRow).toBeVisible({ timeout: 15_000 });
+    await expect(reloadedRow.getByText("user", { exact: true }), "role shown after a reload").toBeVisible();
+    await expect(reloadedRow.getByText("$5.00"), "per-member budget shown after a reload").toBeVisible();
+
+    const after = await teamInfo(page, teamId);
+    expect(after.team_info.models, "model access untouched by a member edit").toEqual(before.team_info.models);
+    expect(otherMembers(after, memberId), "the rest of the roster untouched by a member edit").toEqual(
+      otherMembers(before, memberId),
+    );
+  });
+});
diff --git a/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts b/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts
new file mode 100644
index 00000000000..75fb3be9b64
--- /dev/null
+++ b/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts
@@ -0,0 +1,177 @@
+import { test, expect, type Browser, type BrowserContext, type Page as PlaywrightPage } from "@playwright/test";
+import { Page } from "../../fixtures/pages";
+import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation";
+import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic";
+
+const PASSWORD = "E2e-Member-Perms-Pass-1!";
+
+const auth = () => ({ Authorization: `Bearer ${masterKey()}` });
+
+async function sessionKey(page: PlaywrightPage): Promise<string> {
+  const cookie = (await page.context().cookies()).find((candidate) => candidate.name === "token");
+  expect(cookie?.value, "logged-in session carries a token cookie").toBeTruthy();
+  const payload = JSON.parse(Buffer.from(cookie!.value.split(".")[1], "base64url").toString("utf-8")) as {
+    key?: string;
+  };
+  expect(payload.key, "session JWT carries the virtual key the dashboard calls with").toMatch(/^sk-/);
+  return payload.key!;
+}
+
+async function signIn(browser: Browser, email: string): Promise<BrowserContext> {
+  const context = await browser.newContext({ storageState: { cookies: [], origins: [] } });
+  const page = await context.newPage();
+  await page.goto("/ui/login");
+  await page.getByPlaceholder("Enter your username").fill(email);
+  await page.getByPlaceholder("Enter your password").fill(PASSWORD);
+  await page.getByRole("button", { name: "Login", exact: true }).click();
+  await expect(page.locator("a", { hasText: "Virtual Keys" })).toBeVisible({ timeout: 30_000 });
+  await dismissFeedbackPopup(page);
+  return context;
+}
+
+test.describe("Team Admin - Member permissions", () => {
+  test.use({ storageState: { cookies: [], origins: [] } });
+
+  test("Granting /key/generate lets a plain member create a team key that serves traffic", async ({
+    browser,
+    request,
+  }) => {
+    const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`;
+    const adminId = `e2e-perm-admin-${stamp}`;
+    const memberId = `e2e-perm-member-${stamp}`;
+    const adminEmail = `${adminId}@test.local`;
+    const memberEmail = `${memberId}@test.local`;
+    const teamAlias = `e2e-perm-team-${stamp}`;
+
+    const createUser = async (userId: string, email: string): Promise<void> => {
+      const created = await request.post("/user/new", {
+        headers: auth(),
+        data: { user_id: userId, user_email: email, user_role: "internal_user", auto_create_key: false },
+      });
+      expect(created.ok(), `POST /user/new for ${userId} (${created.status()}): ${await created.text()}`).toBe(true);
+      const password = await request.post("/user/update", {
+        headers: auth(),
+        data: { user_id: userId, password: PASSWORD },
+      });
+      expect(password.ok(), `POST /user/update for ${userId} (${password.status()})`).toBe(true);
+    };
+
+    let teamId = "";
+    const createdKeys: string[] = [];
+    const contexts: BrowserContext[] = [];
+    try {
+      await createUser(adminId, adminEmail);
+      await createUser(memberId, memberEmail);
+
+      const teamRes = await request.post("/team/new", {
+        headers: auth(),
+        data: {
+          team_alias: teamAlias,
+          models: [CHAT_MODEL_A],
+          members_with_roles: [
+            { user_id: adminId, role: "admin" },
+            { user_id: memberId, role: "user" },
+          ],
+        },
+      });
+      expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true);
+      teamId = (await teamRes.json()).team_id as string;
+
+      const memberContext = await signIn(browser, memberEmail);
+      contexts.push(memberContext);
+      const memberPage = memberContext.pages()[0];
+      const memberSessionKey = await sessionKey(memberPage);
+
+      const refused = await memberPage.request.post("/key/generate", {
+        headers: { Authorization: `Bearer ${memberSessionKey}`, "Content-Type": "application/json" },
+        data: { team_id: teamId, key_alias: `e2e-perm-denied-${stamp}` },
+      });
+      expect(refused.status(), "a plain member cannot mint a team key before the grant").toBe(401);
+      expect(await refused.text()).toContain("/key/generate");
+
+      const adminContext = await signIn(browser, adminEmail);
+      contexts.push(adminContext);
+      const adminPage = adminContext.pages()[0];
+      await navigateToPage(adminPage, Page.Teams);
+      await clickTeamId(adminPage, teamId);
+      await adminPage.getByRole("tab", { name: "Member Permissions" }).click();
+
+      for (const route of ["/key/generate", "/key/update"]) {
+        await adminPage.getByRole("row").filter({ hasText: route }).getByRole("checkbox").check();
+      }
+      await adminPage.getByRole("button", { name: "Save Changes" }).click();
+      await expect(adminPage.getByText("Permissions updated successfully").first()).toBeVisible({ timeout: 10_000 });
+
+      await expect
+        .poll(
+          async () => {
+            const res = await request.get(`/team/permissions_list?team_id=${encodeURIComponent(teamId)}`, {
+              headers: auth(),
+            });
+            if (!res.ok()) return [];
+            return ((await res.json()).team_member_permissions ?? []) as string[];
+          },
+          { message: "the granted permissions never landed in /team/permissions_list", timeout: 20_000 },
+        )
+        .toEqual(expect.arrayContaining(["/key/generate", "/key/update"]));
+
+      const keyAlias = `e2e-perm-key-${stamp}`;
+      await navigateToPage(memberPage, Page.ApiKeys);
+      await memberPage.getByRole("button", { name: /Create New Key/i }).click();
+      await expect(memberPage.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 });
+      await memberPage.getByLabel(/Key Name/).fill(keyAlias);
+
+      const teamSelect = memberPage.getByTestId("team-dropdown").getByRole("combobox");
+      await teamSelect.click();
+      await memberPage.keyboard.type(teamAlias);
+      await memberPage.getByRole("option", { name: teamAlias }).first().click();
+
+      await memberPage.getByRole("combobox", { name: "Select models" }).click();
+      await memberPage.getByRole("option", { name: "All Team Models", exact: true }).click();
+      await memberPage.keyboard.press("Escape");
+
+      await memberPage.getByRole("button", { name: "Create Key", exact: true }).click();
+      const saveDialog = memberPage.getByRole("dialog", { name: "Save your Key" });
+      await expect(saveDialog).toBeVisible({ timeout: 15_000 });
+      const apiKey = (await saveDialog.locator("pre").innerText()).trim();
+      expect(apiKey).toMatch(/^sk-/);
+      createdKeys.push(apiKey);
+      await memberPage.keyboard.press("Escape");
+
+      await expect
+        .poll(
+          async () => {
+            const res = await request.get(
+              `/key/list?team_id=${encodeURIComponent(teamId)}&return_full_object=true&size=100`,
+              { headers: auth() },
+            );
+            if (!res.ok()) return null;
+            const row = ((await res.json()).keys as Record<string, unknown>[]).find(
+              (candidate) => candidate.key_alias === keyAlias,
+            );
+            return row ? [row.user_id, row.team_id] : null;
+          },
+          { message: `key ${keyAlias} never appeared on the team with the member as its owner`, timeout: 20_000 },
+        )
+        .toEqual([memberId, teamId]);
+
+      const served = await request.post("/v1/chat/completions", {
+        headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" },
+        data: { model: CHAT_MODEL_A, messages: [{ role: "user", content: `member key ping ${stamp}` }] },
+      });
+      expect(served.status(), "the delegated key is a real key the gateway serves").toBe(200);
+      expect((await served.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT);
+    } finally {
+      for (const context of contexts) {
+        await context.close();
+      }
+      for (const key of createdKeys) {
+        await request.post("/key/delete", { headers: auth(), data: { keys: [key] } });
+      }
+      if (teamId) {
+        await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } });
+      }
+      await request.post("/user/delete", { headers: auth(), data: { user_ids: [adminId, memberId] } });
+    }
+  });
+});

From 5f2b4d27d712847c0f4fef5333f24848983686f5 Mon Sep 17 00:00:00 2001
From: yujonglee <yujong@berri.ai>
Date: Mon, 7 Sep 2026 11:23:12 -0700
Subject: [PATCH 123/164] test(ocr): add SDK callback E2E parity (#40061)

* test(ocr): cover SDK callback parity

* test(ocr): compare callback kwargs
---
 litellm/ocr/main.py                           |  74 +++--
 .../shared/parity/compare.py                  |   1 +
 .../shared/parity/models.py                   |  21 ++
 .../shared/parity/runner.py                   |  10 +-
 .../shared/parity/test_parity.py              |  52 ++-
 .../strategies/e2e_parity/__init__.py         |   4 +-
 .../e2e_parity/sdk/ocr/test_sdk_parity.py     | 295 +++++++++++++++++-
 .../unit_tests_mapping/cases/ocr.py           |   4 +
 tests/test_litellm/ocr/test_rust_bridge.py    |  22 ++
 9 files changed, 453 insertions(+), 30 deletions(-)

diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py
index 56c2292d00c..df3f9d2096b 100644
--- a/litellm/ocr/main.py
+++ b/litellm/ocr/main.py
@@ -29,6 +29,7 @@ from litellm.llms.base_llm.ocr.transformation import (
 )
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
 from litellm.rust_bridge import ocr as rust_ocr_bridge
+from litellm.rust_bridge.bindings import native_exception_types
 from litellm.rust_bridge.configuration import rust_enabled
 from litellm.types.router import GenericLiteLLMParams
 from litellm.utils import ProviderConfigManager, client
@@ -281,6 +282,33 @@ def _prepare_rust_ocr_call(
     )
 
 
+def _map_rust_ocr_error(
+    error: Exception,
+    prepared_request: _PreparedOCRRequest,
+    exception_types: tuple[type[BaseException], type[BaseException]] | None,
+) -> Exception:
+    if exception_types is None:
+        return error
+    _, upstream_error = exception_types
+    if not isinstance(error, upstream_error):
+        return error
+    error_args: Final = cast(  # cast-ok: BaseException.args is typed with Any in the standard library stubs
+        tuple[object, ...], error.args
+    )
+    status_value: Final = error_args[0] if error_args else 0
+    message_value: Final = error_args[1] if len(error_args) > 1 else str(error)
+    status: Final = status_value if isinstance(status_value, int) else 0
+    message: Final = message_value if isinstance(message_value, str) else str(message_value)
+    error_factory: Final = cast(  # cast-ok: the legacy provider interface leaves callable parameters untyped
+        Callable[..., Exception], prepared_request.provider_config.get_error_class
+    )
+    return error_factory(
+        error_message=message,
+        status_code=status or 500,
+        headers={},  # mutable-ok: provider error factories require a concrete header dict
+    )
+
+
 def _run_rust_ocr(
     prepared_request: _PreparedOCRRequest,
     resolve_api_key: Callable[[str], str | None],
@@ -291,16 +319,19 @@ def _run_rust_ocr(
         prepared_request=prepared_request,
         resolve_api_key=resolve_api_key,
     )
-    rust_response: Final = rust_ocr_bridge.ocr(
-        model=prepared_request.model,
-        document=prepared_request.document,
-        api_key=prepared.api_key,
-        api_base=prepared.api_base,
-        custom_llm_provider=prepared_request.custom_llm_provider,
-        extra_headers=prepared.headers,
-        optional_params=prepared.optional_params,
-        timeout=prepared_request.effective_timeout,
-    )
+    try:
+        rust_response: Final = rust_ocr_bridge.ocr(
+            model=prepared_request.model,
+            document=prepared_request.document,
+            api_key=prepared.api_key,
+            api_base=prepared.api_base,
+            custom_llm_provider=prepared_request.custom_llm_provider,
+            extra_headers=prepared.headers,
+            optional_params=prepared.optional_params,
+            timeout=prepared_request.effective_timeout,
+        )
+    except Exception as error:
+        raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error
     if rust_response is None:
         return None
     return OCRResponse.model_validate(rust_response)
@@ -316,16 +347,19 @@ async def _run_rust_aocr(
         prepared_request=prepared_request,
         resolve_api_key=resolve_api_key,
     )
-    rust_response: Final = await rust_ocr_bridge.aocr(
-        model=prepared_request.model,
-        document=prepared_request.document,
-        api_key=prepared.api_key,
-        api_base=prepared.api_base,
-        custom_llm_provider=prepared_request.custom_llm_provider,
-        extra_headers=prepared.headers,
-        optional_params=prepared.optional_params,
-        timeout=prepared_request.effective_timeout,
-    )
+    try:
+        rust_response: Final = await rust_ocr_bridge.aocr(
+            model=prepared_request.model,
+            document=prepared_request.document,
+            api_key=prepared.api_key,
+            api_base=prepared.api_base,
+            custom_llm_provider=prepared_request.custom_llm_provider,
+            extra_headers=prepared.headers,
+            optional_params=prepared.optional_params,
+            timeout=prepared_request.effective_timeout,
+        )
+    except Exception as error:
+        raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error
     if rust_response is None:
         return None
     return OCRResponse.model_validate(rust_response)
diff --git a/tests/rust-python-harness/shared/parity/compare.py b/tests/rust-python-harness/shared/parity/compare.py
index adf85e5c8d7..88239ed042a 100644
--- a/tests/rust-python-harness/shared/parity/compare.py
+++ b/tests/rust-python-harness/shared/parity/compare.py
@@ -78,3 +78,4 @@ def assert_parity(baseline: Execution, candidate: Execution, baseline_user_agent
     validate_harness(baseline, candidate, baseline_user_agent)
     assert_request_parity(baseline.requests, candidate.requests)
     assert_value_parity(baseline.report, candidate.report)
+    assert_value_parity(baseline.callbacks, candidate.callbacks, path="$.callbacks")
diff --git a/tests/rust-python-harness/shared/parity/models.py b/tests/rust-python-harness/shared/parity/models.py
index 898b58d23ee..5612e218824 100644
--- a/tests/rust-python-harness/shared/parity/models.py
+++ b/tests/rust-python-harness/shared/parity/models.py
@@ -37,6 +37,25 @@ class SDKError(BaseModel):
     llm_provider: str | None
 
 
+class CallbackObservation(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    hook: Literal[
+        "log_success_event",
+        "async_log_success_event",
+        "log_failure_event",
+        "async_log_failure_event",
+    ]
+    phase: Literal["success", "failure"]
+    model: str | None
+    call_type: str | None
+    litellm_call_id: str | None
+    metadata: JsonValue
+    kwargs: JsonValue
+    payload: JsonValue
+    error: SDKError | None
+
+
 class SDKJsonChunk(BaseModel):
     model_config = ConfigDict(frozen=True)
 
@@ -119,6 +138,7 @@ class Execution(BaseModel):
 
     requests: tuple[CapturedRequest, ...]
     report: SDKReport
+    callbacks: tuple[CallbackObservation, ...] | None = None
 
 
 class SDKCommand(BaseModel):
@@ -133,6 +153,7 @@ class WorkerSuccess(BaseModel):
 
     status: Literal["ok"] = "ok"
     report: SDKReport
+    callbacks: tuple[CallbackObservation, ...] | None = None
 
 
 class WorkerFailure(BaseModel):
diff --git a/tests/rust-python-harness/shared/parity/runner.py b/tests/rust-python-harness/shared/parity/runner.py
index 43a583382cb..feae0201aee 100644
--- a/tests/rust-python-harness/shared/parity/runner.py
+++ b/tests/rust-python-harness/shared/parity/runner.py
@@ -39,9 +39,7 @@ class SubprocessRunner:
         return (
             sys.executable,
             "-m",
-            ".".join(
-                self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts
-            ),
+            ".".join(self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts),
             "--parity-worker",
             provider_url,
         )
@@ -113,7 +111,11 @@ class SubprocessWorker:
             )
         assert isinstance(result, WorkerSuccess)
         try:
-            return Execution(requests=self.provider.take_requests(len(responses)), report=result.report)
+            return Execution(
+                requests=self.provider.take_requests(len(responses)),
+                report=result.report,
+                callbacks=result.callbacks,
+            )
         except AssertionError:
             self.provider.reset()
             raise
diff --git a/tests/rust-python-harness/shared/parity/test_parity.py b/tests/rust-python-harness/shared/parity/test_parity.py
index 83daccdf8ba..7c8355fbfd8 100644
--- a/tests/rust-python-harness/shared/parity/test_parity.py
+++ b/tests/rust-python-harness/shared/parity/test_parity.py
@@ -7,7 +7,7 @@ import pytest
 from pydantic import BaseModel, ConfigDict, JsonValue, PrivateAttr
 
 from .compare import assert_model_parity, assert_parity
-from .models import CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report
+from .models import CallbackObservation, CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report
 
 SENTINEL: Final = "python-parity-fallback"
 
@@ -69,6 +69,56 @@ def test_parity_rejects_response_difference() -> None:
         assert_parity(python, rust, SENTINEL)
 
 
+def test_parity_distinguishes_unobserved_callbacks_from_zero_events() -> None:
+    python: Final = _execution(user_agent=SENTINEL)
+    rust: Final = _execution(user_agent="litellm-rust").model_copy(update={"callbacks": ()})
+
+    with pytest.raises(AssertionError, match=r"\$\.callbacks"):
+        assert_parity(python, rust, SENTINEL)
+
+
+def test_parity_rejects_callback_payload_difference() -> None:
+    observation: Final = CallbackObservation(
+        hook="log_success_event",
+        phase="success",
+        model="test-model",
+        call_type="ocr",
+        litellm_call_id="test-call",
+        metadata={"profile": "success"},
+        kwargs={"model": "test-model"},
+        payload={"model": "test-model", "pages": []},
+        error=None,
+    )
+    python: Final = _execution(user_agent=SENTINEL).model_copy(update={"callbacks": (observation,)})
+    rust: Final = _execution(user_agent="litellm-rust").model_copy(
+        update={"callbacks": (observation.model_copy(update={"payload": {"model": "changed", "pages": []}}),)}
+    )
+
+    with pytest.raises(AssertionError, match=r"\$\.callbacks"):
+        assert_parity(python, rust, SENTINEL)
+
+
+def test_parity_rejects_callback_kwargs_difference() -> None:
+    observation: Final = CallbackObservation(
+        hook="log_success_event",
+        phase="success",
+        model="test-model",
+        call_type="ocr",
+        litellm_call_id="test-call",
+        metadata={"profile": "success"},
+        kwargs={"model": "test-model"},
+        payload={"model": "test-model", "pages": []},
+        error=None,
+    )
+    python: Final = _execution(user_agent=SENTINEL).model_copy(update={"callbacks": (observation,)})
+    rust: Final = _execution(user_agent="litellm-rust").model_copy(
+        update={"callbacks": (observation.model_copy(update={"kwargs": {"model": "changed"}}),)}
+    )
+
+    with pytest.raises(AssertionError, match=r"\$\.callbacks"):
+        assert_parity(python, rust, SENTINEL)
+
+
 def test_parity_rejects_error_difference() -> None:
     python: Final = Execution(
         requests=(),
diff --git a/tests/rust-python-harness/strategies/e2e_parity/__init__.py b/tests/rust-python-harness/strategies/e2e_parity/__init__.py
index f668e178eef..95346bbd496 100644
--- a/tests/rust-python-harness/strategies/e2e_parity/__init__.py
+++ b/tests/rust-python-harness/strategies/e2e_parity/__init__.py
@@ -18,8 +18,8 @@ CASES: Final[tuple[CaseDefinition, ...]] = (
             coverage=Coverage.PARTIAL,
             module="tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.test_sdk_parity",
             note=(
-                "Recorded sync/async SDK parity; invalid-model provider errors differ, "
-                "and Reducto lacks a Rust contract."
+                "Recorded sync/async SDK parity with focused success/error callback profiles; "
+                "Reducto lacks a Rust contract, and known provider parity gaps remain."
             ),
         ),
         surface="sdk",
diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py
index e72980752f2..5c4abc78081 100644
--- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py
+++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py
@@ -1,23 +1,31 @@
 from __future__ import annotations
 
 import asyncio
+import datetime
+import queue
 import sys
 import tempfile
+import time
 import traceback
-from collections.abc import Callable, Coroutine, Generator
+from collections.abc import Callable, Coroutine, Generator, Mapping
 from contextlib import contextmanager
 from enum import Enum
 from functools import partial
 from pathlib import Path
 from typing import Annotated, Final, Literal, cast
+from urllib.parse import urlsplit, urlunsplit
 
 from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
 
+from litellm.integrations.custom_logger import CustomLogger
 from litellm.llms.base_llm.ocr.transformation import OCRResponse
 
 from .....shared.parity.compare import assert_parity
 from .....shared.parity.fixtures.store import fixture_id, recorded_fixtures
 from .....shared.parity.models import (
+    JSON_VALUE_ADAPTER,
+    CallbackObservation,
+    Execution,
     SDKCommand,
     SDKError,
     SDKReport,
@@ -39,6 +47,9 @@ from .fixtures.config import configured_fixture_directory
 from .fixtures.models import OcrParityCase, OcrSdkInput
 
 API_KEY: Final = "test-key"
+CALLBACK_DELAY_SECONDS: Final = 0.05
+CALLBACK_DRAIN_TIMEOUT_SECONDS: Final = 10.0
+CALLBACK_TERMINALS: Final[tuple[Literal["success", "failure"], ...]] = ("success", "failure")
 PYTHON_HTTP_SENTINEL: Final = "python-ocr-parity-fallback"
 PYTHON_VARIANT: Final = ExecutionVariant(name="Python", environment=(("LITELLM_RUST", "0"),))
 RUST_VARIANT: Final = ExecutionVariant(name="Rust", environment=(("LITELLM_RUST", "1"),))
@@ -75,10 +86,148 @@ class InvalidOcrWorkerCase(BaseModel):
     case: InvalidOcrCase
 
 
-OcrWorkerCase = Annotated[RecordedOcrWorkerCase | InvalidOcrWorkerCase, Field(discriminator="kind")]
+class CallbackOcrWorkerCase(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    kind: Literal["callback"] = "callback"
+    case: OcrParityCase
+    terminal: Literal["success", "failure"]
+
+
+OcrWorkerCase = Annotated[
+    RecordedOcrWorkerCase | InvalidOcrWorkerCase | CallbackOcrWorkerCase,
+    Field(discriminator="kind"),
+]
 OCR_WORKER_CASE_ADAPTER: Final[TypeAdapter[OcrWorkerCase]] = TypeAdapter(OcrWorkerCase)
 
 
+class RecordingCallback(CustomLogger):
+    def __init__(self) -> None:
+        self.message_logging: Final = True
+        self.turn_off_message_logging: Final = False
+        self._observations: Final[queue.SimpleQueue[CallbackObservation]] = queue.SimpleQueue()
+
+    def _normalized_kwargs(self, value: object, key: str | None = None) -> JsonValue:
+        if value is None or isinstance(value, (bool, int, float)):
+            return value
+        if isinstance(value, str):
+            if key != "api_base":
+                return value
+            parsed: Final = urlsplit(value)
+            return urlunsplit(("", "", parsed.path, parsed.query, parsed.fragment))
+        if isinstance(value, datetime.datetime):
+            return "datetime"
+        if isinstance(value, Exception):
+            return sdk_error_report(value).model_dump(mode="json")
+        if isinstance(value, BaseModel):
+            return self._normalized_kwargs(value.model_dump(mode="json"), key)
+        if isinstance(value, Mapping):
+            if any(not isinstance(map_key, str) for map_key in value):
+                raise TypeError("callback kwarg mappings must use string keys")
+            return {
+                map_key: self._normalized_kwargs(map_value, map_key)
+                for map_key, map_value in value.items()
+            }
+        if isinstance(value, (list, tuple)):
+            return [self._normalized_kwargs(item) for item in value]
+        raise TypeError(f"unsupported callback kwarg type: {type(value)}")
+
+    def _record(
+        self,
+        hook: Literal[
+            "log_success_event",
+            "async_log_success_event",
+            "log_failure_event",
+            "async_log_failure_event",
+        ],
+        phase: Literal["success", "failure"],
+        kwargs: dict[str, object],
+        response_obj: object,
+    ) -> None:
+        raw_litellm_params: Final = kwargs.get("litellm_params")
+        litellm_params: Final[Mapping[str, object]] = (
+            cast(Mapping[str, object], raw_litellm_params) if isinstance(raw_litellm_params, Mapping) else {}
+        )
+        raw_metadata: Final = litellm_params.get("metadata")
+        metadata_mapping: Final[Mapping[str, object]] = (
+            cast(Mapping[str, object], raw_metadata) if isinstance(raw_metadata, Mapping) else {}
+        )
+        metadata: Final = JSON_VALUE_ADAPTER.validate_python(
+            {key: metadata_mapping[key] for key in ("callback_profile", "sdk_route") if key in metadata_mapping}
+        )
+        raw_error: Final = kwargs.get("exception")
+        error: Final = sdk_error_report(raw_error) if isinstance(raw_error, Exception) else None
+        normalized_kwargs: Final = self._normalized_kwargs(kwargs)
+        payload_source: Final = (
+            response_obj.model_dump(mode="json") if isinstance(response_obj, BaseModel) else response_obj
+        )
+        payload: Final = JSON_VALUE_ADAPTER.validate_python(payload_source)
+        raw_model: Final = kwargs.get("model")
+        raw_call_type: Final = kwargs.get("call_type")
+        raw_call_id: Final = kwargs.get("litellm_call_id")
+        self._observations.put(
+            CallbackObservation(
+                hook=hook,
+                phase=phase,
+                model=raw_model if isinstance(raw_model, str) else None,
+                call_type=str(raw_call_type) if raw_call_type is not None else None,
+                litellm_call_id=raw_call_id if isinstance(raw_call_id, str) else None,
+                metadata=metadata,
+                kwargs=normalized_kwargs,
+                payload=payload,
+                error=error,
+            )
+        )
+
+    def log_success_event(
+        self,
+        kwargs: dict[str, object],
+        response_obj: object,
+        start_time: datetime.datetime,
+        end_time: datetime.datetime,
+    ) -> None:
+        del start_time, end_time
+        time.sleep(CALLBACK_DELAY_SECONDS)
+        self._record("log_success_event", "success", kwargs, response_obj)
+
+    async def async_log_success_event(
+        self,
+        kwargs: dict[str, object],
+        response_obj: object,
+        start_time: datetime.datetime,
+        end_time: datetime.datetime,
+    ) -> None:
+        del start_time, end_time
+        await asyncio.sleep(CALLBACK_DELAY_SECONDS)
+        self._record("async_log_success_event", "success", kwargs, response_obj)
+
+    def log_failure_event(
+        self,
+        kwargs: dict[str, object],
+        response_obj: object,
+        start_time: datetime.datetime,
+        end_time: datetime.datetime,
+    ) -> None:
+        del start_time, end_time
+        time.sleep(CALLBACK_DELAY_SECONDS)
+        self._record("log_failure_event", "failure", kwargs, response_obj)
+
+    async def async_log_failure_event(
+        self,
+        kwargs: dict[str, object],
+        response_obj: object,
+        start_time: datetime.datetime,
+        end_time: datetime.datetime,
+    ) -> None:
+        del start_time, end_time
+        await asyncio.sleep(CALLBACK_DELAY_SECONDS)
+        self._record("async_log_failure_event", "failure", kwargs, response_obj)
+
+    def observations(self) -> tuple[CallbackObservation, ...]:
+        observations: Final = tuple(self._observations.get_nowait() for _ in range(self._observations.qsize()))
+        return tuple(sorted(observations, key=lambda observation: observation.hook))
+
+
 INVALID_OCR_CASES: Final = (
     InvalidOcrCase(
         name="unsupported_provider",
@@ -229,6 +378,106 @@ def _execute_invalid_sdk_case(
     return _execute_sdk_call(call_kwargs, route, event_loop)
 
 
+def _callback_call_id(route: SDKRoute, terminal: Literal["success", "failure"]) -> str:
+    return f"ocr-callback-{route.value}-{terminal}"
+
+
+def _callback_metadata(route: SDKRoute, terminal: Literal["success", "failure"]) -> dict[str, str]:
+    return {"callback_profile": terminal, "sdk_route": route.value}
+
+
+def _drain_callback_delivery(route: SDKRoute, event_loop: asyncio.AbstractEventLoop) -> None:
+    if route is SDKRoute.AOCR:
+        from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+
+        async def drain_async_callbacks() -> None:
+            await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=CALLBACK_DRAIN_TIMEOUT_SECONDS)
+            await GLOBAL_LOGGING_WORKER.stop()
+
+        event_loop.run_until_complete(drain_async_callbacks())
+
+    from litellm.litellm_core_utils.thread_pool_executor import executor
+
+    executor.shutdown(wait=True, cancel_futures=False)
+
+
+def _execute_callback_sdk_case(
+    case: OcrParityCase,
+    route: SDKRoute,
+    terminal: Literal["success", "failure"],
+    mock_url: str,
+    event_loop: asyncio.AbstractEventLoop,
+) -> WorkerSuccess:
+    callback: Final = RecordingCallback()
+    call_kwargs: Final = {
+        **_call_kwargs(case.litellm_input, mock_url, route),
+        "callbacks": [callback],
+        "litellm_call_id": _callback_call_id(route, terminal),
+        "litellm_trace_id": _callback_call_id(route, terminal),
+        "metadata": _callback_metadata(route, terminal),
+    }
+    report: Final = _execute_sdk_call(call_kwargs, route, event_loop)
+    _drain_callback_delivery(route, event_loop)
+    return WorkerSuccess(report=report, callbacks=callback.observations())
+
+
+def _assert_callback_lifecycle(
+    execution: Execution,
+    route: SDKRoute,
+    terminal: Literal["success", "failure"],
+) -> None:
+    observations: Final = execution.callbacks
+    assert observations is not None, f"{route.value} {terminal} callbacks were not observed"
+    expected_hooks: Final = (
+        (f"log_{terminal}_event",)
+        if route is SDKRoute.OCR
+        else ("async_log_success_event",)
+        if terminal == "success"
+        else ("async_log_failure_event", "log_failure_event")
+    )
+    actual_hooks: Final = tuple(observation.hook for observation in observations)
+    assert actual_hooks == expected_hooks, (
+        f"{route.value} {terminal} expected callback hooks {expected_hooks}, received {actual_hooks}"
+    )
+    expected_call_id: Final = _callback_call_id(route, terminal)
+    expected_metadata: Final = _callback_metadata(route, terminal)
+    for observation in observations:
+        assert observation.phase == terminal
+        assert observation.call_type == route.value
+        assert observation.litellm_call_id == expected_call_id
+        assert observation.metadata == expected_metadata
+        assert observation.model
+        if terminal == "success":
+            assert isinstance(execution.report, SDKSuccess)
+            assert observation.payload == execution.report.response
+            assert observation.error is None
+        else:
+            assert isinstance(execution.report, SDKError)
+            assert observation.payload is None
+            assert observation.error is not None
+            assert observation.error.exception_type
+            assert observation.error.message
+            assert observation.error.status_code is not None
+            assert observation.error.status_code >= 400
+
+
+def _check_callback_ocr_sdk_parity(
+    case: OcrParityCase,
+    route: SDKRoute,
+    terminal: Literal["success", "failure"],
+    case_file: Path,
+    runner: SubprocessRunner,
+) -> None:
+    with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers:
+        python_worker, rust_worker = workers
+        python: Final = python_worker.execute(case_file, route.value, case.provider_responses)
+        rust: Final = rust_worker.execute(case_file, route.value, case.provider_responses)
+
+    _assert_callback_lifecycle(python, route, terminal)
+    _assert_callback_lifecycle(rust, route, terminal)
+    assert_parity(python, rust, PYTHON_HTTP_SENTINEL)
+
+
 def _check_recorded_ocr_sdk_parity(
     ocr_fixture: OcrParityCase,
     route: SDKRoute,
@@ -276,6 +525,25 @@ def _write_worker_case(directory: Path, index: int, case: OcrWorkerCase) -> Path
     return case_file
 
 
+def _callback_fixture(
+    fixtures: tuple[OcrParityCase, ...],
+    terminal: Literal["success", "failure"],
+) -> OcrParityCase:
+    matching: Final = tuple(
+        fixture
+        for fixture in fixtures
+        if fixture.litellm_input.contract == "mistral"
+        and (
+            all(response.status_code < 400 for response in fixture.provider_responses)
+            if terminal == "success"
+            else any(response.status_code >= 400 for response in fixture.provider_responses)
+        )
+    )
+    if not matching:
+        raise AssertionError(f"no recorded Mistral OCR {terminal} fixture is available for callback parity")
+    return min(matching, key=lambda fixture: fixture_id(fixture.litellm_input, fixture.litellm_input.model))
+
+
 @contextmanager
 def parity_checks() -> Generator[tuple[E2ECheck, ...]]:
     fixtures: Final = tuple(
@@ -298,6 +566,17 @@ def parity_checks() -> Generator[tuple[E2ECheck, ...]]:
             _write_worker_case(directory, len(recorded_files) + index, InvalidOcrWorkerCase(case=case))
             for index, case in enumerate(INVALID_OCR_CASES)
         )
+        callback_cases: Final[tuple[tuple[Literal["success", "failure"], OcrParityCase], ...]] = tuple(
+            (terminal, _callback_fixture(fixtures, terminal)) for terminal in CALLBACK_TERMINALS
+        )
+        callback_files: Final = tuple(
+            _write_worker_case(
+                directory,
+                len(recorded_files) + len(invalid_files) + index,
+                CallbackOcrWorkerCase(case=case, terminal=terminal),
+            )
+            for index, (terminal, case) in enumerate(callback_cases)
+        )
         with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers:
             recorded: Final = tuple(
                 E2ECheck(
@@ -315,7 +594,15 @@ def parity_checks() -> Generator[tuple[E2ECheck, ...]]:
                 for case, case_file in zip(INVALID_OCR_CASES, invalid_files, strict=True)
                 for route in SDKRoute
             )
-            yield (*recorded, *invalid)
+            callbacks: Final = tuple(
+                E2ECheck(
+                    f"callback:{route.value}:{terminal}",
+                    partial(_check_callback_ocr_sdk_parity, case, route, terminal, case_file, runner),
+                )
+                for (terminal, case), case_file in zip(callback_cases, callback_files, strict=True)
+                for route in SDKRoute
+            )
+            yield (*recorded, *invalid, *callbacks)
 
 
 def _execute_worker_command(
@@ -333,6 +620,8 @@ def _execute_worker_command(
                 return WorkerSuccess(report=_execute_sdk_case(recorded.litellm_input, route, mock_url, event_loop))
             case InvalidOcrWorkerCase(case=invalid):
                 return WorkerSuccess(report=_execute_invalid_sdk_case(invalid, route, mock_url, event_loop))
+            case CallbackOcrWorkerCase(case=callback_case, terminal=terminal):
+                return _execute_callback_sdk_case(callback_case, route, terminal, mock_url, event_loop)
     except Exception:
         return WorkerFailure(error=traceback.format_exc())
 
diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py
index 3e6c4060134..0e771f0dc17 100644
--- a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py
+++ b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py
@@ -230,6 +230,10 @@ _HOST_ONLY_BRIDGE_EXCLUSIONS: Final = tuple(
             "test_ocr_exception_type_uses_resolved_provider_context",
             "Python wraps bridge exceptions into public errors.",
         ),
+        (
+            "test_rust_upstream_error_uses_ocr_provider_error_mapping",
+            "Python maps native upstream errors through the selected OCR provider config.",
+        ),
         ("test_aocr_routes_to_async_rust_when_enabled", "Python selects and invokes the async native bridge."),
         ("test_aocr_exception_type_uses_resolved_provider_context", "Python wraps async bridge exceptions."),
         ("test_ocr_forwards_timeout_to_rust", "Python converts and forwards explicit timeouts."),
diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py
index d28184b96a6..c34833221cc 100644
--- a/tests/test_litellm/ocr/test_rust_bridge.py
+++ b/tests/test_litellm/ocr/test_rust_bridge.py
@@ -9,6 +9,7 @@ import httpx
 import pytest
 
 import litellm
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.base_llm.ocr.transformation import OCRResponse
 from litellm.rust_bridge import configuration
 
@@ -39,6 +40,10 @@ class CapturedException(Exception):
     pass
 
 
+class RustUpstreamError(Exception):
+    pass
+
+
 class RecordingBridge:
     """A fake ``RustOcr`` callable that records the args it was handed."""
 
@@ -183,6 +188,9 @@ class FakeOCRConfig:
     ) -> str:
         return f"{api_base or 'https://api.mistral.ai/v1'}/ocr"
 
+    def get_error_class(self, error_message: str, status_code: int, headers: dict[str, str]) -> BaseLLMException:
+        return BaseLLMException(status_code=status_code, message=error_message, headers=headers)
+
 
 def build_prepared_request(
     *,
@@ -486,6 +494,20 @@ def test_run_rust_ocr_prepares_request_and_wraps_response():
     }
 
 
+def test_rust_upstream_error_uses_ocr_provider_error_mapping():
+    error = RustUpstreamError(400, '{"message":"invalid model"}')
+
+    mapped = ocr_main._map_rust_ocr_error(
+        error,
+        build_prepared_request(),
+        (RuntimeError, RustUpstreamError),
+    )
+
+    assert isinstance(mapped, BaseLLMException)
+    assert mapped.status_code == 400
+    assert mapped.message == '{"message":"invalid model"}'
+
+
 def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing():
     bridge = RecordingBridge()
     litellm.rust(True)

From 4b3355bdc6b9807d5d1c817094449df03b6af881 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng@berri.ai>
Date: Mon, 7 Sep 2026 11:30:46 -0700
Subject: [PATCH 124/164] test(e2e): prove the virtual key lifecycle on every
 gateway replica (#40023)

* test(e2e): prove the virtual key lifecycle on every replica

Walks one virtual key through create, read, partial update, clear, enforce
and delete against a live proxy and database, reading every write back on
every gateway replica.

The management suite already had single write-then-read tests for keys, but
none of them proved that a partial /key/update leaves the untouched fields
alone, that an explicit null clears a field, or that a write is visible on
more than the one gateway that took it.

Adds read_back_everywhere to the shared ProxyClient: it polls a GET path on
every URL in PROXY_REPLICA_URLS until each replica's parsed body satisfies
the caller's predicate, and fails naming the replica that never converged.
The CLEAR sentinel in the e2e models makes an explicit JSON null expressible
in a body the transport otherwise strips of None fields.

Documents /key/update's merge patch semantics on the endpoint docstring.

* test(e2e): prove key revocation and field preservation on every replica

Applies the findings from an adversarial review of the first commit.

The delete step only checked that chat was refused on the gateway that took
the write, so it would have passed while a sibling gateway kept serving the
deleted key. It now serves one call from every replica first, so each has the
key cached and the delete has something to revoke everywhere, then polls every
replica for the refusal.

The file also carried its own poll loop that tested the deadline before
attempting, so it gave up one attempt early and skipped the attempt landing
exactly on the deadline. It now shares the harness helper, which is generic
over the polled value rather than over a parsed body, so the same loop covers
both the info read-back and the chat refusal.

The model the enforcement step registers now carries a unique marker in its
alias, matching every other deployment this suite creates, so concurrent runs
never share one model group.

The docstring sentence claimed an explicit null clears any field. It does not:
the metadata-backed fields merge into stored metadata, where a null is a silent
no-op, and only the key's own columns clear. Regenerating the dashboard types
picks up the corrected text.

* fix(e2e): delete a deployment that never becomes servable

Registering a model posts /model/new and then waits for every replica to list
it. When that wait timed out the deployment already existed in the database but
its id had never been returned, so no caller could delete it and the row
outlived the run. It is now deleted before the failure propagates.

Found by review on the key lifecycle suite, whose module fixture registers a
deployment this way, but every caller of the shared helper had the same
exposure.

* docs(e2e): drop the duplicated notes from the lifecycle docstrings

The delete method restated what the warm-up helper already explains, and the
module restated the merge patch rule that the endpoint and the request model
both document.
---
 .../key_management_endpoints.py               |   5 +
 tests/e2e/coverage_registry/mgmt.yaml         |   2 +
 .../e2e/management/test_key_lifecycle_e2e.py  | 299 ++++++++++++++++++
 tests/e2e/models.py                           |  36 ++-
 tests/e2e/proxy_client.py                     | 152 ++++++++-
 tests/e2e/test_proxy_client.py                | 114 ++++++-
 ui/litellm-dashboard/src/lib/http/schema.d.ts |   5 +
 7 files changed, 592 insertions(+), 21 deletions(-)
 create mode 100644 tests/e2e/management/test_key_lifecycle_e2e.py

diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index ebcfab090b5..8e51e250319 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -2840,6 +2840,11 @@ async def update_key_fn(
     """
     Update an existing API key's parameters.
 
+    The body is a merge patch: a field left out keeps its stored value, and on the key's own columns
+    an explicit null clears it. The metadata-backed fields below are the exception, merging into the
+    stored metadata instead: passing one as null leaves it unchanged, while `metadata` itself
+    replaces the stored metadata wholesale.
+
     Parameters:
     - key: Optional[str] - The key to update. Either key or key_alias must be provided.
     - key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases)
diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml
index d571fb36546..860d96a50b4 100644
--- a/tests/e2e/coverage_registry/mgmt.yaml
+++ b/tests/e2e/coverage_registry/mgmt.yaml
@@ -5,6 +5,8 @@
 - {id: mgmt.key.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:2462", rationale: "Budget/model changes persist"}
 - {id: mgmt.key.update.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:2462", rationale: "Non-admin cannot escalate perms"}
 - {id: mgmt.key.update.happy_path, module: mgmt, tier: P1, surface: ui, assertions: [happy_path], source: "key_management_endpoints.py:2462", rationale: "Key edit through the dashboard"}
+- {id: mgmt.key.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "key_management_endpoints.py:2829", rationale: "A partial /key/update changes only the field it names; alias, models, limits, budget window, team and metadata read back unchanged on every gateway replica"}
+- {id: mgmt.key.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "key_management_endpoints.py:2829", rationale: "An explicit null on /key/update clears max_budget and budget_duration, and the derived budget_reset_at with it, on every gateway replica"}
 - {id: mgmt.key.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3122", rationale: "Deletion revokes future calls"}
 - {id: mgmt.key.delete.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:3122", rationale: "Non-owner cannot delete"}
 - {id: mgmt.key.info.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3380", rationale: "Info reflects all writes"}
diff --git a/tests/e2e/management/test_key_lifecycle_e2e.py b/tests/e2e/management/test_key_lifecycle_e2e.py
new file mode 100644
index 00000000000..4c8effc4d24
--- /dev/null
+++ b/tests/e2e/management/test_key_lifecycle_e2e.py
@@ -0,0 +1,299 @@
+"""Live e2e: one virtual key walked through its whole lifecycle, read back on every
+gateway replica.
+
+Create, read, partial update, clear, enforce, delete: one method per step, and every
+step creates its own team and key (both deleted on teardown) so a step reruns or skips
+on its own. Writes go through the control plane; read-backs poll every URL in
+PROXY_REPLICA_URLS until each replica converges, because a write that is visible on the
+gateway that took it and stale on its neighbour is exactly the failure this file exists
+to catch. Revocation is the slowest of those: a deleted key stays usable on the other
+replicas until their auth cache entry expires, so the delete step polls each of them
+rather than asserting once.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable, Iterator, Mapping
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Final
+
+import pytest
+
+from e2e_config import unique_marker
+from e2e_http import Result, StreamingResponse, Success, UnknownApiError, unwrap
+from lifecycle import ResourceManager
+from management_client import MODEL_ACCESS_DENIED_MARKER, ManagementClient
+from models import (
+    CLEAR,
+    ChatBody,
+    ChatMessage,
+    KeyGenerateBody,
+    KeyGenerateResponse,
+    KeyInfo,
+    KeyInfoParams,
+    KeyInfoResponse,
+    KeyMetadata,
+    KeyUpdateBody,
+    LiteLLMParamsBody,
+    TeamNewBody,
+)
+from proxy_client import Converged, NotConverged, Poller, await_converged, await_converged_everywhere
+from transport import Transport
+
+pytestmark = pytest.mark.e2e
+
+BACKING_MODEL: Final = "gpt-4o-mini"
+DENIED_MODEL: Final = "gpt-5.5"
+MAX_BUDGET: Final = 25.0
+TPM_LIMIT: Final = 313131
+RPM_LIMIT: Final = 323232
+UPDATED_RPM_LIMIT: Final = 424242
+BUDGET_DURATION: Final = "30d"
+
+
+@dataclass(frozen=True, slots=True)
+class CreatedKey:
+    written: KeyGenerateBody
+    response: KeyGenerateResponse
+
+    @property
+    def key(self) -> str:
+        return self.response.key
+
+
+@pytest.fixture(scope="module")
+def mock_deployment(client: ManagementClient) -> Iterator[str]:
+    """A deployment that answers from a canned response, so the enforcement step needs no
+    provider key. The alias carries a unique marker, like every other model this suite
+    registers, so concurrent runs never share one model group."""
+    model_name: Final = f"e2e-key-lifecycle-{unique_marker()}"
+    model_id: Final = client.proxy.create_model(model_name, LiteLLMParamsBody(model=BACKING_MODEL, mock_response="ok"))
+    try:
+        yield model_name
+    finally:
+        client.proxy.delete_model(model_id)
+
+
+def _await[T](client: ManagementClient, poller: Poller[T], converged: Callable[[T], bool], failure: str) -> T:
+    outcome: Final = await_converged(
+        poller,
+        converged=converged,
+        timeout=client.proxy.poll_timeout,
+        interval=client.proxy.poll_interval,
+        now=time.monotonic,
+        sleep=time.sleep,
+    )
+    match outcome:
+        case Converged(result=result):
+            return result
+        case NotConverged(last_result=last):
+            pytest.fail(f"{failure}; last outcome: {last}")
+
+
+def _chat_poller(transport: Transport, key: str, model: str) -> Poller[StreamingResponse]:
+    return lambda: transport.send(
+        "/chat/completions",
+        headers=transport.bearer(key),
+        json=ChatBody(
+            model=model,
+            messages=[ChatMessage(role="user", content=f"say hi {unique_marker()}")],
+            max_tokens=16,
+        ),
+    )
+
+
+def _create_key(client: ManagementClient, resources: ResourceManager, model: str) -> CreatedKey:
+    marker: Final = unique_marker()
+    team_id: Final = client.create_team(TeamNewBody(team_alias=f"e2e-key-lifecycle-team-{marker}"))
+    resources.defer(lambda: client.delete_team(team_id))
+    written: Final = KeyGenerateBody(
+        key_alias=f"e2e-key-lifecycle-{marker}",
+        models=[model],
+        max_budget=MAX_BUDGET,
+        tpm_limit=TPM_LIMIT,
+        rpm_limit=RPM_LIMIT,
+        budget_duration=BUDGET_DURATION,
+        metadata=KeyMetadata(tag=marker),
+        team_id=team_id,
+    )
+    response: Final = unwrap(client.generate_key(written))
+    resources.defer(lambda: client.proxy.delete_key(response.key))
+    return CreatedKey(written=written, response=response)
+
+
+def _key_info_everywhere(
+    client: ManagementClient, key: str, settled: Callable[[KeyInfo], bool]
+) -> Mapping[str, KeyInfo]:
+    def converged(result: Result[KeyInfoResponse]) -> bool:
+        return isinstance(result, Success) and settled(result.data.info)
+
+    reads: Final = client.proxy.read_back_everywhere(
+        "/key/info", params=KeyInfoParams(key=key), response_type=KeyInfoResponse, converged=converged
+    )
+    return MappingProxyType({replica: unwrap(read).info for replica, read in reads.items()})
+
+
+def _is_key_not_found(result: Result[KeyInfoResponse]) -> bool:
+    return isinstance(result, UnknownApiError) and result.status_code == 404
+
+
+def _assert_reads_back(info: KeyInfo, expected: KeyGenerateBody, replica: str) -> None:
+    for field, observed, wanted in (
+        ("key_alias", info.key_alias, expected.key_alias),
+        ("models", info.models, expected.models),
+        ("max_budget", info.max_budget, expected.max_budget),
+        ("tpm_limit", info.tpm_limit, expected.tpm_limit),
+        ("rpm_limit", info.rpm_limit, expected.rpm_limit),
+        ("budget_duration", info.budget_duration, expected.budget_duration),
+        ("team_id", info.team_id, expected.team_id),
+        ("metadata", info.metadata, expected.metadata),
+    ):
+        assert observed == wanted, f"{replica}: /key/info reports {field}={observed!r}, expected {wanted!r}"
+
+
+def _poll_chat_ok(client: ManagementClient, key: str, model: str) -> None:
+    _ = _await(
+        client,
+        _chat_poller(client.proxy.transport, key, model),
+        lambda outcome: outcome.ok,
+        f"chat on {model} never succeeded for the key before the deadline",
+    )
+
+
+def _warm_every_replica(client: ManagementClient, key: str, model: str) -> None:
+    """Serve one call from every replica, so each has the key in its auth cache. Without
+    this the revocation check below would only prove a replica rejects a key it never
+    knew, which is true of any random string."""
+    for replica, transport in client.proxy.replicas.items():
+        _ = _await(
+            client,
+            _chat_poller(transport, key, model),
+            lambda outcome: outcome.ok,
+            f"{replica}: chat on {model} never succeeded for the key before the deadline",
+        )
+
+
+def _assert_chat_rejected_everywhere(client: ManagementClient, key: str, model: str) -> None:
+    outcomes: Final = await_converged_everywhere(
+        {replica: _chat_poller(transport, key, model) for replica, transport in client.proxy.replicas.items()},
+        converged=lambda outcome: outcome.status_code == 401,
+        timeout=client.proxy.poll_timeout,
+        interval=client.proxy.poll_interval,
+        now=time.monotonic,
+        sleep=time.sleep,
+    )
+    for replica, outcome in outcomes.items():
+        assert isinstance(outcome, Converged), (
+            f"{replica}: the deleted key was still accepted on chat after "
+            f"{client.proxy.poll_timeout}s, last status {outcome.last_result.status_code}"
+        )
+
+
+class TestKeyLifecycle:
+    def test_create_echoes_every_field_written(
+        self, client: ManagementClient, resources: ResourceManager, mock_deployment: str
+    ) -> None:
+        created: Final = _create_key(client, resources, mock_deployment)
+
+        response: Final = created.response
+        for field, observed, wanted in (
+            ("key_alias", response.key_alias, created.written.key_alias),
+            ("models", response.models, created.written.models),
+            ("max_budget", response.max_budget, created.written.max_budget),
+            ("tpm_limit", response.tpm_limit, created.written.tpm_limit),
+            ("rpm_limit", response.rpm_limit, created.written.rpm_limit),
+            ("budget_duration", response.budget_duration, created.written.budget_duration),
+            ("team_id", response.team_id, created.written.team_id),
+            ("metadata", response.metadata, created.written.metadata),
+        ):
+            assert observed == wanted, f"/key/generate echoed {field}={observed!r}, sent {wanted!r}"
+
+    def test_read_reflects_the_create_on_every_replica(
+        self, client: ManagementClient, resources: ResourceManager, mock_deployment: str
+    ) -> None:
+        created: Final = _create_key(client, resources, mock_deployment)
+
+        infos: Final = _key_info_everywhere(
+            client, created.key, lambda info: info.key_alias == created.written.key_alias
+        )
+        for replica, info in infos.items():
+            _assert_reads_back(info, created.written, replica)
+            assert info.budget_reset_at is not None, (
+                f"{replica}: /key/info reports no budget_reset_at for budget_duration={BUDGET_DURATION!r}"
+            )
+
+    @pytest.mark.covers("mgmt.key.update.preserves_unrelated_fields")
+    def test_partial_update_changes_only_the_named_field(
+        self, client: ManagementClient, resources: ResourceManager, mock_deployment: str
+    ) -> None:
+        created: Final = _create_key(client, resources, mock_deployment)
+        before: Final = _key_info_everywhere(client, created.key, lambda info: info.rpm_limit == RPM_LIMIT)
+
+        _ = unwrap(client.update_key(KeyUpdateBody(key=created.key, rpm_limit=UPDATED_RPM_LIMIT)))
+
+        after: Final = _key_info_everywhere(client, created.key, lambda info: info.rpm_limit == UPDATED_RPM_LIMIT)
+        for replica, info in after.items():
+            _assert_reads_back(info, created.written.model_copy(update={"rpm_limit": UPDATED_RPM_LIMIT}), replica)
+            assert info.budget_reset_at == before[replica].budget_reset_at, (
+                f"{replica}: budget_reset_at moved from {before[replica].budget_reset_at!r} to "
+                f"{info.budget_reset_at!r} on a /key/update that did not name budget_duration"
+            )
+
+    @pytest.mark.covers("mgmt.key.update.clear_persists")
+    def test_explicit_null_clears_the_budget_and_its_reset_time(
+        self, client: ManagementClient, resources: ResourceManager, mock_deployment: str
+    ) -> None:
+        created: Final = _create_key(client, resources, mock_deployment)
+        _ = _key_info_everywhere(client, created.key, lambda info: info.max_budget == MAX_BUDGET)
+
+        _ = unwrap(client.update_key(KeyUpdateBody(key=created.key, max_budget=CLEAR, budget_duration=CLEAR)))
+
+        cleared: Final = _key_info_everywhere(client, created.key, lambda info: info.max_budget is None)
+        for replica, info in cleared.items():
+            assert info.budget_duration is None, (
+                f"{replica}: budget_duration={info.budget_duration!r} survived an explicit null"
+            )
+            assert info.budget_reset_at is None, (
+                f"{replica}: clearing budget_duration left budget_reset_at={info.budget_reset_at!r}"
+            )
+            _assert_reads_back(
+                info, created.written.model_copy(update={"max_budget": None, "budget_duration": None}), replica
+            )
+
+    def test_key_serves_its_model_and_is_denied_others(
+        self, client: ManagementClient, resources: ResourceManager, mock_deployment: str
+    ) -> None:
+        created: Final = _create_key(client, resources, mock_deployment)
+
+        _poll_chat_ok(client, created.key, mock_deployment)
+
+        denied: Final = client.chat_status(created.key, DENIED_MODEL, f"say hi {unique_marker()}")
+        assert denied.status_code == 403, (
+            f"chat on {DENIED_MODEL!r} outside the key's model list must be denied 403, got "
+            f"{denied.status_code}: {denied.body[:300]}"
+        )
+        assert MODEL_ACCESS_DENIED_MARKER in denied.body, (
+            f"403 body must be a model-access denial, got: {denied.body[:300]}"
+        )
+
+    def test_delete_revokes_info_and_chat_on_every_replica(
+        self, client: ManagementClient, resources: ResourceManager, mock_deployment: str
+    ) -> None:
+        """The teardown's deferred delete fires again on the already-deleted key by
+        design: the deferred cleanup must survive this test failing before the
+        in-body delete, and a repeat /key/delete is a cheap no-op the warn-only
+        teardown absorbs."""
+        created: Final = _create_key(client, resources, mock_deployment)
+        _warm_every_replica(client, created.key, mock_deployment)
+
+        client.delete_key_strict(created.key)
+
+        _ = client.proxy.read_back_everywhere(
+            "/key/info",
+            params=KeyInfoParams(key=created.key),
+            response_type=KeyInfoResponse,
+            converged=_is_key_not_found,
+        )
+        _assert_chat_rejected_everywhere(client, created.key, mock_deployment)
diff --git a/tests/e2e/models.py b/tests/e2e/models.py
index c01687d0b31..fa124bdf6bd 100644
--- a/tests/e2e/models.py
+++ b/tests/e2e/models.py
@@ -8,9 +8,9 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 from datetime import datetime
-from typing import Literal
+from typing import Final, Literal
 
-from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_validator
+from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator
 
 # ---------- keys ----------
 
@@ -49,6 +49,7 @@ class KeyMetadata(BaseModel):
     logging: list[KeyLoggingCallback] | None = None
     priority: str | None = None
     batch_enqueued_token_limit: int | None = None
+    tag: str | None = None
 
 
 class ObjectPermission(BaseModel):
@@ -81,6 +82,14 @@ class KeyGenerateBody(BaseModel):
 
 class KeyGenerateResponse(BaseModel):
     key: str
+    key_alias: str | None = None
+    models: list[str] = []
+    max_budget: float | None = None
+    tpm_limit: int | None = None
+    rpm_limit: int | None = None
+    budget_duration: str | None = None
+    team_id: str | None = None
+    metadata: KeyMetadata | None = None
 
 
 class KeyRegenerateBody(BaseModel):
@@ -122,6 +131,7 @@ class KeyInfo(BaseModel):
     blocked: bool | None = None
     spend: float | None = None
     max_budget: float | None = None
+    budget_duration: str | None = None
     budget_reset_at: str | None = None
     budget_id: str | None = None
     litellm_budget_table: LiteLLMBudgetTable | None = None
@@ -919,12 +929,34 @@ class CredentialCreateResponse(BaseModel):
 # ---------- key / team / user / organization management ----------
 
 
+class Cleared(BaseModel):
+    """An explicit JSON null in a merge-patch body. The transport drops `None` fields
+    before sending (`exclude_none`), so `None` means "leave the stored value alone"; a
+    field set to `CLEAR` reaches the wire as `null`, which tells the proxy to clear it."""
+
+    model_config = ConfigDict(frozen=True)
+
+    @model_serializer
+    def _as_null(self) -> None:
+        return None
+
+
+CLEAR: Final = Cleared()
+
+
 class KeyUpdateBody(BaseModel):
+    """POST /key/update is a merge patch: a field left `None` is dropped from the body and
+    keeps its stored value, `CLEAR` sends an explicit null that clears it (`budget_duration`
+    clears `budget_reset_at` with it), and `metadata` replaces the stored metadata wholesale."""
+
     key: str
     models: list[str] | None = None
     key_alias: str | None = None
     tpm_limit: int | None = None
     rpm_limit: int | None = None
+    max_budget: float | Cleared | None = None
+    budget_duration: str | Cleared | None = None
+    metadata: KeyMetadata | None = None
 
 
 class KeyBlockBody(BaseModel):
diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py
index 3b15e57dfcf..520cbfde5a9 100644
--- a/tests/e2e/proxy_client.py
+++ b/tests/e2e/proxy_client.py
@@ -16,6 +16,8 @@ from datetime import datetime
 from types import MappingProxyType
 from typing import Final
 
+from pydantic import BaseModel
+
 from e2e_http import (
     AnthropicHeaders,
     AuthHeaders,
@@ -149,9 +151,7 @@ def await_servable(
     last_result: Result[ModelsListResponse] | None = None
     while True:
         t = now()
-        phase_deadline = (
-            started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds
-        )
+        phase_deadline = started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds
         remaining = phase_deadline - t
         if remaining <= 0:
             if (
@@ -164,9 +164,7 @@ def await_servable(
 
         poll_timeout = min(request_timeout, remaining)
         last_result = list_models(poll_timeout)
-        listed = isinstance(last_result, Success) and any(
-            entry.id == model_name for entry in last_result.data.data
-        )
+        listed = isinstance(last_result, Success) and any(entry.id == model_name for entry in last_result.data.data)
         t = now()
         if not listed:
             first_seen_at = None
@@ -179,9 +177,7 @@ def await_servable(
         elif t - first_seen_at >= db_sync_seconds:
             return Servable()
 
-        phase_deadline = (
-            started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds
-        )
+        phase_deadline = started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds
         wait = min(interval, phase_deadline - now())
         if wait > 0:
             sleep(wait)
@@ -239,6 +235,88 @@ def servable_timeout_message(
     )
 
 
+type Poller[T] = Callable[[], T]
+
+
+@dataclass(frozen=True, slots=True)
+class Converged[T]:
+    result: T
+
+
+@dataclass(frozen=True, slots=True)
+class NotConverged[T]:
+    """The deadline passed without a read satisfying the predicate; `last_result` is
+    the final read, so the caller can tell a stale body from a failed request."""
+
+    last_result: T
+
+
+type ConvergeOutcome[T] = Converged[T] | NotConverged[T]
+
+
+def await_converged[T](
+    poll: Poller[T],
+    *,
+    converged: Callable[[T], bool],
+    timeout: float,
+    interval: float,
+    now: Callable[[], float],
+    sleep: Callable[[float], None],
+) -> ConvergeOutcome[T]:
+    """Poll until a read satisfies `converged` or `timeout` elapses.
+
+    Polls before testing the deadline, so a zero or already-spent budget still gets one
+    attempt, and sleeps only min(interval, time left), so the attempt that lands exactly
+    on the deadline is taken rather than skipped. Clock and sleep are injected."""
+    deadline: Final = now() + timeout
+    while True:
+        result = poll()
+        if converged(result):
+            return Converged(result=result)
+        remaining = deadline - now()
+        if remaining <= 0:
+            return NotConverged(last_result=result)
+        sleep(min(interval, remaining))
+
+
+def await_converged_everywhere[T](
+    pollers: Mapping[str, Poller[T]],
+    *,
+    converged: Callable[[T], bool],
+    timeout: float,
+    interval: float,
+    now: Callable[[], float],
+    sleep: Callable[[float], None],
+) -> Mapping[str, ConvergeOutcome[T]]:
+    """`await_converged` against every replica in turn, each with the full budget, so a
+    replica that lags behind the one a write landed on is polled until it catches up
+    rather than failing on its first stale read."""
+    return MappingProxyType(
+        {
+            replica: await_converged(
+                poll, converged=converged, timeout=timeout, interval=interval, now=now, sleep=sleep
+            )
+            for replica, poll in pollers.items()
+        }
+    )
+
+
+def first_lagging_replica[T](
+    outcomes: Mapping[str, ConvergeOutcome[T]],
+) -> tuple[str, NotConverged[T]] | None:
+    return next(
+        ((replica, outcome) for replica, outcome in outcomes.items() if isinstance(outcome, NotConverged)),
+        None,
+    )
+
+
+def converge_timeout_message(*, what: str, replica: str, timeout: float, last_result: object) -> str:
+    return (
+        f"{what} on {replica} never converged within {timeout}s of the write "
+        f"(control/data-plane propagation issue); last read: {last_result}"
+    )
+
+
 @dataclass(frozen=True, slots=True)
 class ProxyClient:
     transport: Transport
@@ -290,6 +368,52 @@ class ProxyClient:
             )
         ).info
 
+    def read_back_everywhere[R: BaseModel](
+        self,
+        path: str,
+        *,
+        params: BaseModel,
+        response_type: type[R],
+        converged: Callable[[Result[R]], bool],
+    ) -> Mapping[str, Result[R]]:
+        """GET `path` under the master key on every replica in PROXY_REPLICA_URLS (the
+        data-plane URL alone when the stack exports no per-gateway addresses), polling
+        each to poll_timeout until its read satisfies `converged`. Returns that read per
+        replica, or fails naming the first replica that never converged and its last
+        read. Behind a load balancer the single address proves one replica converged,
+        not all of them; only per-gateway addresses make this a fleet-wide proof."""
+        outcomes: Final = await_converged_everywhere(
+            {
+                url: self._body_poller(transport, path, params, response_type)
+                for url, transport in self.replicas.items()
+            },
+            converged=converged,
+            timeout=self.poll_timeout,
+            interval=self.poll_interval,
+            now=time.monotonic,
+            sleep=time.sleep,
+        )
+        lagging: Final = first_lagging_replica(outcomes)
+        if lagging is not None:
+            replica, outcome = lagging
+            raise AssertionError(
+                converge_timeout_message(
+                    what=f"GET {path}",
+                    replica=replica,
+                    timeout=self.poll_timeout,
+                    last_result=outcome.last_result,
+                )
+            )
+        return MappingProxyType(
+            {replica: outcome.result for replica, outcome in outcomes.items() if isinstance(outcome, Converged)}
+        )
+
+    @staticmethod
+    def _body_poller[R: BaseModel](
+        transport: Transport, path: str, params: BaseModel, response_type: type[R]
+    ) -> Poller[Result[R]]:
+        return lambda: transport.get(path, headers=transport.master, params=params, response_type=response_type)
+
     def model_info(self) -> list[ModelInfoEntry]:
         """Every configured deployment with the price the proxy resolved for it
         (config override merged over cost-map defaults)."""
@@ -320,9 +444,7 @@ class ProxyClient:
             response_type=FileListResponse,
         )
 
-    def list_fine_tuning_jobs(
-        self, key: str, params: FineTuningJobsParams
-    ) -> Result[FineTuningJobsResponse]:
+    def list_fine_tuning_jobs(self, key: str, params: FineTuningJobsParams) -> Result[FineTuningJobsResponse]:
         return self.transport.get(
             "/v1/fine_tuning/jobs",
             headers=self.transport.bearer(key),
@@ -375,7 +497,11 @@ class ProxyClient:
             )
         ).model_id
         written_at = time.monotonic()
-        self._await_model_servable(body.model_name, listed_for)
+        try:
+            self._await_model_servable(body.model_name, listed_for)
+        except BaseException:
+            self.delete_model(model_id)
+            raise
         settle_propagation(written_at)
         return model_id
 
diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py
index a508c97b9fb..2caac58333f 100644
--- a/tests/e2e/test_proxy_client.py
+++ b/tests/e2e/test_proxy_client.py
@@ -1,9 +1,11 @@
-"""Harness coverage for the model barrier that gates on every replica.
+"""Harness coverage for the barriers that gate on every replica.
 
 No proxy needed and no ``e2e`` marker: this pins that a model registered through
 the control plane only counts as servable once every configured replica lists it
-on /v1/models, which is what keeps a two-gateway stack from handing a test a
-model that one gateway has not reloaded yet. The fakes are plain pollers and an
+on /v1/models, and that a management write only counts as read back once every
+replica's read satisfies the caller's predicate, which is what keeps a two-gateway
+stack from handing a test a model or a key that one gateway has not caught up on
+yet. The fakes are plain pollers standing in for each replica's transport plus an
 injected clock, so nothing here monkeypatches anything.
 """
 
@@ -12,18 +14,33 @@ from __future__ import annotations
 from collections.abc import Iterable, Mapping
 from dataclasses import dataclass
 from itertools import chain, repeat
+from types import MappingProxyType
 from typing import Final
 
 import pytest
 
 from e2e_config import parse_replica_urls
-from e2e_http import Success
-from models import ModelListEntry, ModelsListResponse
-from proxy_client import ModelsPoller, NotServableOn, Servable, await_servable_everywhere
+from e2e_http import Result, Success
+from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse
+from proxy_client import (
+    Poller,
+    ConvergeOutcome,
+    Converged,
+    ModelsPoller,
+    NotConverged,
+    NotServableOn,
+    Servable,
+    await_converged_everywhere,
+    await_servable_everywhere,
+    first_lagging_replica,
+    converge_timeout_message,
+)
 
 MODEL: Final = "gpt-under-test"
 TIMEOUT: Final = 10.0
 INTERVAL: Final = 2.0
+RPM_BEFORE_UPDATE: Final = 100
+RPM_AFTER_UPDATE: Final = 200
 
 
 @dataclass
@@ -78,6 +95,91 @@ class TestAwaitServableEverywhere:
         assert _await(pollers) == Servable()
 
 
+def _key_info(rpm_limit: int) -> Success[KeyInfoResponse]:
+    return Success(status_code=200, data=KeyInfoResponse(info=KeyInfo(rpm_limit=rpm_limit)))
+
+
+def _reads(results: Iterable[Result[KeyInfoResponse]]) -> Poller[Result[KeyInfoResponse]]:
+    it: Final = iter(results)
+    return lambda: next(it)
+
+
+def _updated(result: Result[KeyInfoResponse]) -> bool:
+    return isinstance(result, Success) and result.data.info.rpm_limit == RPM_AFTER_UPDATE
+
+
+def _converge(
+    pollers: Mapping[str, Poller[Result[KeyInfoResponse]]], clock: FakeClock
+) -> Mapping[str, ConvergeOutcome[Result[KeyInfoResponse]]]:
+    return await_converged_everywhere(
+        pollers,
+        converged=_updated,
+        timeout=TIMEOUT,
+        interval=INTERVAL,
+        now=clock.now,
+        sleep=clock.sleep,
+    )
+
+
+class TestAwaitConvergedEverywhere:
+    def test_waits_for_the_replica_that_lags_behind_the_write(self) -> None:
+        clock: Final = FakeClock()
+        pollers: Final = MappingProxyType(
+            {
+                "gateway-1": _reads(repeat(_key_info(RPM_AFTER_UPDATE))),
+                "gateway-2": _reads(
+                    chain(repeat(_key_info(RPM_BEFORE_UPDATE), 2), repeat(_key_info(RPM_AFTER_UPDATE)))
+                ),
+            }
+        )
+        outcomes: Final = _converge(pollers, clock)
+        assert outcomes == {
+            "gateway-1": Converged(result=_key_info(RPM_AFTER_UPDATE)),
+            "gateway-2": Converged(result=_key_info(RPM_AFTER_UPDATE)),
+        }
+        assert first_lagging_replica(outcomes) is None
+        assert clock.elapsed == 2 * INTERVAL
+
+    def test_names_the_replica_that_never_converges_with_its_last_read(self) -> None:
+        clock: Final = FakeClock()
+        pollers: Final = MappingProxyType(
+            {
+                "gateway-1": _reads(repeat(_key_info(RPM_AFTER_UPDATE))),
+                "gateway-2": _reads(repeat(_key_info(RPM_BEFORE_UPDATE))),
+            }
+        )
+        outcomes: Final = _converge(pollers, clock)
+        assert first_lagging_replica(outcomes) == (
+            "gateway-2",
+            NotConverged(last_result=_key_info(RPM_BEFORE_UPDATE)),
+        )
+        assert clock.elapsed == TIMEOUT
+        message: Final = converge_timeout_message(
+            what="GET /key/info",
+            replica="gateway-2",
+            timeout=TIMEOUT,
+            last_result=_key_info(RPM_BEFORE_UPDATE),
+        )
+        assert "gateway-2" in message and "/key/info" in message and str(RPM_BEFORE_UPDATE) in message
+
+    def test_each_replica_gets_its_own_full_budget(self) -> None:
+        """A replica that converges late must not eat into the next replica's budget: both
+        need most of the timeout here, so one shared deadline would starve the second."""
+        clock: Final = FakeClock()
+        slow: Final = chain(repeat(_key_info(RPM_BEFORE_UPDATE), 3), repeat(_key_info(RPM_AFTER_UPDATE)))
+        pollers: Final = MappingProxyType(
+            {
+                "gateway-1": _reads(slow),
+                "gateway-2": _reads(
+                    chain(repeat(_key_info(RPM_BEFORE_UPDATE), 3), repeat(_key_info(RPM_AFTER_UPDATE)))
+                ),
+            }
+        )
+        outcomes: Final = _converge(pollers, clock)
+        assert first_lagging_replica(outcomes) is None
+        assert clock.elapsed == 2 * 3 * INTERVAL
+
+
 class TestParseReplicaUrls:
     def test_splits_and_trims_the_gateway_addresses(self) -> None:
         raw: Final = " http://127.0.0.1:4010/, http://127.0.0.1:4011 "
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 108e50b52a1..36fd744efc0 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -8017,6 +8017,11 @@ export interface paths {
          * Update Key Fn
          * @description Update an existing API key's parameters.
          *
+         *     The body is a merge patch: a field left out keeps its stored value, and on the key's own columns
+         *     an explicit null clears it. The metadata-backed fields below are the exception, merging into the
+         *     stored metadata instead: passing one as null leaves it unchanged, while `metadata` itself
+         *     replaces the stored metadata wholesale.
+         *
          *     Parameters:
          *     - key: Optional[str] - The key to update. Either key or key_alias must be provided.
          *     - key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases)

From 96c698030bdc9978e925e0fe01e93636f9878442 Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Mon, 7 Sep 2026 11:38:53 -0700
Subject: [PATCH 125/164] fix(spend_logs): keep partition DDL transactions
 alive for their statement timeout (#40098)

* fix(spend_logs): keep partition DDL transactions alive for their statement timeout

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* style: ruff format changed files

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(lint): avoid dict-literal kwargs and keep cast-ok on the cast line

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(lint): cast at the call site instead of widening PrismaClient.tx

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(spend_logs): require partition tx timeout to strictly exceed statement bound

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: jesus <jesus@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../spend_logs_partition_manager.py           | 30 ++++++++--
 .../test_spend_logs_partition_manager.py      | 57 ++++++++++++++-----
 2 files changed, 69 insertions(+), 18 deletions(-)

diff --git a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py
index 221c142d9d3..b756bfeb6f6 100644
--- a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py
+++ b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py
@@ -16,7 +16,12 @@ keeps the batched-DELETE path, so existing deployments are untouched.
 import re
 from collections.abc import Callable
 from datetime import date, datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Final, TypeAlias
+from typing import (
+    TYPE_CHECKING,
+    Final,
+    TypeAlias,
+    cast,  # noqa: TID251  # db.tx is reached through untyped __getattr__ delegation
+)
 
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import (
@@ -25,6 +30,8 @@ from litellm.constants import (
 )
 
 if TYPE_CHECKING:
+    from prisma.client import TransactionManager
+
     from litellm.proxy.utils import PrismaClient
 
 SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs"
@@ -116,6 +123,21 @@ def select_partitions_to_drop(partitions: list[tuple[str, datetime | None]], cut
     return [name for name, upper in partitions if upper is not None and upper <= cutoff]
 
 
+_TX_COMMIT_SLACK: Final = timedelta(seconds=5)
+
+
+def _bounded_tx(prisma_client: "PrismaClient", timeout_ms: int) -> "TransactionManager":
+    """
+    Open an interactive transaction that outlives the statement bound it
+    carries. prisma's default 5s transaction timeout would close it mid
+    lock-wait, after which the engine answers the next call with a 422.
+    """
+    return cast(  # cast-ok: PrismaWrapper delegates tx via __getattr__ (untyped)
+        "TransactionManager",
+        prisma_client.db.tx(timeout=timedelta(milliseconds=timeout_ms) + _TX_COMMIT_SLACK),
+    )
+
+
 class SpendLogsPartitionManager:
     def __init__(
         self,
@@ -137,7 +159,7 @@ class SpendLogsPartitionManager:
         if budget_ms is None:
             return False
         try:
-            async with prisma_client.db.tx() as tx:
+            async with _bounded_tx(prisma_client, budget_ms) as tx:
                 await tx.execute_raw(f"SET LOCAL statement_timeout = {budget_ms}")
                 rows: Final = await tx.query_raw(
                     """
@@ -172,7 +194,7 @@ class SpendLogsPartitionManager:
         wait for the lock and statement_timeout bounds the work itself, so a
         partition this run cannot get is simply left for the next one.
         """
-        async with prisma_client.db.tx() as tx:
+        async with _bounded_tx(prisma_client, timeout_ms) as tx:
             await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
             await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}")
             await tx.execute_raw(statement)
@@ -209,7 +231,7 @@ class SpendLogsPartitionManager:
     async def _list_partitions(
         self, prisma_client: "PrismaClient", timeout_ms: int
     ) -> list[tuple[str, datetime | None]]:
-        async with prisma_client.db.tx() as tx:
+        async with _bounded_tx(prisma_client, timeout_ms) as tx:
             await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
             rows: Final = await tx.query_raw(
                 """
diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py
index 4ea655b8871..609dd13afc2 100644
--- a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py
+++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py
@@ -4,7 +4,7 @@ selection, the non-partitioned no-op safety path, and the drop/ensure SQL flow.
 """
 
 from contextlib import asynccontextmanager
-from datetime import date, datetime, timezone
+from datetime import date, datetime, timedelta, timezone
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -19,7 +19,6 @@ from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import (
     upcoming_partitions,
 )
 
-
 DDL_TIMEOUT_MS = 30000
 
 
@@ -28,19 +27,23 @@ def _budget(ms: "int | None" = DDL_TIMEOUT_MS):
     return lambda: ms
 
 
-def _wire_tx(db) -> list[str]:
+def _wire_tx(db) -> "tuple[list[str], list[int | timedelta | None]]":
     """
     Model the prisma seam the partition DDL uses.
 
     Every statement this manager issues, DDL and catalog query alike, runs inside
     db.tx() so it can carry SET LOCAL timeouts. Those SET LOCAL statements are
-    collected in the returned list rather than forwarded, so assertions on
-    db.execute_raw and db.query_raw still see only the real statements.
+    collected in the first returned list rather than forwarded, so assertions on
+    db.execute_raw and db.query_raw still see only the real statements. The
+    second list records the interactive-transaction timeout each tx was opened
+    with.
     """
     session_settings: list[str] = []
+    tx_timeouts: list[int | timedelta | None] = []
 
     @asynccontextmanager
-    async def _tx():
+    async def _tx(*, max_wait: "int | timedelta | None" = None, timeout: "int | timedelta | None" = None):
+        tx_timeouts.append(timeout)
         tx = MagicMock()
 
         async def _execute_raw(sql, *args):
@@ -57,7 +60,7 @@ def _wire_tx(db) -> list[str]:
         yield tx
 
     db.tx = _tx
-    return session_settings
+    return session_settings, tx_timeouts
 
 
 def test_period_start_per_interval():
@@ -227,7 +230,7 @@ async def test_partition_ddl_carries_a_statement_and_lock_timeout():
             }
         ]
     )
-    session_settings = _wire_tx(client.db)
+    session_settings, _ = _wire_tx(client.db)
 
     await mgr.ensure_partitions(client, _budget(7000))
     await mgr.drop_partitions_older_than(client, datetime(2026, 6, 5, tzinfo=timezone.utc), _budget(7000))
@@ -249,7 +252,7 @@ async def test_catalog_queries_carry_a_statement_timeout():
     mgr = SpendLogsPartitionManager()
     client = MagicMock()
     client.db.query_raw = AsyncMock(return_value=[])
-    session_settings = _wire_tx(client.db)
+    session_settings, _ = _wire_tx(client.db)
 
     await mgr.is_partitioned(client, _budget(4000))
     assert session_settings == ["SET LOCAL statement_timeout = 4000"], (
@@ -263,6 +266,34 @@ async def test_catalog_queries_carry_a_statement_timeout():
     )
 
 
+@pytest.mark.asyncio
+async def test_partition_transactions_outlive_their_statement_bound():
+    """
+    prisma's interactive transaction has its own timeout, 5s by default, which
+    keeps ticking while a statement waits on the partition lock. A tx shorter
+    than the SET LOCAL bound it carries is closed mid-lock-wait and the engine
+    then answers the next call with a 422, so the partition is silently not
+    created. A tx equal to the bound is closed too: the statement can consume
+    its whole bound waiting on the lock, then still needs to commit.
+    """
+    mgr = SpendLogsPartitionManager()
+    client = MagicMock()
+    client.db.execute_raw = AsyncMock(return_value=0)
+    client.db.query_raw = AsyncMock(return_value=[])
+    _, tx_timeouts = _wire_tx(client.db)
+
+    await mgr.is_partitioned(client, _budget())
+    await mgr.ensure_partitions(client, _budget())
+    await mgr.drop_partitions_older_than(client, datetime(2026, 6, 5, tzinfo=timezone.utc), _budget())
+
+    assert len(tx_timeouts) > 0
+    for tx_timeout in tx_timeouts:
+        assert isinstance(tx_timeout, timedelta)
+        assert tx_timeout > timedelta(milliseconds=DDL_TIMEOUT_MS), (
+            f"tx timeout {tx_timeout} does not outlive its {DDL_TIMEOUT_MS}ms statement bound"
+        )
+
+
 @pytest.mark.asyncio
 async def test_partition_loops_stop_when_the_budget_runs_out_mid_way():
     """
@@ -308,17 +339,15 @@ async def test_partition_maintenance_issues_nothing_when_the_budget_is_already_s
 
 
 def test_unsupported_interval_raises():
-    with pytest.raises(ValueError, match='Unsupported partition interval: year'):
+    with pytest.raises(ValueError, match="Unsupported partition interval: year"):
         period_start(date(2026, 6, 1), "year")
-    with pytest.raises(ValueError, match='Unsupported partition interval: year'):
+    with pytest.raises(ValueError, match="Unsupported partition interval: year"):
         next_period_start(date(2026, 6, 1), "year")
 
 
 def test_parse_partition_upper_bound_unparseable_to_value_is_none():
     """A TO(...) value that is not a valid timestamp must not raise; return None."""
-    assert (
-        parse_partition_upper_bound("FOR VALUES FROM ('x') TO ('not-a-date')") is None
-    )
+    assert parse_partition_upper_bound("FOR VALUES FROM ('x') TO ('not-a-date')") is None
 
 
 @pytest.mark.asyncio

From a88002c0348cc6b5bc5805dc27fd4d7e87d87db2 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Mon, 7 Sep 2026 11:43:32 -0700
Subject: [PATCH 126/164] test(guardrails): document audit prompt storage
 prerequisite

---
 tests/e2e/CONTRIBUTING.md                         |  9 +++++++++
 tests/e2e/guardrails/test_presidio_masking_e2e.py | 13 ++++++++++---
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md
index b83300cc8d5..1183096b81e 100644
--- a/tests/e2e/CONTRIBUTING.md
+++ b/tests/e2e/CONTRIBUTING.md
@@ -52,6 +52,15 @@ The suites run against a live proxy, so bring one up first by running the litell
 
 Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy. The presidio guardrail tests need a running Presidio analyzer and anonymizer the proxy can reach, addressed by `PRESIDIO_ANALYZER_API_BASE` / `PRESIDIO_ANONYMIZER_API_BASE`
 
+The Presidio spend-log audit test also requires prompt storage on the proxy, because its assertions inspect the detected entities and their scores. Add the following to the proxy config before starting it:
+
+```yaml
+general_settings:
+  store_prompts_in_spend_logs: true
+```
+
+Alternatively, start the proxy with `STORE_PROMPTS_IN_SPEND_LOGS=true litellm --config <your-e2e-config>.yml --port 4000`. Setting the variable only on the pytest process does not configure the proxy. With prompt storage disabled, the proxy correctly redacts `guardrail_response`, so that configuration cannot exercise this test's entity-detail assertions. Enable this only on a test stack using synthetic prompts
+
 A couple of logging destinations are configured on the proxy rather than by the test. The Weave tests scope their callback to the key they create, but litellm builds the `weave_otel` logger from `WANDB_API_KEY` and `WANDB_PROJECT_ID` before it applies the per-key vars, so the proxy needs both in its own environment or the key-scoped callback never initializes and nothing ships
 
 ### The pull request check
diff --git a/tests/e2e/guardrails/test_presidio_masking_e2e.py b/tests/e2e/guardrails/test_presidio_masking_e2e.py
index 4edf6c4ceb6..49d698938ce 100644
--- a/tests/e2e/guardrails/test_presidio_masking_e2e.py
+++ b/tests/e2e/guardrails/test_presidio_masking_e2e.py
@@ -19,6 +19,12 @@ PRESIDIO_ANONYMIZER_API_BASE; missing env is a hard failure, never a skip.
 Each guardrail registers with an explicit presidio_filter_scope so only the
 configured hook's callback exists (the default "both" registers input masking
 AND a post_call output masker), and is deleted on teardown.
+
+The spend-log audit test requires general_settings.store_prompts_in_spend_logs:
+true in the proxy config, or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy
+process environment before startup. Setting it only on pytest has no effect.
+Without this opt-in, redacting guardrail_response is expected proxy behavior;
+this suite deliberately requires the detected-entity details to remain visible.
 """
 
 from __future__ import annotations
@@ -374,9 +380,10 @@ class TestPresidioSpendLogRecord:
         )
 
         assert not isinstance(record.guardrail_response, str), (
-            "guardrail_response must be the detected-entity list, but it is the string "
-            f"{record.guardrail_response!r}. A redaction marker here means the proxy is not storing "
-            "prompts in spend logs, which also strips the per-entity scores from the guardrail panel"
+            "This audit test requires general_settings.store_prompts_in_spend_logs: true in the proxy config "
+            "or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy process environment before startup "
+            "(not just the pytest environment). Default prompt redaction is valid proxy behavior, "
+            f"but prevents entity-detail assertions; got guardrail_response={record.guardrail_response!r}"
         )
         entities = _ENTITY_LIST_ADAPTER.validate_python(record.guardrail_response)
         assert entities, (

From 55fe4a78943711531bcd35fbc30ddd8ff3f30f3b Mon Sep 17 00:00:00 2001
From: Roman D <roman.dovgan@justeattakeaway.com>
Date: Mon, 7 Sep 2026 19:54:59 +0100
Subject: [PATCH 127/164] feat(proxy): resolve root_path per request from a
 configured prefix list (SERVER_ROOT_PATHS) (#35935)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat(proxy): resolve root_path per request from SERVER_ROOT_PATHS

One deployment can encode exactly one client-visible URL path prefix
today: SERVER_ROOT_PATH is a scalar stamped onto the app at startup, so
a pod fronting several ingress prefixes 404s every prefix but one before
any handler runs, and MCP OAuth discovery can emit only one prefix's
URLs (RFC 9728 section 3 exact-match fails for the rest).

Add an opt-in outermost ASGI middleware that matches the request path
against a configured prefix list (SERVER_ROOT_PATHS, comma-separated) on
a segment boundary and sets scope["root_path"] for that request only.
Everything downstream is stock Starlette: route matching strips
root_path so routes stay registered root-relative, and request.base_url
re-includes it, so the discovery documents' resource and the 401
challenges' resource_metadata land under the prefix the client actually
called — with no discovery-builder changes.

LazyFeatureMiddleware now strips the scope root_path (falling back to
the cached SERVER_ROOT_PATH scalar) before feature prefix matching, so
lazily-registered routers — the MCP OAuth discovery router among them —
load under per-request prefixes.

Follow-up to the routing discussion on #35226; composes with, but does
not depend on, #35576.

* fix(proxy): import Sequence from collections.abc (ruff UP035 strict-budget gate)

* review(greptile): trim implementation commentary; fixture-own MCP registry state in tests

Addresses both P2s from the first Greptile pass:
- per_request_root_path_middleware.py (and the related _lazy_features /
  proxy_server comments) cut down to the constraints the code cannot
  express, per repo comment guidance
- the new discovery tests no longer clear/repopulate the shared MCP
  registry inline; a fixture snapshots it, hands the test an empty
  registry, and restores it afterwards so no state leaks between cases

* fix(lint): mutable-ok marker on the prefix accumulator (LIT002 type-discipline gate)

* fix(proxy): tie 401 challenges and get_custom_url to the per-request root_path

The per-request root_path middleware sets scope["root_path"] to the
prefix the client actually called, but the OAuth 401 challenges
(raise_user_oauth_challenge / raise_token_exchange_challenge) still
built their resource_metadata from SERVER_ROOT_PATH. On a pod fronting
several prefixes, the challenge advertised a discovery URL under a
different prefix than the discovery document served — the two
disagreed on where the resource metadata lives, and a strict RFC 9728
client refused the challenge. Route the challenges through a small
ContextVar the middleware populates so they read the same effective
root_path Starlette resolves the request under.

The same accessor fixes get_custom_url: when a request lives under a
SERVER_ROOT_PATHS-matched prefix, request.base_url already carries it,
so appending the SERVER_ROOT_PATH scalar on top produced e.g.
/tenant-a/legacy/sso/callback — a path that does not exist. Reading
the per-request prefix instead (and relying on join_paths's tail-dedup)
keeps SSO login/callback URLs under one prefix — the one the request
actually arrived on.

Fallback: outside a request (module-load-time UI URL builders,
background tasks) the ContextVar is unset and the accessor reads
SERVER_ROOT_PATH, matching get_server_root_path() so scalar-only
deployments are byte-identical.

* fix(mcp): challenge URL under per-request prefix must route, and mock parity

Two follow-ups to the review fix that made the 401 challenge use the
per-request root_path:

1. oauth_protected_resource_path must pick the URL structure that
   actually routes for the mechanism in use:
   - The scalar SERVER_ROOT_PATH deployment registers the well-known
     routes with the prefix INSERTED (via well_known_root_suffix at
     import time), matching RFC 8414 §3. The challenge URL must use the
     same insertion or a client fetching it 404s.
   - The per-request SERVER_ROOT_PATHS deployment can't register routes
     per prefix; PerRequestRootPathMiddleware strips the prefix from
     scope["path"] and the router matches the un-inserted route. The
     URL must place the prefix BEFORE .well-known so the strip leaves a
     matching path.
   The previous fix used the insertion form for both, which 404'd the
   discovery fetch on the per-request path — the discovery doc and the
   challenge would then disagree on where the resource metadata lives,
   the very failure the review flagged. End-to-end verified: the URL
   the challenge advertises routes and the doc's `resource` field
   equals the URL the client originally called (RFC 9728 §3).

2. get_request_root_path now delegates its fallback through
   get_server_root_path() instead of reading the env directly, so every
   existing `monkeypatch.setattr("litellm.proxy.utils.get_server_root_path"`
   test override keeps working. This unstubbed the mock on the /v2/login
   test that failed on the last CI run.

Plus the lint budget: annotate the local accumulator Final, tag the
scope["root_path"] rewrite as an intentional ASGI-contract mutation,
tag the reused `path`/`root_path` rebinds in LazyFeatureMiddleware, and
add reason strings to the two new PLC0415 lazy-import noqas.

* test(mcp): pin the reviewer's expected end-state — challenge URL routes, resource matches called URL

End-to-end regression test that mounts the discoverable router + the
per-request root_path middleware, hits an MCP endpoint that raises
raise_user_oauth_challenge, fetches the resource_metadata URL the
challenge advertises, and checks the returned document's `resource`
equals the URL the client originally called (RFC 9728 §3 exact match).

Covers /tenant-a, /tenant-b, and the unprefixed path on the same app so
a regression on any prefix — challenge URL 404s, or doc emits a
different prefix than the client called — fails at this test rather
than in a strict MCP client's discovery.

---------

Co-authored-by: gym-cmd <186399764+gym-cmd@users.noreply.github.com>
---
 .../mcp_server/mcp_server_manager.py          |  13 +-
 .../outbound_credentials/adapter.py           |  26 +-
 .../proxy/_experimental/mcp_server/server.py  |   8 +-
 litellm/proxy/_lazy_features.py               |  19 +-
 .../per_request_root_path_middleware.py       | 122 +++++++
 litellm/proxy/proxy_server.py                 |  20 ++
 litellm/proxy/utils.py                        |  14 +-
 .../outbound_credentials/test_adapter.py      |  86 ++++-
 .../mcp_server/test_discoverable_endpoints.py | 224 +++++++++++++
 .../test_per_request_root_path_middleware.py  | 302 ++++++++++++++++++
 tests/test_litellm/proxy/test_proxy_server.py |  76 +++++
 .../proxy/utils/helpers/test_url_helpers.py   |  73 +++++
 12 files changed, 948 insertions(+), 35 deletions(-)
 create mode 100644 litellm/proxy/middleware/per_request_root_path_middleware.py
 create mode 100644 tests/test_litellm/proxy/middleware/test_per_request_root_path_middleware.py

diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
index 612596bc803..e7fd650a324 100644
--- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
@@ -163,7 +163,10 @@ from litellm.proxy.common_utils.user_api_key_cache import get_management_object_
 from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import (
     id_jag_assertion_capture_gap_at_startup,
 )
-from litellm.proxy.utils import PrismaClient, ProxyLogging, get_server_root_path
+from litellm.proxy.middleware.per_request_root_path_middleware import (
+    get_request_root_path,
+)
+from litellm.proxy.utils import PrismaClient, ProxyLogging
 from litellm.repositories.table_repositories import MCPServerRepository
 from litellm.types.llms.custom_http import httpxSpecialProvider
 from litellm.types.mcp import (
@@ -3858,7 +3861,7 @@ class MCPServerManager:
                 if err.tag == "unauthorized" and isinstance(spec.config, AuthorizationCodeConfig):
                     # authorization_code's missing per-user token -> the per-server browser-OAuth
                     # challenge, built here where the full MCPServer is in hand.
-                    raise_user_oauth_challenge(server, root_path=get_server_root_path())
+                    raise_user_oauth_challenge(server, root_path=get_request_root_path())
                 if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig):
                     # token_exchange (OBO): a missing/rejected subject token -> the RFC 9728 challenge
                     # pointing at the IdP the client must SSO with to obtain one, rather than an opaque
@@ -3866,7 +3869,7 @@ class MCPServerManager:
                     # Access) threads its claims blob into the challenge for the client to satisfy.
                     raise_token_exchange_challenge(
                         server,
-                        root_path=get_server_root_path(),
+                        root_path=get_request_root_path(),
                         claims=err.unauthorized.claims,
                     )
                 raise_public(err)
@@ -3914,7 +3917,7 @@ class MCPServerManager:
         if spec is None or not isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)):
             return
         if subject_token is None and isinstance(spec.config, TokenExchangeConfig):
-            raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path())
+            raise_token_exchange_challenge(resolved_server, root_path=get_request_root_path())
         match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec):
             case Ok(_):
                 return
@@ -3922,7 +3925,7 @@ class MCPServerManager:
                 if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig):
                     raise_token_exchange_challenge(
                         resolved_server,
-                        root_path=get_server_root_path(),
+                        root_path=get_request_root_path(),
                         claims=err.unauthorized.claims,
                     )
                 raise_public(err)
diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py
index ea2318bd6f1..77979a15199 100644
--- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py
+++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py
@@ -12,6 +12,7 @@ every other mode so the caller defers to v1 (parity-safe); it grows one branch p
 from __future__ import annotations
 
 import base64
+import os
 from typing import TYPE_CHECKING, Final, Literal, NoReturn
 
 from fastapi import HTTPException
@@ -310,13 +311,30 @@ def raise_public(error: CredError) -> NoReturn:
 def oauth_protected_resource_path(root_path: str, server: MCPServer) -> str:
     """The server's RFC 9728 Protected Resource Metadata path, the shared anchor of both challenges.
 
-    ``root_path`` is the proxy's ``SERVER_ROOT_PATH``, resolved by the caller (the imperative shell)
-    so this stays a pure function of its inputs; ``"/"`` and ``""`` both mean no prefix. The path is
-    relative, so it resolves against the caller's own host (correct even behind a reverse proxy).
+    ``root_path`` is the prefix the request was routed under, resolved by the caller (the imperative
+    shell); ``"/"`` and ``""`` both mean no prefix. The path is relative, so it resolves against the
+    caller's own host (correct even behind a reverse proxy).
+
+    URL structure depends on how the prefix is served:
+
+    - The scalar ``SERVER_ROOT_PATH`` deployment registers the well-known routes with the prefix
+      *inserted* into the path (via :func:`well_known_root_suffix` at import time), matching RFC 8414
+      §3 well-known path insertion. When ``root_path`` equals ``SERVER_ROOT_PATH`` the URL must use
+      the same insertion or a client fetching it 404s.
+    - The per-request ``SERVER_ROOT_PATHS`` deployment can't register routes per prefix (the prefix
+      set is dynamic and could contain many entries); the middleware strips the prefix from
+      ``scope["path"]`` and the router matches the un-inserted well-known route. The URL must place
+      the prefix *before* ``.well-known`` so the strip leaves a matching path.
+
+    Picking the wrong form 404s the client's discovery fetch — the discovery document and the 401
+    challenge would then disagree on where the resource metadata lives.
     """
     prefix: Final = "" if root_path == "/" else root_path
     name: Final = server.alias or server.server_name or server.name or server.server_id
-    return f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}"
+    scalar_env: Final = os.getenv("SERVER_ROOT_PATH", "").rstrip("/")
+    if not prefix or (scalar_env and prefix == scalar_env):
+        return f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}"
+    return f"{prefix}/.well-known/oauth-protected-resource/mcp/{name}"
 
 
 def raise_user_oauth_challenge(server: MCPServer, *, root_path: str) -> NoReturn:
diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 26f5d6e7c8c..60a9af89cc3 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -3848,12 +3848,14 @@ if MCP_AVAILABLE:
             # then exchanges. A tool-call-time 401 would be wrapped into a JSON-RPC error and the
             # header lost, so the discovery flow needs this pre-emptive challenge.
             if server and server.auth_type == MCPAuth.oauth2_token_exchange and not oauth2_headers:
-                from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (  # noqa: PLC0415
+                from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (  # noqa: PLC0415  # lazy: adapter pulls MCP subgraph
                     raise_token_exchange_challenge,
                 )
-                from litellm.proxy.utils import get_server_root_path  # noqa: PLC0415
+                from litellm.proxy.middleware.per_request_root_path_middleware import (  # noqa: PLC0415  # lazy: middleware imports proxy utils
+                    get_request_root_path,
+                )
 
-                raise_token_exchange_challenge(server, root_path=get_server_root_path())
+                raise_token_exchange_challenge(server, root_path=get_request_root_path())
 
             # Exchange-backed modes (token_exchange's OBO mint, id_jag's stored-assertion mint): run
             # the exchange here at the transport edge, so a rejected subject raises the RFC 9728
diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py
index 3f90e6c0a7a..50e0a961a49 100644
--- a/litellm/proxy/_lazy_features.py
+++ b/litellm/proxy/_lazy_features.py
@@ -295,15 +295,18 @@ class LazyFeatureMiddleware:
         # Short-circuit once every feature has loaded.
         if scope["type"] in ("http", "websocket") and len(self._loaded) < len(self._features):
             path = scope.get("path", "")
-            # Strip SERVER_ROOT_PATH so prefix matching works under a server
-            # root path. Without this, requests like /api/v1/policies/... never
-            # match the registered prefixes (/policies/...) and lazy features
-            # stay unloaded — every endpoint under them returns 404. The
+            # Strip the request's root_path so prefix matching works under a
+            # server root path. Without this, requests like /api/v1/policies/...
+            # never match the registered prefixes (/policies/...) and lazy
+            # features stay unloaded — every endpoint under them returns 404.
+            # scope["root_path"] wins over the cached env scalar: FastAPI
+            # stamps SERVER_ROOT_PATH there, and PerRequestRootPathMiddleware
+            # resolves SERVER_ROOT_PATHS prefixes there per request. The
             # `+ "/"` boundary prevents false-positive matches (e.g. /apiv2
-            # against root /api). If the path doesn't start with the prefix
-            # (e.g. a reverse proxy already stripped it), we leave it alone.
-            if self._root_path and path.startswith(self._root_path + "/"):
-                path = path[len(self._root_path) :]
+            # against root /api); a pre-stripped path is left alone.
+            root_path: Final = str(scope.get("root_path", "")).rstrip("/") or self._root_path
+            if root_path and path.startswith(root_path + "/"):
+                path = path[len(root_path) :]  # rebind-ok: local strip after the boundary check above
             for feat in self._features:
                 if feat.module_path in self._loaded:
                     continue
diff --git a/litellm/proxy/middleware/per_request_root_path_middleware.py b/litellm/proxy/middleware/per_request_root_path_middleware.py
new file mode 100644
index 00000000000..d5df91458d2
--- /dev/null
+++ b/litellm/proxy/middleware/per_request_root_path_middleware.py
@@ -0,0 +1,122 @@
+"""Per-request ``root_path`` resolution from ``SERVER_ROOT_PATHS``.
+
+``SERVER_ROOT_PATH`` is a startup scalar, so one deployment serves exactly one
+client-visible URL path prefix; a request under any other prefix 404s before a
+handler runs. When the ingress preserves several prefixes into one pod (e.g.
+``/tenant-a/*`` and ``/tenant-b/*``), the matched prefix becomes that
+request's ``scope["root_path"]`` instead: Starlette strips it during route
+matching and rebuilds it into ``request.base_url``, so every emitted URL —
+the MCP OAuth discovery ``resource`` (RFC 9728 §3) and the 401 challenges'
+``resource_metadata`` among them — lands under the prefix the client called.
+Opt-in: with ``SERVER_ROOT_PATHS`` unset the middleware is not added at all.
+"""
+
+import os
+from collections.abc import Sequence
+from contextvars import ContextVar
+from typing import Final
+
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from litellm._logging import verbose_proxy_logger
+
+SERVER_ROOT_PATHS_ENV: Final = "SERVER_ROOT_PATHS"
+
+# The effective ``root_path`` for the currently-handled request. Populated by
+# ``PerRequestRootPathMiddleware`` from the (possibly-mutated) scope so code
+# that emits URLs off the request path — the 401 challenges' resource_metadata
+# and ``get_custom_url``'s SSO callbacks among them — can pick up the prefix
+# the client actually called without threading scope through every call site.
+# ``None`` means "middleware did not run" (the ``SERVER_ROOT_PATHS`` env is
+# unset, so no per-request prefix exists); readers fall back to the scalar
+# ``SERVER_ROOT_PATH`` in that case, which matches the pre-middleware behavior.
+_request_root_path_var: Final[ContextVar[str | None]] = ContextVar("_request_root_path_var", default=None)
+
+
+def get_request_root_path() -> str:
+    """Return the effective ``root_path`` for the current request.
+
+    Reads the value ``PerRequestRootPathMiddleware`` stashed for this request;
+    falls back through :func:`~litellm.proxy.utils.get_server_root_path` (i.e.
+    the ``SERVER_ROOT_PATH`` env) when the middleware did not run — the
+    scalar-only deployment. Delegating to the existing helper keeps every
+    existing ``monkeypatch.setattr("litellm.proxy.utils.get_server_root_path"``
+    test override working, and keeps a single source of truth for the scalar.
+    """
+    value: Final = _request_root_path_var.get()
+    if value is not None:
+        return value
+    # Lazy import: utils.py imports this module (via the lazy import inside
+    # get_custom_url), so a top-level import would build a cycle at load time.
+    from litellm.proxy.utils import get_server_root_path  # noqa: PLC0415  # lazy import breaks a two-way dep
+
+    return get_server_root_path()
+
+
+def normalize_root_paths(raw_paths: Sequence[str]) -> tuple[str, ...]:
+    """Strip whitespace and trailing slashes, dedupe, order longest-first;
+    warn and drop entries missing a leading ``/`` and the bare root."""
+    kept: Final[list[str]] = []  # mutable-ok: local accumulator; escapes only as a tuple
+    for entry in raw_paths:
+        candidate = entry.strip()
+        if not candidate:
+            continue
+        if not candidate.startswith("/"):
+            verbose_proxy_logger.warning(
+                "%s entry %r does not start with '/' and will be ignored.",
+                SERVER_ROOT_PATHS_ENV,
+                entry,
+            )
+            continue
+        candidate = candidate.rstrip("/")
+        if not candidate:
+            verbose_proxy_logger.warning(
+                "%s entry %r is the bare root and will be ignored; a root-mounted deployment needs no entry.",
+                SERVER_ROOT_PATHS_ENV,
+                entry,
+            )
+            continue
+        if candidate not in kept:
+            kept.append(candidate)
+    return tuple(sorted(kept, key=len, reverse=True))
+
+
+def get_server_root_paths() -> tuple[str, ...]:
+    """The normalized ``SERVER_ROOT_PATHS`` prefixes, empty when unset."""
+    configured: Final = os.getenv(SERVER_ROOT_PATHS_ENV, "")
+    if not configured.strip():
+        return ()
+    return normalize_root_paths(configured.split(","))
+
+
+class PerRequestRootPathMiddleware:
+    """Sets ``scope["root_path"]`` to the configured prefix matching the
+    request path on a whole-segment boundary. ``scope["path"]`` is left
+    untouched (Starlette strips ``root_path`` at route-match time). Must be
+    the outermost middleware so inner middlewares and the router see the
+    resolved value; a matched prefix overrides a scalar ``SERVER_ROOT_PATH``
+    for that request.
+    """
+
+    def __init__(self, app: ASGIApp, root_paths: Sequence[str]) -> None:
+        self.app = app
+        self.root_paths: Final = normalize_root_paths(root_paths)
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] in ("http", "websocket"):
+            path: Final = scope.get("path", "")
+            for prefix in self.root_paths:
+                if path == prefix or path.startswith(prefix + "/"):
+                    scope["root_path"] = prefix  # rebind-ok: ASGI middleware contract; Router and base_url read it
+                    break
+            # Stash the effective root_path (matched prefix, or the scope's
+            # existing value when nothing matched — i.e. FastAPI's scalar
+            # SERVER_ROOT_PATH) so code that emits URLs off the request path
+            # picks the same prefix the router will resolve the request under.
+            token: Final = _request_root_path_var.set(str(scope.get("root_path", "")))
+            try:
+                await self.app(scope, receive, send)
+            finally:
+                _request_root_path_var.reset(token)
+            return
+        await self.app(scope, receive, send)
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index ba5714fe950..0915b8dd1b9 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -598,6 +598,10 @@ from litellm.proxy.middleware.admission_control_middleware import (
 from litellm.proxy.middleware.in_flight_requests_middleware import (
     InFlightRequestsMiddleware,
 )
+from litellm.proxy.middleware.per_request_root_path_middleware import (
+    PerRequestRootPathMiddleware,
+    get_server_root_paths,
+)
 from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware
 from litellm.proxy.middleware.request_size_limit_middleware import (
     RequestSizeLimitMiddleware,
@@ -18384,6 +18388,22 @@ app.add_middleware(
     get_settings=lambda: get_admission_control_settings(general_settings),
     state=admission_control_state,
 )
+# Added last on purpose - last-added is outermost, and the client-visible URL
+# prefix must be resolved into scope["root_path"] before any inner middleware
+# or the router inspects the path. Only added when SERVER_ROOT_PATHS is
+# configured, so the default deployment's middleware stack is unchanged.
+_server_root_paths: Final = get_server_root_paths()
+if _server_root_paths:
+    if server_root_path and server_root_path != "/":
+        verbose_proxy_logger.warning(
+            "Both SERVER_ROOT_PATH=%r and SERVER_ROOT_PATHS=%r are set. A request "
+            "matching a SERVER_ROOT_PATHS prefix overrides the scalar root_path for "
+            "that request; unmatched requests keep SERVER_ROOT_PATH. Configure one "
+            "mechanism or the other.",
+            server_root_path,
+            _server_root_paths,
+        )
+    app.add_middleware(PerRequestRootPathMiddleware, root_paths=_server_root_paths)
 
 
 async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "StreamingResponse":
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index f1103ce6a29..ddf31cb1d8a 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -7212,7 +7212,19 @@ def get_custom_url(request_base_url: str, route: str | None = None) -> str:
     else:
         base_url = request_base_url
 
-    server_root_path: Final = get_server_root_path()
+    # get_request_root_path() returns the prefix the router is actually
+    # resolving this request under: the matched SERVER_ROOT_PATHS entry when
+    # PerRequestRootPathMiddleware ran, otherwise the SERVER_ROOT_PATH scalar.
+    # This keeps the emitted URL under one prefix — the one the client called —
+    # instead of stacking the scalar onto a request already living under a
+    # dynamic prefix (which would produce /tenant-a/legacy/... — a path that
+    # doesn't exist). join_paths()'s tail-dedup then collapses the append when
+    # base_url (i.e. request.base_url) already ends in the same prefix.
+    from litellm.proxy.middleware.per_request_root_path_middleware import (  # noqa: PLC0415  # lazy: middleware imports utils
+        get_request_root_path,
+    )
+
+    server_root_path: Final = get_request_root_path()
     if route is not None:
         if server_root_path != "":
             # First join base_url with server_root_path, then with route
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py
index d67d0df4d0e..1b003e11993 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py
@@ -476,16 +476,47 @@ def test_raise_public_plain_unauthorized_has_no_challenge():
 
 
 @pytest.mark.parametrize(
-    "root_path, expected_prefix",
+    "root_path",
     [
-        ("/", ""),  # "/" means no prefix
-        ("", ""),  # empty means no prefix
-        ("/api/v1", "/api/v1"),  # a real root path is prepended verbatim
+        "/",  # "/" means no prefix
+        "",  # empty means no prefix
     ],
 )
-def test_oauth_protected_resource_path_honors_root_path(root_path, expected_prefix):
+def test_oauth_protected_resource_path_no_prefix(root_path, monkeypatch):
+    monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
     path = oauth_protected_resource_path(root_path, _server(alias="my-srv"))
-    assert path == f"/.well-known/oauth-protected-resource{expected_prefix}/mcp/my-srv"
+    assert path == "/.well-known/oauth-protected-resource/mcp/my-srv"
+
+
+def test_oauth_protected_resource_path_scalar_prefix_uses_rfc8414_insertion(monkeypatch):
+    # A scalar SERVER_ROOT_PATH deployment registers the well-known routes with
+    # the prefix inserted (via well_known_root_suffix at import time). The URL
+    # must match that insertion or a client fetching it 404s.
+    monkeypatch.setenv("SERVER_ROOT_PATH", "/api/v1")
+    path = oauth_protected_resource_path("/api/v1", _server(alias="my-srv"))
+    assert path == "/.well-known/oauth-protected-resource/api/v1/mcp/my-srv"
+
+
+def test_oauth_protected_resource_path_per_request_prefix_goes_before_wellknown(monkeypatch):
+    # Per-request deployment: SERVER_ROOT_PATHS matched /tenant-a for this
+    # request but the scalar SERVER_ROOT_PATH is unset. Routes were registered
+    # without the well-known insertion, so the URL must place the prefix
+    # *before* .well-known — PerRequestRootPathMiddleware strips it and the
+    # router matches the un-inserted route.
+    monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
+    path = oauth_protected_resource_path("/tenant-a", _server(alias="my-srv"))
+    assert path == "/tenant-a/.well-known/oauth-protected-resource/mcp/my-srv"
+
+
+def test_oauth_protected_resource_path_dynamic_prefix_wins_over_scalar(monkeypatch):
+    # Both env vars configured: the middleware matched a SERVER_ROOT_PATHS
+    # prefix (/tenant-a) that differs from the scalar (/legacy). The URL must
+    # advertise /tenant-a — the prefix the client called — with no /legacy
+    # segment stacked onto it. Same review-fix invariant get_custom_url pins.
+    monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy")
+    path = oauth_protected_resource_path("/tenant-a", _server(alias="my-srv"))
+    assert path == "/tenant-a/.well-known/oauth-protected-resource/mcp/my-srv"
+    assert "/legacy" not in path
 
 
 @pytest.mark.parametrize(
@@ -510,7 +541,11 @@ def test_raise_user_oauth_challenge_points_at_per_server_prm():
     )
 
 
-def test_raise_user_oauth_challenge_includes_server_root_path():
+def test_raise_user_oauth_challenge_includes_server_root_path(monkeypatch):
+    # The scalar deployment: routes are registered with the prefix inserted
+    # (via well_known_root_suffix at import time), so the challenge URL uses
+    # the RFC 8414 §3 insertion form.
+    monkeypatch.setenv("SERVER_ROOT_PATH", "/api/v1")
     with pytest.raises(HTTPException) as exc_info:
         raise_user_oauth_challenge(_server(alias="my-srv"), root_path="/api/v1")
     assert (
@@ -519,6 +554,20 @@ def test_raise_user_oauth_challenge_includes_server_root_path():
     )
 
 
+def test_raise_user_oauth_challenge_per_request_prefix_is_routable(monkeypatch):
+    # Per-request deployment (SERVER_ROOT_PATHS matched /tenant-a): the
+    # challenge URL must place /tenant-a before .well-known so the client's
+    # discovery fetch routes through the same middleware strip the original
+    # request went through. The scalar-inserted form would 404 here.
+    monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
+    with pytest.raises(HTTPException) as exc_info:
+        raise_user_oauth_challenge(_server(alias="my-srv"), root_path="/tenant-a")
+    assert (
+        exc_info.value.headers["WWW-Authenticate"]
+        == 'Bearer resource_metadata="/tenant-a/.well-known/oauth-protected-resource/mcp/my-srv"'
+    )
+
+
 def test_raise_token_exchange_challenge_is_rfc9728_invalid_token():
     from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
         raise_token_exchange_challenge,
@@ -535,17 +584,30 @@ def test_raise_token_exchange_challenge_is_rfc9728_invalid_token():
     assert "error_description=" in www
 
 
-def test_raise_token_exchange_challenge_includes_server_root_path():
+def test_raise_token_exchange_challenge_includes_server_root_path(monkeypatch):
     from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
         raise_token_exchange_challenge,
     )
 
+    monkeypatch.setenv("SERVER_ROOT_PATH", "/api/v1")
     with pytest.raises(HTTPException) as exc_info:
         raise_token_exchange_challenge(_server(alias="obo-srv"), root_path="/api/v1")
     www = exc_info.value.headers["WWW-Authenticate"]
     assert 'resource_metadata="/.well-known/oauth-protected-resource/api/v1/mcp/obo-srv"' in www
 
 
+def test_raise_token_exchange_challenge_per_request_prefix_is_routable(monkeypatch):
+    from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
+        raise_token_exchange_challenge,
+    )
+
+    monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
+    with pytest.raises(HTTPException) as exc_info:
+        raise_token_exchange_challenge(_server(alias="obo-srv"), root_path="/tenant-a")
+    www = exc_info.value.headers["WWW-Authenticate"]
+    assert 'resource_metadata="/tenant-a/.well-known/oauth-protected-resource/mcp/obo-srv"' in www
+
+
 def test_raise_token_exchange_challenge_static_form_is_unchanged_without_step_up():
     from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
         raise_token_exchange_challenge,
@@ -589,9 +651,7 @@ def test_id_jag_client_secret_maps_to_config():
     # ID-JAG asserts the user's id_token; the access_token default maps to id_token.
     assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:id_token"
     assert isinstance(spec.config.client_auth, ClientSecretAuth)
-    assert spec.config.client_auth.client_secret.get_secret_value() == (
-        "litellm-client-secret"
-    )
+    assert spec.config.client_auth.client_secret.get_secret_value() == ("litellm-client-secret")
 
 
 def test_id_jag_private_key_maps_to_private_key_jwt_auth():
@@ -617,9 +677,7 @@ def test_id_jag_private_key_wins_over_client_secret():
 
 
 def test_id_jag_honors_explicit_subject_token_type():
-    spec = to_server_spec(
-        _id_jag_server(subject_token_type="urn:ietf:params:oauth:token-type:saml2")
-    )
+    spec = to_server_spec(_id_jag_server(subject_token_type="urn:ietf:params:oauth:token-type:saml2"))
     assert spec is not None and isinstance(spec.config, IdJagConfig)
     assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:saml2"
 
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
index a9bb24bbef9..be4206a1faf 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
@@ -10342,6 +10342,230 @@ async def test_upstream_resource_sent_on_dcr_bridge_relay_authorize():
     assert query["client_id"] == ["caller-client"]
 
 
+# ---------------------------------------------------------------------------
+# Per-request root_path (SERVER_ROOT_PATHS / PerRequestRootPathMiddleware):
+# one app fronting several client-visible URL path prefixes, each prefix's
+# discovery documents emitting URLs under the prefix the client called
+# (RFC 9728 §3 exact-match). A scalar PROXY_BASE_URL / SERVER_ROOT_PATH can
+# encode at most one prefix per pod; these tests pin the N-prefix case.
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def _no_proxy_base_url(monkeypatch):
+    """Discovery must derive URLs from the request in these tests, so the
+    scalar env overrides are cleared."""
+    monkeypatch.delenv("PROXY_BASE_URL", raising=False)
+    monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
+
+
+@pytest.fixture
+def _isolated_mcp_registry():
+    """Fixture-owned registry state: snapshot the shared registry, hand the
+    test an empty one, restore afterwards so nothing leaks between cases."""
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+        global_mcp_server_manager,
+    )
+
+    saved = dict(global_mcp_server_manager.registry)
+    global_mcp_server_manager.registry.clear()
+    try:
+        yield global_mcp_server_manager.registry
+    finally:
+        global_mcp_server_manager.registry.clear()
+        global_mcp_server_manager.registry.update(saved)
+
+
+def _prefixed_discovery_client(prefixes):
+    from fastapi import FastAPI
+    from fastapi.testclient import TestClient
+
+    from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router
+    from litellm.proxy.middleware.per_request_root_path_middleware import (
+        PerRequestRootPathMiddleware,
+    )
+
+    app = FastAPI()
+    app.include_router(router)
+    app.add_middleware(PerRequestRootPathMiddleware, root_paths=prefixes)
+    return TestClient(app)
+
+
+class TestPerRequestRootPathDiscovery:
+    def test_prefixed_wellknown_not_routable_without_middleware(self, _isolated_mcp_registry):
+        """Control: on a plain app (the only shape a scalar root_path can
+        express), a prefixed well-known request 404s before any discovery
+        builder runs — the routing gap this feature exists to close."""
+        from fastapi import FastAPI
+        from fastapi.testclient import TestClient
+
+        from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router
+
+        server = _create_oauth2_server(server_id="srv_a", name="server_a", server_name="server_a", alias="server_a")
+        _isolated_mcp_registry[server.server_id] = server
+        app = FastAPI()
+        app.include_router(router)
+        client = TestClient(app)
+        resp = client.get("/tenant-a/.well-known/oauth-protected-resource/mcp/server_a")
+        assert resp.status_code == 404
+
+    def test_two_prefixes_one_app_each_resource_matches_the_called_url(
+        self, _no_proxy_base_url, _isolated_mcp_registry
+    ):
+        """The multi-origin case itself: two prefixes served by the same app,
+        each per-server document's ``resource`` equal to the URL its client
+        called — including the prefix."""
+        for sid, name in (("srv_a", "server_a"), ("srv_b", "server_b")):
+            _isolated_mcp_registry[sid] = _create_oauth2_server(server_id=sid, name=name, server_name=name, alias=name)
+        client = _prefixed_discovery_client(["/tenant-a", "/tenant-b"])
+
+        resp_a = client.get("/tenant-a/.well-known/oauth-protected-resource/mcp/server_a")
+        resp_b = client.get("/tenant-b/.well-known/oauth-protected-resource/mcp/server_b")
+
+        assert resp_a.status_code == 200
+        assert resp_a.json()["resource"] == "http://testserver/tenant-a/mcp/server_a"
+        assert resp_b.status_code == 200
+        assert resp_b.json()["resource"] == "http://testserver/tenant-b/mcp/server_b"
+
+        # Every URL the document advertises stays under the request's
+        # prefix, so it resolves on this same app.
+        for auth_server in resp_a.json()["authorization_servers"]:
+            assert auth_server.startswith("http://testserver/tenant-a/")
+
+    def test_unprefixed_requests_unchanged_on_the_same_app(self, _no_proxy_base_url, _isolated_mcp_registry):
+        """Backward compat on the very same app: a root request emits the
+        document byte-identical to a deployment without the middleware."""
+        server = _create_oauth2_server(server_id="srv_a", name="server_a", server_name="server_a", alias="server_a")
+        _isolated_mcp_registry[server.server_id] = server
+        client = _prefixed_discovery_client(["/tenant-a"])
+        resp = client.get("/.well-known/oauth-protected-resource/mcp/server_a")
+        assert resp.status_code == 200
+        assert resp.json()["resource"] == "http://testserver/mcp/server_a"
+
+    def test_unlisted_prefix_404s(self, _no_proxy_base_url):
+        client = _prefixed_discovery_client(["/tenant-a"])
+        assert client.get("/tenant-c/.well-known/oauth-protected-resource/mcp/server_a").status_code == 404
+
+    def test_aggregate_documents_and_as_endpoints_under_prefix(self, _no_proxy_base_url, _isolated_mcp_registry):
+        """Aggregate PRM/AS documents carry the prefix, and the advertised
+        authorize endpoint actually resolves under it — the 404 trap that
+        invalidated prefixing discovery URLs without per-request routing
+        (#35226 review round 1)."""
+        client = _prefixed_discovery_client(["/tenant-a"])
+
+        prm = client.get("/tenant-a/.well-known/oauth-protected-resource/mcp")
+        asm = client.get("/tenant-a/.well-known/oauth-authorization-server/mcp")
+
+        assert prm.status_code == 200
+        assert prm.json()["resource"] == "http://testserver/tenant-a/mcp"
+        assert prm.json()["authorization_servers"] == ["http://testserver/tenant-a/mcp"]
+
+        assert asm.status_code == 200
+        assert asm.json()["issuer"] == "http://testserver/tenant-a/mcp"
+        assert asm.json()["authorization_endpoint"] == "http://testserver/tenant-a/authorize"
+
+        # The prefixed authorize URL routes to the real handler (not 404):
+        # under per-request root_path the whole app is reachable per-prefix,
+        # so discovery may advertise prefixed AS endpoints safely.
+        assert client.get("/tenant-a/authorize").status_code != 404
+
+    def test_passthrough_challenge_metadata_url_carries_prefix(self, _no_proxy_base_url):
+        """The WWW-Authenticate resource_metadata URL a 401 advertises must
+        land under the request's prefix, or the client is bounced to a
+        document whose ``resource`` cannot match the URL it called."""
+        from litellm.proxy._experimental.mcp_server.oauth_utils import (
+            get_passthrough_resource_metadata_url,
+        )
+
+        def _scope(path, root_path=None):
+            scope = {
+                "type": "http",
+                "method": "POST",
+                "path": path,
+                "headers": [],
+                "scheme": "http",
+                "server": ("testserver", 80),
+                "client": ("1.2.3.4", 4444),
+            }
+            if root_path is not None:
+                scope["root_path"] = root_path
+            return scope
+
+        prefixed = get_passthrough_resource_metadata_url(
+            scope=_scope("/tenant-a/mcp/github", root_path="/tenant-a"),
+            server_name="github",
+        )
+        assert prefixed == "http://testserver/tenant-a/.well-known/oauth-protected-resource/mcp/github"
+
+        # Regression guard: no root_path → today's URL, unchanged.
+        bare = get_passthrough_resource_metadata_url(
+            scope=_scope("/mcp/github"),
+            server_name="github",
+        )
+        assert bare == "http://testserver/.well-known/oauth-protected-resource/mcp/github"
+
+    def test_user_oauth_challenge_url_routes_and_resource_matches_client_url(
+        self, _no_proxy_base_url, _isolated_mcp_registry
+    ):
+        """The reviewer's expected end-state, pinned end-to-end: an MCP endpoint
+        raising ``raise_user_oauth_challenge`` under a per-request prefix must
+        emit a resource_metadata URL the client can actually fetch, and the
+        document it returns must carry the same prefix the client originally
+        called. If either half breaks the client's discovery is dead."""
+        import re
+
+        from fastapi import FastAPI, HTTPException, Request
+        from fastapi.testclient import TestClient
+
+        from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
+            raise_user_oauth_challenge,
+        )
+        from litellm.proxy.middleware.per_request_root_path_middleware import (
+            PerRequestRootPathMiddleware,
+            get_request_root_path,
+        )
+
+        server = _create_oauth2_server(
+            server_id="srv_a", name="server_a", server_name="server_a", alias="server_a"
+        )
+        _isolated_mcp_registry[server.server_id] = server
+
+        app = FastAPI()
+        app.include_router(router)
+
+        @app.post("/mcp/{name}")
+        def _mcp(name: str, request: Request):
+            try:
+                raise_user_oauth_challenge(server, root_path=get_request_root_path())
+            except HTTPException as exc:
+                return {"www_authenticate": exc.headers["WWW-Authenticate"]}
+
+        app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a", "/tenant-b"])
+        client = TestClient(app)
+
+        for prefix, mcp_url in (
+            ("/tenant-a", "http://testserver/tenant-a/mcp/server_a"),
+            ("/tenant-b", "http://testserver/tenant-b/mcp/server_a"),
+            ("", "http://testserver/mcp/server_a"),
+        ):
+            call = client.post(f"{prefix}/mcp/server_a")
+            assert call.status_code == 200, call.text
+            www = call.json()["www_authenticate"]
+            match = re.search(r'resource_metadata="([^"]+)"', www)
+            assert match, www
+            discovery = client.get(match.group(1))
+            # The challenge URL must route (a client that can't fetch it has
+            # no way to reach the resource metadata).
+            assert discovery.status_code == 200, (
+                f"challenge URL {match.group(1)} for prefix {prefix!r} 404s; "
+                "the client can't reach the resource metadata."
+            )
+            # And the doc's `resource` must equal the URL the client called
+            # (RFC 9728 §3 exact match): a mismatch bounces a strict client.
+            assert discovery.json()["resource"] == mcp_url, discovery.json()
+
+
 def _s256(verifier: str) -> str:
     return urlsafe_b64encode(hashlib.sha256(verifier.encode("ascii")).digest()).rstrip(b"=").decode("ascii")
 
diff --git a/tests/test_litellm/proxy/middleware/test_per_request_root_path_middleware.py b/tests/test_litellm/proxy/middleware/test_per_request_root_path_middleware.py
new file mode 100644
index 00000000000..61eea00bf7d
--- /dev/null
+++ b/tests/test_litellm/proxy/middleware/test_per_request_root_path_middleware.py
@@ -0,0 +1,302 @@
+"""Tests for PerRequestRootPathMiddleware (``SERVER_ROOT_PATHS``).
+
+One deployment fronting several client-visible URL path prefixes: the matched
+prefix becomes that request's ``root_path``, so Starlette route matching and
+``request.base_url`` — and therefore every URL the proxy emits, the MCP OAuth
+discovery documents among them — resolve under the prefix the client called.
+"""
+
+import pytest
+from fastapi import FastAPI, Request
+from fastapi.testclient import TestClient
+
+from litellm.proxy.middleware.per_request_root_path_middleware import (
+    PerRequestRootPathMiddleware,
+    get_request_root_path,
+    get_server_root_paths,
+    normalize_root_paths,
+)
+
+
+class TestNormalizeRootPaths:
+    def test_strips_whitespace_and_trailing_slash(self):
+        assert normalize_root_paths([" /tenant-a/ ", "/tenant-b"]) == (
+            "/tenant-a",
+            "/tenant-b",
+        )
+
+    def test_drops_empty_entries(self):
+        assert normalize_root_paths(["", "  ", "/tenant-a"]) == ("/tenant-a",)
+
+    def test_drops_entries_without_leading_slash(self):
+        # A typo'd entry must not silently match nothing at request time.
+        assert normalize_root_paths(["tenant-a", "/tenant-b"]) == ("/tenant-b",)
+
+    def test_drops_bare_root(self):
+        # "/" would turn every request into a root_path rewrite; a
+        # root-mounted deployment needs no entry at all.
+        assert normalize_root_paths(["/", "/tenant-a"]) == ("/tenant-a",)
+
+    def test_dedupes(self):
+        assert normalize_root_paths(["/t", "/t/", " /t "]) == ("/t",)
+
+    def test_longest_first_for_nested_prefixes(self):
+        # Longest-first ordering is what makes the most-specific nested
+        # prefix win at match time.
+        assert normalize_root_paths(["/t", "/t/deep"]) == ("/t/deep", "/t")
+
+
+class TestGetServerRootPaths:
+    def test_unset_env_is_empty(self, monkeypatch):
+        monkeypatch.delenv("SERVER_ROOT_PATHS", raising=False)
+        assert get_server_root_paths() == ()
+
+    def test_empty_env_is_empty(self, monkeypatch):
+        monkeypatch.setenv("SERVER_ROOT_PATHS", "")
+        assert get_server_root_paths() == ()
+
+    def test_comma_separated_entries(self, monkeypatch):
+        monkeypatch.setenv("SERVER_ROOT_PATHS", "/tenant-a, /tenant-b/")
+        assert get_server_root_paths() == ("/tenant-a", "/tenant-b")
+
+
+def _capture_scope_middleware(root_paths):
+    """Middleware wired to a downstream that records the scope it received."""
+    captured = {}
+
+    async def downstream(scope, receive, send):
+        captured.update(scope)
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b""})
+
+    return PerRequestRootPathMiddleware(downstream, root_paths=root_paths), captured
+
+
+async def _run(mw, scope):
+    async def receive():
+        return {"type": "http.request", "body": b"", "more_body": False}
+
+    async def send(message):
+        pass
+
+    await mw(scope, receive, send)
+
+
+class TestPerRequestRootPathMiddleware:
+    @pytest.mark.asyncio
+    async def test_matched_prefix_becomes_root_path_path_untouched(self):
+        # Starlette's router strips root_path from the (unmodified) path at
+        # match time, so the middleware must NOT rewrite scope["path"].
+        mw, captured = _capture_scope_middleware(["/tenant-a"])
+        await _run(mw, {"type": "http", "path": "/tenant-a/mcp/x", "method": "GET", "headers": []})
+        assert captured["root_path"] == "/tenant-a"
+        assert captured["path"] == "/tenant-a/mcp/x"
+
+    @pytest.mark.asyncio
+    async def test_exact_prefix_matches(self):
+        mw, captured = _capture_scope_middleware(["/tenant-a"])
+        await _run(mw, {"type": "http", "path": "/tenant-a", "method": "GET", "headers": []})
+        assert captured["root_path"] == "/tenant-a"
+
+    @pytest.mark.asyncio
+    async def test_segment_boundary_prevents_sibling_match(self):
+        # /tenant-ab must not match the /tenant-a prefix.
+        mw, captured = _capture_scope_middleware(["/tenant-a"])
+        await _run(mw, {"type": "http", "path": "/tenant-ab/mcp", "method": "GET", "headers": []})
+        assert "root_path" not in captured
+
+    @pytest.mark.asyncio
+    async def test_unmatched_path_untouched(self):
+        mw, captured = _capture_scope_middleware(["/tenant-a"])
+        await _run(mw, {"type": "http", "path": "/chat/completions", "method": "GET", "headers": []})
+        assert "root_path" not in captured
+        assert captured["path"] == "/chat/completions"
+
+    @pytest.mark.asyncio
+    async def test_longest_nested_prefix_wins(self):
+        mw, captured = _capture_scope_middleware(["/t", "/t/deep"])
+        await _run(mw, {"type": "http", "path": "/t/deep/mcp", "method": "GET", "headers": []})
+        assert captured["root_path"] == "/t/deep"
+
+    @pytest.mark.asyncio
+    async def test_matched_prefix_overrides_scalar_root_path(self):
+        # FastAPI(root_path=SERVER_ROOT_PATH) stamps the scalar before the
+        # middleware stack runs; a matched dynamic prefix wins for that
+        # request (combining both mechanisms is warned about at startup).
+        mw, captured = _capture_scope_middleware(["/tenant-a"])
+        await _run(
+            mw,
+            {"type": "http", "path": "/tenant-a/mcp", "root_path": "/legacy", "method": "GET", "headers": []},
+        )
+        assert captured["root_path"] == "/tenant-a"
+
+    @pytest.mark.asyncio
+    async def test_unmatched_request_keeps_scalar_root_path(self):
+        mw, captured = _capture_scope_middleware(["/tenant-a"])
+        await _run(
+            mw,
+            {"type": "http", "path": "/legacy/mcp", "root_path": "/legacy", "method": "GET", "headers": []},
+        )
+        assert captured["root_path"] == "/legacy"
+
+    @pytest.mark.asyncio
+    async def test_websocket_scope_matched(self):
+        mw, captured = _capture_scope_middleware(["/tenant-a"])
+        await _run(mw, {"type": "websocket", "path": "/tenant-a/ws", "headers": []})
+        assert captured["root_path"] == "/tenant-a"
+
+    @pytest.mark.asyncio
+    async def test_lifespan_scope_passes_through(self):
+        called = {}
+
+        async def downstream(scope, receive, send):
+            called["scope"] = scope
+
+        mw = PerRequestRootPathMiddleware(downstream, root_paths=["/tenant-a"])
+        await _run(mw, {"type": "lifespan"})
+        assert called["scope"] == {"type": "lifespan"}
+
+
+def _routed_client(prefixes):
+    app = FastAPI()
+
+    @app.get("/where")
+    def where(request: Request):
+        return {
+            "base_url": str(request.base_url),
+            "root_path": request.scope.get("root_path", ""),
+        }
+
+    app.add_middleware(PerRequestRootPathMiddleware, root_paths=prefixes)
+    return TestClient(app)
+
+
+class TestEndToEndRouting:
+    def test_two_prefixes_route_on_one_app(self):
+        # The property a scalar SERVER_ROOT_PATH cannot provide: two
+        # client-visible prefixes served by the same app, each request
+        # reconstructing its own base URL.
+        client = _routed_client(["/tenant-a", "/tenant-b"])
+
+        resp_a = client.get("/tenant-a/where")
+        resp_b = client.get("/tenant-b/where")
+
+        assert resp_a.status_code == 200
+        assert resp_a.json() == {
+            "base_url": "http://testserver/tenant-a/",
+            "root_path": "/tenant-a",
+        }
+        assert resp_b.status_code == 200
+        assert resp_b.json() == {
+            "base_url": "http://testserver/tenant-b/",
+            "root_path": "/tenant-b",
+        }
+
+    def test_unprefixed_route_still_served(self):
+        client = _routed_client(["/tenant-a"])
+        resp = client.get("/where")
+        assert resp.status_code == 200
+        assert resp.json()["base_url"] == "http://testserver/"
+
+    def test_unlisted_prefix_404s(self):
+        client = _routed_client(["/tenant-a"])
+        assert client.get("/tenant-c/where").status_code == 404
+
+
+class TestGetRequestRootPath:
+    """``get_request_root_path`` is the accessor that plumbs the middleware's
+    resolved prefix to code that doesn't have scope in hand — the 401 challenge
+    builders in ``mcp_server_manager`` / ``server`` and ``get_custom_url`` on the
+    SSO callback path. Reading the SERVER_ROOT_PATH scalar there would emit URLs
+    under a prefix the client didn't call, and stack a second prefix onto ones it
+    did (the two review points this fixture pins)."""
+
+    def test_falls_back_to_server_root_path_env_outside_a_request(self, monkeypatch):
+        # Outside a request the middleware's ContextVar is unset. The scalar
+        # env still owns the answer, so pre-middleware call sites (module-load
+        # UI URL builders, background tasks) behave exactly as they did.
+        monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy")
+        assert get_request_root_path() == "/legacy"
+
+    def test_returns_empty_string_when_no_env_and_no_request(self, monkeypatch):
+        monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
+        assert get_request_root_path() == ""
+
+    def test_returns_matched_prefix_inside_a_request(self, monkeypatch):
+        # With both env vars set, a request matching a SERVER_ROOT_PATHS prefix
+        # must see that prefix — not the SERVER_ROOT_PATH scalar — so the URL
+        # it emits stays under the prefix the router will resolve it against.
+        monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy")
+
+        seen: list[str] = []
+        app = FastAPI()
+
+        @app.get("/where")
+        def where():
+            seen.append(get_request_root_path())
+            return {}
+
+        app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a"])
+        client = TestClient(app)
+
+        assert client.get("/tenant-a/where").status_code == 200
+        assert seen == ["/tenant-a"]
+
+    def test_unmatched_request_falls_through_to_scope_scalar(self, monkeypatch):
+        # A request the middleware saw but did not match keeps whatever
+        # scope["root_path"] the app was mounted under (the scalar). The
+        # ContextVar still reflects the effective per-request answer, so
+        # emitted URLs and the router agree even on the fallback path.
+        monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy")
+        seen: list[str] = []
+
+        app = FastAPI(root_path="/legacy")
+
+        @app.get("/where")
+        def where():
+            seen.append(get_request_root_path())
+            return {}
+
+        app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a"])
+        client = TestClient(app)
+
+        assert client.get("/legacy/where").status_code == 200
+        assert seen == ["/legacy"]
+
+    def test_each_request_sees_its_own_prefix(self, monkeypatch):
+        # Two sequential requests through the same app must each see the
+        # prefix they arrived under, so one tenant's client is never sent
+        # the URL of another tenant's origin.
+        monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
+        seen: list[tuple[str, str]] = []
+        app = FastAPI()
+
+        @app.get("/where")
+        def where(tag: str):
+            seen.append((tag, get_request_root_path()))
+            return {}
+
+        app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a", "/tenant-b"])
+        client = TestClient(app)
+        assert client.get("/tenant-a/where?tag=a").status_code == 200
+        assert client.get("/tenant-b/where?tag=b").status_code == 200
+        assert seen == [("a", "/tenant-a"), ("b", "/tenant-b")]
+
+    def test_context_var_reset_after_request(self, monkeypatch):
+        # A ContextVar left set after the request finishes would poison the
+        # module-load-time callers that read it lazily (they'd think they were
+        # inside a request under the last-seen prefix).
+        monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy")
+
+        app = FastAPI()
+
+        @app.get("/where")
+        def where():
+            return {"prefix": get_request_root_path()}
+
+        app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a"])
+        client = TestClient(app)
+        assert client.get("/tenant-a/where").json() == {"prefix": "/tenant-a"}
+        # After the request finishes, the scalar-env fallback owns the answer
+        # again — nothing was left stashed from the last request's scope.
+        assert get_request_root_path() == "/legacy"
diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py
index 9393ec0f8e6..b7bb58378d4 100644
--- a/tests/test_litellm/proxy/test_proxy_server.py
+++ b/tests/test_litellm/proxy/test_proxy_server.py
@@ -9224,6 +9224,82 @@ class TestLazyFeatureMiddleware:
         else:
             assert loads == [], f"{case}: feature must not load"
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "env_root_path,scope_root_path,request_path,should_load,case",
+        [
+            # Per-request root_path (PerRequestRootPathMiddleware under
+            # SERVER_ROOT_PATHS) with no scalar env: strip and match.
+            ("", "/tenant-a", "/tenant-a/dummy/x", True, "per-request root_path strip"),
+            # scope root_path is authoritative over the cached env scalar.
+            ("/api/v1", "/tenant-a", "/tenant-a/dummy/x", True, "scope wins over env scalar"),
+            # Boundary check still applies to the per-request value.
+            ("", "/tenant-a", "/tenant-ab/dummy/x", False, "boundary check on scope root_path"),
+            # Empty scope root_path falls back to the env scalar.
+            ("/api/v1", "", "/api/v1/dummy/x", True, "empty scope falls back to env"),
+        ],
+    )
+    async def test_per_request_root_path_handling(
+        self, monkeypatch, env_root_path, scope_root_path, request_path, should_load, case
+    ):
+        """
+        ``scope["root_path"]`` must be stripped before prefix matching when
+        set — the scalar SERVER_ROOT_PATH lands there via
+        ``FastAPI(root_path=...)``, and PerRequestRootPathMiddleware
+        (SERVER_ROOT_PATHS) resolves a per-request prefix there. Otherwise
+        lazily-registered features — the MCP OAuth discovery router among
+        them — stay unloaded under a client-visible prefix and 404.
+        """
+        from fastapi import FastAPI
+
+        from litellm.proxy._lazy_features import (
+            LazyFeature,
+            LazyFeatureMiddleware,
+        )
+
+        monkeypatch.setenv("SERVER_ROOT_PATH", env_root_path)
+
+        loads = []
+
+        def fake_register(app, module):
+            loads.append(getattr(module, "__name__", "?"))
+
+        feat = LazyFeature(
+            name=f"dummy_prr_{case}",
+            module_path="json",
+            path_prefixes=("/dummy",),
+            register_fn=fake_register,
+        )
+
+        async def downstream(scope, receive, send):
+            await send({"type": "http.response.start", "status": 200, "headers": []})
+            await send({"type": "http.response.body", "body": b""})
+
+        target_app = FastAPI()
+        mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,))
+
+        async def receive():
+            return {"type": "http.request", "body": b"", "more_body": False}
+
+        async def send(message):
+            pass
+
+        await mw(
+            {
+                "type": "http",
+                "path": request_path,
+                "root_path": scope_root_path,
+                "method": "GET",
+                "headers": [],
+            },
+            receive,
+            send,
+        )
+        if should_load:
+            assert loads == ["json"], f"{case}: expected feature to load"
+        else:
+            assert loads == [], f"{case}: feature must not load"
+
     @pytest.mark.asyncio
     async def test_concurrent_first_requests_only_register_once(self):
         """
diff --git a/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py b/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py
index 31ea1bdce74..5f23c5fc20e 100644
--- a/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py
+++ b/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py
@@ -314,3 +314,76 @@ def test_normalize_route_for_root_path_error_path_when_route_not_under_root(
     _clear_url_env(monkeypatch)
     monkeypatch.setenv("SERVER_ROOT_PATH", "/proxy")
     assert normalize_route_for_root_path("/other/v1/chat") is None
+
+
+# ---------------------------------------------------------------------------
+# get_custom_url under a per-request prefix (SERVER_ROOT_PATHS):
+# when a request lives under a dynamic prefix, request.base_url already
+# contains it. Appending the SERVER_ROOT_PATH scalar on top produced
+# double-prefixed SSO callback / login URLs (a path that doesn't exist on
+# the deployment). These pin that the emitted URL now stays under one
+# prefix — the one the request actually arrived on.
+# ---------------------------------------------------------------------------
+
+
+def test_get_custom_url_uses_per_request_prefix_when_middleware_ran(monkeypatch):
+    """The middleware stashes the effective per-request prefix in a ContextVar.
+    ``get_custom_url`` reads that in preference to the SERVER_ROOT_PATH scalar,
+    and ``join_paths``'s tail-dedup collapses the append so a request whose
+    ``base_url`` already ends in ``/tenant-a`` does not become
+    ``/tenant-a/legacy/route``."""
+    _clear_url_env(monkeypatch)
+    monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy")
+
+    from litellm.proxy.middleware.per_request_root_path_middleware import (
+        _request_root_path_var,
+    )
+
+    token = _request_root_path_var.set("/tenant-a")
+    try:
+        # request.base_url already carries the tenant prefix; the scalar
+        # SERVER_ROOT_PATH must not be re-appended on top.
+        result = get_custom_url(
+            request_base_url="https://request.example.com/tenant-a/",
+            route="/v1/chat",
+        )
+    finally:
+        _request_root_path_var.reset(token)
+
+    assert result == "https://request.example.com/tenant-a/v1/chat"
+
+
+def test_get_custom_url_no_double_prefix_when_both_env_vars_configured(monkeypatch):
+    """Regression guard for the review point: with SERVER_ROOT_PATH also set
+    (scalar-legacy) and the request matched by SERVER_ROOT_PATHS (per-request),
+    the emitted URL is under one prefix — the per-request one — never both."""
+    _clear_url_env(monkeypatch)
+    monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy")
+    monkeypatch.setenv("SERVER_ROOT_PATHS", "/tenant-a")
+
+    from litellm.proxy.middleware.per_request_root_path_middleware import (
+        _request_root_path_var,
+    )
+
+    token = _request_root_path_var.set("/tenant-a")
+    try:
+        # No "/legacy" ever appears — the fix pins the reviewer's expected
+        # behavior: only one prefix should apply per request.
+        result = get_custom_url("https://api.example.com/tenant-a", "/sso/callback")
+    finally:
+        _request_root_path_var.reset(token)
+
+    assert result == "https://api.example.com/tenant-a/sso/callback"
+    assert "/legacy" not in result
+
+
+def test_get_custom_url_scalar_only_still_stamps_root_path(monkeypatch):
+    """Pre-middleware deployments have not opted into SERVER_ROOT_PATHS at all;
+    the ContextVar stays unset and the SERVER_ROOT_PATH scalar owns the answer
+    — the behavior every existing scalar-only deployment relies on."""
+    _clear_url_env(monkeypatch)
+    monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy")
+
+    result = get_custom_url("https://request.example.com", "/v1/chat")
+
+    assert result == "https://request.example.com/legacy/v1/chat"

From 1ae3216120a49fd2ac1ea6a2f43e42e8a33c276b Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Mon, 7 Sep 2026 12:05:43 -0700
Subject: [PATCH 128/164] fix(router): preserve default heuristic updates
 (#40007)

---
 .../auto_router_tuning_baseline.py            | 14 ++++--
 .../test_model_management_endpoints.py        | 10 ++--
 .../proxy/proxy_server/test_lifecycle.py      | 21 ++++++++
 .../test_auto_router_tuning_baseline.py       | 48 +++++++++++++++++--
 4 files changed, 80 insertions(+), 13 deletions(-)

diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_utils/auto_router_tuning_baseline.py
index 1e6412856a6..b82269d5824 100644
--- a/litellm/router_utils/auto_router_tuning_baseline.py
+++ b/litellm/router_utils/auto_router_tuning_baseline.py
@@ -10,7 +10,7 @@ from pydantic import ValidationError
 
 from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig
 
-TUNING_BASELINE_PARAM_NAME: Final = "auto_router_tuning_baseline"
+TUNING_BASELINE_PARAM_NAME: Final = "auto_router_tuning_baseline_v2"
 
 HEURISTIC_V1_TUNING_FIELDS: Final = (
     "tiers",
@@ -29,6 +29,8 @@ HEURISTIC_V1_TUNING_FIELDS: Final = (
     "keyword_tier_rules",
 )
 
+_TUNING_FIELD_SET: Final = frozenset(HEURISTIC_V1_TUNING_FIELDS)
+
 _V1_SCORING_CLASSIFIER_TYPES: Final = frozenset({"heuristic", "heuristic_first", "hybrid"})
 _AUTO_ROUTER_COMPLEXITY_PREFIX: Final = "auto_router/complexity_router"
 _EMPTY: Final[Mapping[str, object]] = MappingProxyType({})
@@ -40,12 +42,16 @@ def _mapping(value: object) -> Mapping[str, object]:
 
 
 def tuning_fingerprint(complexity_router_config: object) -> str | None:
-    """Digest of normalized heuristic-v1 tuning fields, or None when the config is invalid."""
+    """Digest of explicitly supplied heuristic-v1 tuning fields, normalized, or None when invalid."""
+    raw: Final = _mapping(complexity_router_config)
     try:
-        validated: Final = ComplexityRouterConfig.model_validate(_mapping(complexity_router_config))
+        validated: Final = ComplexityRouterConfig.model_validate(raw)
     except ValidationError:
         return None
-    payload: Final = validated.model_dump(mode="json", include=frozenset(HEURISTIC_V1_TUNING_FIELDS))
+    supplied: Final = ((_TUNING_FIELD_SET - frozenset(("tier_model_configs",))) & frozenset(raw)) | (
+        frozenset(("tier_model_configs",)) if validated.tier_model_configs else frozenset()
+    )
+    payload: Final = validated.model_dump(mode="json", include=supplied)
     return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
 
 
diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
index 63c08529b4d..d90338f8480 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
@@ -4447,13 +4447,13 @@ class TestStrategyRouterWriteValidation:
             self.litellm_proxymodeltable = MagicMock(
                 create=AsyncMock(),
                 update=AsyncMock(),
-                find_many=AsyncMock(
-                    return_value=tuple(
-                        LiteLLM_ProxyModelTable.model_validate(row) for row in self.tuning_rows
-                    )
-                ),
+                find_many=AsyncMock(side_effect=self._find_many),
             )
 
+        async def _find_many(self, where: object = None) -> tuple[LiteLLM_ProxyModelTable, ...]:
+            json.dumps(where)  # prisma serializes the filter with json.dumps and rejects a mappingproxy
+            return tuple(LiteLLM_ProxyModelTable.model_validate(row) for row in self.tuning_rows)
+
         @property
         def db(self) -> "TestStrategyRouterWriteValidation._FakeTx":
             return self
diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py
index a518c9289ae..7db85dc6943 100644
--- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py
+++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py
@@ -821,6 +821,27 @@ def test_proxy_startup_event_warns_for_global_budget_without_database():
     )
 
 
+@pytest.mark.asyncio
+async def test_tuning_baseline_v2_is_created_alongside_the_legacy_row():
+    from litellm.router_utils.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT
+
+    prisma_client = MagicMock()
+    prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None)
+    prisma_client.db.litellm_config.create = AsyncMock()
+    deployment = {
+        "model_name": "a",
+        "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": {}},
+    }
+
+    result = await ProxyStartupEvent._load_heuristic_v1_tuning_baselines(prisma_client, [deployment])
+
+    assert result == {'yaml:["a",[]]': DEFAULT_TUNING_FINGERPRINT}
+    assert prisma_client.db.litellm_config.create.await_args.kwargs["data"] == {
+        "param_name": "auto_router_tuning_baseline_v2",
+        "param_value": json.dumps(dict(result)),
+    }
+
+
 @pytest.mark.asyncio
 async def test_tuning_baseline_waits_for_a_complete_db_model_census(monkeypatch):
     prisma_client = MagicMock()
diff --git a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py
index c3395ab641f..fa7a96adb20 100644
--- a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py
+++ b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py
@@ -72,6 +72,9 @@ class TestTuningFingerprint:
             config["classifier_llm_config"] = {"model": "judge"}
         assert tuning_fingerprint(config) != DEFAULT_TUNING_FINGERPRINT
 
+    def test_explicit_empty_tier_model_configs_follow_omission(self) -> None:
+        assert tuning_fingerprint({"tier_model_configs": {}}) == DEFAULT_TUNING_FINGERPRINT
+
     def test_tier_model_overrides_change_the_fingerprint(self) -> None:
         plain = tuning_fingerprint({"tiers": {"SIMPLE": "x"}})
         with_override = tuning_fingerprint(
@@ -80,11 +83,39 @@ class TestTuningFingerprint:
         assert plain != with_override
 
     def test_non_tuning_fields_do_not_change_the_fingerprint(self) -> None:
-        assert tuning_fingerprint({"return_raw_model_name": True, "session_affinity": True}) == DEFAULT_TUNING_FINGERPRINT
+        assert (
+            tuning_fingerprint({"return_raw_model_name": True, "session_affinity": True}) == DEFAULT_TUNING_FINGERPRINT
+        )
 
     def test_invalid_config_has_no_fingerprint(self) -> None:
         assert tuning_fingerprint({"tier_boundaries": "not-a-mapping"}) is None
 
+    def test_a_release_changing_a_shipped_default_is_not_an_operator_edit(self, monkeypatch) -> None:
+        """An omitted setting follows the shipped default and stays off the quota when that default moves;
+        only what the operator wrote is fingerprinted, so an explicit value equal to the old default still counts."""
+        import litellm.router_strategy.complexity_router.config as config_module
+
+        untouched = _router("a", {})
+        tiers_only = _router("b", {"tiers": _TIERS})
+        pinned = _router("c", {"dimension_weights": dict(config_module.DEFAULT_DIMENSION_WEIGHTS)})
+        baselines = snapshot_tuning_baselines([untouched, tiers_only, pinned])
+        assert tuning_fingerprint(pinned["litellm_params"]["complexity_router_config"]) != DEFAULT_TUNING_FINGERPRINT
+
+        monkeypatch.setattr(
+            config_module,
+            "DEFAULT_DIMENSION_WEIGHTS",
+            {**config_module.DEFAULT_DIMENSION_WEIGHTS, "codePresence": 0.99},
+        )
+        monkeypatch.setattr(
+            config_module,
+            "DEFAULT_TIER_BOUNDARIES",
+            {**config_module.DEFAULT_TIER_BOUNDARIES, "simple_medium": 0.42},
+        )
+
+        assert tuning_fingerprint({}) == DEFAULT_TUNING_FINGERPRINT
+        assert mutable_tuned_identities([untouched, tiers_only, pinned], baselines) == frozenset()
+        assert mutable_tuned_identities([untouched], snapshot_tuning_baselines([])) == frozenset()
+
 
 class TestRouterIdentity:
     def test_db_rows_key_on_model_id_and_yaml_rows_on_name_and_tags(self) -> None:
@@ -161,12 +192,18 @@ class TestQuota:
         new_c = _router("c", {"tiers": _TIERS})
 
         assert tuning_quota_violation(candidate=edited_a, others=[legacy_b], baselines=baselines, limit=1) is None
-        assert tuning_quota_violation(candidate=edited_a, others=[edited_a, legacy_b], baselines=baselines, limit=1) is None
+        assert (
+            tuning_quota_violation(candidate=edited_a, others=[edited_a, legacy_b], baselines=baselines, limit=1)
+            is None
+        )
         assert tuning_quota_violation(candidate=legacy_a, others=[edited_b], baselines=baselines, limit=1) is None
         assert tuning_quota_violation(candidate=edited_b, others=[edited_a], baselines=baselines, limit=1) is not None
         assert tuning_quota_violation(candidate=new_c, others=[edited_a], baselines=baselines, limit=1) is not None
         assert tuning_quota_violation(candidate=new_c, others=[edited_a], baselines=baselines, limit=None) is None
-        assert tuning_quota_violation(candidate=legacy_a, others=[edited_a, edited_b], baselines=baselines, limit=1) is None
+        assert (
+            tuning_quota_violation(candidate=legacy_a, others=[edited_a, edited_b], baselines=baselines, limit=1)
+            is None
+        )
 
     def test_reverting_to_baseline_frees_the_quota(self) -> None:
         legacy_a = _router("a", {"tiers": _TIERS})
@@ -174,7 +211,10 @@ class TestQuota:
         baselines = snapshot_tuning_baselines([legacy_a, legacy_b])
         edited_b = _router("b", {"tiers": _TIERS})
         assert tuning_quota_violation(candidate=edited_b, others=[legacy_a], baselines=baselines, limit=1) is None
-        assert tuning_quota_violation(candidate=edited_b, others=[legacy_a, edited_b], baselines=baselines, limit=1) is None
+        assert (
+            tuning_quota_violation(candidate=edited_b, others=[legacy_a, edited_b], baselines=baselines, limit=1)
+            is None
+        )
 
     def test_violation_message_names_the_limit_and_remedy(self) -> None:
         message = tuning_limit_violation(held=2, limit=1)

From eea4c860f76ab5eadbe082ff36d081c00ec7346b Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Mon, 7 Sep 2026 12:12:22 -0700
Subject: [PATCH 129/164] feat(ui): add key-scoped auto-router usage tab
 (#39999)

* feat(ui): add key-scoped auto-router usage tab

GET /auto_router/benchmarks takes an optional api_key filter, applied in the
rollup aggregate on the primary key's leading column. Proxy admins get a
separate Auto-router usage tab on key detail pages with spend, baseline,
savings, tier routing, cache metrics and the existing router selector

* fix(ui): share key analytics date range
---
 litellm/proxy/db/autorouter_session_rollup.py |  4 +-
 .../auto_router_endpoints.py                  |  2 +
 .../spend/test_autorouter_session_rollup.py   | 32 ++++++
 .../test_auto_router_endpoints.py             |  5 +-
 .../AutoRouterBenchmarksTab.test.tsx          | 28 +++++-
 .../_components/AutoRouterBenchmarksTab.tsx   |  9 +-
 .../useAutoRouterBenchmarks.test.ts           | 31 +++++-
 .../_components/useAutoRouterBenchmarks.ts    |  4 +-
 ...useDailyActivityRange.integration.test.tsx |  8 +-
 .../useDailyActivityRange.test.tsx            | 10 +-
 .../_components/useDailyActivityRange.ts      | 24 +++--
 ...KeyAutoRouterUsageTab.integration.test.tsx | 98 +++++++++++++++++++
 .../templates/KeyAutoRouterUsageTab.tsx       | 18 ++++
 .../KeySavingsTab.integration.test.tsx        | 12 ++-
 .../components/templates/KeySavingsTab.tsx    | 19 ++--
 .../templates/key_info_view.test.tsx          | 56 +++++++++++
 .../components/templates/key_info_view.tsx    | 26 ++++-
 ui/litellm-dashboard/src/lib/http/schema.d.ts |  2 +
 18 files changed, 354 insertions(+), 34 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx
 create mode 100644 ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.tsx

diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py
index 9c637a62dc1..b33bffbaaa1 100644
--- a/litellm/proxy/db/autorouter_session_rollup.py
+++ b/litellm/proxy/db/autorouter_session_rollup.py
@@ -37,7 +37,9 @@ CACHE_TTL_1H_SECONDS: Final = 3600
 AUTOROUTER_BENCHMARKS_SQL: Final = """
 WITH windowed AS (
     SELECT * FROM "LiteLLM_AutoRouterSession"
-    WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp
+    WHERE last_turn_at >= $1::timestamp
+      AND first_turn_at < $2::timestamp
+      AND ($3::text IS NULL OR api_key = $3::text)
 ),
 tier_maps AS (
     SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns
diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py
index 2a7813bc140..0f0323b45f8 100644
--- a/litellm/proxy/management_endpoints/auto_router_endpoints.py
+++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py
@@ -645,6 +645,7 @@ async def get_auto_router_benchmarks(
         str | None, Query(description="YYYY-MM-DD UTC, inclusive (defaults to 30 days before end_date)")
     ] = None,
     end_date: Annotated[str | None, Query(description="YYYY-MM-DD UTC, inclusive (defaults to today)")] = None,
+    api_key: Annotated[str | None, Query(description="Filter to one virtual key token hash")] = None,
 ) -> AutoRouterBenchmarksResponse:
     """
     Benchmarks for the auto-router dashboard: session shape, savings against the configured
@@ -681,6 +682,7 @@ async def get_auto_router_benchmarks(
         AUTOROUTER_BENCHMARKS_SQL,
         start_day.isoformat(),
         (end_day + timedelta(days=1)).isoformat(),
+        api_key,
     )
     rows: Final = _SESSION_AGG_ROWS.validate_python(raw_rows or ())
     groups: Final = (
diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py
index 7bc61c40ea0..c2272f3d20d 100644
--- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py
+++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py
@@ -169,6 +169,7 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db):
         AUTOROUTER_BENCHMARKS_SQL,
         (T0 - timedelta(days=1)).isoformat(),
         (T0 + timedelta(days=1)).isoformat(),
+        None,
     )
     matching = [row for row in rows if row["router_name"] == router]
     assert len(matching) == 1
@@ -181,6 +182,33 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db):
     assert grouped["session_seconds"] == pytest.approx(60.0)
 
 
+async def test_the_benchmarks_aggregate_can_filter_to_one_key(db):
+    router = f"r-{uuid.uuid4()}"
+    first_key = f"k-{uuid.uuid4()}"
+    second_key = f"k-{uuid.uuid4()}"
+    await _turn(db, first_key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, saved=0.5)
+    await _turn(db, second_key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, saved=9.0)
+
+    rows = await db.query_raw(
+        AUTOROUTER_BENCHMARKS_SQL,
+        (T0 - timedelta(days=1)).isoformat(),
+        (T0 + timedelta(days=1)).isoformat(),
+        first_key,
+    )
+    matching = [row for row in rows if row["router_name"] == router]
+    assert len(matching) == 1
+    assert matching[0]["sessions"] == 1
+    assert matching[0]["saved_spend"] == pytest.approx(0.5)
+
+    unknown_key_rows = await db.query_raw(
+        AUTOROUTER_BENCHMARKS_SQL,
+        (T0 - timedelta(days=1)).isoformat(),
+        (T0 + timedelta(days=1)).isoformat(),
+        f"k-{uuid.uuid4()}",
+    )
+    assert [row for row in unknown_key_rows if row["router_name"] == router] == []
+
+
 async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db):
     key = f"k-{uuid.uuid4()}"
     router = f"r-{uuid.uuid4()}"
@@ -191,6 +219,7 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db
         AUTOROUTER_BENCHMARKS_SQL,
         (T0 - timedelta(days=1)).isoformat(),
         (T0 + timedelta(days=1)).isoformat(),
+        None,
     )
     matching = sorted(
         (row for row in rows if row["router_name"] == router),
@@ -253,6 +282,7 @@ async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db):
         AUTOROUTER_BENCHMARKS_SQL,
         (T0 - timedelta(days=1)).isoformat(),
         (T0 + timedelta(days=1)).isoformat(),
+        None,
     )
     grouped = next(row for row in rows if row["router_name"] == router)
     assert grouped["tier_turns"] == {"simple": 2, "complex": 1}
@@ -280,6 +310,7 @@ async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(d
         AUTOROUTER_BENCHMARKS_SQL,
         (T0 - timedelta(days=1)).isoformat(),
         (T0 + timedelta(days=1)).isoformat(),
+        None,
     )
     by_type = {row["router_type"]: row["tier_turns"] for row in rows if row["router_name"] == router}
     assert by_type == {"complexity": {"medium": 1}, "quality": {"2": 1}}
@@ -294,6 +325,7 @@ async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db):
         AUTOROUTER_BENCHMARKS_SQL,
         (T0 - timedelta(days=1)).isoformat(),
         (T0 + timedelta(days=1)).isoformat(),
+        None,
     )
     grouped = next(row for row in rows if row["router_name"] == router)
     assert grouped["tier_turns"] == {}
diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py
index 21f8d985f22..35e76c96c14 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py
@@ -489,6 +489,7 @@ class TestAutoRouterBenchmarks:
         monkeypatch: pytest.MonkeyPatch,
         rows: Sequence[Mapping[str, object]],
         model_list: Sequence[object],
+        api_key: str | None = None,
     ) -> AutoRouterBenchmarksResponse:
         from litellm.proxy import proxy_server
         from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks
@@ -503,6 +504,7 @@ class TestAutoRouterBenchmarks:
             user_api_key_dict=ADMIN,
             start_date="2026-07-01",
             end_date="2026-08-01",
+            api_key=api_key,
         )
 
     ROW = _SessionAggRow(
@@ -652,8 +654,9 @@ class TestAutoRouterBenchmarks:
             user_api_key_dict=ADMIN,
             start_date="2026-07-01",
             end_date="2026-08-01",
+            api_key="key-hash",
         )
-        assert captured["params"] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00")
+        assert captured["params"] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00", "key-hash")
         assert response.routers_in_scope == 1
         assert response.groups[0].router_name == "live-auto"
         assert response.groups[0].saved_pct == response.totals.saved_pct == 75.0
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx
index ce0cd75cd36..e7c6adc478c 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx
@@ -22,7 +22,7 @@ vi.mock("@/components/shared/advanced_date_picker", () => ({
 
 import { useAutoRouters } from "@/app/(dashboard)/hooks/models/useModels";
 
-import AutoRouterBenchmarksTab from "./AutoRouterBenchmarksTab";
+import AutoRouterBenchmarksTab, { AutoRouterUsageView } from "./AutoRouterBenchmarksTab";
 import type {
   AutoRouterBenchmarkGroup,
   AutoRouterBenchmarksResponse,
@@ -359,13 +359,37 @@ describe("AutoRouterBenchmarksTab", () => {
     mockHook({ data: response([group()]) });
     const { dateValue, onDateChange } = renderTab();
 
-    expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue);
+    expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue, undefined);
     expect(screen.getByText("Jul 6 – Aug 5 (UTC)")).toBeInTheDocument();
 
     fireEvent.click(screen.getByTestId("date-picker"));
     expect(onDateChange).toHaveBeenCalledWith({ from: new Date(2026, 7, 1), to: new Date(2026, 7, 5) });
   });
 
+  it("scopes the query to one key when the usage view is mounted for a key", () => {
+    mockHook({ data: response([group()]) });
+    const dateValue = { from: new Date(2026, 6, 6), to: new Date(2026, 7, 5) };
+    const activity = {
+      dateValue,
+      onDateChange: vi.fn(),
+      results: [],
+      loading: false,
+      isFetchingMore: false,
+      progress: { currentPage: 1, totalPages: 1 },
+      cancelled: false,
+      cancel: vi.fn(),
+    };
+    render(
+      <QueryClientProvider client={new QueryClient()}>
+        <AutoRouterUsageView accessToken="sk-test" activity={activity} apiKey="key-hash-1" />
+      </QueryClientProvider>,
+    );
+
+    expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue, "key-hash-1");
+    expect(screen.getByText("Total estimated savings")).toBeInTheDocument();
+    expect(screen.queryByRole("tab", { name: "Shadow Evals" })).not.toBeInTheDocument();
+  });
+
   it("shows usage by default and mounts shadow evals only when its sub-tab is selected", () => {
     mockHook({ data: response([group()]) });
     renderTab();
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx
index 09a0cf0242b..39e6b0fd390 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx
@@ -273,12 +273,13 @@ const BenchmarksBody: React.FC<BenchmarksBodyProps> = ({ isPending, error, data,
 
 interface AutoRouterBenchmarksTabProps {
   accessToken: string | null;
-  activity: DailyActivityRange;
+  activity: Pick<DailyActivityRange, "dateValue" | "onDateChange">;
+  apiKey?: string;
 }
 
-const UsageView: React.FC<AutoRouterBenchmarksTabProps> = ({ accessToken, activity }) => {
+export const AutoRouterUsageView: React.FC<AutoRouterBenchmarksTabProps> = ({ accessToken, activity, apiKey }) => {
   const { dateValue, onDateChange } = activity;
-  const { data, isPending, error } = useAutoRouterBenchmarks(accessToken, dateValue);
+  const { data, isPending, error } = useAutoRouterBenchmarks(accessToken, dateValue, apiKey);
   const [selectedKey, setSelectedKey] = useState<string>(ALL_ROUTERS);
   const { data: autoRouters } = useAutoRouters();
 
@@ -347,7 +348,7 @@ const AutoRouterBenchmarksTab: React.FC<AutoRouterBenchmarksTabProps> = ({ acces
       </TabsList>
 
       <TabsContent value="usage" keepMounted={visitedTabs.includes("usage")}>
-        <UsageView accessToken={accessToken} activity={activity} />
+        <AutoRouterUsageView accessToken={accessToken} activity={activity} />
       </TabsContent>
       <TabsContent value="shadow-evals" keepMounted={visitedTabs.includes("shadow-evals")}>
         <ShadowEvalSection />
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.test.ts
index e8a858ef3a5..1258c967ec9 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.test.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.test.ts
@@ -3,7 +3,9 @@ import { describe, expect, it, vi } from "vitest";
 vi.mock("@/components/networking", () => ({ formatDate: vi.fn() }));
 vi.mock("@/lib/http/api", () => ({ $api: { useQuery: vi.fn() } }));
 
-import { benchmarksWindow } from "./useAutoRouterBenchmarks";
+import { $api } from "@/lib/http/api";
+
+import { benchmarksWindow, useAutoRouterBenchmarks } from "./useAutoRouterBenchmarks";
 
 const localDay =
   (offsetHours: number) =>
@@ -44,3 +46,30 @@ describe("benchmarksWindow", () => {
     expect(benchmarksWindow({ to: now }, now, pacific)).toEqual({});
   });
 });
+
+describe("useAutoRouterBenchmarks", () => {
+  const range = { from: new Date("2026-07-06T19:00:00Z"), to: new Date("2026-08-05T19:00:00Z") };
+
+  it("forwards the key hash as the endpoint's api_key filter", () => {
+    useAutoRouterBenchmarks("sk-test", range, "key-hash-1");
+
+    const [, path, init] = vi.mocked($api.useQuery).mock.calls.at(-1) as unknown as [
+      string,
+      string,
+      { params: { query: Record<string, string | undefined> } },
+    ];
+    expect(path).toBe("/auto_router/benchmarks");
+    expect(init.params.query.api_key).toBe("key-hash-1");
+  });
+
+  it("leaves the read deployment-wide when no key is given", () => {
+    useAutoRouterBenchmarks("sk-test", range);
+
+    const [, , init] = vi.mocked($api.useQuery).mock.calls.at(-1) as unknown as [
+      string,
+      string,
+      { params: { query: Record<string, string | undefined> } },
+    ];
+    expect(init.params.query.api_key).toBeUndefined();
+  });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.ts
index 155c9d0c696..284342f0aec 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.ts
@@ -23,10 +23,10 @@ export const benchmarksWindow = (
   };
 };
 
-export const useAutoRouterBenchmarks = (accessToken: string | null, range: DateRange) =>
+export const useAutoRouterBenchmarks = (accessToken: string | null, range: DateRange, apiKey?: string) =>
   $api.useQuery(
     "get",
     "/auto_router/benchmarks",
-    { params: { query: benchmarksWindow(range, new Date()) } },
+    { params: { query: { ...benchmarksWindow(range, new Date()), api_key: apiKey } } },
     { enabled: Boolean(accessToken && range.from && range.to), retry: false },
   );
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.integration.test.tsx
index 78280e32eed..51dfe9fda90 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.integration.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.integration.test.tsx
@@ -24,12 +24,18 @@ describe("useScopedDailyActivityRange wiring", () => {
     );
     global.fetch = mockFetch;
 
-    renderHook(() => useScopedDailyActivityRange("sk-token", { userId: "u1", apiKey: "hash-abc" }));
+    const activity = {
+      dateValue: { from: new Date(2026, 7, 1), to: new Date(2026, 7, 10) },
+      onDateChange: vi.fn(),
+    };
+    renderHook(() => useScopedDailyActivityRange("sk-token", { userId: "u1", apiKey: "hash-abc" }, activity));
 
     await waitFor(() => expect(mockFetch).toHaveBeenCalled());
 
     const url = String(mockFetch.mock.calls[0][0]);
     expect(url).toContain("user_id=u1");
     expect(url).toContain("api_key=hash-abc");
+    expect(url).toContain("start_date=2026-08-01");
+    expect(url).toContain("end_date=2026-08-10");
   });
 });
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx
index 6c9281060a4..00902aa9fdd 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx
@@ -25,11 +25,19 @@ vi.mock("@/components/networking", () => ({
 }));
 
 import { userDailyActivityAggregatedCall } from "@/components/networking";
-import { useDailyActivityRange } from "./useDailyActivityRange";
+import { useActivityDateRange, useDailyActivityRange } from "./useDailyActivityRange";
 
 const argsOfLastCall = () => mockUsePaginatedDailyActivity.mock.calls.at(-1)?.[0].args as unknown[];
 
 describe("useDailyActivityRange", () => {
+  it("offers date-range state without starting a daily-activity query", () => {
+    const { result } = renderHook(() => useActivityDateRange());
+
+    expect(result.current.dateValue.from).toBeInstanceOf(Date);
+    expect(result.current.dateValue.to).toBeInstanceOf(Date);
+    expect(mockUsePaginatedDailyActivity).not.toHaveBeenCalled();
+  });
+
   it("queries every user's activity for an admin", () => {
     renderHook(() => useDailyActivityRange("test-token", "u1", "proxy_admin"));
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts
index 8ed57e36a94..3435b57dbc8 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts
@@ -37,14 +37,20 @@ export interface DailyActivityScope {
   apiKey?: string | null;
 }
 
-export const useScopedDailyActivityRange = (
-  accessToken: string | null,
-  scope: DailyActivityScope,
-): DailyActivityRange => {
+export type ActivityDateRange = Pick<DailyActivityRange, "dateValue" | "onDateChange">;
+
+export const useActivityDateRange = (): ActivityDateRange => {
   const initialFrom = useMemo(() => new Date(new Date().getTime() - THIRTY_DAYS_MS), []);
   const initialTo = useMemo(() => new Date(), []);
   const [dateValue, setDateValue] = useState<DateRange>({ from: initialFrom, to: initialTo });
+  return { dateValue, onDateChange: setDateValue };
+};
 
+export const useScopedDailyActivityRange = (
+  accessToken: string | null,
+  scope: DailyActivityScope,
+  { dateValue, onDateChange }: ActivityDateRange,
+): DailyActivityRange => {
   const startTime = dateValue.from ?? null;
   const endTime = dateValue.to ?? null;
   const { userId, apiKey = null } = scope;
@@ -63,7 +69,7 @@ export const useScopedDailyActivityRange = (
 
   return {
     dateValue,
-    onDateChange: setDateValue,
+    onDateChange,
     results: data.results as DailyData[],
     loading,
     isFetchingMore,
@@ -77,7 +83,7 @@ export const useDailyActivityRange = (
   accessToken: string | null,
   userId: string | null,
   userRole: string,
-): DailyActivityRange =>
-  useScopedDailyActivityRange(accessToken, {
-    userId: spendScopeUserId(userRole, userId),
-  });
+): DailyActivityRange => {
+  const dateRange = useActivityDateRange();
+  return useScopedDailyActivityRange(accessToken, { userId: spendScopeUserId(userRole, userId) }, dateRange);
+};
diff --git a/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx
new file mode 100644
index 00000000000..8cfd9d1941e
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx
@@ -0,0 +1,98 @@
+import { screen } from "@testing-library/react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+import { renderWithProviders, testQueryClient } from "../../../tests/test-utils";
+import KeyAutoRouterUsageTab from "./KeyAutoRouterUsageTab";
+
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+  default: () => ({ accessToken: "test-token", userId: "admin-123", userRole: "Admin" }),
+}));
+
+const jsonResponse = (body: unknown) =>
+  new Response(JSON.stringify(body), { status: 200, headers: { "content-type": "application/json" } });
+
+const cache = {
+  coverage_pct: 100,
+  hit_rate_pct: 50,
+  same_model: { turns: 2, hits: 1, hit_rate_pct: 50 },
+  first_visit: { turns: 1, hits: 0, hit_rate_pct: 0 },
+  return_to_tier: { turns: 1, hits: 1, hit_rate_pct: 100 },
+  unordered_turns: 0,
+  return_misses_expired: 0,
+  return_misses_within_ttl: 0,
+  return_misses_unknown: 0,
+  ttl_5m_turns: 0,
+  ttl_1h_turns: 0,
+};
+
+const stats = {
+  sessions: 2,
+  turns: 4,
+  avg_turns_per_session: 2,
+  avg_session_seconds: 30,
+  avg_tokens_per_session: 100,
+  spend: 1.25,
+  saved_spend: 8.75,
+  baseline_spend: 10,
+  saved_pct: 87.5,
+  saved_per_session: 4.375,
+  cache,
+};
+
+const benchmarks = {
+  start_date: "2025-01-01",
+  end_date: "2025-01-31",
+  routers_in_scope: 2,
+  totals: stats,
+  groups: [
+    { router_name: "router-one", router_type: "complexity", tier_turns: { SIMPLE: 4 }, ...stats },
+    {
+      router_name: "router-two",
+      router_type: "complexity",
+      tier_turns: { SIMPLE: 1 },
+      ...stats,
+      spend: 0.25,
+      saved_spend: 0.75,
+      baseline_spend: 1,
+    },
+  ],
+};
+
+const noDeployments = { data: [], total_count: 0, current_page: 1, total_pages: 1, size: 1000 };
+const fetchMock = vi.fn(async (request: Request | string) => {
+  const url = typeof request === "string" ? request : request.url;
+  if (url.includes("/auto_router/benchmarks")) return jsonResponse(benchmarks);
+  return jsonResponse(noDeployments);
+});
+const requestedUrls = () =>
+  fetchMock.mock.calls.map(([request]) => (typeof request === "string" ? request : request.url));
+
+describe("KeyAutoRouterUsageTab", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    testQueryClient.clear();
+    vi.stubGlobal("fetch", fetchMock);
+  });
+
+  it("renders this key's spend, baseline, savings and per-router filter", async () => {
+    const activity = {
+      dateValue: { from: new Date(2025, 0, 1), to: new Date(2025, 0, 31) },
+      onDateChange: vi.fn(),
+    };
+    renderWithProviders(<KeyAutoRouterUsageTab accessToken="test-token" keyToken="key-hash-1" activity={activity} />);
+
+    expect(await screen.findByText("$8.75")).toBeInTheDocument();
+    expect(screen.getByText("Actual auto-router spend")).toBeInTheDocument();
+    expect(screen.getByText("$1.25")).toBeInTheDocument();
+    expect(screen.getByText("Estimated spend at highest-tier model")).toBeInTheDocument();
+    expect(screen.getByText("$10.00")).toBeInTheDocument();
+    expect(screen.getByText("Auto-router prompt caching")).toBeInTheDocument();
+    expect(screen.getAllByText("50.0%").length).toBeGreaterThan(0);
+    expect(screen.getByText("All auto-routers")).toBeInTheDocument();
+
+    const benchmarkUrl = new URL(requestedUrls().find((url) => url.includes("/auto_router/benchmarks")) ?? "");
+    expect(benchmarkUrl.searchParams.get("api_key")).toBe("key-hash-1");
+    expect(benchmarkUrl.searchParams.get("start_date")).toBe("2025-01-01");
+    expect(benchmarkUrl.searchParams.get("end_date")).toBe("2025-01-31");
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.tsx b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.tsx
new file mode 100644
index 00000000000..9e1cee1ec2e
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.tsx
@@ -0,0 +1,18 @@
+"use client";
+
+import React from "react";
+
+import { AutoRouterUsageView } from "@/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab";
+import type { ActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange";
+
+interface KeyAutoRouterUsageTabProps {
+  accessToken: string | null;
+  keyToken: string;
+  activity: ActivityDateRange;
+}
+
+const KeyAutoRouterUsageTab: React.FC<KeyAutoRouterUsageTabProps> = ({ accessToken, keyToken, activity }) => (
+  <AutoRouterUsageView accessToken={accessToken} activity={activity} apiKey={keyToken} />
+);
+
+export default KeyAutoRouterUsageTab;
diff --git a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx
index 6c7d65ee566..385c3967d02 100644
--- a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx
+++ b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx
@@ -46,6 +46,11 @@ const mockActivity = (
 
 const scopedRange = () => vi.spyOn(useScopedDailyActivityRangeModule, "useScopedDailyActivityRange");
 
+const activity = {
+  dateValue: { from: new Date(2025, 0, 1), to: new Date(2025, 0, 31) },
+  onDateChange: vi.fn(),
+};
+
 const renderTab = (props: Partial<React.ComponentProps<typeof KeySavingsTab>> = {}) =>
   render(
     <KeySavingsTab
@@ -53,6 +58,7 @@ const renderTab = (props: Partial<React.ComponentProps<typeof KeySavingsTab>> =
       keyToken="key-abc123"
       userId="user-123"
       userRole="Internal User"
+      activity={activity}
       {...props}
     />,
   );
@@ -114,7 +120,7 @@ describe("KeySavingsTab", () => {
 
     renderTab({ userId: "user-456", userRole: "Internal User" });
 
-    expect(hook).toHaveBeenCalledWith("test-token", { userId: "user-456", apiKey: "key-abc123" });
+    expect(hook).toHaveBeenCalledWith("test-token", { userId: "user-456", apiKey: "key-abc123" }, activity);
     expect(screen.getByTestId("key-savings-scope-note")).toHaveTextContent("Showing your own requests");
   });
 
@@ -123,7 +129,7 @@ describe("KeySavingsTab", () => {
 
     renderTab({ userId: "admin-123", userRole: "Admin" });
 
-    expect(hook).toHaveBeenCalledWith("test-token", { userId: null, apiKey: "key-abc123" });
+    expect(hook).toHaveBeenCalledWith("test-token", { userId: null, apiKey: "key-abc123" }, activity);
     expect(screen.queryByTestId("key-savings-scope-note")).not.toBeInTheDocument();
   });
 
@@ -132,7 +138,7 @@ describe("KeySavingsTab", () => {
 
     renderTab({ userId: "org-admin-1", userRole: "Org Admin" });
 
-    expect(hook).toHaveBeenCalledWith("test-token", { userId: "org-admin-1", apiKey: "key-abc123" });
+    expect(hook).toHaveBeenCalledWith("test-token", { userId: "org-admin-1", apiKey: "key-abc123" }, activity);
     expect(screen.getByTestId("key-savings-scope-note")).toHaveTextContent("Showing your own requests");
   });
 });
diff --git a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx
index 967c3c7eff0..c33529eb042 100644
--- a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx
+++ b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx
@@ -22,7 +22,10 @@ import {
   usd,
   withStartAnchor,
 } from "@/app/(dashboard)/cost-optimization/_components/costOptimizationUtils";
-import { useScopedDailyActivityRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange";
+import {
+  useScopedDailyActivityRange,
+  type ActivityDateRange,
+} from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange";
 
 interface KeySavingsTabProps {
   accessToken: string | null;
@@ -30,19 +33,19 @@ interface KeySavingsTabProps {
   keyToken: string;
   userId: string | null;
   userRole: string;
+  activity: ActivityDateRange;
 }
 
-const KeySavingsTab: React.FC<KeySavingsTabProps> = ({ accessToken, keyToken, userId, userRole }) => {
+const KeySavingsTab: React.FC<KeySavingsTabProps> = ({ accessToken, keyToken, userId, userRole, activity }) => {
   // Proxy admins read the whole key. For anyone else the endpoint applies the caller's own user_id
   // alongside the key filter, so the figures cover only that viewer's requests on this key -- said
   // plainly in the scope note below rather than left to be misread as the key's total.
   const readsWholeKey = hasProxyWideSpendView(userRole);
-  const activity = useScopedDailyActivityRange(accessToken, {
-    userId: spendScopeUserId(userRole, userId),
-    apiKey: keyToken,
-  });
-
-  const { dateValue, onDateChange, results, loading, isFetchingMore } = activity;
+  const { dateValue, onDateChange, results, loading, isFetchingMore } = useScopedDailyActivityRange(
+    accessToken,
+    { userId: spendScopeUserId(userRole, userId), apiKey: keyToken },
+    activity,
+  );
   const startTime = dateValue.from ?? null;
   const endTime = dateValue.to ?? null;
 
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 7a27cc41e5e..b403255b329 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
@@ -26,6 +26,30 @@ vi.mock("./key_edit_view", () => ({
   },
 }));
 
+import type { ActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange";
+
+const AnalyticsDateControl = ({ activity, keyToken }: { activity: ActivityDateRange; keyToken: string }) => (
+  <div>
+    <span data-testid="key-auto-router-usage">{keyToken}</span>
+    <output aria-label="Selected dates">{activity.dateValue.from?.toISOString()}</output>
+    {[1, 10].map((day) => (
+      <button
+        key={day}
+        onClick={() => activity.onDateChange({ from: new Date(Date.UTC(2026, 7, day)), to: new Date(2026, 7, 20) })}
+      >
+        Select August {day}
+      </button>
+    ))}
+  </div>
+);
+
+vi.mock("./KeyAutoRouterUsageTab", () => ({
+  default: (props: React.ComponentProps<typeof AnalyticsDateControl>) => <AnalyticsDateControl {...props} />,
+}));
+vi.mock("./KeySavingsTab", () => ({
+  default: (props: React.ComponentProps<typeof AnalyticsDateControl>) => <AnalyticsDateControl {...props} />,
+}));
+
 vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({
   default: vi.fn(),
 }));
@@ -176,6 +200,38 @@ describe("KeyInfoView", () => {
     await userEvent.click(await screen.findByRole("button", { name: /more key actions/i }));
   };
 
+  it("shows key-scoped auto-router usage as its own admin tab", async () => {
+    vi.mocked(useAuthorized).mockReturnValue({ ...baseUseAuthorizedMock, userRole: "Admin" });
+    renderWithProviders(<KeyInfoView keyData={MOCK_KEY_DATA} onClose={() => {}} keyId="test-key-id" teams={[]} />);
+
+    await userEvent.click(screen.getByRole("tab", { name: "Auto-router usage" }));
+
+    expect(screen.getByTestId("key-auto-router-usage")).toHaveTextContent("test-token-123");
+  });
+
+  it("preserves dates in both directions across unmounted analytics panels", async () => {
+    vi.mocked(useAuthorized).mockReturnValue({ ...baseUseAuthorizedMock, userRole: "Admin" });
+    renderWithProviders(<KeyInfoView keyData={MOCK_KEY_DATA} onClose={() => {}} keyId="test-key-id" teams={[]} />);
+
+    expect(screen.queryByLabelText("Selected dates")).not.toBeInTheDocument();
+    await userEvent.click(screen.getByRole("tab", { name: "Savings" }));
+    await userEvent.click(screen.getByRole("button", { name: "Select August 1" }));
+    await userEvent.click(screen.getByRole("tab", { name: "Auto-router usage" }));
+    expect(screen.getByLabelText("Selected dates")).toHaveTextContent("2026-08-01T00:00:00.000Z");
+    await userEvent.click(screen.getByRole("button", { name: "Select August 10" }));
+    await userEvent.click(screen.getByRole("tab", { name: "Settings" }));
+    expect(screen.queryByLabelText("Selected dates")).not.toBeInTheDocument();
+    await userEvent.click(screen.getByRole("tab", { name: "Savings" }));
+    expect(screen.getByLabelText("Selected dates")).toHaveTextContent("2026-08-10T00:00:00.000Z");
+  });
+
+  it("does not offer the admin-only auto-router usage tab to an internal user", () => {
+    vi.mocked(useAuthorized).mockReturnValue({ ...baseUseAuthorizedMock, userRole: "Internal User" });
+    renderWithProviders(<KeyInfoView keyData={MOCK_KEY_DATA} onClose={() => {}} keyId="test-key-id" teams={[]} />);
+
+    expect(screen.queryByRole("tab", { name: "Auto-router usage" })).not.toBeInTheDocument();
+  });
+
   describe("last updated", () => {
     const renderWithTimestamps = (overrides: Partial<KeyResponse>) => {
       vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock);
diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx
index 0dd0dd6d6af..f5c682a2ee0 100644
--- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx
+++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx
@@ -16,8 +16,15 @@ import { modelGroupHref, teamDetailHref } from "@/utils/entityLinks";
 import { BadgeLink } from "@/components/shared/BadgeLink";
 import { KeyInfoHeader } from "./KeyInfoHeader";
 import KeySavingsTab from "./KeySavingsTab";
+import KeyAutoRouterUsageTab from "./KeyAutoRouterUsageTab";
+import { useActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange";
 import { useEffect, useState } from "react";
-import { isProxyAdminRole, isUserTeamAdminForSingleTeam, rolesWithWriteAccess } from "../../utils/roles";
+import {
+  hasProxyWideSpendView,
+  isProxyAdminRole,
+  isUserTeamAdminForSingleTeam,
+  rolesWithWriteAccess,
+} from "../../utils/roles";
 import { mapDisplayToInternalNames, mapInternalToDisplayNames } from "../callback_info_helpers";
 import AutoRotationView from "../common_components/AutoRotationView";
 import DeleteResourceModal from "../common_components/DeleteResourceModal";
@@ -81,6 +88,7 @@ export default function KeyInfoView({
   backButtonText = "Back to Keys",
 }: KeyInfoViewProps) {
   const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized();
+  const activityDateRange = useActivityDateRange();
   const queryClient = useQueryClient();
   const canEditGuardrails = premiumUser || (userRole != null && rolesWithWriteAccess.includes(userRole));
   const { teams: teamsData } = useTeams();
@@ -618,6 +626,11 @@ export default function KeyInfoView({
           <TabsTrigger value="savings" className="flex-none rounded-none px-4 py-2">
             Savings
           </TabsTrigger>
+          {hasProxyWideSpendView(userRole) && (
+            <TabsTrigger value="auto-router-usage" className="flex-none rounded-none px-4 py-2">
+              Auto-router usage
+            </TabsTrigger>
+          )}
           <TabsTrigger value="settings" className="flex-none rounded-none px-4 py-2">
             Settings
           </TabsTrigger>
@@ -761,9 +774,20 @@ export default function KeyInfoView({
               keyToken={currentKeyData.token}
               userId={userID}
               userRole={userRole}
+              activity={activityDateRange}
             />
           </TabsContent>
 
+          {hasProxyWideSpendView(userRole) && (
+            <TabsContent value="auto-router-usage">
+              <KeyAutoRouterUsageTab
+                accessToken={accessToken}
+                keyToken={currentKeyData.token}
+                activity={activityDateRange}
+              />
+            </TabsContent>
+          )}
+
           {/* Settings Panel */}
           <TabsContent value="settings" keepMounted>
             <Card className="block p-6">
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 36fd744efc0..842a3da4122 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -41393,6 +41393,8 @@ export interface operations {
                 start_date?: string | null;
                 /** @description YYYY-MM-DD UTC, inclusive (defaults to today) */
                 end_date?: string | null;
+                /** @description Filter to one virtual key token hash */
+                api_key?: string | null;
             };
             header?: never;
             path?: never;

From cc287a7d8f484dc3eb063aff33999068f733b00c Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Mon, 7 Sep 2026 12:23:47 -0700
Subject: [PATCH 130/164] fix(ui): hide the Create Vector Store flow from non
 proxy admins (#40148)

* fix(ui): hide the Create Vector Store flow from non proxy admins

The vector stores page rendered the Create Vector Store tab, the
+ Add Vector Store button and a GET /credentials call for every role,
while the proxy only lets proxy admins call POST /vector_store/new and
GET /credentials. Internal users landed on the create form and got an
Only proxy admin error toast. Gate all three on isProxyAdminRole and
default everyone else to the Manage tab, matching the Indexes tab and
the Add Model gating.

Resolves LIT-7131

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(ui): exclude view-only admin sessions from the vector store create flow

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../vector-stores/_components/index.test.tsx  | 51 ++++++++++++++++---
 .../vector-stores/_components/index.tsx       | 35 ++++++++-----
 .../app/(dashboard)/vector-stores/page.tsx    |  6 ++-
 3 files changed, 70 insertions(+), 22 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx
index 7137da201d8..3224496b13a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx
@@ -45,7 +45,7 @@ describe("VectorStoreManagement loading state", () => {
 
   it("should resolve the loading state when accessToken is null instead of showing the skeleton forever", async () => {
     const user = userEvent.setup();
-    render(<VectorStoreManagement accessToken={null} userID={null} userRole={null} />);
+    render(<VectorStoreManagement accessToken={null} userID={null} userRole={null} isViewOnly={false} />);
     await openManageTab(user);
     expect(await screen.findByText("table-loaded")).toBeInTheDocument();
     expect(mockVectorStoreListCall).not.toHaveBeenCalled();
@@ -59,7 +59,7 @@ describe("VectorStoreManagement loading state", () => {
         resolveFetch = resolve;
       }),
     );
-    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin" />);
+    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin" isViewOnly={false} />);
     await openManageTab(user);
     expect(screen.getByText("table-loading")).toBeInTheDocument();
 
@@ -69,6 +69,43 @@ describe("VectorStoreManagement loading state", () => {
   });
 });
 
+describe("VectorStoreManagement create flow visibility", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockVectorStoreListCall.mockResolvedValue({ data: [] });
+    mockCredentialListCall.mockResolvedValue({ credentials: [] });
+  });
+
+  it.each([
+    { label: "Internal User", userRole: "Internal User", isViewOnly: false },
+    { label: "Internal Viewer", userRole: "Internal Viewer", isViewOnly: true },
+    { label: "proxy_admin_viewer session (userRole Admin, isViewOnly)", userRole: "Admin", isViewOnly: true },
+    { label: "Org Admin", userRole: "Org Admin", isViewOnly: false },
+  ])(
+    "should hide the Create Vector Store tab and button and skip /credentials for $label",
+    async ({ userRole, isViewOnly }) => {
+      render(
+        <VectorStoreManagement accessToken="sk-test" userID="user-1" userRole={userRole} isViewOnly={isViewOnly} />,
+      );
+      await waitFor(() => expect(mockVectorStoreListCall).toHaveBeenCalledWith("sk-test"));
+      expect(screen.queryByRole("tab", { name: "Create Vector Store" })).not.toBeInTheDocument();
+      expect(screen.getByRole("tab", { name: "Manage Vector Stores" })).toHaveAttribute("aria-selected", "true");
+      expect(await screen.findByText("table-loaded")).toBeInTheDocument();
+      expect(screen.queryByRole("button", { name: "+ Add Vector Store" })).not.toBeInTheDocument();
+      expect(mockCredentialListCall).not.toHaveBeenCalled();
+    },
+  );
+
+  it("should keep the Create Vector Store tab and button and fetch /credentials for a proxy admin", async () => {
+    const user = userEvent.setup();
+    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin" isViewOnly={false} />);
+    await waitFor(() => expect(mockCredentialListCall).toHaveBeenCalledWith("sk-test"));
+    expect(screen.getByRole("tab", { name: "Create Vector Store" })).toHaveAttribute("aria-selected", "true");
+    await openManageTab(user);
+    expect(screen.getByRole("button", { name: "+ Add Vector Store" })).toBeInTheDocument();
+  });
+});
+
 describe("VectorStoreManagement Indexes tab", () => {
   beforeEach(() => {
     vi.clearAllMocks();
@@ -88,7 +125,7 @@ describe("VectorStoreManagement Indexes tab", () => {
         },
       ],
     });
-    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin" />);
+    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin" isViewOnly={false} />);
     await user.click(screen.getByRole("tab", { name: "Indexes" }));
     expect(await screen.findByText("support-docs-index")).toBeInTheDocument();
     expect(screen.getByText("support-docs-store")).toBeInTheDocument();
@@ -96,7 +133,7 @@ describe("VectorStoreManagement Indexes tab", () => {
   });
 
   it("should not render the Indexes tab for an Admin Viewer", async () => {
-    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin Viewer" />);
+    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin Viewer" isViewOnly={true} />);
     await waitFor(() => expect(mockVectorStoreListCall).toHaveBeenCalledWith("sk-test"));
     expect(screen.getByRole("tab", { name: "Manage Vector Stores" })).toBeInTheDocument();
     expect(screen.queryByRole("tab", { name: "Indexes" })).not.toBeInTheDocument();
@@ -125,7 +162,7 @@ describe("VectorStoreManagement Indexes tab", () => {
         },
       ],
     });
-    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin" />);
+    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin" isViewOnly={false} />);
     await user.click(screen.getByRole("tab", { name: "Indexes" }));
     await user.click(await screen.findByRole("button", { name: "support-docs-store" }));
     expect(await screen.findByTestId("vector-store-info-view")).toHaveTextContent("vs-1");
@@ -135,7 +172,7 @@ describe("VectorStoreManagement Indexes tab", () => {
   it("should link to the feature docs and a GitHub issue for unsupported providers on the Indexes tab", async () => {
     const user = userEvent.setup();
     mockIndexesListCall.mockResolvedValue({ object: "list", data: [] });
-    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin" />);
+    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin" isViewOnly={false} />);
     await user.click(screen.getByRole("tab", { name: "Indexes" }));
     expect(screen.getByRole("link", { name: "vector store index docs" })).toHaveAttribute(
       "href",
@@ -149,7 +186,7 @@ describe("VectorStoreManagement Indexes tab", () => {
   });
 
   it("should not call indexesListCall until the Indexes tab is clicked", async () => {
-    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin" />);
+    render(<VectorStoreManagement accessToken="sk-test" userID="user-1" userRole="Admin" isViewOnly={false} />);
     await waitFor(() => expect(mockVectorStoreListCall).toHaveBeenCalledWith("sk-test"));
     expect(screen.getByRole("tab", { name: "Indexes" })).toBeInTheDocument();
     expect(mockIndexesListCall).not.toHaveBeenCalled();
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx
index 1745c710c51..285a96520bd 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx
@@ -24,9 +24,10 @@ interface VectorStoreProps {
   accessToken: string | null;
   userID: string | null;
   userRole: string | null;
+  isViewOnly: boolean;
 }
 
-const VectorStoreManagement: React.FC<VectorStoreProps> = ({ accessToken, userID, userRole }) => {
+const VectorStoreManagement: React.FC<VectorStoreProps> = ({ accessToken, userID, userRole, isViewOnly }) => {
   const [vectorStores, setVectorStores] = useState<VectorStore[]>([]);
   const [isLoadingVectorStores, setIsLoadingVectorStores] = useState(true);
   const [isCreateModalVisible, setIsCreateModalVisible] = useState(false);
@@ -37,7 +38,9 @@ const VectorStoreManagement: React.FC<VectorStoreProps> = ({ accessToken, userID
   const [selectedVectorStoreId, setSelectedVectorStoreId] = useState<string | null>(null);
   const [editVectorStore, setEditVectorStore] = useState(false);
   const [isDeleting, setIsDeleting] = useState(false);
-  const { onTabChange, hasVisited } = useVisitedTabs("create");
+  const canCreateVectorStores = isProxyAdminRole(userRole || "") && !isViewOnly;
+  const defaultTab = canCreateVectorStores ? "create" : "manage";
+  const { onTabChange, hasVisited } = useVisitedTabs(defaultTab);
 
   const fetchVectorStores = async () => {
     if (!accessToken) {
@@ -56,7 +59,7 @@ const VectorStoreManagement: React.FC<VectorStoreProps> = ({ accessToken, userID
   };
 
   const fetchCredentials = async () => {
-    if (!accessToken) return;
+    if (!accessToken || !canCreateVectorStores) return;
     try {
       const response = await credentialListCall(accessToken);
       setCredentials(response.credentials || []);
@@ -153,11 +156,13 @@ const VectorStoreManagement: React.FC<VectorStoreProps> = ({ accessToken, userID
           You can use vector stores to store and retrieve LLM embeddings.
         </p>
 
-        <Tabs defaultValue="create" onValueChange={onTabChange}>
+        <Tabs defaultValue={defaultTab} onValueChange={onTabChange}>
           <TabsList variant="line" className="mb-6 h-auto w-full justify-start rounded-none p-0">
-            <TabsTrigger value="create" className="flex-none rounded-none px-4 py-2">
-              Create Vector Store
-            </TabsTrigger>
+            {canCreateVectorStores && (
+              <TabsTrigger value="create" className="flex-none rounded-none px-4 py-2">
+                Create Vector Store
+              </TabsTrigger>
+            )}
             <TabsTrigger value="manage" className="flex-none rounded-none px-4 py-2">
               Manage Vector Stores
             </TabsTrigger>
@@ -171,14 +176,18 @@ const VectorStoreManagement: React.FC<VectorStoreProps> = ({ accessToken, userID
             )}
           </TabsList>
 
-          <TabsContent keepMounted={hasVisited("create")} value="create">
-            <CreateVectorStore accessToken={accessToken} onSuccess={handleVectorStoreCreated} />
-          </TabsContent>
+          {canCreateVectorStores && (
+            <TabsContent keepMounted={hasVisited("create")} value="create">
+              <CreateVectorStore accessToken={accessToken} onSuccess={handleVectorStoreCreated} />
+            </TabsContent>
+          )}
 
           <TabsContent keepMounted={hasVisited("manage")} value="manage">
-            <Button className="mb-4" onClick={() => setIsCreateModalVisible(true)}>
-              + Add Vector Store
-            </Button>
+            {canCreateVectorStores && (
+              <Button className="mb-4" onClick={() => setIsCreateModalVisible(true)}>
+                + Add Vector Store
+              </Button>
+            )}
 
             <div className="grid grid-cols-1 gap-2 pt-2 pb-2 w-full mt-2">
               <VectorStoreTable
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/page.tsx
index 34759152d11..94fedb557a7 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/page.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/page.tsx
@@ -4,6 +4,8 @@ import VectorStoreManagement from "./_components";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 
 export default function VectorStores() {
-  const { accessToken, userRole, userId } = useAuthorized();
-  return <VectorStoreManagement accessToken={accessToken} userRole={userRole} userID={userId} />;
+  const { accessToken, userRole, userId, isViewOnly } = useAuthorized();
+  return (
+    <VectorStoreManagement accessToken={accessToken} userRole={userRole} userID={userId} isViewOnly={isViewOnly} />
+  );
 }

From c949843157f890ea253b51af4fa825e78543feb8 Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Mon, 7 Sep 2026 12:24:04 -0700
Subject: [PATCH 131/164] fix(ui): send empty vector_stores when the last team
 vector store is removed (#40144)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../src/components/team/TeamInfo.test.tsx           | 13 +++++++++++++
 .../src/components/team/TeamInfo.tsx                |  2 +-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
index 041ba4b548d..b286c8c1303 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
@@ -2052,6 +2052,19 @@ describe("TeamInfoView - the exact bytes the update call sends", () => {
     expect(objectPermission.agent_access_groups).toStrictEqual([]);
   });
 
+  it("sends an empty vector_stores array after the last vector store chip is removed", async () => {
+    const user = userEvent.setup({ delay: null });
+    await openEditor(user);
+
+    await user.click(within(screen.getByLabelText("vs-1")).getByRole("button"));
+    expect(screen.queryByLabelText("vs-1")).not.toBeInTheDocument();
+
+    const payload = await save(user);
+
+    const objectPermission = wireBody(payload).object_permission as Record<string, unknown>;
+    expect(objectPermission.vector_stores).toStrictEqual([]);
+  });
+
   it("resends every stored value once both sections are opened", async () => {
     const user = userEvent.setup({ delay: null });
     await openEditor(user);
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
index efd69a79fc0..4947cec1441 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
@@ -1039,7 +1039,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
       delete values.agents_and_groups;
 
       // Handle vector stores permissions
-      if (values.vector_stores && values.vector_stores.length > 0) {
+      if (values.vector_stores) {
         updateData.object_permission.vector_stores = values.vector_stores;
       }
 

From 192e38fa7ba2f529cfaad3bcfc28d2613aca28d5 Mon Sep 17 00:00:00 2001
From: Yassin Kortam <yassin@berri.ai>
Date: Mon, 7 Sep 2026 12:28:38 -0700
Subject: [PATCH 132/164] feat(skills): semantic search over the LiteLLM-hosted
 skill registry (#39401)

* feat(skills): semantic search over the LiteLLM-hosted skill registry

Adds GET /v1/skills?query= (custom_llm_provider=litellm_proxy) and a
skill_search MCP virtual tool, ranking the caller's accessible skills by
semantic similarity, mirroring the A2A agent registry search (LIT-6309).

Also fixes a pre-existing bug where create_skill() dropped description and
instructions for the litellm_proxy provider, which left every LiteLLM-hosted
skill with no searchable text.

* fix(mcp): coerce skill_search top_k instead of raising 500 on malformed input

The MCP-REST skill_search dispatch validated raw tool arguments through a
pydantic model directly, so a non-numeric top_k raised a ValidationError
that the endpoint's catch-all turned into an HTTP 500. Mirrors the
agent_search branch's tolerant coerce_top_k handling instead.

* fix(skills): enforce key limits on search embeddings and bound the semantic index

Semantic search embeddings now run the same pre_call_hook the /embeddings
route runs, so key rate limits, budgets and guardrails apply before the
embedding model is called. The shared SemanticTextIndex caps cached vectors
and evicts the least recently searched entries, and each skill's embedded
text is capped so one skill cannot inflate the embedding batch

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(skills): surface proxy 429s from search embeddings instead of a 503

ProxyRateLimitError is also an OpenAIError, so the search engine was folding
a key rate limit into skill_search_unavailable. Proxy HTTPExceptions now
propagate so the caller gets the same 429 the /embeddings route returns

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(skills): import assert_never from typing_extensions for Python 3.10

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(skills): embed the request as the pre-call hooks returned it, not the original text

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(skills): keep the litellm_proxy provider check for GET /v1/skills?query= inside llms/

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(skills): move the GET /v1/skills?query= endpoint tests under tests/test_litellm/proxy

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .github/workflows/test-unit.yml               |   1 +
 litellm/__init__.py                           |   1 +
 .../llms/litellm_proxy/skills/constants.py    |   5 +
 litellm/llms/litellm_proxy/skills/handler.py  |  19 +-
 .../llms/litellm_proxy/skills/skill_search.py | 161 +++++++
 .../litellm_proxy/skills/transformation.py    |  11 +-
 .../mcp_server/rest_endpoints.py              |  13 +
 .../proxy/_experimental/mcp_server/server.py  |   9 +
 .../_experimental/mcp_server/tool_search.py   |  68 ++-
 litellm/proxy/_lazy_openapi_snapshot.json     |  57 ++-
 litellm/proxy/agent_endpoints/agent_search.py |   7 +-
 litellm/proxy/agent_endpoints/endpoints.py    |   3 +-
 .../anthropic_endpoints/skills_endpoints.py   |  97 +++-
 .../proxy/common_utils/semantic_text_index.py |  79 +++-
 litellm/skills/main.py                        |   6 +-
 litellm/types/llms/anthropic_skills.py        |   9 +
 .../litellm_proxy/skills/test_skill_search.py | 436 ++++++++++++++++++
 .../mcp_server/test_mcp_tool_search.py        |  50 +-
 .../agent_endpoints/test_agent_search.py      |  27 +-
 .../test_skills_endpoints.py                  | 175 +++++++
 tests/test_litellm/skills/test_skills_main.py |  57 +++
 ui/litellm-dashboard/src/lib/http/schema.d.ts |  14 +
 22 files changed, 1259 insertions(+), 46 deletions(-)
 create mode 100644 litellm/llms/litellm_proxy/skills/skill_search.py
 create mode 100644 tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py
 create mode 100644 tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py
 create mode 100644 tests/test_litellm/skills/test_skills_main.py

diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml
index 33245ec5b5f..cc606339a20 100644
--- a/.github/workflows/test-unit.yml
+++ b/.github/workflows/test-unit.yml
@@ -116,6 +116,7 @@ jobs:
               tests/test_litellm/rerank_api
               tests/test_litellm/rust_bridge
               tests/test_litellm/sandbox
+              tests/test_litellm/skills
               tests/test_litellm/test_router
               tests/test_litellm/vector_stores
               tests/test_litellm/videos
diff --git a/litellm/__init__.py b/litellm/__init__.py
index 42c0ea881fd..dc2f40af46e 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -495,6 +495,7 @@ public_model_groups: Optional[List[str]] = None
 public_agent_groups: Optional[List[str]] = None
 agent_search_embedding_model: Optional[str] = None
 mcp_tool_search: Optional[Mapping[str, object]] = None
+skill_search_embedding_model: Optional[str] = None
 # Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]])
 # New format: { "displayName": { "url": "...", "index": 0 } }
 # Old format: { "displayName": "url" } (for backward compatibility)
diff --git a/litellm/llms/litellm_proxy/skills/constants.py b/litellm/llms/litellm_proxy/skills/constants.py
index a6c88718f11..04a3a7dbc91 100644
--- a/litellm/llms/litellm_proxy/skills/constants.py
+++ b/litellm/llms/litellm_proxy/skills/constants.py
@@ -16,3 +16,8 @@ DEFAULT_MAX_ITERATIONS: Final[int] = 10
 
 DEFAULT_SANDBOX_TIMEOUT: Final[int] = 120
 """Default timeout in seconds for sandbox code execution."""
+
+MAX_SKILLS_PER_SEARCH: Final[int] = 5000
+"""Upper bound on how many of the caller's accessible skills a single semantic
+search embeds. Ranking runs in memory over this candidate set (no tsvector/DB-side
+filtering yet), so this caps worst-case embedding cost per search request."""
diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py
index 73f6ed23092..9b625cb0571 100644
--- a/litellm/llms/litellm_proxy/skills/handler.py
+++ b/litellm/llms/litellm_proxy/skills/handler.py
@@ -6,11 +6,15 @@ Used by the transformation layer and skills injection hook.
 """
 
 import uuid
+from collections.abc import Sequence
 from typing import Final
 
 from litellm._logging import verbose_logger
 from litellm.caching.in_memory_cache import InMemoryCache
-from litellm.llms.litellm_proxy.skills.constants import LITELLM_SKILL_ID_PREFIX
+from litellm.llms.litellm_proxy.skills.constants import (
+    LITELLM_SKILL_ID_PREFIX,
+    MAX_SKILLS_PER_SEARCH,
+)
 from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth
 from litellm.proxy.common_utils.resource_ownership import (
     get_primary_resource_owner_scope,
@@ -131,6 +135,19 @@ class LiteLLMSkillsHandler:
         )
         return [_prisma_skill_to_litellm(s) for s in skills]
 
+    @staticmethod
+    async def list_skills_for_search(
+        user_api_key_dict: UserAPIKeyAuth | None = None,
+    ) -> Sequence[LiteLLM_SkillsTable]:
+        """Every skill the caller can access, for ranking. Same owner-scope filter as
+        ``list_skills``, but unpaginated (up to ``MAX_SKILLS_PER_SEARCH``) since a query
+        must be scored against the whole accessible set, not one page of it."""
+        return await LiteLLMSkillsHandler.list_skills(
+            limit=MAX_SKILLS_PER_SEARCH,
+            offset=0,
+            user_api_key_dict=user_api_key_dict,
+        )
+
     @staticmethod
     async def _load_skill(skill_id: str) -> object | None:
         """Cache-first read of the Prisma skill row. Owner-scope filtering
diff --git a/litellm/llms/litellm_proxy/skills/skill_search.py b/litellm/llms/litellm_proxy/skills/skill_search.py
new file mode 100644
index 00000000000..f975c6c4cab
--- /dev/null
+++ b/litellm/llms/litellm_proxy/skills/skill_search.py
@@ -0,0 +1,161 @@
+"""Semantic ranking over the LiteLLM-hosted skill registry, shared by GET /v1/skills?query= and the skill_search MCP tool."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Final, TypeAlias
+
+from pydantic import BaseModel, ConfigDict
+
+from litellm.llms.litellm_proxy.skills.constants import MAX_SKILLS_PER_SEARCH
+from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler
+from litellm.proxy.common_utils.semantic_text_index import (
+    Embedder,
+    EmbeddingFailed,
+    SemanticTextIndex,
+    router_embedder,
+)
+from litellm.types.utils import LlmProviders
+
+if TYPE_CHECKING:
+    from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth
+    from litellm.proxy.utils import ProxyLogging
+    from litellm.router import Router
+
+DEFAULT_SKILL_SEARCH_TOP_K: Final = 5
+MAX_SKILL_SEARCH_TOP_K: Final = 100
+"""Matches the ``le=100`` bound GET /v1/skills?query= enforces via FastAPI's Query
+validation, so the MCP tool can't return a larger payload than the REST endpoint allows."""
+MAX_SKILL_SEARCH_TEXT_CHARS: Final = 4000
+"""Per-skill cap on the title + description + instructions text that gets embedded, so one
+search embeds at most ``MAX_SKILLS_PER_SEARCH * MAX_SKILL_SEARCH_TEXT_CHARS`` characters no
+matter how long the stored instructions are."""
+
+
+@dataclass(frozen=True, slots=True)
+class SkillSearchHit:
+    skill: LiteLLM_SkillsTable
+    score: float
+
+
+@dataclass(frozen=True, slots=True)
+class SkillSearchHits:
+    hits: tuple[SkillSearchHit, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class SkillSearchNotConfigured:
+    reason: str
+
+
+@dataclass(frozen=True, slots=True)
+class SkillSearchEmbeddingFailed:
+    reason: str
+
+
+@dataclass(frozen=True, slots=True)
+class SkillSearchUnsupportedProvider:
+    reason: str
+
+
+SkillSearchOutcome: TypeAlias = SkillSearchHits | SkillSearchNotConfigured | SkillSearchEmbeddingFailed
+HostedSkillSearchOutcome: TypeAlias = SkillSearchOutcome | SkillSearchUnsupportedProvider
+
+
+class SkillSearchResult(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    skill_id: str
+    display_title: str | None
+    description: str | None
+    score: float
+
+
+def skill_search_text(skill: LiteLLM_SkillsTable) -> str:
+    joined: Final = "\n".join(part for part in (skill.display_title, skill.description, skill.instructions) if part)
+    return joined[:MAX_SKILL_SEARCH_TEXT_CHARS]
+
+
+def skill_search_result(hit: SkillSearchHit) -> SkillSearchResult:
+    return SkillSearchResult(
+        skill_id=hit.skill.skill_id,
+        display_title=hit.skill.display_title,
+        description=hit.skill.description,
+        score=hit.score,
+    )
+
+
+class SkillSearchIndex:
+    """Caches one vector per distinct skill text per embedding model, so repeat searches only embed the query."""
+
+    def __init__(self, max_entries: int = MAX_SKILLS_PER_SEARCH) -> None:
+        self._index: Final = SemanticTextIndex(max_entries=max_entries)
+
+    async def search(
+        self,
+        query: str,
+        skills: Sequence[LiteLLM_SkillsTable],
+        top_k: int,
+        embed: Embedder,
+        embedding_model: str,
+    ) -> SkillSearchHits | SkillSearchEmbeddingFailed:
+        texts: Final = tuple(skill_search_text(skill) for skill in skills)
+        scores: Final = await self._index.scores(query, texts, embed, embedding_model)
+        if isinstance(scores, EmbeddingFailed):
+            return SkillSearchEmbeddingFailed(reason=scores.reason)
+        ranked: Final = sorted(
+            (SkillSearchHit(skill=skill, score=score) for skill, score in zip(skills, scores, strict=True)),
+            key=lambda hit: hit.score,
+            reverse=True,
+        )
+        return SkillSearchHits(hits=tuple(ranked[:top_k]))
+
+
+global_skill_search_index: Final = SkillSearchIndex()
+
+
+async def search_skills(
+    query: str,
+    skills: Sequence[LiteLLM_SkillsTable],
+    top_k: int,
+    router: Router | None,
+    embedding_model: str | None,
+    index: SkillSearchIndex,
+    user_api_key_dict: UserAPIKeyAuth,
+    proxy_logging_obj: ProxyLogging,
+) -> SkillSearchOutcome:
+    if embedding_model is None:
+        return SkillSearchNotConfigured(
+            reason="skill search needs litellm_settings.skill_search_embedding_model set to an embedding model from model_list"
+        )
+    if router is None:
+        return SkillSearchNotConfigured(reason="skill search needs a model_list so the embedding model can be called")
+    embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj)
+    return await index.search(query, skills, top_k, embed, embedding_model)
+
+
+async def search_hosted_skills(
+    custom_llm_provider: str | None,
+    query: str,
+    top_k: int,
+    router: Router | None,
+    embedding_model: str | None,
+    index: SkillSearchIndex,
+    user_api_key_dict: UserAPIKeyAuth,
+    proxy_logging_obj: ProxyLogging,
+) -> HostedSkillSearchOutcome:
+    """GET /v1/skills?query= for the skills LiteLLM hosts itself: only ``litellm_proxy`` has a registry to rank."""
+    if custom_llm_provider != LlmProviders.LITELLM_PROXY.value:
+        return SkillSearchUnsupportedProvider(reason="query is only supported for custom_llm_provider=litellm_proxy")
+    skills: Final = await LiteLLMSkillsHandler.list_skills_for_search(user_api_key_dict=user_api_key_dict)
+    return await search_skills(
+        query=query,
+        skills=skills,
+        top_k=top_k,
+        router=router,
+        embedding_model=embedding_model,
+        index=index,
+        user_api_key_dict=user_api_key_dict,
+        proxy_logging_obj=proxy_logging_obj,
+    )
diff --git a/litellm/llms/litellm_proxy/skills/transformation.py b/litellm/llms/litellm_proxy/skills/transformation.py
index c972dc349c9..9fc2d2cbb45 100644
--- a/litellm/llms/litellm_proxy/skills/transformation.py
+++ b/litellm/llms/litellm_proxy/skills/transformation.py
@@ -154,7 +154,7 @@ class LiteLLMSkillsTransformationHandler:
             user_api_key_dict=user_api_key_dict,
         )
 
-        return self._db_skill_to_response(db_skill)
+        return self.db_skill_to_response(db_skill)
 
     def list_skills_handler(
         self,
@@ -222,7 +222,9 @@ class LiteLLMSkillsTransformationHandler:
             user_api_key_dict=user_api_key_dict,
         )
 
-        skills: Final = [self._db_skill_to_response(s) for s in db_skills]
+        skills: Final = [  # mutable-ok: ListSkillsResponse.data needs list[Skill]; never mutated after
+            self.db_skill_to_response(s) for s in db_skills
+        ]
         return ListSkillsResponse(
             data=skills,
             has_more=len(skills) >= limit,
@@ -288,7 +290,7 @@ class LiteLLMSkillsTransformationHandler:
             skill_id=skill_id,
             user_api_key_dict=user_api_key_dict,
         )
-        return self._db_skill_to_response(db_skill)
+        return self.db_skill_to_response(db_skill)
 
     def delete_skill_handler(
         self,
@@ -354,7 +356,7 @@ class LiteLLMSkillsTransformationHandler:
             type=result.get("type", "skill_deleted"),
         )
 
-    def _db_skill_to_response(self, db_skill: "LiteLLM_SkillsTable") -> Skill:
+    def db_skill_to_response(self, db_skill: "LiteLLM_SkillsTable") -> Skill:
         """
         Convert a database skill record to Anthropic-compatible Skill response.
 
@@ -375,4 +377,5 @@ class LiteLLMSkillsTransformationHandler:
             latest_version=db_skill.latest_version,
             source=db_skill.source or "custom",
             type="skill",
+            description=db_skill.description,
         )
diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
index 5fbfad54a39..102129ffdd0 100644
--- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
+++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
@@ -104,6 +104,9 @@ if MCP_AVAILABLE:
     from mcp.types import Tool as MCPTool
 
     from litellm.experimental_mcp_client.client import MCPClient
+    from litellm.llms.litellm_proxy.skills.skill_search import (
+        DEFAULT_SKILL_SEARCH_TOP_K,
+    )
     from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
         _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES,
         global_mcp_server_manager,
@@ -188,10 +191,12 @@ if MCP_AVAILABLE:
             AGENT_SEARCH_TOOL_NAME,
             DEFAULT_AGENT_SEARCH_TOP_K,
             MCP_TOOL_SEARCH_TOOL_NAME,
+            SKILL_SEARCH_TOOL_NAME,
             coerce_top_k,
             handle_agent_search,
             handle_mcp_tool_call,
             handle_mcp_tool_search,
+            handle_skill_search,
         )
         from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
         from litellm.proxy.proxy_server import general_settings, proxy_config, proxy_logging_obj
@@ -210,6 +215,14 @@ if MCP_AVAILABLE:
                 ),
                 user_api_key_dict=user_api_key_dict,
             )
+        if tool_name == SKILL_SEARCH_TOOL_NAME:
+            return await handle_skill_search(
+                query=str(tool_arguments.get("query", "")),
+                top_k=coerce_top_k(
+                    tool_arguments.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K
+                ),
+                user_api_key_dict=user_api_key_dict,
+            )
         rest_client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
         (
             virtual_mcp_auth_header,
diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 60a9af89cc3..0b424c31c4b 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -911,15 +911,18 @@ if MCP_AVAILABLE:
         Returns a CallToolResult when ``name`` is a virtual tool, else ``None`` so
         the caller falls through to normal tool routing.
         """
+        from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K
         from litellm.proxy._experimental.mcp_server.tool_search import (
             AGENT_SEARCH_TOOL_NAME,
             DEFAULT_AGENT_SEARCH_TOP_K,
             MCP_TOOL_SEARCH_TOOL_NAME,
+            SKILL_SEARCH_TOOL_NAME,
             VIRTUAL_TOOL_NAMES,
             coerce_top_k,
             handle_agent_search,
             handle_mcp_tool_call,
             handle_mcp_tool_search,
+            handle_skill_search,
         )
 
         if name not in VIRTUAL_TOOL_NAMES:
@@ -961,6 +964,12 @@ if MCP_AVAILABLE:
                 top_k=coerce_top_k(args.get("top_k", DEFAULT_AGENT_SEARCH_TOP_K), default=DEFAULT_AGENT_SEARCH_TOP_K),
                 user_api_key_dict=user_api_key_auth,
             )
+        if name == SKILL_SEARCH_TOOL_NAME:
+            return await handle_skill_search(
+                query=str(args.get("query", "")),
+                top_k=coerce_top_k(args.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K),
+                user_api_key_dict=user_api_key_auth,
+            )
         virtual_logging_obj: Final = await _build_virtual_call_logging_obj(
             name=name,
             arguments=args,
diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py
index af02c11ad86..f19340d30cb 100644
--- a/litellm/proxy/_experimental/mcp_server/tool_search.py
+++ b/litellm/proxy/_experimental/mcp_server/tool_search.py
@@ -11,6 +11,7 @@ from pydantic import ValidationError
 from typing_extensions import ReadOnly, Required, assert_never
 
 import litellm
+from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K
 from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K
 from litellm.proxy.common_utils.semantic_text_index import (
     Embedder,
@@ -30,7 +31,10 @@ MCP_TOOL_SEARCH_SETTINGS_KEY: Final[str] = "mcp_tool_search"
 MCP_TOOL_SEARCH_TOOL_NAME: Final[str] = "mcp_tool_search"
 MCP_TOOL_CALL_TOOL_NAME: Final[str] = "mcp_tool_call"
 AGENT_SEARCH_TOOL_NAME: Final[str] = "agent_search"
-VIRTUAL_TOOL_NAMES: Final = frozenset((MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME))
+SKILL_SEARCH_TOOL_NAME: Final[str] = "skill_search"
+VIRTUAL_TOOL_NAMES: Final = frozenset(
+    (MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, SKILL_SEARCH_TOOL_NAME)
+)
 
 
 def coerce_top_k(value: Any, default: int = 5) -> int:
@@ -199,8 +203,28 @@ _AGENT_SEARCH_DEFINITION: Final[VirtualToolDefinition] = {
 }
 
 
+_SKILL_SEARCH_DEFINITION: Final[VirtualToolDefinition] = {
+    "name": SKILL_SEARCH_TOOL_NAME,
+    "description": "Find registered skills by describing what you need in natural language. Returns the best "
+    "matching skills you can access, ranked by semantic similarity, each with its skill_id, display_title, "
+    "description, and score.",
+    "inputSchema": {
+        "type": "object",
+        "properties": {
+            "query": {"type": "string", "description": "What you need the skill to do, in natural language."},
+            "top_k": {
+                "type": "integer",
+                "description": "Maximum number of skills to return.",
+                "default": DEFAULT_SKILL_SEARCH_TOP_K,
+            },
+        },
+        "required": _json_array("query"),
+    },
+}
+
+
 def get_virtual_tool_definitions() -> tuple[VirtualToolDefinition, ...]:
-    return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION)
+    return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION, _SKILL_SEARCH_DEFINITION)
 
 
 def _text_tool_result(text: str, is_error: bool) -> CallToolResult:
@@ -223,7 +247,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI
     )
     from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents
     from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
-    from litellm.proxy.proxy_server import llm_router
+    from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
 
     await check_feature_access_for_user(user_api_key_dict, "agents")
     outcome: Final = await search_agents(
@@ -234,6 +258,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI
         embedding_model=litellm.agent_search_embedding_model,
         index=global_agent_search_index,
         user_api_key_dict=user_api_key_dict,
+        proxy_logging_obj=proxy_logging_obj,
     )
     match outcome:
         case AgentSearchHits(hits):
@@ -245,6 +270,39 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI
             assert_never(outcome)
 
 
+async def handle_skill_search(query: str, top_k: int, user_api_key_dict: UserAPIKeyAuth) -> CallToolResult:
+    from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler
+    from litellm.llms.litellm_proxy.skills.skill_search import (
+        MAX_SKILL_SEARCH_TOP_K,
+        SkillSearchEmbeddingFailed,
+        SkillSearchHits,
+        SkillSearchNotConfigured,
+        global_skill_search_index,
+        search_skills,
+        skill_search_result,
+    )
+    from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
+
+    outcome: Final = await search_skills(
+        query=query,
+        skills=await LiteLLMSkillsHandler.list_skills_for_search(user_api_key_dict),
+        top_k=min(max(top_k, 1), MAX_SKILL_SEARCH_TOP_K),
+        router=llm_router,
+        embedding_model=litellm.skill_search_embedding_model,
+        index=global_skill_search_index,
+        user_api_key_dict=user_api_key_dict,
+        proxy_logging_obj=proxy_logging_obj,
+    )
+    match outcome:
+        case SkillSearchHits(hits):
+            results: Final = tuple(skill_search_result(hit).model_dump() for hit in hits)
+            return _text_tool_result(json.dumps(results), is_error=False)
+        case SkillSearchNotConfigured(reason) | SkillSearchEmbeddingFailed(reason):
+            return _text_tool_result(reason, is_error=True)
+        case _:
+            assert_never(outcome)
+
+
 async def handle_mcp_tool_search(
     query: str,
     top_k: int,
@@ -257,7 +315,7 @@ async def handle_mcp_tool_search(
     raw_headers: dict[str, str] | None = None,
 ) -> CallToolResult:
     from litellm.proxy._experimental.mcp_server.server import _list_mcp_tools
-    from litellm.proxy.proxy_server import llm_router
+    from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
 
     settings: Final = mcp_tool_search_settings()
     if isinstance(settings, ValidationError):
@@ -271,7 +329,7 @@ async def handle_mcp_tool_search(
         )
     ranker: Final = (
         SemanticToolRanker(
-            embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict),
+            embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict, proxy_logging_obj),
             embedding_model=settings.embedding_model,
             index=global_mcp_tool_search_index,
         )
diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json
index 91a97ad6544..4093f2c5248 100644
--- a/litellm/proxy/_lazy_openapi_snapshot.json
+++ b/litellm/proxy/_lazy_openapi_snapshot.json
@@ -4687,6 +4687,17 @@
               "title": "Created At",
               "type": "string"
             },
+            "description": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Description"
+            },
             "display_title": {
               "anyOf": [
                 {
@@ -4713,6 +4724,17 @@
               ],
               "title": "Latest Version"
             },
+            "search_score": {
+              "anyOf": [
+                {
+                  "type": "number"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Search Score"
+            },
             "source": {
               "title": "Source",
               "type": "string"
@@ -4781,7 +4803,7 @@
     "paths": {
       "/v1/skills": {
         "get": {
-          "description": "List skills on Anthropic.\n\nRequires `?beta=true` query parameter.\n\nModel-based routing (for multi-account support):\n- Pass model via header: `x-litellm-model: claude-account-1`\n- Pass model via query: `?model=claude-account-1`\n- Pass model via body: `{\"model\": \"claude-account-1\"}`\n\nExample usage:\n```bash\n# Basic usage\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\"       -H \"Authorization: Bearer your-key\"\n\n# With model-based routing\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\"       -H \"Authorization: Bearer your-key\"       -H \"x-litellm-model: claude-account-1\"\n```\n\nReturns: ListSkillsResponse with list of skills",
+          "description": "List skills on Anthropic.\n\nRequires `?beta=true` query parameter.\n\nModel-based routing (for multi-account support):\n- Pass model via header: `x-litellm-model: claude-account-1`\n- Pass model via query: `?model=claude-account-1`\n- Pass model via body: `{\"model\": \"claude-account-1\"}`\n\nExample usage:\n```bash\n# Basic usage\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\"       -H \"Authorization: Bearer your-key\"\n\n# With model-based routing\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\"       -H \"Authorization: Bearer your-key\"       -H \"x-litellm-model: claude-account-1\"\n```\n\nPass `?custom_llm_provider=litellm_proxy&query=<task>` to rank the LiteLLM-hosted skills you can\naccess by semantic similarity instead of paging through the whole registry:\n```bash\ncurl \"http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5\"       -H \"Authorization: Bearer your-key\"\n```\n\nReturns: ListSkillsResponse with list of skills",
           "operationId": "list_skills_v1_skills_get",
           "parameters": [
             {
@@ -4849,6 +4871,39 @@
                 "default": "anthropic",
                 "title": "Custom Llm Provider"
               }
+            },
+            {
+              "description": "Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model.",
+              "in": "query",
+              "name": "query",
+              "required": false,
+              "schema": {
+                "anyOf": [
+                  {
+                    "minLength": 1,
+                    "type": "string"
+                  },
+                  {
+                    "type": "null"
+                  }
+                ],
+                "description": "Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model.",
+                "title": "Query"
+              }
+            },
+            {
+              "description": "With query: the maximum number of ranked skills to return.",
+              "in": "query",
+              "name": "top_k",
+              "required": false,
+              "schema": {
+                "default": 5,
+                "description": "With query: the maximum number of ranked skills to return.",
+                "maximum": 100,
+                "minimum": 1,
+                "title": "Top K",
+                "type": "integer"
+              }
             }
           ],
           "responses": {
diff --git a/litellm/proxy/agent_endpoints/agent_search.py b/litellm/proxy/agent_endpoints/agent_search.py
index 76e3fe6c5ad..65a89bb2c7a 100644
--- a/litellm/proxy/agent_endpoints/agent_search.py
+++ b/litellm/proxy/agent_endpoints/agent_search.py
@@ -18,6 +18,7 @@ from litellm.types.agents import AgentResponse
 
 if TYPE_CHECKING:
     from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.utils import ProxyLogging
     from litellm.router import Router
 
 DEFAULT_AGENT_SEARCH_TOP_K: Final = 5
@@ -132,6 +133,7 @@ async def search_agents(
     embedding_model: str | None,
     index: AgentSearchIndex,
     user_api_key_dict: UserAPIKeyAuth,
+    proxy_logging_obj: ProxyLogging,
 ) -> AgentSearchOutcome:
     if embedding_model is None:
         return AgentSearchNotConfigured(
@@ -139,6 +141,5 @@ async def search_agents(
         )
     if router is None:
         return AgentSearchNotConfigured(reason="agent search needs a model_list so the embedding model can be called")
-    return await index.search(
-        query, agents, top_k, router_embedder(router, embedding_model, user_api_key_dict), embedding_model
-    )
+    embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj)
+    return await index.search(query, agents, top_k, embed, embedding_model)
diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py
index cc17672553b..aa8979a73c6 100644
--- a/litellm/proxy/agent_endpoints/endpoints.py
+++ b/litellm/proxy/agent_endpoints/endpoints.py
@@ -249,7 +249,7 @@ def _agent_search_error(status_code: int, error: str, message: str) -> HTTPExcep
 async def _rank_agents_by_query(
     query: str, agents: Sequence[AgentResponse], top_k: int, user_api_key_dict: UserAPIKeyAuth
 ) -> tuple[AgentResponse, ...]:
-    from litellm.proxy.proxy_server import llm_router
+    from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
 
     outcome: Final = await search_agents(
         query=query,
@@ -259,6 +259,7 @@ async def _rank_agents_by_query(
         embedding_model=litellm.agent_search_embedding_model,
         index=global_agent_search_index,
         user_api_key_dict=user_api_key_dict,
+        proxy_logging_obj=proxy_logging_obj,
     )
     match outcome:
         case AgentSearchHits(hits):
diff --git a/litellm/proxy/anthropic_endpoints/skills_endpoints.py b/litellm/proxy/anthropic_endpoints/skills_endpoints.py
index 9390bf4c537..4426c0b547a 100644
--- a/litellm/proxy/anthropic_endpoints/skills_endpoints.py
+++ b/litellm/proxy/anthropic_endpoints/skills_endpoints.py
@@ -2,11 +2,23 @@
 Anthropic Skills API endpoints - /v1/skills
 """
 
-from typing import Final
+from types import MappingProxyType
+from typing import Annotated, Final
 
 import orjson
-from fastapi import APIRouter, Depends, Request, Response
+from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response
+from typing_extensions import ReadOnly, TypedDict, assert_never
 
+import litellm
+from litellm.llms.litellm_proxy.skills.skill_search import (
+    DEFAULT_SKILL_SEARCH_TOP_K,
+    SkillSearchEmbeddingFailed,
+    SkillSearchHits,
+    SkillSearchNotConfigured,
+    SkillSearchUnsupportedProvider,
+    global_skill_search_index,
+    search_hosted_skills,
+)
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
@@ -23,6 +35,51 @@ from litellm.types.llms.anthropic_skills import (
 router: Final = APIRouter()
 
 
+class _SkillSearchErrorDetail(TypedDict):
+    error: ReadOnly[str]
+    message: ReadOnly[str]
+
+
+def _skill_search_error(status_code: int, error: str, message: str) -> HTTPException:
+    detail: Final[_SkillSearchErrorDetail] = {"error": error, "message": message}
+    return HTTPException(status_code=status_code, detail=detail)
+
+
+async def _search_skills(
+    custom_llm_provider: str | None, query: str, top_k: int, user_api_key_dict: UserAPIKeyAuth
+) -> ListSkillsResponse:
+    from litellm.llms.litellm_proxy.skills.transformation import (
+        LiteLLMSkillsTransformationHandler,
+    )
+    from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
+
+    outcome: Final = await search_hosted_skills(
+        custom_llm_provider=custom_llm_provider,
+        query=query,
+        top_k=top_k,
+        router=llm_router,
+        embedding_model=litellm.skill_search_embedding_model,
+        index=global_skill_search_index,
+        user_api_key_dict=user_api_key_dict,
+        proxy_logging_obj=proxy_logging_obj,
+    )
+    to_response: Final = LiteLLMSkillsTransformationHandler().db_skill_to_response
+    match outcome:
+        case SkillSearchHits(hits):
+            skills: Final = [  # mutable-ok: ListSkillsResponse.data requires list[Skill]; never mutated after
+                to_response(hit.skill).model_copy(update=MappingProxyType({"search_score": hit.score})) for hit in hits
+            ]
+            return ListSkillsResponse(data=skills, has_more=False, next_page=None)
+        case SkillSearchUnsupportedProvider(reason):
+            raise _skill_search_error(400, "skill_search_unsupported_provider", reason)
+        case SkillSearchNotConfigured(reason):
+            raise _skill_search_error(400, "skill_search_not_configured", reason)
+        case SkillSearchEmbeddingFailed(reason):
+            raise _skill_search_error(503, "skill_search_unavailable", reason)
+        case _:
+            assert_never(outcome)
+
+
 @router.post(
     "/v1/skills",
     tags=["[beta] Anthropic Skills API"],
@@ -134,32 +191,58 @@ async def list_skills(
     after_id: str | None = None,
     before_id: str | None = None,
     custom_llm_provider: str | None = "anthropic",
+    query: Annotated[
+        str | None,
+        Query(
+            min_length=1,
+            description="Describe what you need in natural language to rank the skills you can access by "
+            "semantic similarity over their title and description. Each result carries a search_score. "
+            "Only supported for custom_llm_provider=litellm_proxy. Requires "
+            "litellm_settings.skill_search_embedding_model.",
+        ),
+    ] = None,
+    top_k: Annotated[
+        int,
+        Query(ge=1, le=100, description="With query: the maximum number of ranked skills to return."),
+    ] = DEFAULT_SKILL_SEARCH_TOP_K,
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
 ):
     """
     List skills on Anthropic.
-    
+
     Requires `?beta=true` query parameter.
-    
+
     Model-based routing (for multi-account support):
     - Pass model via header: `x-litellm-model: claude-account-1`
     - Pass model via query: `?model=claude-account-1`
     - Pass model via body: `{"model": "claude-account-1"}`
-    
+
     Example usage:
     ```bash
     # Basic usage
     curl "http://localhost:4000/v1/skills?beta=true&limit=10" \
       -H "Authorization: Bearer your-key"
-    
+
     # With model-based routing
     curl "http://localhost:4000/v1/skills?beta=true&limit=10" \
       -H "Authorization: Bearer your-key" \
       -H "x-litellm-model: claude-account-1"
     ```
-    
+
+    Pass `?custom_llm_provider=litellm_proxy&query=<task>` to rank the LiteLLM-hosted skills you can
+    access by semantic similarity instead of paging through the whole registry:
+    ```bash
+    curl "http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5" \
+      -H "Authorization: Bearer your-key"
+    ```
+
     Returns: ListSkillsResponse with list of skills
     """
+    if query is not None:
+        return await _search_skills(
+            custom_llm_provider=custom_llm_provider, query=query, top_k=top_k, user_api_key_dict=user_api_key_dict
+        )
+
     from litellm.proxy.proxy_server import (
         general_settings,
         llm_router,
diff --git a/litellm/proxy/common_utils/semantic_text_index.py b/litellm/proxy/common_utils/semantic_text_index.py
index 0820459af49..b8d3595163e 100644
--- a/litellm/proxy/common_utils/semantic_text_index.py
+++ b/litellm/proxy/common_utils/semantic_text_index.py
@@ -5,10 +5,11 @@ from __future__ import annotations
 import math
 from collections.abc import Awaitable, Mapping, Sequence
 from dataclasses import dataclass
-from itertools import chain
+from itertools import chain, islice
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
 
+from fastapi import HTTPException
 from openai import OpenAIError
 from pydantic import BaseModel, ConfigDict
 
@@ -16,10 +17,14 @@ from litellm.exceptions import BudgetExceededError
 
 if TYPE_CHECKING:
     from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.utils import ProxyLogging
     from litellm.router import Router
 
 Vector: TypeAlias = tuple[float, ...]
 
+DEFAULT_MAX_CACHED_VECTORS: Final = 5000
+"""Ceiling on how many (embedding model, text) vectors one index keeps; the least recently searched are evicted first."""
+
 
 class Embedder(Protocol):
     def __call__(self, texts: Sequence[str]) -> Awaitable[Sequence[Vector]]: ...
@@ -42,6 +47,16 @@ class _EmbeddingData(BaseModel):
     data: tuple[_EmbeddingItem, ...]
 
 
+class _EmbeddingRequest(BaseModel):
+    """The /embeddings-shaped request as the pre-call hooks (rate limits, budgets, guardrails) hand it back."""
+
+    model_config = ConfigDict(frozen=True, extra="ignore")
+
+    model: str
+    input: tuple[str, ...]
+    metadata: dict[str, object]  # mutable-ok: the router mutates the metadata dict it is handed
+
+
 def cosine_similarity(left: Vector, right: Vector) -> float:
     dot: Final = sum(a * b for a, b in zip(left, right, strict=True))
     norms: Final = math.sqrt(sum(a * a for a in left)) * math.sqrt(sum(b * b for b in right))
@@ -57,23 +72,40 @@ def embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, obj
     }
 
 
-def router_embedder(router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth) -> Embedder:
+def router_embedder(
+    router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging
+) -> Embedder:
+    """Embeds through the router after the same key rate-limit, budget and guardrail pre-call hooks /embeddings runs."""
+
     async def embed(texts: Sequence[str]) -> Sequence[Vector]:
-        batch: Final = list(texts)  # mutable-ok: Router.aembedding accepts only str | list input
+        request: Final = {  # mutable-ok: pre_call_hook mutates the request dict in place
+            "model": embedding_model,
+            "input": list(texts),  # mutable-ok: Router.aembedding accepts only str | list input
+            "metadata": embedding_spend_metadata(user_api_key_dict),
+        }
+        processed: Final = _EmbeddingRequest.model_validate(
+            await proxy_logging_obj.pre_call_hook(
+                user_api_key_dict=user_api_key_dict, data=request, call_type="aembedding"
+            )
+        )
         response: Final = await router.aembedding(
-            model=embedding_model, input=batch, metadata=embedding_spend_metadata(user_api_key_dict)
+            model=processed.model,
+            input=list(processed.input),  # mutable-ok: Router.aembedding accepts only str | list input
+            metadata=processed.metadata,
         )
         return tuple(item.embedding for item in _EmbeddingData.model_validate(response.model_dump()).data)
 
     return embed
 
 
-_NO_VECTORS: Final[Mapping[str, Vector]] = MappingProxyType({})
+_CacheKey: TypeAlias = tuple[str, str]
 
 
 async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed:
     try:
         vectors: Final = tuple(await embed(texts))
+    except HTTPException:
+        raise
     except (OpenAIError, ValueError, BudgetExceededError) as exc:
         return EmbeddingFailed(reason=f"embedding the search query failed: {exc}")
     if len(vectors) != len(texts):
@@ -111,20 +143,34 @@ async def _embed_query_and_texts(
 
 
 class SemanticTextIndex:
-    """Caches one vector per distinct text per embedding model, so repeat searches only embed the query."""
+    """Caches one vector per distinct text per embedding model, so repeat searches only embed the query.
 
-    def __init__(self) -> None:
-        self._vectors: Mapping[str, Mapping[str, Vector]] = MappingProxyType({})
+    Holds at most ``max_entries`` vectors across all models: once full, the texts no recent search touched go first."""
 
-    def _merged(self, embedding_model: str, embedded: _Embedded) -> Mapping[str, Vector]:
-        kept: Final = MappingProxyType(
+    def __init__(self, max_entries: int = DEFAULT_MAX_CACHED_VECTORS) -> None:
+        self._max_entries: Final = max_entries
+        self._vectors: Mapping[_CacheKey, Vector] = MappingProxyType({})
+
+    def _cached(self, embedding_model: str) -> Mapping[str, Vector]:
+        return MappingProxyType(
+            {text: vector for (model, text), vector in self._vectors.items() if model == embedding_model}
+        )
+
+    def _merged(self, embedding_model: str, embedded: _Embedded, texts: Sequence[str]) -> Mapping[_CacheKey, Vector]:
+        dimension: Final = len(embedded.query_vector)
+        touched: Final = MappingProxyType({(embedding_model, text): embedded.vectors[text] for text in texts})
+        untouched: Final = MappingProxyType(
             {
-                text: vector
-                for text, vector in self._vectors.get(embedding_model, _NO_VECTORS).items()
-                if len(vector) == len(embedded.query_vector)
+                key: vector
+                for key, vector in chain(
+                    self._vectors.items(),
+                    (((embedding_model, text), vector) for text, vector in embedded.vectors.items()),
+                )
+                if key not in touched and (key[0] != embedding_model or len(vector) == dimension)
             }
         )
-        return MappingProxyType({**kept, **embedded.vectors})
+        ordered: Final = MappingProxyType({**untouched, **touched})
+        return MappingProxyType(dict(islice(ordered.items(), max(len(ordered) - self._max_entries, 0), None)))
 
     async def scores(
         self, query: str, texts: Sequence[str], embed: Embedder, embedding_model: str
@@ -132,11 +178,10 @@ class SemanticTextIndex:
         """Cosine similarity of `query` to each entry of `texts`, in the same order."""
         if not texts:
             return ()
-        cached: Final = self._vectors.get(embedding_model, _NO_VECTORS)
-        embedded: Final = await _embed_query_and_texts(embed, query, texts, cached)
+        embedded: Final = await _embed_query_and_texts(embed, query, texts, self._cached(embedding_model))
         if isinstance(embedded, EmbeddingFailed):
             return embedded
         if not _same_dimension(embedded.query_vector, embedded.vectors, texts):
             return EmbeddingFailed(reason=f"embedding model {embedding_model} returned vectors of mixed dimensions")
-        self._vectors = MappingProxyType({**self._vectors, embedding_model: self._merged(embedding_model, embedded)})
+        self._vectors = self._merged(embedding_model, embedded, texts)
         return tuple(cosine_similarity(embedded.query_vector, embedded.vectors[text]) for text in texts)
diff --git a/litellm/skills/main.py b/litellm/skills/main.py
index 002419dbad4..71fd78f11a3 100644
--- a/litellm/skills/main.py
+++ b/litellm/skills/main.py
@@ -182,10 +182,14 @@ def create_skill(
         if extra_body:
             create_request.update(extra_body)
 
-        # Route to LiteLLM DB if custom_llm_provider="litellm_proxy"
+        # Route to LiteLLM DB if custom_llm_provider="litellm_proxy". description/instructions
+        # arrive as top-level kwargs from the REST form endpoint, or nested in extra_body from
+        # the SDK convention used by other providers' create_request above.
         if custom_llm_provider == LlmProviders.LITELLM_PROXY.value:
             return _get_litellm_skills_handler().create_skill_handler(
                 display_title=display_title,
+                description=kwargs.get("description") or (extra_body.get("description") if extra_body else None),
+                instructions=kwargs.get("instructions") or (extra_body.get("instructions") if extra_body else None),
                 files=files,
                 metadata=_get_skill_request_metadata(kwargs, extra_body),
                 user_id=kwargs.get("user_id"),
diff --git a/litellm/types/llms/anthropic_skills.py b/litellm/types/llms/anthropic_skills.py
index 51eefe7154f..4b27f9b17ef 100644
--- a/litellm/types/llms/anthropic_skills.py
+++ b/litellm/types/llms/anthropic_skills.py
@@ -57,6 +57,15 @@ class Skill(BaseModel):
     updated_at: str
     """ISO 8601 timestamp of when the skill was last updated"""
 
+    description: str | None = None
+    """Description of the skill. Populated for the LiteLLM-hosted registry
+    (custom_llm_provider="litellm_proxy"); Anthropic's list endpoint does not
+    return a description, so this is None there."""
+
+    search_score: float | None = None
+    """Semantic similarity to the ``query`` passed to ``GET /v1/skills``. None
+    unless a query was given."""
+
 
 class ListSkillsResponse(BaseModel):
     """Response from listing skills"""
diff --git a/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py
new file mode 100644
index 00000000000..a0f22a59f0c
--- /dev/null
+++ b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py
@@ -0,0 +1,436 @@
+import asyncio
+import json
+from collections.abc import Sequence
+from types import MappingProxyType
+from typing import Final
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from openai import APIConnectionError
+
+import litellm
+from litellm.llms.litellm_proxy.skills.skill_search import (
+    MAX_SKILL_SEARCH_TEXT_CHARS,
+    SkillSearchEmbeddingFailed,
+    SkillSearchHits,
+    SkillSearchIndex,
+    SkillSearchNotConfigured,
+    search_skills,
+    skill_search_text,
+)
+from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth
+from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
+from litellm.proxy.common_utils.semantic_text_index import Vector, cosine_similarity
+
+CALLER: Final = UserAPIKeyAuth(api_key="hashed-caller-key", team_id="team-1", user_id="user-1")
+
+TRANSLATOR: Final = LiteLLM_SkillsTable(
+    skill_id="translate-file",
+    display_title="Document Translator",
+    description="Converts files from one language into another",
+    instructions="Take an uploaded document and produce it in the target language",
+)
+SQL_ANALYST: Final = LiteLLM_SkillsTable(
+    skill_id="warehouse-sql-analyst",
+    display_title="Warehouse SQL Analyst",
+    description="Runs SQL against the inventory database",
+)
+TRIP_PLANNER: Final = LiteLLM_SkillsTable(
+    skill_id="trip-planner",
+    display_title="Trip Planner",
+    description="Books flights and hotels",
+)
+SKILLS: Final = (TRANSLATOR, SQL_ANALYST, TRIP_PLANNER)
+
+VECTORS: Final = MappingProxyType(
+    {
+        "language translation": (1.0, 0.0, 0.0),
+        skill_search_text(TRANSLATOR): (0.9, 0.1, 0.0),
+        skill_search_text(SQL_ANALYST): (0.0, 1.0, 0.0),
+        skill_search_text(TRIP_PLANNER): (0.3, 0.0, 1.0),
+    }
+)
+
+
+def _pass_through_key_limits() -> MagicMock:
+    limits = MagicMock()
+    limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
+    return limits
+
+
+def _embedding_router() -> MagicMock:
+    router = MagicMock()
+    router.aembedding = AsyncMock(
+        side_effect=lambda model, input, metadata: litellm.EmbeddingResponse(
+            model=model,
+            data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)],
+        )
+    )
+    return router
+
+
+class FakeEmbedder:
+    def __init__(self) -> None:
+        self.calls: list[tuple[str, ...]] = []  # mutable-ok: test spy recording embed inputs
+
+    async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
+        self.calls.append(tuple(texts))
+        return tuple(VECTORS[text] for text in texts)
+
+
+class FixedDimensionEmbedder:
+    def __init__(self, dimensions: int) -> None:
+        self.dimensions: Final = dimensions
+        self.calls: list[tuple[str, ...]] = []  # mutable-ok: test spy recording embed inputs
+
+    async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
+        self.calls.append(tuple(texts))
+        await asyncio.sleep(0)
+        return tuple((1.0,) * self.dimensions for _ in texts)
+
+
+class TestSkillSearchText:
+    def test_joins_title_description_and_instructions(self) -> None:
+        assert skill_search_text(TRANSLATOR) == (
+            "Document Translator\n"
+            "Converts files from one language into another\n"
+            "Take an uploaded document and produce it in the target language"
+        )
+
+    def test_missing_fields_fall_back_to_whatever_is_present(self) -> None:
+        assert skill_search_text(LiteLLM_SkillsTable(skill_id="bare", display_title="bare")) == "bare"
+
+    def test_all_fields_absent_is_an_empty_string(self) -> None:
+        assert skill_search_text(LiteLLM_SkillsTable(skill_id="empty")) == ""
+
+    def test_oversized_instructions_are_cut_so_one_skill_cannot_blow_up_the_embedding_batch(self) -> None:
+        bloated = LiteLLM_SkillsTable(
+            skill_id="bloated", display_title="Bloated", instructions="x" * (MAX_SKILL_SEARCH_TEXT_CHARS * 3)
+        )
+        text = skill_search_text(bloated)
+        assert len(text) == MAX_SKILL_SEARCH_TEXT_CHARS
+        assert text.startswith("Bloated\n")
+
+
+class TestCosineSimilarity:
+    def test_identical_direction_scores_one(self) -> None:
+        assert cosine_similarity((2.0, 0.0), (1.0, 0.0)) == pytest.approx(1.0)
+
+    def test_orthogonal_scores_zero(self) -> None:
+        assert cosine_similarity((1.0, 0.0), (0.0, 1.0)) == pytest.approx(0.0)
+
+    def test_zero_vector_scores_zero_instead_of_dividing(self) -> None:
+        assert cosine_similarity((0.0, 0.0), (1.0, 0.0)) == 0.0
+
+
+class TestSkillSearchIndex:
+    @pytest.mark.asyncio
+    async def test_ranks_by_similarity_and_truncates_to_top_k(self) -> None:
+        outcome = await SkillSearchIndex().search(
+            "language translation", SKILLS, top_k=2, embed=FakeEmbedder(), embedding_model="m"
+        )
+        assert isinstance(outcome, SkillSearchHits)
+        assert [hit.skill.skill_id for hit in outcome.hits] == ["translate-file", "trip-planner"]
+        assert outcome.hits[0].score > outcome.hits[1].score
+
+    @pytest.mark.asyncio
+    async def test_second_search_only_embeds_the_query(self) -> None:
+        index = SkillSearchIndex()
+        embedder = FakeEmbedder()
+        await index.search("language translation", SKILLS, top_k=5, embed=embedder, embedding_model="m")
+        await index.search("language translation", SKILLS, top_k=5, embed=embedder, embedding_model="m")
+        assert len(embedder.calls[0]) == 1 + len(SKILLS)
+        assert embedder.calls[1] == ("language translation",)
+
+    @pytest.mark.asyncio
+    async def test_switching_embedding_models_does_not_reuse_cached_vectors(self) -> None:
+        index = SkillSearchIndex()
+        await index.search("language translation", SKILLS, top_k=5, embed=FakeEmbedder(), embedding_model="small")
+        wide = FixedDimensionEmbedder(2)
+        outcome = await index.search("language translation", SKILLS, top_k=5, embed=wide, embedding_model="wide")
+        assert isinstance(outcome, SkillSearchHits)
+        assert len(wide.calls[0]) == 1 + len(SKILLS)
+
+    @pytest.mark.asyncio
+    async def test_cached_vectors_of_another_dimension_are_re_embedded(self) -> None:
+        index = SkillSearchIndex()
+        await index.search("language translation", SKILLS, top_k=5, embed=FakeEmbedder(), embedding_model="m")
+        fallback = FixedDimensionEmbedder(2)
+        outcome = await index.search("language translation", SKILLS, top_k=5, embed=fallback, embedding_model="m")
+        assert isinstance(outcome, SkillSearchHits)
+        assert fallback.calls == [
+            ("language translation",),
+            ("language translation", *(skill_search_text(skill) for skill in SKILLS)),
+        ]
+
+    @pytest.mark.asyncio
+    async def test_re_embedding_a_subset_drops_the_other_skills_old_vectors(self) -> None:
+        index = SkillSearchIndex()
+        await index.search("language translation", SKILLS, top_k=5, embed=FakeEmbedder(), embedding_model="m")
+        wide = FixedDimensionEmbedder(2)
+        await index.search("language translation", SKILLS[:1], top_k=5, embed=wide, embedding_model="m")
+        await index.search("language translation", SKILLS, top_k=5, embed=wide, embedding_model="m")
+        assert wide.calls[-1] == ("language translation", *(skill_search_text(skill) for skill in SKILLS[1:]))
+
+    @pytest.mark.asyncio
+    async def test_concurrent_searches_keep_each_others_vectors(self) -> None:
+        index = SkillSearchIndex()
+        embedder = FixedDimensionEmbedder(3)
+        await asyncio.gather(
+            index.search("q", SKILLS[:1], top_k=5, embed=embedder, embedding_model="m"),
+            index.search("q", SKILLS[1:], top_k=5, embed=embedder, embedding_model="m"),
+        )
+        await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m")
+        assert embedder.calls[-1] == ("q",)
+
+    @pytest.mark.asyncio
+    async def test_least_recently_searched_skills_are_evicted_once_the_index_is_full(self) -> None:
+        index = SkillSearchIndex(max_entries=len(SKILLS))
+        embedder = FixedDimensionEmbedder(3)
+        newcomer = LiteLLM_SkillsTable(skill_id="newcomer", display_title="Newcomer")
+        await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m")
+        await index.search("q", SKILLS[:1], top_k=5, embed=embedder, embedding_model="m")
+        await index.search("q", (newcomer,), top_k=5, embed=embedder, embedding_model="m")
+        await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m")
+        assert embedder.calls[-1] == ("q", skill_search_text(SKILLS[1]))
+
+    @pytest.mark.asyncio
+    async def test_deleted_skills_stop_occupying_the_index_after_enough_new_ones(self) -> None:
+        index = SkillSearchIndex(max_entries=2)
+        embedder = FixedDimensionEmbedder(3)
+        for generation in range(50):
+            skill = LiteLLM_SkillsTable(skill_id=f"gen-{generation}", display_title=f"Generation {generation}")
+            await index.search("q", (skill,), top_k=5, embed=embedder, embedding_model="m")
+        await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m")
+        await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m")
+        assert embedder.calls[-1] == ("q", skill_search_text(SKILLS[0]))
+
+    @pytest.mark.asyncio
+    async def test_mixed_dimensions_in_one_batch_become_embedding_failed(self) -> None:
+        async def mixed(texts: Sequence[str]) -> Sequence[Vector]:
+            return ((1.0, 0.0), *((1.0, 0.0, 0.0) for _ in texts[1:]))
+
+        outcome = await SkillSearchIndex().search("q", SKILLS, top_k=5, embed=mixed, embedding_model="m")
+        assert isinstance(outcome, SkillSearchEmbeddingFailed)
+        assert "mixed dimensions" in outcome.reason
+
+    @pytest.mark.asyncio
+    async def test_no_accessible_skills_returns_no_hits_without_embedding(self) -> None:
+        embedder = FakeEmbedder()
+        outcome = await SkillSearchIndex().search("anything", (), top_k=5, embed=embedder, embedding_model="m")
+        assert outcome == SkillSearchHits(hits=())
+        assert embedder.calls == []
+
+    @pytest.mark.asyncio
+    async def test_provider_error_becomes_embedding_failed(self) -> None:
+        async def failing(texts: Sequence[str]) -> Sequence[Vector]:
+            raise APIConnectionError(request=MagicMock())
+
+        outcome = await SkillSearchIndex().search("q", SKILLS, top_k=5, embed=failing, embedding_model="m")
+        assert isinstance(outcome, SkillSearchEmbeddingFailed)
+        assert "embedding the search query failed" in outcome.reason
+
+    @pytest.mark.asyncio
+    async def test_wrong_vector_count_becomes_embedding_failed(self) -> None:
+        async def short(texts: Sequence[str]) -> Sequence[Vector]:
+            return ((1.0, 0.0, 0.0),)
+
+        outcome = await SkillSearchIndex().search("q", SKILLS, top_k=5, embed=short, embedding_model="m")
+        assert isinstance(outcome, SkillSearchEmbeddingFailed)
+
+
+class TestSearchSkills:
+    @pytest.mark.asyncio
+    async def test_no_embedding_model_is_not_configured(self) -> None:
+        outcome = await search_skills(
+            "q",
+            SKILLS,
+            5,
+            router=MagicMock(),
+            embedding_model=None,
+            index=SkillSearchIndex(),
+            user_api_key_dict=CALLER,
+            proxy_logging_obj=_pass_through_key_limits(),
+        )
+        assert isinstance(outcome, SkillSearchNotConfigured)
+        assert "skill_search_embedding_model" in outcome.reason
+
+    @pytest.mark.asyncio
+    async def test_no_router_is_not_configured(self) -> None:
+        outcome = await search_skills(
+            "q",
+            SKILLS,
+            5,
+            router=None,
+            embedding_model="m",
+            index=SkillSearchIndex(),
+            user_api_key_dict=CALLER,
+            proxy_logging_obj=_pass_through_key_limits(),
+        )
+        assert isinstance(outcome, SkillSearchNotConfigured)
+
+    @pytest.mark.asyncio
+    async def test_router_embeddings_are_read_from_the_response(self) -> None:
+        router = _embedding_router()
+        outcome = await search_skills(
+            "language translation",
+            SKILLS,
+            1,
+            router=router,
+            embedding_model="text-embedding-3-small",
+            index=SkillSearchIndex(),
+            user_api_key_dict=CALLER,
+            proxy_logging_obj=_pass_through_key_limits(),
+        )
+        assert isinstance(outcome, SkillSearchHits)
+        assert [hit.skill.skill_id for hit in outcome.hits] == ["translate-file"]
+        assert router.aembedding.await_args.kwargs["model"] == "text-embedding-3-small"
+
+    @pytest.mark.asyncio
+    async def test_embedding_spend_is_attributed_to_the_calling_key(self) -> None:
+        router = _embedding_router()
+        await search_skills(
+            "language translation",
+            SKILLS,
+            1,
+            router=router,
+            embedding_model="text-embedding-3-small",
+            index=SkillSearchIndex(),
+            user_api_key_dict=CALLER,
+            proxy_logging_obj=_pass_through_key_limits(),
+        )
+        metadata = router.aembedding.await_args.kwargs["metadata"]
+        assert metadata["user_api_key"] == "hashed-caller-key"
+        assert metadata["user_api_key_team_id"] == "team-1"
+        assert metadata["user_api_key_user_id"] == "user-1"
+
+    @pytest.mark.asyncio
+    async def test_key_limits_are_checked_against_the_real_embedding_call_before_it_runs(self) -> None:
+        router = _embedding_router()
+        key_limits = _pass_through_key_limits()
+        await search_skills(
+            "language translation",
+            SKILLS,
+            1,
+            router=router,
+            embedding_model="text-embedding-3-small",
+            index=SkillSearchIndex(),
+            user_api_key_dict=CALLER,
+            proxy_logging_obj=key_limits,
+        )
+        checked = key_limits.pre_call_hook.await_args.kwargs
+        assert checked["user_api_key_dict"] is CALLER
+        assert checked["call_type"] == "aembedding"
+        assert checked["data"]["model"] == "text-embedding-3-small"
+        assert checked["data"]["input"] == router.aembedding.await_args.kwargs["input"]
+        assert checked["data"]["metadata"]["user_api_key"] == "hashed-caller-key"
+
+    @pytest.mark.asyncio
+    async def test_the_embedding_model_sees_the_request_as_the_guardrails_rewrote_it(self) -> None:
+        router = MagicMock()
+        router.aembedding = AsyncMock(
+            side_effect=lambda model, input, metadata: litellm.EmbeddingResponse(
+                model=model,
+                data=[{"object": "embedding", "index": i, "embedding": [1.0, 0.0, 0.0]} for i in range(len(input))],
+            )
+        )
+        key_limits = MagicMock()
+        key_limits.pre_call_hook = AsyncMock(
+            side_effect=lambda user_api_key_dict, data, call_type: {
+                **data,
+                "input": ["[MASKED]" for _ in data["input"]],
+                "metadata": {**data["metadata"], "guardrail": "masked"},
+            }
+        )
+        await search_skills(
+            "language translation",
+            SKILLS,
+            1,
+            router=router,
+            embedding_model="text-embedding-3-small",
+            index=SkillSearchIndex(),
+            user_api_key_dict=CALLER,
+            proxy_logging_obj=key_limits,
+        )
+        sent = tuple(call.kwargs for call in router.aembedding.await_args_list)
+        assert sent
+        assert all(set(call["input"]) == {"[MASKED]"} for call in sent)
+        assert all(call["metadata"]["guardrail"] == "masked" for call in sent)
+
+    @pytest.mark.asyncio
+    async def test_a_key_over_its_limit_never_reaches_the_embedding_model(self) -> None:
+        router = _embedding_router()
+        key_limits = MagicMock()
+        key_limits.pre_call_hook = AsyncMock(side_effect=ProxyRateLimitError(detail="rpm exceeded"))
+        with pytest.raises(ProxyRateLimitError) as raised:
+            await search_skills(
+                "language translation",
+                SKILLS,
+                1,
+                router=router,
+                embedding_model="text-embedding-3-small",
+                index=SkillSearchIndex(),
+                user_api_key_dict=CALLER,
+                proxy_logging_obj=key_limits,
+            )
+        assert raised.value.status_code == 429
+        router.aembedding.assert_not_awaited()
+
+
+@pytest.fixture
+def accessible_skills(monkeypatch: pytest.MonkeyPatch) -> AsyncMock:
+    list_for_search = AsyncMock(return_value=list(SKILLS))
+    monkeypatch.setattr(
+        "litellm.llms.litellm_proxy.skills.handler.LiteLLMSkillsHandler.list_skills_for_search", list_for_search
+    )
+    return list_for_search
+
+
+@pytest.fixture
+def key_limits(monkeypatch: pytest.MonkeyPatch) -> MagicMock:
+    limits = _pass_through_key_limits()
+    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", limits)
+    return limits
+
+
+@pytest.fixture
+def embedding_router(monkeypatch: pytest.MonkeyPatch, key_limits: MagicMock) -> MagicMock:
+    router = _embedding_router()
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
+    monkeypatch.setattr(litellm, "skill_search_embedding_model", "text-embedding-3-small")
+    return router
+
+
+class TestHandleSkillSearchMCP:
+    @pytest.mark.asyncio
+    async def test_top_k_is_clamped_to_the_same_ceiling_as_rest(
+        self, accessible_skills: AsyncMock, embedding_router: MagicMock
+    ) -> None:
+        from litellm.llms.litellm_proxy.skills.skill_search import MAX_SKILL_SEARCH_TOP_K
+        from litellm.proxy._experimental.mcp_server.tool_search import handle_skill_search
+
+        many_skills: Final = tuple(
+            LiteLLM_SkillsTable(
+                skill_id=f"skill-{i}", display_title=TRIP_PLANNER.display_title, description=TRIP_PLANNER.description
+            )
+            for i in range(MAX_SKILL_SEARCH_TOP_K + 50)
+        )
+        accessible_skills.return_value = list(many_skills)
+
+        result = await handle_skill_search(
+            query="language translation", top_k=10_000, user_api_key_dict=UserAPIKeyAuth(user_id="u")
+        )
+        assert result.isError is False
+        assert len(json.loads(result.content[0].text)) == MAX_SKILL_SEARCH_TOP_K
+
+    @pytest.mark.asyncio
+    async def test_top_k_below_one_is_raised_to_one(
+        self, accessible_skills: AsyncMock, embedding_router: MagicMock
+    ) -> None:
+        from litellm.proxy._experimental.mcp_server.tool_search import handle_skill_search
+
+        result = await handle_skill_search(
+            query="language translation", top_k=0, user_api_key_dict=UserAPIKeyAuth(user_id="u")
+        )
+        assert result.isError is False
+        assert len(json.loads(result.content[0].text)) == 1
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py
index 239f89ebd90..798b0001af1 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py
@@ -24,6 +24,7 @@ from litellm.proxy._experimental.mcp_server.tool_search import (
     AGENT_SEARCH_TOOL_NAME,
     MCP_TOOL_CALL_TOOL_NAME,
     MCP_TOOL_SEARCH_TOOL_NAME,
+    SKILL_SEARCH_TOOL_NAME,
     SemanticToolRanker,
     ToolSearchResult,
     coerce_top_k,
@@ -272,8 +273,8 @@ class TestSearchTools:
 
 
 class TestGetVirtualToolDefinitions:
-    def test_returns_three_tools(self) -> None:
-        assert len(get_virtual_tool_definitions()) == 3
+    def test_returns_four_tools(self) -> None:
+        assert len(get_virtual_tool_definitions()) == 4
 
     def test_agent_search_schema_requires_query(self) -> None:
         tools = get_virtual_tool_definitions()
@@ -330,6 +331,7 @@ class TestGetVirtualToolDefinitions:
             MCP_TOOL_SEARCH_TOOL_NAME,
             MCP_TOOL_CALL_TOOL_NAME,
             AGENT_SEARCH_TOOL_NAME,
+            SKILL_SEARCH_TOOL_NAME,
         }
 
 
@@ -364,7 +366,12 @@ class TestListToolRestApiWithToolSearch:
 
         assert result["error"] is None
         tool_names = [t["name"] for t in result["tools"]]
-        assert set(tool_names) == {MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME}
+        assert set(tool_names) == {
+            MCP_TOOL_SEARCH_TOOL_NAME,
+            MCP_TOOL_CALL_TOOL_NAME,
+            AGENT_SEARCH_TOOL_NAME,
+            SKILL_SEARCH_TOOL_NAME,
+        }
 
     @pytest.mark.asyncio
     async def test_returns_full_catalog_when_flag_disabled(self) -> None:
@@ -737,6 +744,31 @@ class TestCallToolRestApiVirtualTools:
         assert mock_search.await_args.kwargs["top_k"] == 1
         assert mock_search.await_args.kwargs["agents"] == (translator,)
 
+    @pytest.mark.asyncio
+    async def test_skill_search_call_tolerates_malformed_top_k(self) -> None:
+        """Regression: a caller-supplied non-numeric top_k must be coerced to the default,
+        the same as agent_search, instead of raising a pydantic ValidationError that the
+        endpoint's catch-all turns into an HTTP 500."""
+        from mcp.types import CallToolResult, TextContent
+
+        from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K
+
+        user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
+        request = self._make_request(
+            {"name": SKILL_SEARCH_TOOL_NAME, "arguments": {"query": "translate a document", "top_k": "not-a-number"}}
+        )
+        fake_result = CallToolResult(content=[TextContent(type="text", text="[]")], isError=False)
+        with patch(  # test-quality-ok: the embedding router only resolves via proxy_server globals, no injection seam
+            "litellm.proxy._experimental.mcp_server.tool_search.handle_skill_search",
+            new_callable=AsyncMock,
+            return_value=fake_result,
+        ) as mock_search:
+            result = await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict)
+
+        assert result.isError is False
+        assert mock_search.await_args.kwargs["top_k"] == DEFAULT_SKILL_SEARCH_TOP_K
+        assert mock_search.await_args.kwargs["query"] == "translate a document"
+
     @pytest.mark.asyncio
     async def test_agent_search_call_reports_missing_embedding_model_as_tool_error(self) -> None:
         from litellm.proxy.agent_endpoints.agent_search import AgentSearchNotConfigured
@@ -782,10 +814,15 @@ class TestCallToolRestApiVirtualTools:
 
         router = MagicMock()
         router.aembedding = AsyncMock(side_effect=fake_aembedding)
+        key_limits = MagicMock()
+        key_limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
         with (
             patch(  # test-quality-ok: the proxy's router is a module global; the handler reaches it the way production does
                 "litellm.proxy.proxy_server.llm_router", router
             ),
+            patch(  # test-quality-ok: the proxy's key-limit hooks are a module global; the embedding call runs them like /embeddings does
+                "litellm.proxy.proxy_server.proxy_logging_obj", key_limits
+            ),
             patch(  # test-quality-ok: the authorized catalog is the seam every virtual tool shares; the ranking under test stays real
                 "litellm.proxy._experimental.mcp_server.server._list_mcp_tools",
                 new_callable=AsyncMock,
@@ -795,6 +832,8 @@ class TestCallToolRestApiVirtualTools:
             result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict)
 
         assert mock_list.await_args.kwargs["user_api_key_auth"] is user_api_key_dict
+        assert key_limits.pre_call_hook.await_args.kwargs["call_type"] == "aembedding"
+        assert key_limits.pre_call_hook.await_args.kwargs["data"]["model"] == "emb"
         assert result.isError is False
         assert [t["name"] for t in json.loads(result.content[0].text)] == [FX_TOOL.name]
 
@@ -810,7 +849,9 @@ class TestCallToolRestApiVirtualTools:
         assert "mcp_tool_search.embedding_model" in result.content[0].text
 
     @pytest.mark.asyncio
-    async def test_mcp_tool_search_reports_invalid_settings_as_tool_error(self, monkeypatch: pytest.MonkeyPatch) -> None:
+    async def test_mcp_tool_search_reports_invalid_settings_as_tool_error(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
         monkeypatch.setattr(litellm, "mcp_tool_search", {"top_k": 0})
         user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
         result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict)
@@ -1196,6 +1237,7 @@ class TestHandleListToolsVirtual:
             MCP_TOOL_SEARCH_TOOL_NAME,
             MCP_TOOL_CALL_TOOL_NAME,
             AGENT_SEARCH_TOOL_NAME,
+            SKILL_SEARCH_TOOL_NAME,
         }
 
 
diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py
index daca244c0a1..c02ed1f37e5 100644
--- a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py
+++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py
@@ -211,11 +211,24 @@ class TestAgentSearchIndex:
         assert isinstance(outcome, AgentSearchEmbeddingFailed)
 
 
+def _pass_through_key_limits() -> MagicMock:
+    limits = MagicMock()
+    limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
+    return limits
+
+
 class TestSearchAgents:
     @pytest.mark.asyncio
     async def test_no_embedding_model_is_not_configured(self) -> None:
         outcome = await search_agents(
-            "q", AGENTS, 5, router=MagicMock(), embedding_model=None, index=AgentSearchIndex(), user_api_key_dict=CALLER
+            "q",
+            AGENTS,
+            5,
+            router=MagicMock(),
+            embedding_model=None,
+            index=AgentSearchIndex(),
+            user_api_key_dict=CALLER,
+            proxy_logging_obj=_pass_through_key_limits(),
         )
         assert isinstance(outcome, AgentSearchNotConfigured)
         assert "agent_search_embedding_model" in outcome.reason
@@ -223,7 +236,14 @@ class TestSearchAgents:
     @pytest.mark.asyncio
     async def test_no_router_is_not_configured(self) -> None:
         outcome = await search_agents(
-            "q", AGENTS, 5, router=None, embedding_model="m", index=AgentSearchIndex(), user_api_key_dict=CALLER
+            "q",
+            AGENTS,
+            5,
+            router=None,
+            embedding_model="m",
+            index=AgentSearchIndex(),
+            user_api_key_dict=CALLER,
+            proxy_logging_obj=_pass_through_key_limits(),
         )
         assert isinstance(outcome, AgentSearchNotConfigured)
 
@@ -244,6 +264,7 @@ class TestSearchAgents:
             embedding_model="text-embedding-3-small",
             index=AgentSearchIndex(),
             user_api_key_dict=CALLER,
+            proxy_logging_obj=_pass_through_key_limits(),
         )
         assert isinstance(outcome, AgentSearchHits)
         assert [hit.agent.agent_id for hit in outcome.hits] == ["translator"]
@@ -266,6 +287,7 @@ class TestSearchAgents:
             embedding_model="text-embedding-3-small",
             index=AgentSearchIndex(),
             user_api_key_dict=CALLER,
+            proxy_logging_obj=_pass_through_key_limits(),
         )
         metadata = router.aembedding.await_args.kwargs["metadata"]
         assert metadata["user_api_key"] == "hashed-caller-key"
@@ -302,6 +324,7 @@ def embedding_router(monkeypatch: pytest.MonkeyPatch) -> MagicMock:
         )
     )
     monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
+    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", _pass_through_key_limits())
     monkeypatch.setattr(litellm, "agent_search_embedding_model", "text-embedding-3-small")
     return router
 
diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py
new file mode 100644
index 00000000000..ae6b1471c93
--- /dev/null
+++ b/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py
@@ -0,0 +1,175 @@
+from types import MappingProxyType
+from typing import Final
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from openai import APIConnectionError
+
+import litellm
+from litellm.llms.litellm_proxy.skills.skill_search import skill_search_text
+from litellm.proxy._types import LiteLLM_SkillsTable, LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy.anthropic_endpoints.skills_endpoints import router
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
+
+TRANSLATOR: Final = LiteLLM_SkillsTable(
+    skill_id="translate-file",
+    display_title="Document Translator",
+    description="Converts files from one language into another",
+    instructions="Take an uploaded document and produce it in the target language",
+)
+SQL_ANALYST: Final = LiteLLM_SkillsTable(
+    skill_id="warehouse-sql-analyst",
+    display_title="Warehouse SQL Analyst",
+    description="Runs SQL against the inventory database",
+)
+TRIP_PLANNER: Final = LiteLLM_SkillsTable(
+    skill_id="trip-planner",
+    display_title="Trip Planner",
+    description="Books flights and hotels",
+)
+SKILLS: Final = (TRANSLATOR, SQL_ANALYST, TRIP_PLANNER)
+
+VECTORS: Final = MappingProxyType(
+    {
+        "language translation": (1.0, 0.0, 0.0),
+        skill_search_text(TRANSLATOR): (0.9, 0.1, 0.0),
+        skill_search_text(SQL_ANALYST): (0.0, 1.0, 0.0),
+        skill_search_text(TRIP_PLANNER): (0.3, 0.0, 1.0),
+    }
+)
+
+
+def _client(role: LitellmUserRoles) -> TestClient:
+    app = FastAPI()
+    app.include_router(router)
+    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="u", user_role=role)
+    return TestClient(app)
+
+
+@pytest.fixture
+def accessible_skills(monkeypatch: pytest.MonkeyPatch) -> AsyncMock:
+    list_for_search = AsyncMock(return_value=list(SKILLS))
+    monkeypatch.setattr(
+        "litellm.llms.litellm_proxy.skills.handler.LiteLLMSkillsHandler.list_skills_for_search", list_for_search
+    )
+    return list_for_search
+
+
+@pytest.fixture
+def key_limits(monkeypatch: pytest.MonkeyPatch) -> MagicMock:
+    limits = MagicMock()
+    limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
+    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", limits)
+    return limits
+
+
+@pytest.fixture
+def embedding_router(monkeypatch: pytest.MonkeyPatch, key_limits: MagicMock) -> MagicMock:
+    embedding_router = MagicMock()
+    embedding_router.aembedding = AsyncMock(
+        side_effect=lambda model, input, metadata: litellm.EmbeddingResponse(
+            model=model,
+            data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)],
+        )
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", embedding_router)
+    monkeypatch.setattr(litellm, "skill_search_embedding_model", "text-embedding-3-small")
+    return embedding_router
+
+
+class TestGetSkillsQuery:
+    def test_query_ranks_and_scores_and_truncates(
+        self, accessible_skills: AsyncMock, embedding_router: MagicMock
+    ) -> None:
+        response = _client(LitellmUserRoles.PROXY_ADMIN).get(
+            "/v1/skills",
+            params={"custom_llm_provider": "litellm_proxy", "query": "language translation", "top_k": 2},
+            headers={"Authorization": "Bearer k"},
+        )
+        assert response.status_code == 200
+        body = response.json()["data"]
+        assert [skill["id"] for skill in body] == ["translate-file", "trip-planner"]
+        assert body[0]["search_score"] > body[1]["search_score"]
+        assert embedding_router.aembedding.await_args.kwargs["metadata"]["user_api_key_user_id"] == "u"
+
+    def test_restricted_key_only_ranks_the_skills_it_can_access(
+        self, accessible_skills: AsyncMock, embedding_router: MagicMock
+    ) -> None:
+        accessible_skills.return_value = [SQL_ANALYST]
+        response = _client(LitellmUserRoles.INTERNAL_USER).get(
+            "/v1/skills",
+            params={"custom_llm_provider": "litellm_proxy", "query": "language translation"},
+            headers={"Authorization": "Bearer k"},
+        )
+        assert response.status_code == 200
+        assert [skill["id"] for skill in response.json()["data"]] == ["warehouse-sql-analyst"]
+
+    def test_no_accessible_skills_is_a_no_match_empty_result(
+        self, accessible_skills: AsyncMock, embedding_router: MagicMock
+    ) -> None:
+        accessible_skills.return_value = []
+        response = _client(LitellmUserRoles.INTERNAL_USER).get(
+            "/v1/skills",
+            params={"custom_llm_provider": "litellm_proxy", "query": "anything"},
+            headers={"Authorization": "Bearer k"},
+        )
+        assert response.status_code == 200
+        assert response.json()["data"] == []
+        embedding_router.aembedding.assert_not_awaited()
+
+    def test_query_is_unsupported_for_the_anthropic_passthrough_provider(
+        self, accessible_skills: AsyncMock, embedding_router: MagicMock
+    ) -> None:
+        response = _client(LitellmUserRoles.PROXY_ADMIN).get(
+            "/v1/skills", params={"query": "anything"}, headers={"Authorization": "Bearer k"}
+        )
+        assert response.status_code == 400
+        assert response.json()["detail"]["error"] == "skill_search_unsupported_provider"
+        accessible_skills.assert_not_awaited()
+
+    def test_missing_embedding_model_is_a_400(
+        self, accessible_skills: AsyncMock, embedding_router: MagicMock, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setattr(litellm, "skill_search_embedding_model", None)
+        response = _client(LitellmUserRoles.PROXY_ADMIN).get(
+            "/v1/skills",
+            params={"custom_llm_provider": "litellm_proxy", "query": "anything"},
+            headers={"Authorization": "Bearer k"},
+        )
+        assert response.status_code == 400
+        assert response.json()["detail"]["error"] == "skill_search_not_configured"
+
+    def test_embedding_provider_failure_is_a_503(
+        self, accessible_skills: AsyncMock, embedding_router: MagicMock
+    ) -> None:
+        embedding_router.aembedding = AsyncMock(side_effect=APIConnectionError(request=MagicMock()))
+        response = _client(LitellmUserRoles.PROXY_ADMIN).get(
+            "/v1/skills",
+            params={"custom_llm_provider": "litellm_proxy", "query": "anything"},
+            headers={"Authorization": "Bearer k"},
+        )
+        assert response.status_code == 503
+        assert response.json()["detail"]["error"] == "skill_search_unavailable"
+
+    def test_a_key_over_its_rate_limit_gets_a_429_without_embedding(
+        self, accessible_skills: AsyncMock, embedding_router: MagicMock, key_limits: MagicMock
+    ) -> None:
+        key_limits.pre_call_hook = AsyncMock(side_effect=ProxyRateLimitError(detail="rpm exceeded"))
+        response = _client(LitellmUserRoles.PROXY_ADMIN).get(
+            "/v1/skills",
+            params={"custom_llm_provider": "litellm_proxy", "query": "language translation"},
+            headers={"Authorization": "Bearer k"},
+        )
+        assert response.status_code == 429
+        embedding_router.aembedding.assert_not_awaited()
+
+    def test_top_k_is_validated(self, accessible_skills: AsyncMock, embedding_router: MagicMock) -> None:
+        response = _client(LitellmUserRoles.PROXY_ADMIN).get(
+            "/v1/skills",
+            params={"custom_llm_provider": "litellm_proxy", "query": "anything", "top_k": 0},
+            headers={"Authorization": "Bearer k"},
+        )
+        assert response.status_code == 422
diff --git a/tests/test_litellm/skills/test_skills_main.py b/tests/test_litellm/skills/test_skills_main.py
new file mode 100644
index 00000000000..e1c66c8d9ea
--- /dev/null
+++ b/tests/test_litellm/skills/test_skills_main.py
@@ -0,0 +1,57 @@
+from unittest.mock import MagicMock
+
+import litellm.skills.main as skills_main
+from litellm.types.utils import LlmProviders
+
+
+def test_create_skill_forwards_description_and_instructions_from_top_level_kwargs(
+    monkeypatch,
+) -> None:
+    """The REST /v1/skills form endpoint passes description/instructions as top-level
+    kwargs (not extra_body). Regression for a bug where the litellm_proxy dispatch
+    branch of create_skill() dropped both, so every LiteLLM-hosted skill was created
+    with description=None and instructions=None regardless of what the caller sent."""
+    handler = MagicMock()
+    monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler)
+
+    skills_main.create_skill(
+        display_title="Document Translator",
+        description="Converts files from one language into another",
+        instructions="Take an uploaded document and produce it in the target language",
+        custom_llm_provider=LlmProviders.LITELLM_PROXY.value,
+    )
+
+    assert handler.create_skill_handler.call_args.kwargs["description"] == (
+        "Converts files from one language into another"
+    )
+    assert handler.create_skill_handler.call_args.kwargs["instructions"] == (
+        "Take an uploaded document and produce it in the target language"
+    )
+
+
+def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None:
+    """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under
+    extra_body instead of passing them as top-level kwargs; both paths must reach the DB."""
+    handler = MagicMock()
+    monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler)
+
+    skills_main.create_skill(
+        display_title="Warehouse SQL Analyst",
+        extra_body={"description": "Runs SQL against the inventory database", "instructions": "Summarize results"},
+        custom_llm_provider=LlmProviders.LITELLM_PROXY.value,
+    )
+
+    assert handler.create_skill_handler.call_args.kwargs["description"] == (
+        "Runs SQL against the inventory database"
+    )
+    assert handler.create_skill_handler.call_args.kwargs["instructions"] == "Summarize results"
+
+
+def test_create_skill_without_description_or_instructions_passes_none(monkeypatch) -> None:
+    handler = MagicMock()
+    monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler)
+
+    skills_main.create_skill(display_title="Bare Skill", custom_llm_provider=LlmProviders.LITELLM_PROXY.value)
+
+    assert handler.create_skill_handler.call_args.kwargs["description"] is None
+    assert handler.create_skill_handler.call_args.kwargs["instructions"] is None
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 842a3da4122..12e7eb801e6 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -19872,6 +19872,12 @@ export interface paths {
          *     curl "http://localhost:4000/v1/skills?beta=true&limit=10"       -H "Authorization: Bearer your-key"       -H "x-litellm-model: claude-account-1"
          *     ```
          *
+         *     Pass `?custom_llm_provider=litellm_proxy&query=<task>` to rank the LiteLLM-hosted skills you can
+         *     access by semantic similarity instead of paging through the whole registry:
+         *     ```bash
+         *     curl "http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5"       -H "Authorization: Bearer your-key"
+         *     ```
+         *
          *     Returns: ListSkillsResponse with list of skills
          */
         get: operations["list_skills_v1_skills_get"];
@@ -36050,12 +36056,16 @@ export interface components {
         Skill: {
             /** Created At */
             created_at: string;
+            /** Description */
+            description?: string | null;
             /** Display Title */
             display_title?: string | null;
             /** Id */
             id: string;
             /** Latest Version */
             latest_version?: string | null;
+            /** Search Score */
+            search_score?: number | null;
             /** Source */
             source: string;
             /**
@@ -64556,6 +64566,10 @@ export interface operations {
                 after_id?: string | null;
                 before_id?: string | null;
                 custom_llm_provider?: string | null;
+                /** @description Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model. */
+                query?: string | null;
+                /** @description With query: the maximum number of ranked skills to return. */
+                top_k?: number;
             };
             header?: never;
             path?: never;

From e11a8c59ff0497b56e0023e642647be1b338fde8 Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Mon, 7 Sep 2026 19:34:43 +0000
Subject: [PATCH 133/164] fix(ui): show inherited MCP servers on the internal
 user editor and flag access groups with no members (#40036)

* fix(ui): show inherited MCP servers on the internal-user editor and flag access groups with no members

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(ui): consult the unfiltered access group registry before calling a group empty

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../users/_components/user_edit_view.test.tsx | 47 ++++++++++++++++++-
 .../users/_components/user_edit_view.tsx      |  2 +
 .../MCPToolPermissions.test.tsx               | 40 ++++++++++++++++
 .../MCPToolPermissions.tsx                    | 22 ++++++++-
 .../effectiveMcpServers.test.ts               | 24 ++++++++++
 .../effectiveMcpServers.ts                    | 13 +++++
 6 files changed, 146 insertions(+), 2 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx
index 8fb94ce477e..2571eb344f5 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx
@@ -1,8 +1,11 @@
 import { cleanup, fireEvent, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import { renderWithProviders } from "../../../../../tests/test-utils";
+import { renderWithProviders, testQueryClient } from "../../../../../tests/test-utils";
 import { UserEditView } from "./user_edit_view";
+import * as networking from "@/components/networking";
+
+vi.mock("@/components/networking");
 
 vi.mock("@/components/key_team_helpers/fetch_available_models_team_key", () => ({
   getModelDisplayName: vi.fn((model: string) => model),
@@ -59,6 +62,10 @@ describe("UserEditView", () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
+    testQueryClient.clear();
+    vi.mocked(networking.fetchMCPServers).mockResolvedValue([]);
+    vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]);
+    vi.mocked(networking.fetchMCPToolsets).mockResolvedValue([]);
   });
 
   afterEach(() => {
@@ -579,6 +586,44 @@ describe("UserEditView", () => {
       expect(budgetInput.closest("form")).not.toHaveAttribute("novalidate");
     });
 
+    it("shows the tool matrix for servers the user reaches only through an access group or toolset", async () => {
+      vi.mocked(networking.fetchMCPServers).mockResolvedValue([
+        { server_id: "srv-group", server_name: "Group Server", alias: "Group Server", mcp_access_groups: ["group-a"] },
+        { server_id: "srv-toolset", server_name: "Toolset Server", alias: "Toolset Server" },
+      ]);
+      vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue(["group-a"]);
+      vi.mocked(networking.fetchMCPToolsets).mockResolvedValue([
+        {
+          toolset_id: "toolset-a",
+          toolset_name: "Toolset A",
+          tools: [{ server_id: "srv-toolset", tool_name: "list_issues" }],
+        } as never,
+      ]);
+      vi.mocked(networking.listMCPTools).mockResolvedValue({
+        tools: [{ name: "list_issues", description: "List issues" }],
+        error: false,
+      });
+
+      renderWithProviders(
+        <UserEditView
+          {...defaultProps}
+          objectPermission={
+            {
+              mcp_servers: [],
+              mcp_access_groups: ["group-a"],
+              mcp_toolsets: ["toolset-a"],
+              mcp_tool_permissions: {},
+            } as never
+          }
+        />,
+      );
+
+      expect(await screen.findByText("Via access group: group-a")).toBeInTheDocument();
+      expect(await screen.findByText("Via toolset: Toolset A")).toBeInTheDocument();
+      expect(networking.listMCPTools).toHaveBeenCalledWith("test-token", "srv-group");
+      expect(networking.listMCPTools).toHaveBeenCalledWith("test-token", "srv-toolset");
+    });
+
     it("should send objects for the mcp keys seeded from objectPermission", async () => {
       const payload = await submittedPayload({
         objectPermission: {
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx
index 96771dc6dc4..b7a3486c78e 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx
@@ -336,6 +336,8 @@ export function UserEditView({
               <MCPToolPermissions
                 accessToken={accessToken || ""}
                 selectedServers={form.watch("mcp_servers_and_groups")?.servers || []}
+                selectedAccessGroups={form.watch("mcp_servers_and_groups")?.accessGroups || []}
+                selectedToolsets={form.watch("mcp_servers_and_groups")?.toolsets || []}
                 toolPermissions={form.watch("mcp_tool_permissions") || {}}
                 onChange={(toolPerms) => form.setValue("mcp_tool_permissions", toolPerms)}
               />
diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx
index 69a761b4723..149d231fff6 100644
--- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx
+++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx
@@ -19,6 +19,7 @@ describe("MCPToolPermissions", () => {
     vi.clearAllMocks();
     testQueryClient.clear();
     vi.mocked(networking.fetchMCPToolsets).mockResolvedValue([]);
+    vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]);
   });
 
   it("should update tool permissions when user selects a tool", async () => {
@@ -621,6 +622,45 @@ describe("MCPToolPermissions", () => {
       );
 
       expect(await screen.findByText("Unable to load MCP servers")).toBeInTheDocument();
+      expect(screen.queryByText(/has 0 servers/)).not.toBeInTheDocument();
+    });
+
+    it("tells the admin when a loaded access group has no member servers", async () => {
+      vi.mocked(networking.fetchMCPServers).mockResolvedValue([groupServer]);
+      vi.mocked(networking.fetchMCPToolsets).mockResolvedValue([]);
+      vi.mocked(networking.listMCPTools).mockResolvedValue({ tools: groupTools, error: false });
+
+      renderWithProviders(
+        <MCPToolPermissions
+          accessToken={mockAccessToken}
+          selectedServers={[]}
+          selectedAccessGroups={["production-group", "ops_readonly"]}
+          toolPermissions={{}}
+          onChange={vi.fn()}
+        />,
+      );
+
+      expect(await screen.findByText('Access group "ops_readonly" has 0 servers')).toBeInTheDocument();
+      expect(screen.getByText("Group Server")).toBeInTheDocument();
+      expect(screen.queryByText('Access group "production-group" has 0 servers')).not.toBeInTheDocument();
+    });
+
+    it("does not call a group empty when its servers are only hidden from the caller's catalog", async () => {
+      vi.mocked(networking.fetchMCPServers).mockResolvedValue([]);
+      vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue(["production-group"]);
+
+      renderWithProviders(
+        <MCPToolPermissions
+          accessToken={mockAccessToken}
+          selectedServers={[]}
+          selectedAccessGroups={["production-group", "ops_readonly"]}
+          toolPermissions={{}}
+          onChange={vi.fn()}
+        />,
+      );
+
+      expect(await screen.findByText('Access group "ops_readonly" has 0 servers')).toBeInTheDocument();
+      expect(screen.queryByText('Access group "production-group" has 0 servers')).not.toBeInTheDocument();
     });
 
     it("warns when the selected toolsets cannot be resolved to servers", async () => {
diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx
index c866e9cc011..e26f1a6f511 100644
--- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx
+++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx
@@ -4,6 +4,7 @@ import { MCPTool } from "../mcp_tools/types";
 import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group";
 import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
 import { useMCPServers } from "../../app/(dashboard)/hooks/mcpServers/useMCPServers";
+import { useMCPAccessGroups } from "../../app/(dashboard)/hooks/mcpServers/useMCPAccessGroups";
 import { useMCPToolsets } from "../../app/(dashboard)/hooks/mcpServers/useMCPToolsets";
 import McpCrudPermissionPanel from "../mcp_tools/McpCrudPermissionPanel";
 import { classifyToolOp } from "../../utils/mcpToolCrudClassification";
@@ -12,6 +13,7 @@ import {
   EffectiveMcpServer,
   McpGrantSource,
   applyToolPermissionWrite,
+  emptyMcpAccessGroups,
   mcpAllowedToolsFor,
   resolveEffectiveMcpServers,
 } from "./effectiveMcpServers";
@@ -55,7 +57,13 @@ const MCPToolPermissions: React.FC<MCPToolPermissionsProps> = ({
   onChange,
   disabled = false,
 }) => {
-  const { data: allServers = [], isError: serversFailed, isLoading: serversLoading } = useMCPServers();
+  const {
+    data: allServers = [],
+    isError: serversFailed,
+    isLoading: serversLoading,
+    isSuccess: serversLoaded,
+  } = useMCPServers();
+  const { data: populatedAccessGroups = [], isSuccess: accessGroupsLoaded } = useMCPAccessGroups();
   const { data: toolsets = [], isError: toolsetsFailed, isLoading: toolsetsLoading } = useMCPToolsets();
   const [serverTools, setServerTools] = useState<Record<string, MCPTool[]>>({});
   const [loadingTools, setLoadingTools] = useState<Record<string, boolean>>({});
@@ -181,6 +189,18 @@ const MCPToolPermissions: React.FC<MCPToolPermissionsProps> = ({
         </div>
       )}
 
+      {serversLoaded &&
+        accessGroupsLoaded &&
+        emptyMcpAccessGroups(allServers, populatedAccessGroups, selectedAccessGroups).map((group) => (
+          <div key={group} className="p-4 bg-yellow-50 border border-yellow-200 rounded-lg">
+            <p className="text-sm text-yellow-800 font-medium">Access group &quot;{group}&quot; has 0 servers</p>
+            <p className="text-sm text-yellow-700 mt-1">
+              No MCP server lists this group, so it grants nothing. A server defined in config.yaml joins a group
+              through its <code>access_groups</code> key; <code>mcp_access_groups</code> is ignored there
+            </p>
+          </div>
+        ))}
+
       {toolsetsFailed && selectedToolsets.length > 0 && (
         <div className="p-4 bg-yellow-50 border border-yellow-200 rounded-lg">
           <p className="text-sm text-yellow-800 font-medium">Unable to load toolsets</p>
diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.test.ts b/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.test.ts
index b2ac5337245..487c6f9f55e 100644
--- a/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.test.ts
+++ b/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.test.ts
@@ -2,6 +2,7 @@ import { describe, it, expect } from "vitest";
 import { MCPServer, MCPToolset } from "../mcp_tools/types";
 import {
   applyToolPermissionWrite,
+  emptyMcpAccessGroups,
   mcpAllowedToolsFor,
   mcpServersForIdentifier,
   mcpToolPermissionKeyFor,
@@ -101,6 +102,29 @@ describe("mcpToolPermissionKeyFor", () => {
   });
 });
 
+describe("emptyMcpAccessGroups", () => {
+  const grouped = server({ server_id: "srv-group", server_name: "Grouped", mcp_access_groups: ["prod"] });
+  const objectGrouped = {
+    ...grouped,
+    server_id: "srv-obj",
+    mcp_access_groups: [{ name: "legacy" }],
+  } as unknown as MCPServer;
+
+  it("names only the selected groups no loaded server belongs to", () => {
+    expect(emptyMcpAccessGroups([grouped, objectGrouped], [], ["prod", "legacy", "ops_readonly"])).toEqual([
+      "ops_readonly",
+    ]);
+  });
+
+  it("names every selected group when no server is loaded and the registry is empty", () => {
+    expect(emptyMcpAccessGroups([], [], ["prod"])).toEqual(["prod"]);
+  });
+
+  it("trusts the group registry when the caller's catalog hides the member servers", () => {
+    expect(emptyMcpAccessGroups([], ["prod"], ["prod", "ops_readonly"])).toEqual(["ops_readonly"]);
+  });
+});
+
 describe("resolveEffectiveMcpServers", () => {
   const direct = server({ server_id: "srv-direct", server_name: "Direct" });
   const grouped = server({ server_id: "srv-group", server_name: "Grouped", mcp_access_groups: ["prod"] });
diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.ts b/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.ts
index 3320dc135c9..b3ba24f3c59 100644
--- a/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.ts
+++ b/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.ts
@@ -54,6 +54,19 @@ const accessGroupNamesOf = (server: MCPServer): readonly string[] =>
     return [typeof parsed.data === "string" ? parsed.data : parsed.data.name];
   });
 
+// The server catalog can be trimmed to the caller's grants, so the unfiltered group registry
+// (GET /v1/mcp/access_groups) has to agree before a group is called empty.
+export const emptyMcpAccessGroups = (
+  allServers: readonly MCPServer[],
+  populatedAccessGroups: readonly string[],
+  selectedAccessGroups: readonly string[],
+): readonly string[] =>
+  selectedAccessGroups.filter(
+    (group) =>
+      !populatedAccessGroups.includes(group) &&
+      !allServers.some((server) => accessGroupNamesOf(server).includes(group)),
+  );
+
 // Which servers an identifier names, with the same precedence the backend's expand_permission_list
 // applies: a string that is a registry server id names exactly that server, and only a string that
 // is not falls back to server_name/alias, which can name several. Matching all three fields at once

From 6908318c168cf50d750b3535fc3a2e4def0abf8e Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Mon, 7 Sep 2026 19:37:44 +0000
Subject: [PATCH 134/164] feat(keys): allow editing soft budget on existing
 keys (#39002)

* feat(keys): allow editing soft budget on existing keys

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(ui): extract KeyBudgetNumberField to keep key_edit_view under max-lines

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* style(ui): format keyEditFormValues with prettier

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(keys): cover soft budget validation and update adapter

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(ui): reject non-finite soft budget values instead of clearing

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(keys): assert soft budget validation returns None for valid values

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(keys): write soft budget and key row in one transaction

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: yassin <yassin@berri.ai>
---
 litellm/proxy/_types.py                       |   1 +
 .../key_management_endpoints.py               | 144 +++++++++-
 .../test_key_management_endpoints.py          | 247 ++++++++++++++++++
 .../templates/KeyEditViewControls.tsx         |  28 ++
 .../KeyInfoView.handleKeyUpdate.test.tsx      | 106 ++++++++
 .../components/templates/keyEditFormValues.ts |   5 +
 .../templates/key_edit_view.test.tsx          |   1 +
 .../components/templates/key_edit_view.tsx    |  26 +-
 .../components/templates/key_info_view.tsx    |  17 ++
 ui/litellm-dashboard/src/lib/http/schema.d.ts |   4 +-
 10 files changed, 557 insertions(+), 22 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index c28ac8848ba..d79b753b5ff 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -1276,6 +1276,7 @@ class UpdateKeyRequest(KeyRequestBase):
     # else they will get overwritten
     duration: str | None = None
     spend: float | None = None
+    soft_budget: float | None = None
     metadata: dict | None = None
     temp_budget_increase: float | None = None
     temp_budget_expiry: datetime | None = None
diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index 8e51e250319..f46c4170071 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -196,6 +196,38 @@ class _ModelRowWhere(TypedDict):
     model_id: ReadOnly[str]
 
 
+class _KeyUpdateResult(TypedDict):
+    token: ReadOnly[str]
+    data: ReadOnly[Mapping[str, object]]
+
+
+class _KeyRowWhere(TypedDict):
+    token: ReadOnly[str]
+
+
+class _BudgetRowWhere(TypedDict):
+    budget_id: ReadOnly[str]
+
+
+class _BudgetRowSoftBudgetUpdate(TypedDict):
+    soft_budget: ReadOnly[float | None]
+    updated_by: ReadOnly[str]
+
+
+class _BudgetRowSoftBudgetCreate(TypedDict):
+    soft_budget: ReadOnly[float]
+    created_by: ReadOnly[str]
+    updated_by: ReadOnly[str]
+
+
+class _KeyUpdateTx(Protocol):
+    @property
+    def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ...
+
+    @property
+    def litellm_budgettable(self) -> "TableActions[prisma_models.LiteLLM_BudgetTable]": ...
+
+
 class _ConfigTableActions(Protocol):
     """Config table surface this module needs; the shared repository seam exposes no ``update``."""
 
@@ -1812,11 +1844,7 @@ async def generate_key_fn(
                 status_code=400,
                 detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"},
             )
-        if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0):
-            raise HTTPException(
-                status_code=400,
-                detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"},
-            )
+        _validate_soft_budget_value(data.soft_budget)
 
         custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = (
             _custom_key_generate_hook(proxy_server)
@@ -2121,6 +2149,88 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_
     return non_default_values
 
 
+def _validate_soft_budget_value(soft_budget: float | None) -> None:
+    if soft_budget is not None and (not math.isfinite(soft_budget) or soft_budget < 0):
+        raise HTTPException(
+            status_code=400,
+            detail={"error": f"soft_budget must be a non-negative finite number. Received: {soft_budget}"},
+        )
+
+
+async def _update_key_soft_budget(
+    db: _KeyUpdateTx,
+    existing_key_row: LiteLLM_VerificationToken,
+    soft_budget: float | None,
+    changed_by: str,
+) -> str | None:
+    existing_budget_id: Final = existing_key_row.budget_id
+    if existing_budget_id is not None:
+        budget_update: Final[_BudgetRowSoftBudgetUpdate] = {"soft_budget": soft_budget, "updated_by": changed_by}
+        budget_where: Final[_BudgetRowWhere] = {"budget_id": existing_budget_id}
+        await db.litellm_budgettable.update(where=budget_where, data=budget_update)
+        return existing_budget_id
+    if soft_budget is None:
+        return None
+    budget_create: Final[_BudgetRowSoftBudgetCreate] = {
+        "soft_budget": soft_budget,
+        "created_by": changed_by,
+        "updated_by": changed_by,
+    }
+    created_budget: Final = await db.litellm_budgettable.create(data=budget_create)
+    return created_budget.budget_id
+
+
+async def _apply_soft_budget_update(
+    data: UpdateKeyRequest,
+    non_default_values: Mapping[str, object],
+    db: _KeyUpdateTx,
+    existing_key_row: LiteLLM_VerificationToken,
+    changed_by: str,
+) -> Mapping[str, object]:
+    remaining: Final = MappingProxyType({k: v for k, v in non_default_values.items() if k != "soft_budget"})
+    updated_budget_id: Final = await _update_key_soft_budget(
+        db=db,
+        existing_key_row=existing_key_row,
+        soft_budget=data.soft_budget,
+        changed_by=changed_by,
+    )
+    if updated_budget_id is not None and existing_key_row.budget_id is None:
+        return MappingProxyType({**remaining, "budget_id": updated_budget_id})
+    return remaining
+
+
+async def _update_key_row_with_soft_budget(
+    prisma_client: PrismaClient,
+    key: str,
+    data: UpdateKeyRequest,
+    non_default_values: Mapping[str, object],
+    existing_key_row: LiteLLM_VerificationToken,
+    changed_by: str,
+) -> _KeyUpdateResult:
+    hashed_token: Final = _hash_token_if_needed(key)
+    key_where: Final[_KeyRowWhere] = {"token": hashed_token}
+    tx: _KeyUpdateTx
+    async with prisma_client.tx() as tx:
+        update_values: Final = await _apply_soft_budget_update(
+            data=data,
+            non_default_values=non_default_values,
+            db=tx,
+            existing_key_row=existing_key_row,
+            changed_by=changed_by,
+        )
+        updated_row: Final = await tx.litellm_verificationtoken.update(
+            where=key_where,
+            data=with_settings_updated_at(
+                prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token}))
+            ),
+        )
+    updated_data: Final[Mapping[str, object]] = (
+        updated_row.model_dump() if updated_row is not None else MappingProxyType({})
+    )
+    result: Final[_KeyUpdateResult] = {"token": hashed_token, "data": updated_data}
+    return result
+
+
 async def prepare_key_update_data(
     data: UpdateKeyRequest | RegenerateKeyRequest,
     existing_key_row: LiteLLM_VerificationToken,
@@ -2659,6 +2769,7 @@ async def _validate_update_key_data(
         (data.max_budget is not None and data.max_budget != existing_key_row.max_budget)
         or data.spend is not None
         or "budget_limits" in data.model_fields_set
+        or "soft_budget" in data.model_fields_set
     )
 
     _existing_metadata: Final = getattr(existing_key_row, "metadata", None)
@@ -2862,7 +2973,7 @@ async def update_key_fn(
     - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
     - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
     - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
-    - soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
+    - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget.
     - max_parallel_requests: Optional[int] - Rate limit for parallel requests
     - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"}
     - tpm_limit: Optional[int] - Tokens per minute limit
@@ -2918,6 +3029,7 @@ async def update_key_fn(
     """
     from litellm.proxy import proxy_server
     from litellm.proxy.proxy_server import (
+        litellm_proxy_admin_name,
         llm_router,
         premium_user,
         prisma_client,
@@ -2933,6 +3045,8 @@ async def update_key_fn(
                 detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"},
             )
 
+        _validate_soft_budget_value(data.soft_budget)
+
         # get the row from db
         existing_key_row: Final = await _get_and_validate_existing_key(
             token=data.key,
@@ -2989,10 +3103,22 @@ async def update_key_fn(
             existing_key_alias=existing_key_row.key_alias,
         )
 
-        _data: Final = {**non_default_values, "token": key}
         if prisma_client is None:
             raise Exception("Not connected to DB!")
-        response: Final = await prisma_client.update_data(token=key, data=_data)
+
+        changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name
+        response: Final = (
+            await _update_key_row_with_soft_budget(
+                prisma_client=prisma_client,
+                key=key,
+                data=data,
+                non_default_values=non_default_values,
+                existing_key_row=existing_key_row,
+                changed_by=changed_by,
+            )
+            if "soft_budget" in data.model_fields_set
+            else await prisma_client.update_data(token=key, data=MappingProxyType({**non_default_values, "token": key}))
+        )
 
         # Delete - key from cache, since it's been updated!
         # key updated - a new model could have been added to this key. it should not block requests after this is done
@@ -6432,7 +6558,7 @@ async def _list_key_helper(
                     {"token": "desc"},  # fallback sort
                 ]
             ),
-            include={"object_permission": True},
+            include={"object_permission": True, "litellm_budget_table": True},
         )
 
     verbose_proxy_logger.debug("Fetched %s keys", len(keys))
diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
index 8766b1a1868..a873a367eab 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
@@ -17692,6 +17692,253 @@ async def test_check_project_key_limits_still_rejects_real_model_outside_project
     assert "Model 'gpt-5.4-mini' not in project's allowed models" in exc_info.value.detail["error"]
 
 
+@pytest.mark.asyncio
+async def test_update_key_soft_budget_updates_existing_budget_row():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _update_key_soft_budget,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id="budget-123")
+    mock_db = MagicMock()
+    mock_db.litellm_budgettable.update = AsyncMock()
+    mock_db.litellm_budgettable.create = AsyncMock()
+
+    result = await _update_key_soft_budget(
+        db=mock_db,
+        existing_key_row=existing_key,
+        soft_budget=25.0,
+        changed_by="user-1",
+    )
+
+    assert result == "budget-123"
+    mock_db.litellm_budgettable.update.assert_awaited_once_with(
+        where={"budget_id": "budget-123"},
+        data={"soft_budget": 25.0, "updated_by": "user-1"},
+    )
+
+
+@pytest.mark.asyncio
+async def test_update_key_soft_budget_clears_existing_budget_row():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _update_key_soft_budget,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id="budget-123")
+    mock_db = MagicMock()
+    mock_db.litellm_budgettable.update = AsyncMock()
+    mock_db.litellm_budgettable.create = AsyncMock()
+
+    result = await _update_key_soft_budget(
+        db=mock_db,
+        existing_key_row=existing_key,
+        soft_budget=None,
+        changed_by="user-1",
+    )
+
+    assert result == "budget-123"
+    mock_db.litellm_budgettable.update.assert_awaited_once_with(
+        where={"budget_id": "budget-123"},
+        data={"soft_budget": None, "updated_by": "user-1"},
+    )
+
+
+@pytest.mark.asyncio
+async def test_update_key_soft_budget_creates_budget_row_when_key_has_none():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _update_key_soft_budget,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None)
+    created_row = MagicMock()
+    created_row.budget_id = "budget-new"
+    mock_db = MagicMock()
+    mock_db.litellm_budgettable.create = AsyncMock(return_value=created_row)
+    mock_db.litellm_budgettable.update = AsyncMock()
+
+    result = await _update_key_soft_budget(
+        db=mock_db,
+        existing_key_row=existing_key,
+        soft_budget=10.5,
+        changed_by="user-1",
+    )
+
+    assert result == "budget-new"
+    mock_db.litellm_budgettable.create.assert_awaited_once_with(
+        data={"soft_budget": 10.5, "created_by": "user-1", "updated_by": "user-1"}
+    )
+
+
+@pytest.mark.asyncio
+async def test_update_key_soft_budget_noop_when_clearing_without_budget_row():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _update_key_soft_budget,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None)
+    mock_db = MagicMock()
+    mock_db.litellm_budgettable.create = AsyncMock()
+    mock_db.litellm_budgettable.update = AsyncMock()
+
+    result = await _update_key_soft_budget(
+        db=mock_db,
+        existing_key_row=existing_key,
+        soft_budget=None,
+        changed_by="user-1",
+    )
+
+    assert result is None
+    mock_db.litellm_budgettable.create.assert_not_awaited()
+    mock_db.litellm_budgettable.update.assert_not_awaited()
+
+
+def test_update_key_request_accepts_soft_budget():
+    request = UpdateKeyRequest(key="sk-test", soft_budget=42.0)
+    assert request.soft_budget == 42.0
+    assert "soft_budget" in request.model_fields_set
+
+
+@pytest.mark.parametrize("valid_value", [None, 0.0, 25.0])
+def test_validate_soft_budget_value_accepts_valid_values(valid_value):
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _validate_soft_budget_value,
+    )
+
+    assert _validate_soft_budget_value(valid_value) is None
+
+
+@pytest.mark.parametrize("invalid_value", [-5.0, float("nan"), float("inf")])
+def test_validate_soft_budget_value_rejects_invalid_values(invalid_value):
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _validate_soft_budget_value,
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        _validate_soft_budget_value(invalid_value)
+
+    assert exc_info.value.status_code == 400
+    assert "soft_budget must be a non-negative finite number" in str(exc_info.value.detail)
+
+
+@pytest.mark.asyncio
+async def test_apply_soft_budget_update_adds_budget_id_for_new_budget_row():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _apply_soft_budget_update,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None)
+    created_row = MagicMock()
+    created_row.budget_id = "budget-created-456"
+    mock_db = MagicMock()
+    mock_db.litellm_budgettable.create = AsyncMock(return_value=created_row)
+    mock_db.litellm_budgettable.update = AsyncMock()
+
+    result = await _apply_soft_budget_update(
+        data=UpdateKeyRequest(key="sk-test", soft_budget=25.0),
+        non_default_values={"soft_budget": 25.0},
+        db=mock_db,
+        existing_key_row=existing_key,
+        changed_by="user-1",
+    )
+
+    assert dict(result) == {"budget_id": "budget-created-456"}
+
+
+@pytest.mark.asyncio
+async def test_apply_soft_budget_update_keeps_existing_budget_id_out_of_token_update():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _apply_soft_budget_update,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id="budget-123")
+    mock_db = MagicMock()
+    mock_db.litellm_budgettable.update = AsyncMock()
+    mock_db.litellm_budgettable.create = AsyncMock()
+
+    result = await _apply_soft_budget_update(
+        data=UpdateKeyRequest(key="sk-test", soft_budget=40.0),
+        non_default_values={"soft_budget": 40.0, "max_budget": 100.0},
+        db=mock_db,
+        existing_key_row=existing_key,
+        changed_by="user-1",
+    )
+
+    assert dict(result) == {"max_budget": 100.0}
+    mock_db.litellm_budgettable.update.assert_awaited_once_with(
+        where={"budget_id": "budget-123"},
+        data={"soft_budget": 40.0, "updated_by": "user-1"},
+    )
+
+
+@pytest.mark.asyncio
+async def test_update_key_row_with_soft_budget_updates_budget_and_key_in_transaction():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _update_key_row_with_soft_budget,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None)
+    created_row = MagicMock(budget_id="budget-new")
+    updated_row = MagicMock()
+    updated_row.model_dump.return_value = {"token": "hashed", "budget_id": "budget-new"}
+    tx = MagicMock()
+    tx.litellm_budgettable.create = AsyncMock(return_value=created_row)
+    tx.litellm_verificationtoken.update = AsyncMock(return_value=updated_row)
+    tx_context = MagicMock()
+    tx_context.__aenter__ = AsyncMock(return_value=tx)
+    tx_context.__aexit__ = AsyncMock(return_value=None)
+    prisma_client = MagicMock()
+    prisma_client.tx.return_value = tx_context
+    prisma_client.jsonify_object = lambda data: dict(data)
+
+    result = await _update_key_row_with_soft_budget(
+        prisma_client=prisma_client,
+        key="sk-test",
+        data=UpdateKeyRequest(key="sk-test", soft_budget=25.0),
+        non_default_values={"soft_budget": 25.0},
+        existing_key_row=existing_key,
+        changed_by="user-1",
+    )
+
+    assert set(result) == {"token", "data"}
+    assert result["data"] == {"token": "hashed", "budget_id": "budget-new"}
+    tx.litellm_verificationtoken.update.assert_awaited_once()
+    update_call = tx.litellm_verificationtoken.update.await_args
+    assert update_call.kwargs["where"] == {"token": result["token"]}
+    assert update_call.kwargs["data"]["budget_id"] == "budget-new"
+    assert "soft_budget" not in update_call.kwargs["data"]
+
+
+@pytest.mark.asyncio
+async def test_update_key_row_with_soft_budget_propagates_transaction_error():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _update_key_row_with_soft_budget,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None)
+    created_row = MagicMock(budget_id="budget-new")
+    tx = MagicMock()
+    tx.litellm_budgettable.create = AsyncMock(return_value=created_row)
+    tx.litellm_verificationtoken.update = AsyncMock(side_effect=RuntimeError("update failed"))
+    tx_context = MagicMock()
+    tx_context.__aenter__ = AsyncMock(return_value=tx)
+    tx_context.__aexit__ = AsyncMock(return_value=None)
+    prisma_client = MagicMock()
+    prisma_client.tx.return_value = tx_context
+    prisma_client.jsonify_object = lambda data: dict(data)
+
+    with pytest.raises(RuntimeError, match="update failed"):
+        await _update_key_row_with_soft_budget(
+            prisma_client=prisma_client,
+            key="sk-test",
+            data=UpdateKeyRequest(key="sk-test", soft_budget=25.0),
+            non_default_values={"soft_budget": 25.0},
+            existing_key_row=existing_key,
+            changed_by="user-1",
+        )
+
+    tx_context.__aexit__.assert_awaited_once()
+    assert tx_context.__aexit__.await_args.args[0] is RuntimeError
+
+
 def test_generate_key_request_blank_team_id_is_personal():
     """The UI Team-field clear submits team_id=""; it must count as no team (LIT-3925)."""
     from litellm.proxy._types import RegenerateKeyRequest
diff --git a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx
index 1ab7b9be52c..2bbbc2bd48b 100644
--- a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx
+++ b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx
@@ -1,7 +1,11 @@
 import React from "react";
+import { Control } from "react-hook-form";
 import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
 import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 import { CircleHelp } from "lucide-react";
+import { FormField } from "@/components/shared/form/FormField";
+import NumericalInput from "../shared/numerical_input";
+import { KeyEditFormValues } from "./keyEditFormValues";
 
 export const labelWithHint = (label: React.ReactNode, hint: string): React.ReactNode => (
   <>
@@ -48,3 +52,27 @@ export const KeyTypeSelect = ({
     </SelectContent>
   </Select>
 );
+
+export const KeyBudgetNumberField = ({
+  control,
+  name,
+  label,
+  placeholder,
+}: {
+  control: Control<KeyEditFormValues>;
+  name: "max_budget" | "soft_budget";
+  label: string;
+  placeholder: string;
+}) => (
+  <FormField control={control} name={name} label={label}>
+    {({ ref: _ref, ...field }) => (
+      <NumericalInput
+        {...field}
+        value={field.value ?? ""}
+        step={0.01}
+        style={{ width: "100%" }}
+        placeholder={placeholder}
+      />
+    )}
+  </FormField>
+);
diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx
index 705679fe2e6..8a42ecff50c 100644
--- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx
+++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx
@@ -1,5 +1,10 @@
 import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { toast } from "@/lib/toast";
+
+vi.mock("@/lib/toast", () => ({
+  toast: { error: vi.fn(), success: vi.fn() },
+}));
 
 // ---- Hoisted shared mocks (safe to use inside vi.mock factories) ----
 const { keyUpdateCallMock, keyDeleteCallMock, mockUseAuthorized } = vi.hoisted(() => {
@@ -483,3 +488,104 @@ describe("KeyInfoView handleKeyUpdate empty strings", () => {
     });
   });
 });
+
+describe("KeyInfoView handleKeyUpdate soft_budget", () => {
+  const premiumAdminAuth = {
+    accessToken: "access_abc",
+    userId: "user_1",
+    userRole: "Admin",
+    premiumUser: true,
+    token: "token_123",
+    userEmail: "test@example.com",
+    disabledPersonalKeyCreation: false,
+    showSSOBanner: false,
+  };
+
+  const renderWithSoftBudget = (softBudget: number | null) => {
+    mockUseAuthorized.mockReturnValue(premiumAdminAuth);
+
+    return render(
+      <KeyInfoView
+        keyId="tok_123"
+        onClose={() => {}}
+        keyData={
+          { ...baseKeyData, litellm_budget_table: softBudget === null ? null : { soft_budget: softBudget } } as any
+        }
+        onKeyDataUpdate={() => {}}
+        teams={[]}
+      />,
+    );
+  };
+
+  it("should send a changed soft_budget as a number", async () => {
+    renderWithSoftBudget(null);
+
+    fireEvent.click(screen.getByText("Settings"));
+    fireEvent.click(screen.getByText("Edit Settings"));
+    (globalThis as any).__TEST_FORM_VALUES = {
+      token: "tok_123",
+      soft_budget: "25",
+    };
+
+    fireEvent.click(screen.getByText("Mock Submit"));
+
+    await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled());
+
+    const [, sentPayload] = keyUpdateCallMock.mock.calls[0];
+    expect(sentPayload.soft_budget).toBe(25);
+  });
+
+  it("should omit an unchanged soft_budget so unrelated edits skip the budget gate", async () => {
+    renderWithSoftBudget(25);
+
+    fireEvent.click(screen.getByText("Settings"));
+    fireEvent.click(screen.getByText("Edit Settings"));
+    (globalThis as any).__TEST_FORM_VALUES = {
+      token: "tok_123",
+      soft_budget: 25,
+      key_alias: "renamed",
+    };
+
+    fireEvent.click(screen.getByText("Mock Submit"));
+
+    await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled());
+
+    const [, sentPayload] = keyUpdateCallMock.mock.calls[0];
+    expect("soft_budget" in sentPayload).toBe(false);
+  });
+
+  it("should forward a cleared soft_budget as an explicit null the JSON body keeps", async () => {
+    renderWithSoftBudget(25);
+
+    fireEvent.click(screen.getByText("Settings"));
+    fireEvent.click(screen.getByText("Edit Settings"));
+    (globalThis as any).__TEST_FORM_VALUES = {
+      token: "tok_123",
+      soft_budget: "",
+    };
+
+    fireEvent.click(screen.getByText("Mock Submit"));
+
+    await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled());
+
+    const [, sentPayload] = keyUpdateCallMock.mock.calls[0];
+    expect(sentPayload.soft_budget).toBeNull();
+    expect(JSON.stringify({ ...sentPayload })).toContain('"soft_budget":null');
+  });
+
+  it("should reject an overflowing soft_budget instead of silently clearing it", async () => {
+    renderWithSoftBudget(25);
+
+    fireEvent.click(screen.getByText("Settings"));
+    fireEvent.click(screen.getByText("Edit Settings"));
+    (globalThis as any).__TEST_FORM_VALUES = {
+      token: "tok_123",
+      soft_budget: "1e309",
+    };
+
+    fireEvent.click(screen.getByText("Mock Submit"));
+
+    await waitFor(() => expect(toast.error).toHaveBeenCalled());
+    expect(keyUpdateCallMock).not.toHaveBeenCalled();
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts
index ad90610b732..f24fb6e5a86 100644
--- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts
+++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts
@@ -22,6 +22,7 @@ export interface KeyEditFormValues {
   models?: string[];
   allowed_routes?: string;
   max_budget?: number | string | null;
+  soft_budget?: number | string | null;
   budget_duration?: string | null;
   tpm_limit?: number | string | null;
   tpm_limit_type?: string | null;
@@ -67,6 +68,8 @@ export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues =>
   allowed_routes:
     Array.isArray(keyData.allowed_routes) && keyData.allowed_routes.length > 0 ? keyData.allowed_routes.join(", ") : "",
   max_budget: keyData.max_budget,
+  soft_budget:
+    (keyData.litellm_budget_table as { soft_budget?: number | null } | null | undefined)?.soft_budget ?? null,
   budget_duration: canonicalBudgetDuration(keyData.budget_duration),
   tpm_limit: keyData.tpm_limit,
   tpm_limit_type: (keyData as { tpm_limit_type?: string | null }).tpm_limit_type ?? null,
@@ -117,6 +120,7 @@ export const keyEditFormSchema = z.object({
   models: z.custom<string[] | undefined>(),
   allowed_routes: z.custom<string | undefined>(),
   max_budget: z.custom<number | string | null | undefined>(),
+  soft_budget: z.custom<number | string | null | undefined>(),
   budget_duration: z.custom<string | null | undefined>(),
   tpm_limit: z.custom<number | string | null | undefined>(),
   tpm_limit_type: z.custom<string | null | undefined>(),
@@ -168,6 +172,7 @@ export const toSubmittedValues = (
   models: values.models,
   allowed_routes: values.allowed_routes,
   max_budget: values.max_budget,
+  soft_budget: values.soft_budget,
   budget_duration: values.budget_duration,
   tpm_limit: values.tpm_limit,
   tpm_limit_type: values.tpm_limit_type,
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 4f732083c3e..35e9268e1c2 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
@@ -1885,6 +1885,7 @@ describe("KeyEditView", () => {
     key_alias: "asdasdas",
     models: [],
     max_budget: 0,
+    soft_budget: null,
     budget_duration: "30d",
     tpm_limit: 10,
     tpm_limit_type: null,
diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
index 0f955fee895..b1bbbeed6f9 100644
--- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
+++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
@@ -32,7 +32,7 @@ import {
   modelSentinelOptions,
   parseAllowedRoutes,
 } from "./keyEditFieldNormalizers";
-import { KeyTypeSelect, labelWithHint } from "./KeyEditViewControls";
+import { KeyBudgetNumberField, KeyTypeSelect, labelWithHint } from "./KeyEditViewControls";
 import {
   AgentsAndGroups,
   KeyEditFormValues,
@@ -417,17 +417,19 @@ export function KeyEditView({
             )}
           </FormField>
 
-          <FormField control={form.control} name="max_budget" label="Max Budget (USD)">
-            {({ ref: _ref, ...field }) => (
-              <NumericalInput
-                {...field}
-                value={field.value ?? ""}
-                step={0.01}
-                style={{ width: "100%" }}
-                placeholder="Enter a numerical value"
-              />
-            )}
-          </FormField>
+          <KeyBudgetNumberField
+            control={form.control}
+            name="max_budget"
+            label="Max Budget (USD)"
+            placeholder="Enter a numerical value"
+          />
+
+          <KeyBudgetNumberField
+            control={form.control}
+            name="soft_budget"
+            label="Soft Budget (USD)"
+            placeholder="Get alerts when spend crosses this value, without blocking requests"
+          />
 
           <FormField control={form.control} name="budget_duration" label="Reset Budget">
             {({ value, onChange, id }) => (
diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx
index f5c682a2ee0..969decb6613 100644
--- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx
+++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx
@@ -218,6 +218,23 @@ export default function KeyInfoView({
       // Handle max budget empty string
       formValues.max_budget = mapEmptyStringToNull(formValues.max_budget);
 
+      // soft_budget is a budget change server-side (admin-gated); only send it when it changed
+      // so a non-admin edit of unrelated fields isn't blocked by that gate.
+      const previousSoftBudget =
+        (currentKeyData.litellm_budget_table as { soft_budget?: number | null } | null | undefined)?.soft_budget ??
+        null;
+      const nextSoftBudget =
+        formValues.soft_budget === "" || formValues.soft_budget == null ? null : Number(formValues.soft_budget);
+      if (nextSoftBudget !== null && !Number.isFinite(nextSoftBudget)) {
+        toast.error("Soft Budget must be a finite number");
+        return;
+      }
+      if (nextSoftBudget === previousSoftBudget) {
+        delete formValues.soft_budget;
+      } else {
+        formValues.soft_budget = nextSoftBudget;
+      }
+
       // Handle object_permission updates
       if (formValues.vector_stores !== undefined) {
         formValues.object_permission = {
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 12e7eb801e6..9de988fde4c 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -8039,7 +8039,7 @@ export interface paths {
          *     - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
          *     - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
          *     - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
-         *     - soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
+         *     - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget.
          *     - max_parallel_requests: Optional[int] - Rate limit for parallel requests
          *     - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"}
          *     - tpm_limit: Optional[int] - Tokens per minute limit
@@ -37761,6 +37761,8 @@ export interface components {
             rpm_limit?: number | null;
             /** Rpm Limit Type */
             rpm_limit_type?: ("guaranteed_throughput" | "best_effort_throughput" | "dynamic") | null;
+            /** Soft Budget */
+            soft_budget?: number | null;
             /** Spend */
             spend?: number | null;
             /** Tag Rpm Limit */

From edeb93e727135b43bd5c50c25630279ea34a1502 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Mon, 7 Sep 2026 12:50:17 -0700
Subject: [PATCH 135/164] ci: prepare workflows for main default branch

---
 .github/workflows/publish-basedpyright-base-counts.yml | 6 +++---
 .github/workflows/sync-together-ai-models.yml          | 6 ++++--
 .github/workflows/test-litellm-ui-unit.yml             | 1 +
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/publish-basedpyright-base-counts.yml b/.github/workflows/publish-basedpyright-base-counts.yml
index cd443a8e9db..27d4682dbd9 100644
--- a/.github/workflows/publish-basedpyright-base-counts.yml
+++ b/.github/workflows/publish-basedpyright-base-counts.yml
@@ -1,6 +1,6 @@
 name: Publish basedpyright base counts
 
-# Every commit on litellm_internal_staging is some branch's future merge-base.
+# Every commit on main or litellm_internal_staging can become a future merge-base.
 # Publishing its per-rule basedpyright counts as an artifact lets
 # scripts/type_check_gate.py download them in seconds instead of paying a
 # 60-110s second basedpyright pass on every fresh worktree or moved merge-base.
@@ -10,13 +10,13 @@ name: Publish basedpyright base counts
 on:
   push:
     branches:
+      - main
       - litellm_internal_staging
   workflow_dispatch:
     inputs:
       ref:
-        description: "Ref to compute and publish base counts for"
+        description: "Ref to compute and publish base counts for (defaults to the workflow run's commit)"
         required: false
-        default: litellm_internal_staging
 
 permissions:
   contents: read
diff --git a/.github/workflows/sync-together-ai-models.yml b/.github/workflows/sync-together-ai-models.yml
index 1daaadeabe2..18d5e3de1eb 100644
--- a/.github/workflows/sync-together-ai-models.yml
+++ b/.github/workflows/sync-together-ai-models.yml
@@ -13,10 +13,12 @@ jobs:
   sync_together_ai_models:
     if: github.repository == 'BerriAI/litellm'
     runs-on: ubuntu-latest
+    env:
+      BASE_BRANCH: ${{ github.event.repository.default_branch }}
     steps:
       - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
         with:
-          ref: litellm_internal_staging
+          ref: ${{ env.BASE_BRANCH }}
           persist-credentials: false
       - name: Set up uv
         uses: ./.github/actions/setup-uv-with-retries
@@ -63,6 +65,6 @@ jobs:
           gh pr create --title "feat(models): sync together_ai model registry" \
             --body-file "$RUNNER_TEMP/pr_body.md" \
             --head "$branch" \
-            --base litellm_internal_staging
+            --base "$BASE_BRANCH"
         env:
           GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}
diff --git a/.github/workflows/test-litellm-ui-unit.yml b/.github/workflows/test-litellm-ui-unit.yml
index 314efcc49d5..cd58f861a87 100644
--- a/.github/workflows/test-litellm-ui-unit.yml
+++ b/.github/workflows/test-litellm-ui-unit.yml
@@ -12,6 +12,7 @@ on:
       - "litellm_**"
   push:
     branches:
+      - main
       - litellm_internal_staging
 
 concurrency:

From e04e5d71138ea1356aacea1ea22ef060855dd596 Mon Sep 17 00:00:00 2001
From: yucheng-berri <yucheng@berri.ai>
Date: Mon, 7 Sep 2026 12:59:57 -0700
Subject: [PATCH 136/164] fix(router): keep provider response headers on
 streaming chat completions (#40091)

* fix(router): keep provider response headers on streaming chat completions

The Router re-wraps a deployment's CustomStreamWrapper in FallbackStreamWrapper
(and its sync twin) so a mid-stream failure can fail over. Neither wrapper
forwarded `_response_headers`, so every streaming chat completion handed the
proxy's callbacks and its response-header builder a wrapper with no provider
headers, and a successful mid-stream fallback still published the failed
deployment's identity, `x-request-id` and rate limit counters.

Forward `_response_headers` into both wrappers, repoint the wrapper at the
deployment that served the stream once a fallback takes over, and rebuild the
proxy's response headers from that deployment while `create_response` still has
the first chunk buffered.

* fix(router): follow a nested fallback to the deployment that served the stream

A fallback the router picks is itself a fallback-aware wrapper, and it only
repoints at its own fallback once it yields, so reading its hidden params at
selection time named a deployment that produced no output. Re-read them when
the first fallback item arrives, which is still before the proxy commits
response headers.

Also addresses review feedback: the streaming header builder reads self.data
instead of taking a coarse request_data parameter, and the new test recorder
local is Final.

* test(router): cover the fallback header adoption helper directly

The router_code_coverage gate wants every router.py function named in a
router test, and this also pins the weak-reference behavior: a wrapper
collected mid-stream must not break the generator still draining it.

* refactor(proxy): take a read-only mapping for the model-id lookup

_get_model_id_from_response only reads its request payload, so a Mapping
says what it needs and the two metadata hops are narrowed instead of
assumed to be dicts.

* test: drop mutable recorder locals and routine comments from the new tests

An AsyncMock await_count and an asyncio.Event say the same thing as a
list and a dict that the test mutates.

* chore(router): justify the two rebinds in the fallback loops

Both are the one-shot re-read that follows a nested fallback, so they get
the repo's rebind-ok note like the rest of the file.
---
 litellm/proxy/common_request_processing.py    | 126 ++--
 .../pass_through_endpoints.py                 |   2 +-
 litellm/router.py                             |  83 ++-
 .../proxy/test_common_request_processing.py   | 247 +++++++-
 tests/test_litellm/test_router.py             | 574 ++++++++++++++++++
 5 files changed, 985 insertions(+), 47 deletions(-)

diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py
index 5e6c9b34332..9720e4b1cf8 100644
--- a/litellm/proxy/common_request_processing.py
+++ b/litellm/proxy/common_request_processing.py
@@ -756,7 +756,7 @@ class _UpstreamClosingStreamingResponse(StreamingResponse):
         content: AsyncGenerator[str, None],
         *,
         media_type: str | None = None,
-        headers: dict | None = None,
+        headers: Mapping[str, str] | None = None,
         status_code: int = status.HTTP_200_OK,
         upstream_generator: AsyncGenerator[str, None] | None = None,
     ) -> None:
@@ -888,25 +888,39 @@ def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]:
     return f"data: {json.dumps({'error': error_obj})}\n\n", "data: [DONE]\n\n"
 
 
+def _sse_stream_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
+    """`headers` plus the two that stop reverse proxies from buffering SSE (issue #28384)."""
+    return MappingProxyType({**headers, **_TTFT_KEEPALIVE_HEADERS})
+
+
+async def _resolve_stream_headers(
+    headers: Mapping[str, str], refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None
+) -> Mapping[str, str]:
+    if refresh_headers is None:
+        return headers
+    try:
+        return await refresh_headers()
+    except Exception as e:  # noqa: BLE001  # a stream whose first chunk is already paid for must not fail over its headers
+        verbose_proxy_logger.exception("Error refreshing streaming response headers: %s", e)
+        return headers
+
+
 async def create_response(
     generator: AsyncGenerator[str, None],
     media_type: str,
-    headers: dict,
+    headers: Mapping[str, str],
     default_status_code: int = status.HTTP_200_OK,
     request: Request | None = None,
+    refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None = None,
 ) -> StreamingResponse | JSONResponse:
     """
     Create streaming response, checking if the first chunk is an error.
     If the first chunk is an error, return a standard JSON error response.
     Otherwise, return StreamingResponse and stream all content.
+
+    ``refresh_headers`` is consulted once the first chunk has been buffered, for
+    callers whose headers can only be known then.
     """
-    # Tell buffering reverse proxies (nginx, ingress-nginx, Envoy) to flush SSE
-    # immediately instead of releasing the whole stream in one batch (issue #28384).
-    streaming_headers: Final = {
-        **headers,
-        "Cache-Control": "no-cache",
-        "X-Accel-Buffering": "no",
-    }
     first_chunk_value: str | None = None
     final_status_code = default_status_code
 
@@ -917,6 +931,7 @@ async def create_response(
 
         # Now get the first chunk from the actual generator
         first_chunk_value = await _buffer_first_chunk_honoring_disconnect(generator, request)
+        resolved_headers: Final = await _resolve_stream_headers(headers, refresh_headers)
 
         if first_chunk_value is not None:
             try:
@@ -943,7 +958,7 @@ async def create_response(
                     return JSONResponse(
                         status_code=final_status_code,
                         content={"error": error_dict},
-                        headers=headers,
+                        headers=resolved_headers,
                     )
             except Exception as e:
                 verbose_proxy_logger.debug("Error parsing first chunk value: %s", e)
@@ -972,7 +987,7 @@ async def create_response(
         return StreamingResponse(
             empty_gen(),
             media_type=media_type,
-            headers=streaming_headers,
+            headers=_sse_stream_headers(await _resolve_stream_headers(headers, refresh_headers)),
             status_code=default_status_code,
         )
     except Exception as e:
@@ -988,7 +1003,7 @@ async def create_response(
         return StreamingResponse(
             error_gen_message(),
             media_type=media_type,
-            headers=streaming_headers,
+            headers=_sse_stream_headers(await _resolve_stream_headers(headers, refresh_headers)),
             status_code=error_status,
         )
 
@@ -1010,7 +1025,7 @@ async def create_response(
     return _UpstreamClosingStreamingResponse(
         combined_generator(),
         media_type=media_type,
-        headers=streaming_headers,
+        headers=_sse_stream_headers(resolved_headers),
         status_code=final_status_code,
         upstream_generator=generator,
     )
@@ -1535,7 +1550,7 @@ class ProxyBaseLLMRequestProcessing:
     @staticmethod
     def _merge_passthrough_streaming_headers(
         response_headers: httpx.Headers | dict | None,
-        custom_headers: dict,
+        custom_headers: Mapping[str, str],
     ) -> dict:
         """
         Merge upstream passthrough headers with proxy/custom headers.
@@ -2143,14 +2158,45 @@ class ProxyBaseLLMRequestProcessing:
         return fallback_model_group
 
     @staticmethod
-    def _get_model_id_from_response(hidden_params: dict, data: dict) -> str:
+    def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str:
         """Extract model_id from hidden_params with fallback to litellm_metadata."""
         model_id = hidden_params.get("model_id", None) or ""
         if not model_id:
-            litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
-            model_info: Final = litellm_metadata.get("model_info", {}) or {}
-            model_id = model_info.get("id", "") or ""
-        return model_id
+            litellm_metadata: Final = data.get("litellm_metadata")
+            model_info: Final = litellm_metadata.get("model_info") if isinstance(litellm_metadata, Mapping) else None
+            model_id = (model_info.get("id") or "") if isinstance(model_info, Mapping) else ""
+        return str(model_id) if model_id else ""
+
+    def _stream_response_headers(
+        self,
+        *,
+        hidden_params: Mapping[str, object],
+        user_api_key_dict: UserAPIKeyAuth,
+        logging_obj: LiteLLMLoggingObj,
+        version: str | None,
+        callback_headers: Mapping[str, str],
+    ) -> Mapping[str, str]:
+        """The streaming response headers describing `hidden_params`' deployment."""
+        return MappingProxyType(
+            {
+                **ProxyBaseLLMRequestProcessing.get_custom_headers(
+                    user_api_key_dict=user_api_key_dict,
+                    call_id=logging_obj.litellm_call_id,
+                    model_id=self._get_model_id_from_response(hidden_params, self.data),
+                    cache_key=hidden_params.get("cache_key") or "",
+                    api_base=hidden_params.get("api_base") or "",
+                    version=version,
+                    response_cost=hidden_params.get("response_cost") or "",
+                    model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
+                    fastest_response_batch_completion=hidden_params.get("fastest_response_batch_completion"),
+                    request_data=self.data,
+                    hidden_params=hidden_params,
+                    litellm_logging_obj=logging_obj,
+                    **(hidden_params.get("additional_headers") or MappingProxyType({})),
+                ),
+                **callback_headers,
+            }
+        )
 
     @staticmethod
     def _get_deployment_model_name(
@@ -2419,31 +2465,32 @@ class ProxyBaseLLMRequestProcessing:
             if self._is_streaming_request(
                 data=self.data, is_streaming_request=is_streaming_request
             ) or self._is_streaming_response(response):  # use generate_responses to stream responses
-                custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
-                    user_api_key_dict=user_api_key_dict,
-                    call_id=logging_obj.litellm_call_id,
-                    model_id=model_id,
-                    cache_key=cache_key,
-                    api_base=api_base,
-                    version=version,
-                    response_cost=response_cost,
-                    model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
-                    fastest_response_batch_completion=fastest_response_batch_completion,
-                    request_data=self.data,
-                    hidden_params=hidden_params,
-                    litellm_logging_obj=logging_obj,
-                    **additional_headers,
-                )
-
                 # Call response headers hook for streaming success
-                callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
+                stream_callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook(
                     data=self.data,
                     user_api_key_dict=user_api_key_dict,
                     response=response,
                     request_headers=dict(request.headers),
                 )
-                if callback_headers:
-                    custom_headers.update(callback_headers)
+                custom_headers: Final = self._stream_response_headers(
+                    hidden_params=hidden_params,
+                    user_api_key_dict=user_api_key_dict,
+                    logging_obj=logging_obj,
+                    version=version,
+                    callback_headers=stream_callback_headers or MappingProxyType({}),
+                )
+
+                async def refresh_stream_headers() -> Mapping[str, str]:
+                    """`custom_headers` rebuilt for whichever deployment served the stream."""
+                    if not getattr(response, "fallback_headers_adopted", False):
+                        return custom_headers
+                    return self._stream_response_headers(
+                        hidden_params=get_hidden_params_dict(response),
+                        user_api_key_dict=user_api_key_dict,
+                        logging_obj=logging_obj,
+                        version=version,
+                        callback_headers=stream_callback_headers or MappingProxyType({}),
+                    )
 
                 # Preserve the original client-requested model (pre-alias mapping) for downstream
                 # streaming generators. Pre-call processing can rewrite `self.data["model"]` for
@@ -2581,6 +2628,7 @@ class ProxyBaseLLMRequestProcessing:
                         media_type="text/event-stream",
                         headers=custom_headers,
                         request=request,
+                        refresh_headers=refresh_stream_headers,
                     )
 
             ### CALL HOOKS ### - modify outgoing data
@@ -3032,7 +3080,7 @@ class ProxyBaseLLMRequestProcessing:
         response: Any,
         proxy_logging_obj: "ProxyLogging",
         user_api_key_dict: "UserAPIKeyAuth",
-        custom_headers: dict,
+        custom_headers: Mapping[str, str],
         request_headers: dict[str, str],
     ) -> Response | None:
         if not self._has_post_call_guardrails_for_passthrough():
diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
index 3cb9acc6110..f1f823e59b5 100644
--- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
@@ -322,7 +322,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
     def get_response_headers(
         headers: httpx.Headers,
         litellm_call_id: str | None = None,
-        custom_headers: dict | None = None,
+        custom_headers: Mapping[str, str] | None = None,
     ) -> dict:
         # Exclude headers that uvicorn writes itself (server, date) and
         # encoding/length headers that don't survive re-serialization.
diff --git a/litellm/router.py b/litellm/router.py
index c7a254cbc02..95cabfad4bd 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -616,6 +616,38 @@ RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset(
 RETRY_BREADCRUMB_LIMIT: Final = 4
 
 
+class FallbackAwareStreamWrapper(CustomStreamWrapper):
+    """Base for the Router's chat-completion stream wrappers, which are built around the
+    attempt the Router picked first and have to repoint themselves when a fallback takes over."""
+
+    fallback_headers_adopted: bool = False
+
+    def adopt_fallback_response_headers(
+        self,
+        fallback_response: object,
+        prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]],
+    ) -> None:
+        """Repoint this wrapper at the deployment that served the stream.
+
+        Replaces rather than merges, so the failed attempt's `x-request-id`, rate limit
+        counters, `model_id` and `api_base` cannot reach the proxy's response headers or
+        its callbacks.
+        """
+        self._response_headers = getattr(fallback_response, "_response_headers", None)
+        fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params
+        if fallback_hidden_params:
+            self._hidden_params = {  # mutable-ok: the rest of litellm writes into _hidden_params
+                **fallback_hidden_params,
+                # dict() because add_retry_fallback_headers mutates additional_headers in place
+                "additional_headers": dict(fallback_headers),  # mutable-ok: see above
+            }
+            self._base_hidden_params = {  # mutable-ok: CustomStreamWrapper keeps this snapshot as a dict
+                **self._hidden_params,
+                "response_cost": None,
+            }
+        self.fallback_headers_adopted = True
+
+
 class Router:
     model_names: set = set()
     cache_responses: bool | None = False
@@ -2576,6 +2608,18 @@ class Router:
             return fallback_hidden_params, {}
         return fallback_hidden_params, cast("dict[str, object]", fallback_headers)
 
+    @staticmethod
+    def _adopt_fallback_response_headers(
+        wrapper_ref: "weakref.ref[FallbackAwareStreamWrapper]",
+        fallback_response: object,
+    ) -> tuple[dict[str, object], dict[str, object]]:
+        """Repoint the wrapper at `fallback_response`, returning its prepared hidden params."""
+        prepared: Final = Router._prepare_fallback_hidden_params(fallback_response)
+        adopting_wrapper: Final = wrapper_ref()
+        if adopting_wrapper is not None:
+            adopting_wrapper.adopt_fallback_response_headers(fallback_response, prepared)
+        return prepared
+
     @staticmethod
     def _apply_fallback_hidden_params_to_item(
         fallback_item: object,
@@ -2615,7 +2659,7 @@ class Router:
 
         held_slot: Final = deployment_slot if deployment_slot is not None else contextlib.AsyncExitStack()
 
-        class FallbackStreamWrapper(CustomStreamWrapper):
+        class FallbackStreamWrapper(FallbackAwareStreamWrapper):
             def __init__(self, async_generator: AsyncGenerator):
                 # Copy attributes from the original model_response
                 super().__init__(
@@ -2623,6 +2667,7 @@ class Router:
                     model=model_response.model,
                     custom_llm_provider=model_response.custom_llm_provider,
                     logging_obj=model_response.logging_obj,
+                    _response_headers=getattr(model_response, "_response_headers", None),
                 )
                 self._async_generator = async_generator
                 inner_chunks: Final[object] = getattr(model_response, "chunks", None)
@@ -2699,8 +2744,17 @@ class Router:
 
                     # If fallback returns a streaming response, iterate over it
                     if hasattr(fallback_response, "__aiter__"):
-                        prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
+                        prepared_fallback_hidden_params = Router._adopt_fallback_response_headers(
+                            wrapper_ref, fallback_response
+                        )
+                        fallback_headers_are_settled = False
                         async for fallback_item in fallback_response:
+                            if not fallback_headers_are_settled:
+                                fallback_headers_are_settled = True  # rebind-ok: one-shot latch
+                                # a fallback that failed over again only repoints itself once it yields
+                                prepared_fallback_hidden_params = (  # rebind-ok: re-read once the fallback yields
+                                    Router._adopt_fallback_response_headers(wrapper_ref, fallback_response)
+                                )
                             Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params)
                             if (
                                 fallback_item
@@ -2742,7 +2796,11 @@ class Router:
                                 e,
                             )
 
-        return FallbackStreamWrapper(stream_with_fallbacks())
+        wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks())
+        # weak, so the generator closing over it does not keep the wrapper out of
+        # refcount teardown and delay the `finally` that releases the deployment slot
+        wrapper_ref: Final = weakref.ref(wrapped_response)
+        return wrapped_response
 
     @staticmethod
     def _extract_partial_responses_usage(
@@ -3171,13 +3229,14 @@ class Router:
         """
         from litellm.exceptions import MidStreamFallbackError
 
-        class SyncFallbackStreamWrapper(CustomStreamWrapper):
+        class SyncFallbackStreamWrapper(FallbackAwareStreamWrapper):
             def __init__(self, sync_generator: Generator):
                 super().__init__(
                     completion_stream=sync_generator,
                     model=model_response.model,
                     custom_llm_provider=model_response.custom_llm_provider,
                     logging_obj=model_response.logging_obj,
+                    _response_headers=getattr(model_response, "_response_headers", None),
                 )
                 self._sync_generator = sync_generator
                 if hasattr(model_response, "_hidden_params"):
@@ -3233,8 +3292,17 @@ class Router:
                     )
 
                     if hasattr(fallback_response, "__iter__"):
-                        prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
+                        prepared_fallback_hidden_params = Router._adopt_fallback_response_headers(
+                            wrapper_ref, fallback_response
+                        )
+                        fallback_headers_are_settled = False
                         for fallback_item in fallback_response:
+                            if not fallback_headers_are_settled:
+                                fallback_headers_are_settled = True  # rebind-ok: one-shot latch
+                                # a fallback that failed over again only repoints itself once it yields
+                                prepared_fallback_hidden_params = (  # rebind-ok: re-read once the fallback yields
+                                    Router._adopt_fallback_response_headers(wrapper_ref, fallback_response)
+                                )
                             Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params)
                             if (
                                 fallback_item
@@ -3272,7 +3340,10 @@ class Router:
                             close_err,
                         )
 
-        return SyncFallbackStreamWrapper(stream_with_fallbacks())
+        wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks())
+        # weak, for the same reason as the async twin
+        wrapper_ref: Final = weakref.ref(wrapped_response)
+        return wrapped_response
 
     async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs):
         """
diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py
index 96d9b0c5a26..6acd9d7258e 100644
--- a/tests/test_litellm/proxy/test_common_request_processing.py
+++ b/tests/test_litellm/proxy/test_common_request_processing.py
@@ -2,7 +2,7 @@ import asyncio
 import copy
 import datetime
 import json
-from types import SimpleNamespace
+from types import MappingProxyType, SimpleNamespace
 from typing import AsyncGenerator, Callable, Final, Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -1809,6 +1809,146 @@ class TestCommonRequestProcessingHelpers:
         response = await create_response(mock_generator(), "text/event-stream", custom_headers)
         assert response.headers["x-custom-header"] == "TestValue"
 
+    async def test_create_streaming_response_refresh_headers_after_first_chunk(self):
+        """LIT-6767: headers a caller can only resolve once the first chunk exists.
+
+        A pre-first-chunk fallback replaces the deployment while the response
+        headers are still uncommitted, so ``refresh_headers`` is consulted after
+        the first chunk is buffered and its result wins.
+        """
+
+        async def mock_generator():
+            yield 'data: {"content": "data"}\n\n'
+            yield "data: [DONE]\n\n"
+
+        refresh_headers: Final = AsyncMock(
+            return_value={"x-litellm-model-id": "fallback-deployment", "llm_provider-x-request-id": "req-FALLBACK"}
+        )
+
+        response = await create_response(
+            mock_generator(),
+            "text/event-stream",
+            {"x-litellm-model-id": "failed-deployment", "llm_provider-x-request-id": "req-FAILED"},
+            refresh_headers=refresh_headers,
+        )
+        assert isinstance(response, StreamingResponse)
+        assert refresh_headers.await_count == 1
+        assert response.headers["x-litellm-model-id"] == "fallback-deployment"
+        assert response.headers["llm_provider-x-request-id"] == "req-FALLBACK"
+        # the buffering headers are still applied on top of the refreshed set
+        assert response.headers["x-accel-buffering"] == "no"
+        assert response.headers["cache-control"] == "no-cache"
+
+    async def test_create_streaming_response_refreshes_only_after_the_first_chunk(self):
+        """LIT-6767: the refresh has to be consulted after the generator produced a chunk.
+
+        A pre-first-chunk fallback only repoints the response while that first chunk is
+        being produced, so a refresh consulted any earlier still describes the attempt
+        that failed and the headers go out wrong.
+        """
+        first_chunk_produced: Final = asyncio.Event()
+
+        async def mock_generator():
+            first_chunk_produced.set()
+            yield 'data: {"content": "data"}\n\n'
+            yield "data: [DONE]\n\n"
+
+        async def refresh_headers():
+            served = "fallback-deployment" if first_chunk_produced.is_set() else "failed-deployment"
+            return {"x-litellm-model-id": served}
+
+        response = await create_response(
+            mock_generator(),
+            "text/event-stream",
+            {"x-litellm-model-id": "failed-deployment"},
+            refresh_headers=refresh_headers,
+        )
+        assert response.headers["x-litellm-model-id"] == "fallback-deployment"
+
+    async def test_create_streaming_response_empty_stream_uses_refreshed_headers(self):
+        """LIT-6767: a fallback that served nothing still gets to name itself.
+
+        The empty-generator branch returns its own StreamingResponse, so it needs the
+        refreshed headers too or the client is told the failed deployment answered.
+        """
+
+        async def mock_generator():
+            return
+            yield  # make it an async generator
+
+        async def refresh_headers():
+            return {"x-litellm-model-id": "fallback-deployment"}
+
+        response = await create_response(
+            mock_generator(),
+            "text/event-stream",
+            {"x-litellm-model-id": "failed-deployment"},
+            refresh_headers=refresh_headers,
+        )
+        assert isinstance(response, StreamingResponse)
+        assert response.headers["x-litellm-model-id"] == "fallback-deployment"
+        assert response.headers["x-accel-buffering"] == "no"
+
+    async def test_create_streaming_response_without_refresh_headers_is_unchanged(self):
+        """LIT-6767: the default keeps the caller-supplied headers verbatim."""
+
+        async def mock_generator():
+            yield 'data: {"content": "data"}\n\n'
+            yield "data: [DONE]\n\n"
+
+        response = await create_response(
+            mock_generator(),
+            "text/event-stream",
+            {"x-litellm-model-id": "failed-deployment"},
+        )
+        assert response.headers["x-litellm-model-id"] == "failed-deployment"
+
+    async def test_create_streaming_response_refresh_headers_failure_keeps_stream(self):
+        """LIT-6767: the first chunk is already paid for, so a failing refresh
+        falls back to the caller's headers instead of erroring the stream."""
+
+        async def mock_generator():
+            yield 'data: {"content": "data"}\n\n'
+            yield "data: [DONE]\n\n"
+
+        async def refresh_headers():
+            raise RuntimeError("boom")
+
+        response = await create_response(
+            mock_generator(),
+            "text/event-stream",
+            {"x-litellm-model-id": "failed-deployment"},
+            refresh_headers=refresh_headers,
+        )
+        assert isinstance(response, StreamingResponse)
+        assert response.status_code == status.HTTP_200_OK
+        assert response.headers["x-litellm-model-id"] == "failed-deployment"
+        assert await self.consume_stream(response) == [
+            'data: {"content": "data"}\n\n',
+            "data: [DONE]\n\n",
+        ]
+
+    async def test_create_response_first_chunk_error_uses_refreshed_headers(self):
+        """LIT-6767: the JSON error response built from a bad first chunk carries
+        the refreshed headers too, so it cannot describe a deployment that no
+        longer served the request."""
+
+        async def mock_generator():
+            yield 'data: {"error": {"code": 403, "message": "forbidden"}}\n\n'
+            yield "data: [DONE]\n\n"
+
+        async def refresh_headers():
+            return {"x-litellm-model-id": "fallback-deployment"}
+
+        response = await create_response(
+            mock_generator(),
+            "text/event-stream",
+            {"x-litellm-model-id": "failed-deployment"},
+            refresh_headers=refresh_headers,
+        )
+        assert isinstance(response, JSONResponse)
+        assert response.headers["x-litellm-model-id"] == "fallback-deployment"
+
     async def test_create_streaming_response_disables_proxy_buffering(self):
         """Regression for #28384: every StreamingResponse create_response returns
         must carry the headers that stop nginx/ingress/Envoy from buffering the
@@ -8014,3 +8154,108 @@ class TestDetachedStreamFailureHook:
         await logging_obj._on_detached_stream_failure(failure)
 
         assert [call["original_exception"] for call in recorder.calls] == [failure]
+
+
+class TestStreamingResponseHeadersFollowFallback:
+    """LIT-6767: the streaming branch has to publish the deployment that served the stream."""
+
+    @staticmethod
+    def _fallback_adopting_stream():
+        class _Stream:
+            def __init__(self) -> None:
+                self._hidden_params = {
+                    "model_id": "failed-deployment",
+                    "api_base": "http://127.0.0.1:20769/v1",
+                    "additional_headers": {"llm_provider-stale-marker": "failed-deployment"},
+                }
+                self.fallback_headers_adopted = False
+
+            def adopt(self) -> None:
+                self._hidden_params = {
+                    "model_id": "served-deployment",
+                    "api_base": "https://api.openai.com",
+                    "additional_headers": {"llm_provider-x-request-id": "req-SERVED"},
+                }
+                self.fallback_headers_adopted = True
+
+        return _Stream()
+
+    @pytest.mark.asyncio
+    async def test_streaming_headers_name_the_deployment_that_served(self, monkeypatch):
+        """A pre-first-chunk fallback repoints the stream while the headers are still
+        uncommitted, so the published headers must describe the fallback, not the attempt
+        the Router picked first."""
+        stream = self._fallback_adopting_stream()
+
+        def select_data_generator(**kwargs):
+            async def generator():
+                stream.adopt()
+                yield 'data: {"choices": [{"delta": {"content": "OK"}}]}\n\n'
+                yield "data: [DONE]\n\n"
+
+            return generator()
+
+        logging_obj = MagicMock()
+        logging_obj.litellm_call_id = "lit-6767-call"
+        logging_obj._defer_async_logging = False
+        logging_obj._on_deferred_stream_complete = None
+        logging_obj.cost_breakdown = None
+
+        processor = ProxyBaseLLMRequestProcessing(
+            data={"model": "oa-midfail", "stream": True, "litellm_logging_obj": logging_obj}
+        )
+
+        proxy_logging_obj = MagicMock(spec=ProxyLogging)
+        proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
+        proxy_logging_obj.update_request_status = AsyncMock(return_value=None)
+        proxy_logging_obj.post_call_success_hook = AsyncMock(
+            side_effect=lambda data, user_api_key_dict, response: response
+        )
+        proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
+            return_value={"x-callback-header": "kept"}
+        )
+
+        async def fake_route_request(**kwargs):
+            async def call():
+                return stream
+
+            return call()
+
+        monkeypatch.setattr(
+            litellm.proxy.common_request_processing, "route_request", fake_route_request
+        )
+
+        result = await processor.base_process_llm_request(
+            request=Request(scope={"type": "http", "headers": []}),
+            fastapi_response=Response(),
+            user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"),
+            route_type="acompletion",
+            proxy_logging_obj=proxy_logging_obj,
+            general_settings={},
+            proxy_config=MagicMock(spec=ProxyConfig),
+            select_data_generator=select_data_generator,
+            is_streaming_request=True,
+            skip_pre_call_logic=True,
+        )
+
+        assert isinstance(result, StreamingResponse)
+        assert result.headers["x-litellm-model-id"] == "served-deployment"
+        assert result.headers["x-litellm-model-api-base"] == "https://api.openai.com"
+        assert result.headers["llm_provider-x-request-id"] == "req-SERVED"
+        assert "llm_provider-stale-marker" not in result.headers
+        assert result.headers["x-callback-header"] == "kept"
+
+
+class TestPassthroughHeadersAcceptImmutableMappings:
+    """LIT-6767: the streaming branch now hands the passthrough helpers an immutable mapping."""
+
+    def test_merge_passthrough_streaming_headers_accepts_a_read_only_mapping(self):
+        merged = ProxyBaseLLMRequestProcessing._merge_passthrough_streaming_headers(
+            response_headers=httpx.Headers({"content-type": "text/event-stream", "transfer-encoding": "chunked"}),
+            custom_headers=MappingProxyType({"x-litellm-model-id": "served-deployment"}),
+        )
+
+        assert merged["x-litellm-model-id"] == "served-deployment"
+        assert merged["content-type"] == "text/event-stream"
+        # the excluded hop-by-hop header is still dropped
+        assert "transfer-encoding" not in merged
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index e9b77076448..eed34c79a06 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -2387,6 +2387,580 @@ async def test_acompletion_streaming_iterator_preserves_hidden_params():
     assert result._hidden_params.get("_response_ms") == 500.0
 
 
+@pytest.mark.asyncio
+async def test_acompletion_streaming_iterator_preserves_response_headers():
+    """LIT-6767: the returned wrapper must carry the provider's raw response headers.
+
+    Proxy callbacks read ``_response_headers`` off the object the router hands
+    back. The wrapper used to be built without it, so every streaming chat
+    completion reported zero raw provider headers while the non-streaming path
+    reported the full set.
+    """
+    from unittest.mock import MagicMock
+
+    from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gpt-4",
+                "litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
+            }
+        ],
+    )
+
+    async def _empty():
+        return
+        yield  # make it an async generator
+
+    provider_headers = {
+        "x-request-id": "req-provider-123",
+        "x-ratelimit-remaining-requests": "42",
+        # a provider must never be able to spoof an internal header
+        "x-litellm-model-id": "spoofed",
+    }
+    source = CustomStreamWrapper(
+        completion_stream=_empty(),
+        model="gpt-4",
+        custom_llm_provider="openai",
+        logging_obj=MagicMock(),
+        _response_headers=provider_headers,
+    )
+
+    result = await router._acompletion_streaming_iterator(
+        model_response=source,
+        messages=[{"role": "user", "content": "hi"}],
+        initial_kwargs={"model": "gpt-4", "stream": True},
+    )
+
+    assert result._response_headers == provider_headers
+    additional_headers = result._hidden_params["additional_headers"]
+    assert additional_headers["llm_provider-x-request-id"] == "req-provider-123"
+    assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42"
+    # internal-header protection survives: the provider value is namespaced, never promoted
+    assert additional_headers["llm_provider-x-litellm-model-id"] == "spoofed"
+    assert "x-litellm-model-id" not in additional_headers
+
+
+def test_completion_streaming_iterator_preserves_response_headers():
+    """LIT-6767, sync counterpart of the async header-preservation test."""
+    from unittest.mock import MagicMock
+
+    from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gpt-4",
+                "litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
+            }
+        ],
+    )
+
+    provider_headers = {"x-request-id": "req-provider-sync", "openai-organization": "org-real"}
+    source = CustomStreamWrapper(
+        completion_stream=iter([]),
+        model="gpt-4",
+        custom_llm_provider="openai",
+        logging_obj=MagicMock(),
+        _response_headers=provider_headers,
+    )
+
+    result = router._completion_streaming_iterator(
+        model_response=source,
+        messages=[{"role": "user", "content": "hi"}],
+        initial_kwargs={"model": "gpt-4", "stream": True},
+    )
+
+    assert result._response_headers == provider_headers
+    assert result._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req-provider-sync"
+
+
+def test_adopt_fallback_response_headers_replaces_rather_than_merges():
+    """LIT-6767: direct unit for FallbackAwareStreamWrapper.adopt_fallback_response_headers.
+
+    Values from the failed attempt must not survive, so the wrapper replaces both
+    ``_response_headers`` and ``_hidden_params`` instead of merging them.
+    """
+    from unittest.mock import MagicMock
+
+    from litellm.router import FallbackAwareStreamWrapper, Router
+
+    wrapper = FallbackAwareStreamWrapper(
+        completion_stream=iter([]),
+        model="gpt-4",
+        custom_llm_provider="openai",
+        logging_obj=MagicMock(),
+        _response_headers={"x-request-id": "req-FAILED"},
+    )
+    wrapper._hidden_params = {
+        "model_id": "failed-deployment",
+        "only_on_failed_attempt": "stale",
+        "additional_headers": {"llm_provider-x-request-id": "req-FAILED"},
+    }
+
+    fallback = MagicMock()
+    fallback._response_headers = {"x-request-id": "req-FALLBACK"}
+    fallback._hidden_params = {
+        "model_id": "fallback-deployment",
+        "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"},
+    }
+
+    wrapper.adopt_fallback_response_headers(
+        fallback, Router._prepare_fallback_hidden_params(fallback)
+    )
+
+    assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"}
+    assert wrapper._hidden_params["model_id"] == "fallback-deployment"
+    assert "only_on_failed_attempt" not in wrapper._hidden_params
+    assert wrapper._hidden_params is not fallback._hidden_params
+    # the snapshot CustomStreamWrapper caches at init has to follow, or a chunk built
+    # from it would still be stamped with the deployment that failed
+    assert wrapper._base_hidden_params["model_id"] == "fallback-deployment"
+    # the nested header dict is copied too, so a later mutation on the fallback
+    # response cannot reach headers the proxy has already published
+    assert wrapper._hidden_params["additional_headers"] is not fallback._hidden_params["additional_headers"]
+    fallback._hidden_params["additional_headers"]["llm_provider-x-request-id"] = "req-MUTATED"
+    assert wrapper._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"}
+
+
+def test_adopt_fallback_response_headers_survives_a_collected_wrapper():
+    """LIT-6767: adoption still returns the fallback's params once the wrapper is gone."""
+    import weakref
+    from unittest.mock import MagicMock
+
+    from litellm.router import FallbackAwareStreamWrapper, Router
+
+    fallback: Final = MagicMock()
+    fallback._response_headers = {"x-request-id": "req-FALLBACK"}
+    fallback._hidden_params = {
+        "model_id": "fallback-deployment",
+        "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"},
+    }
+
+    wrapper = FallbackAwareStreamWrapper(
+        completion_stream=iter([]),
+        model="gpt-4",
+        custom_llm_provider="openai",
+        logging_obj=MagicMock(),
+    )
+    live_ref: Final = weakref.ref(wrapper)
+    prepared: Final = Router._adopt_fallback_response_headers(live_ref, fallback)
+    assert prepared == (fallback._hidden_params, fallback._hidden_params["additional_headers"])
+    assert wrapper.fallback_headers_adopted is True
+    assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"}
+
+    dead_ref: Final = weakref.ref(wrapper)
+    del wrapper
+    assert dead_ref() is None
+    assert Router._adopt_fallback_response_headers(dead_ref, fallback) == prepared
+
+
+def test_adopt_fallback_response_headers_drops_headers_the_fallback_cannot_replace():
+    """LIT-6767: a fallback that carries no raw provider headers publishes none.
+
+    Keeping the failed attempt's raw headers would hand the client and the callbacks a
+    provider ``x-request-id`` for a request that deployment never served, which is the
+    leak this fix exists to close.
+    """
+    from unittest.mock import MagicMock
+
+    from litellm.router import FallbackAwareStreamWrapper, Router
+
+    wrapper = FallbackAwareStreamWrapper(
+        completion_stream=iter([]),
+        model="gpt-4",
+        custom_llm_provider="openai",
+        logging_obj=MagicMock(),
+        _response_headers={"x-request-id": "req-FAILED"},
+    )
+    wrapper._hidden_params = {"model_id": "failed-deployment", "additional_headers": {}}
+
+    fallback = MagicMock()
+    fallback._response_headers = None
+    fallback._hidden_params = {"model_id": "fallback-deployment"}
+
+    wrapper.adopt_fallback_response_headers(
+        fallback, Router._prepare_fallback_hidden_params(fallback)
+    )
+
+    assert wrapper._response_headers is None
+    assert wrapper._hidden_params["model_id"] == "fallback-deployment"
+    assert wrapper.fallback_headers_adopted is True
+
+
+def test_adopt_fallback_response_headers_keeps_identity_when_fallback_has_none():
+    """A fallback response carrying no hidden params keeps the identity headers.
+
+    Publishing no ``x-litellm-*`` header at all for a request the fallback served is
+    worse than keeping what is there, so only the raw provider headers are dropped.
+    """
+    from unittest.mock import MagicMock
+
+    from litellm.router import FallbackAwareStreamWrapper, Router
+
+    wrapper = FallbackAwareStreamWrapper(
+        completion_stream=iter([]),
+        model="gpt-4",
+        custom_llm_provider="openai",
+        logging_obj=MagicMock(),
+        _response_headers={"x-request-id": "req-FAILED"},
+    )
+    hidden_params_before = wrapper._hidden_params
+
+    fallback = object()
+    wrapper.adopt_fallback_response_headers(
+        fallback, Router._prepare_fallback_hidden_params(fallback)
+    )
+
+    assert wrapper._response_headers is None
+    assert wrapper._hidden_params is hidden_params_before
+    assert wrapper.fallback_headers_adopted is True
+
+
+@pytest.mark.asyncio
+async def test_acompletion_streaming_iterator_adopts_fallback_response_headers():
+    """LIT-6767: after a successful pre-first-chunk fallback, the wrapper must
+    describe the deployment that served the stream, with no value left over
+    from the attempt that failed."""
+    from unittest.mock import MagicMock, patch
+
+    from litellm.exceptions import MidStreamFallbackError
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gpt-4",
+                "litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
+            }
+        ],
+    )
+
+    failed_error = MidStreamFallbackError(
+        message="upstream died before the first chunk",
+        model="gpt-4",
+        llm_provider="openai",
+        generated_content="",
+        is_pre_first_chunk=True,
+    )
+
+    class FailedStream:
+        def __init__(self):
+            self.model = "gpt-4"
+            self.custom_llm_provider = "openai"
+            self.logging_obj = MagicMock()
+            self.chunks = []
+            self._response_headers = {"x-request-id": "req-FAILED"}
+            self._hidden_params = {
+                "model_id": "failed-deployment",
+                "api_base": "https://failed.example",
+                "additional_headers": {"llm_provider-x-request-id": "req-FAILED"},
+                "only_on_failed_attempt": "stale",
+            }
+
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            raise failed_error
+
+    class FallbackStream:
+        def __init__(self):
+            self._response_headers = {"x-request-id": "req-FALLBACK"}
+            self._hidden_params = {
+                "model_id": "fallback-deployment",
+                "api_base": "https://fallback.example",
+                "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"},
+            }
+            self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])])
+
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            try:
+                return next(self._chunks)
+            except StopIteration:
+                raise StopAsyncIteration from None
+
+    fallback_stream = FallbackStream()
+    with patch.object(
+        router,
+        "async_function_with_fallbacks_common_utils",
+        return_value=fallback_stream,
+    ):
+        result = await router._acompletion_streaming_iterator(
+            model_response=FailedStream(),
+            messages=[{"role": "user", "content": "hi"}],
+            initial_kwargs={"model": "gpt-4", "stream": True},
+        )
+        # the failed attempt is what the wrapper is built from
+        assert result._response_headers == {"x-request-id": "req-FAILED"}
+        # the very first chunk the fallback produces must already be published under
+        # the fallback's identity: the proxy commits response headers once that chunk
+        # is buffered, so adopting any later is adopting too late
+        await result.__anext__()
+        assert result._response_headers == {"x-request-id": "req-FALLBACK"}
+        assert result._hidden_params["model_id"] == "fallback-deployment"
+        async for _ in result:
+            pass
+
+    assert result._response_headers == {"x-request-id": "req-FALLBACK"}
+    assert result._hidden_params["model_id"] == "fallback-deployment"
+    assert result._hidden_params["api_base"] == "https://fallback.example"
+    assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"}
+    # stale values are removed, not merged over
+    assert "only_on_failed_attempt" not in result._hidden_params
+    # and the wrapper holds its own copy, so later fallback mutations cannot leak in
+    assert result._hidden_params is not fallback_stream._hidden_params
+
+
+@pytest.mark.asyncio
+async def test_acompletion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback():
+    """LIT-6767: a fallback that itself fails over before its first chunk.
+
+    The selected fallback still describes its own failed attempt at selection time, so
+    the wrapper has to re-read it once a chunk exists or it publishes a deployment that
+    produced no output.
+    """
+    from unittest.mock import MagicMock, patch
+
+    from litellm.exceptions import MidStreamFallbackError
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gpt-4",
+                "litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
+            }
+        ],
+    )
+
+    failed_error: Final = MidStreamFallbackError(
+        message="upstream died before the first chunk",
+        model="gpt-4",
+        llm_provider="openai",
+        generated_content="",
+        is_pre_first_chunk=True,
+    )
+
+    class FailedStream:
+        def __init__(self):
+            self.model = "gpt-4"
+            self.custom_llm_provider = "openai"
+            self.logging_obj = MagicMock()
+            self.chunks = []
+            self._response_headers = {"x-request-id": "req-FAILED"}
+            self._hidden_params = {
+                "model_id": "failed-deployment",
+                "additional_headers": {"llm_provider-x-request-id": "req-FAILED"},
+            }
+
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            raise failed_error
+
+    class NestedFallbackStream:
+        """A fallback that repoints itself at a third deployment as it yields."""
+
+        def __init__(self):
+            self._response_headers = {"x-request-id": "req-MIDDLE"}
+            self._hidden_params = {
+                "model_id": "middle-deployment",
+                "additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"},
+            }
+            self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])])
+
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            try:
+                chunk = next(self._chunks)
+            except StopIteration:
+                raise StopAsyncIteration from None
+            self._response_headers = {"x-request-id": "req-SERVED"}
+            self._hidden_params = {
+                "model_id": "served-deployment",
+                "additional_headers": {"llm_provider-x-request-id": "req-SERVED"},
+            }
+            return chunk
+
+    with patch.object(
+        router,
+        "async_function_with_fallbacks_common_utils",
+        return_value=NestedFallbackStream(),
+    ):
+        result = await router._acompletion_streaming_iterator(
+            model_response=FailedStream(),
+            messages=[{"role": "user", "content": "hi"}],
+            initial_kwargs={"model": "gpt-4", "stream": True},
+        )
+        first_chunk: Final = await result.__anext__()
+        # the proxy commits response headers once this chunk is buffered
+        assert result._response_headers == {"x-request-id": "req-SERVED"}
+        assert result._hidden_params["model_id"] == "served-deployment"
+        assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-SERVED"}
+        # and the chunk itself carries the same deployment
+        assert first_chunk._hidden_params["model_id"] == "served-deployment"
+        async for _ in result:
+            pass
+
+    assert result._response_headers == {"x-request-id": "req-SERVED"}
+    assert result._hidden_params["model_id"] == "served-deployment"
+
+
+def test_completion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback():
+    """LIT-6767, sync counterpart of the nested-fallback adoption test."""
+    from unittest.mock import MagicMock, patch
+
+    from litellm.exceptions import MidStreamFallbackError
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gpt-4",
+                "litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
+            }
+        ],
+    )
+
+    failed_error: Final = MidStreamFallbackError(
+        message="upstream died before the first chunk",
+        model="gpt-4",
+        llm_provider="openai",
+        generated_content="",
+        is_pre_first_chunk=True,
+    )
+
+    class FailedStream:
+        def __init__(self):
+            self.model = "gpt-4"
+            self.custom_llm_provider = "openai"
+            self.logging_obj = MagicMock()
+            self.chunks = []
+            self._response_headers = {"x-request-id": "req-FAILED"}
+            self._hidden_params = {
+                "model_id": "failed-deployment",
+                "additional_headers": {"llm_provider-x-request-id": "req-FAILED"},
+            }
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            raise failed_error
+
+    class NestedFallbackStream:
+        """A fallback that repoints itself at a third deployment as it yields."""
+
+        def __init__(self):
+            self._response_headers = {"x-request-id": "req-MIDDLE"}
+            self._hidden_params = {
+                "model_id": "middle-deployment",
+                "additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"},
+            }
+            self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])])
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            chunk = next(self._chunks)
+            self._response_headers = {"x-request-id": "req-SERVED"}
+            self._hidden_params = {
+                "model_id": "served-deployment",
+                "additional_headers": {"llm_provider-x-request-id": "req-SERVED"},
+            }
+            return chunk
+
+    with patch.object(router, "function_with_fallbacks", return_value=NestedFallbackStream()):
+        result = router._completion_streaming_iterator(
+            model_response=FailedStream(),
+            messages=[{"role": "user", "content": "hi"}],
+            initial_kwargs={"model": "gpt-4", "stream": True},
+        )
+        first_chunk: Final = next(result)
+        assert result._response_headers == {"x-request-id": "req-SERVED"}
+        assert result._hidden_params["model_id"] == "served-deployment"
+        assert first_chunk._hidden_params["model_id"] == "served-deployment"
+        for _ in result:
+            pass
+
+    assert result._response_headers == {"x-request-id": "req-SERVED"}
+    assert result._hidden_params["model_id"] == "served-deployment"
+
+
+def test_completion_streaming_iterator_adopts_fallback_response_headers():
+    """LIT-6767, sync counterpart of the fallback-adoption test."""
+    from unittest.mock import MagicMock, patch
+
+    from litellm.exceptions import MidStreamFallbackError
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gpt-4",
+                "litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
+            }
+        ],
+    )
+
+    failed_error = MidStreamFallbackError(
+        message="upstream died before the first chunk",
+        model="gpt-4",
+        llm_provider="openai",
+        generated_content="",
+        is_pre_first_chunk=True,
+    )
+
+    class FailedStream:
+        def __init__(self):
+            self.model = "gpt-4"
+            self.custom_llm_provider = "openai"
+            self.logging_obj = MagicMock()
+            self.chunks = []
+            self._response_headers = {"x-request-id": "req-FAILED"}
+            self._hidden_params = {
+                "model_id": "failed-deployment",
+                "additional_headers": {"llm_provider-x-request-id": "req-FAILED"},
+                "only_on_failed_attempt": "stale",
+            }
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            raise failed_error
+
+    class FallbackStream:
+        def __init__(self):
+            self._response_headers = {"x-request-id": "req-FALLBACK"}
+            self._hidden_params = {
+                "model_id": "fallback-deployment",
+                "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"},
+            }
+
+        def __iter__(self):
+            return iter([])
+
+    with patch.object(router, "function_with_fallbacks", return_value=FallbackStream()):
+        result = router._completion_streaming_iterator(
+            model_response=FailedStream(),
+            messages=[{"role": "user", "content": "hi"}],
+            initial_kwargs={"model": "gpt-4", "stream": True},
+        )
+        assert result._response_headers == {"x-request-id": "req-FAILED"}
+        for _ in result:
+            pass
+
+    assert result._response_headers == {"x-request-id": "req-FALLBACK"}
+    assert result._hidden_params["model_id"] == "fallback-deployment"
+    assert "only_on_failed_attempt" not in result._hidden_params
+
+
 def test_completion_streaming_iterator_fallback_on_429():
     """Sync streaming: MidStreamFallbackError (429 pre-first-chunk) triggers fallback.
 

From b618c7ad8607dcdd7507cb1d69e980893c96732f Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Mon, 7 Sep 2026 13:10:09 -0700
Subject: [PATCH 137/164] fix(proxy): let authorized internal users open vector
 store details (#40150)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm/proxy/_types.py                       |  1 +
 litellm/proxy/vector_store_endpoints/utils.py |  9 ++-
 .../proxy/auth/test_route_checks.py           | 59 ++++++++++++++++++
 .../test_vector_store_access_control.py       | 62 +++++++++++++++++++
 4 files changed, 128 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index d79b753b5ff..abce10690e5 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -848,6 +848,7 @@ class LiteLLMRoutes(enum.Enum):
         "/model/{model_id}/update",
         "/prompt/list",
         "/prompt/info",
+        "/vector_store/info",
         # Project read routes - endpoint scopes results to caller's teams (non-admin)
         "/project/list",
         "/project/info",
diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py
index f2070e6604c..8363aaee99a 100644
--- a/litellm/proxy/vector_store_endpoints/utils.py
+++ b/litellm/proxy/vector_store_endpoints/utils.py
@@ -161,6 +161,8 @@ async def can_user_access_vector_store(
        this vector store id.
     5. The caller's team_id matches the vector store's team_id.
 
+    A dashboard session credential is evaluated against the same effective
+    contexts as listing (its own grants plus each real team of the user).
     Otherwise access is denied.
     """
     if _is_proxy_admin(user_api_key_dict):
@@ -169,7 +171,8 @@ async def can_user_access_vector_store(
     if vector_store.get("team_id") is None:
         return True
 
-    return await _is_vector_store_granted(vector_store, user_api_key_dict)
+    auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict)
+    return await _is_vector_store_granted_to_any(vector_store, auth_contexts)
 
 
 async def _is_vector_store_granted(
@@ -219,7 +222,7 @@ async def _team_auth_context(team_id: str, user_api_key_dict: UserAPIKeyAuth) ->
     )
 
 
-async def _vector_store_listing_auth_contexts(
+async def _vector_store_auth_contexts(
     user_api_key_dict: UserAPIKeyAuth,
 ) -> tuple[UserAPIKeyAuth, ...]:
     if not is_ui_session_credential(user_api_key_dict):
@@ -250,7 +253,7 @@ async def filter_listable_vector_stores(
     if _is_proxy_admin(user_api_key_dict):
         return tuple(vector_stores)
 
-    auth_contexts: Final = await _vector_store_listing_auth_contexts(user_api_key_dict)
+    auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict)
     return tuple([vs for vs in vector_stores if await _is_vector_store_granted_to_any(vs, auth_contexts)])
 
 
diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py
index 48926cb7bc2..5b15d4a7d5e 100644
--- a/tests/test_litellm/proxy/auth/test_route_checks.py
+++ b/tests/test_litellm/proxy/auth/test_route_checks.py
@@ -3217,6 +3217,65 @@ def test_internal_user_blocked_from_search_tool_writes(route):
     assert "Your role=internal_user" in str(exc_info.value)
 
 
+@pytest.mark.parametrize(
+    "user_role",
+    [
+        LitellmUserRoles.INTERNAL_USER.value,
+        LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value,
+    ],
+)
+def test_non_admin_can_open_vector_store_details(user_role):
+    """Regression for LIT-7132: the dashboard lists a vector store via /vector_store/list
+    (an LLM API route) but opened it via /vector_store/info, which no non-admin allowlist
+    granted, so the route gate 401'd before the handler's per-store access check ran."""
+    user_obj = LiteLLM_UserTable(
+        user_id="test_user",
+        user_email="user@example.com",
+        user_role=user_role,
+    )
+    valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role)
+    request = MagicMock(spec=Request)
+    request.query_params = {}
+
+    granted = RouteChecks.non_proxy_admin_allowed_routes_check(
+        user_obj=user_obj,
+        _user_role=user_role,
+        route="/vector_store/info",
+        request=request,
+        valid_token=valid_token,
+        request_data={},
+    )
+    assert granted is None
+
+
+@pytest.mark.parametrize(
+    "route",
+    ["/vector_store/new", "/vector_store/update", "/vector_store/delete"],
+)
+def test_internal_user_blocked_from_vector_store_writes(route):
+    user_obj = LiteLLM_UserTable(
+        user_id="test_user",
+        user_email="user@example.com",
+        user_role=LitellmUserRoles.INTERNAL_USER.value,
+    )
+    valid_token = UserAPIKeyAuth(
+        user_id="test_user",
+        user_role=LitellmUserRoles.INTERNAL_USER.value,
+    )
+    request = MagicMock(spec=Request)
+    request.query_params = {}
+
+    with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update"):
+        RouteChecks.non_proxy_admin_allowed_routes_check(
+            user_obj=user_obj,
+            _user_role=LitellmUserRoles.INTERNAL_USER.value,
+            route=route,
+            request=request,
+            valid_token=valid_token,
+            request_data={},
+        )
+
+
 def test_proxy_admin_viewer_can_read_another_users_info():
     """Admin Viewer has read parity with Proxy Admin, so the /user/info
     key-ownership gate must not apply to it — the Users page reads every row."""
diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py
index 93049b21460..dc4f2900038 100644
--- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py
+++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py
@@ -244,3 +244,65 @@ async def test_list_vector_stores_dashboard_session_resolves_real_teams(
         ),
     ):
         assert await _listed_ids(alice) == expected
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("user_team_ids", "expected_status"),
+    [
+        (["team_a"], 200),
+        (["team_b"], 403),
+        ([], 403),
+    ],
+)
+async def test_get_vector_store_info_dashboard_session_resolves_real_teams(
+    user_team_ids: list[str], expected_status: int
+):
+    """Regression for LIT-7132: /vector_store/info must grant a dashboard session the same team-owned stores
+    /vector_store/list shows it, instead of judging the session's reserved litellm-dashboard team id."""
+    from litellm.models.team import LiteLLM_TeamTableCachedObj
+    from litellm.proxy.vector_store_endpoints.management_endpoints import (
+        get_vector_store_info,
+    )
+    from litellm.types.vector_stores import VectorStoreInfoRequest
+
+    alice = UserAPIKeyAuth(
+        team_id="litellm-dashboard",
+        user_id="alice",
+        user_role=LitellmUserRoles.INTERNAL_USER,
+    )
+
+    async def fake_get_team_object(team_id: str, **_kwargs: object) -> LiteLLM_TeamTableCachedObj:
+        return LiteLLM_TeamTableCachedObj(team_id=team_id)
+
+    mock_prisma = MagicMock()
+    mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock(
+        return_value=MagicMock(model_dump=lambda: dict(_TEAM_A_OWNED))
+    )
+
+    async def outcome() -> int:
+        try:
+            response = await get_vector_store_info(
+                data=VectorStoreInfoRequest(vector_store_id="vs_team_a"), user_api_key_dict=alice
+            )
+        except HTTPException as exc:
+            return exc.status_code
+        assert response["vector_store"]["vector_store_id"] == "vs_team_a"
+        return 200
+
+    with (
+        patch(  # test-quality-ok: the endpoint reads the store row through the module-level prisma client, no injection seam
+            "litellm.proxy.proxy_server.prisma_client", mock_prisma
+        ),
+        patch(  # test-quality-ok: the endpoint consults the module-level registry before the DB, no injection seam
+            "litellm.vector_store_registry", None
+        ),
+        patch(  # test-quality-ok: team rows come from the module-level prisma client, no injection seam
+            "litellm.proxy.auth.auth_checks.get_team_object", new=fake_get_team_object
+        ),
+        patch(  # test-quality-ok: the user row comes from the module-level prisma client, no injection seam
+            "litellm.proxy.vector_store_endpoints.utils.resolve_ui_session_team_ids",
+            new=AsyncMock(return_value=user_team_ids),
+        ),
+    ):
+        assert await outcome() == expected_status

From 7c1745cc7178df28434ddd39b4e6861f64826331 Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Mon, 7 Sep 2026 13:55:27 -0700
Subject: [PATCH 138/164] feat(ui): make automatic auto-router setup
 discoverable and show what it configured (#40146)

* fix(ui): expand Detailed Configuration after automatic auto-router setup

* feat(ui): promote automatic auto-router setup to a callout banner

* test(ui): read tier chips through testing-library queries to stay in lint budget

* style(ui): drop explanatory comments per repo convention

* test(ui): reject unexpected automatic tier models
---
 .../add_model/add_auto_router_tab.test.tsx    | 43 +++++++++++++++----
 .../add_model/add_auto_router_tab.tsx         | 23 +++++-----
 2 files changed, 45 insertions(+), 21 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
index 5605a993ded..014854ac712 100644
--- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
@@ -60,6 +60,20 @@ const optionByLabel = (label: string): HTMLElement | undefined =>
 
 const isOptionDisabled = (option: HTMLElement): boolean => option.getAttribute("aria-disabled") === "true";
 
+const tierChips = (tier: string): HTMLElement => {
+  const placeholder = `Select model(s) for ${tier.toLowerCase()} queries`;
+  const chips = screen
+    .getAllByRole("toolbar")
+    .find((candidate) => within(candidate).queryByLabelText(placeholder) !== null);
+  if (!chips) throw new Error(`No tier row found for "${tier}"`);
+  return chips;
+};
+
+const expectTierModel = (tier: string, model: string): void => {
+  const chips = within(tierChips(tier)).getAllByLabelText(/.+/, { selector: '[data-slot="combobox-chip"]' });
+  expect(chips.map((chip) => chip.getAttribute("aria-label"))).toEqual([model]);
+};
+
 const selectTemplate = async (label: string): Promise<void> => {
   await userEvent.click(optionByLabel(label)!);
 };
@@ -188,11 +202,10 @@ describe("AddAutoRouterTab", () => {
     const button = await screen.findByTestId("configure-automatically-button");
     await userEvent.click(button);
 
-    expect(
-      screen.getByText(
-        /Simple: gpt-5.6-luna.*Medium: claude-sonnet-5.*Complex: claude-opus-5.*Reasoning: claude-opus-5/,
-      ),
-    ).toBeInTheDocument();
+    expectTierModel("Simple", "gpt-5.6-luna");
+    expectTierModel("Medium", "claude-sonnet-5");
+    expectTierModel("Complex", "claude-opus-5");
+    expectTierModel("Reasoning", "claude-opus-5");
     expect(toast.success).not.toHaveBeenCalledWith(expect.stringContaining("Configured with"));
   });
 
@@ -209,9 +222,23 @@ describe("AddAutoRouterTab", () => {
     const button = await screen.findByTestId("configure-automatically-button");
     await userEvent.click(button);
 
-    expect(
-      screen.getByText(/Simple: gpt-5.6-luna.*Medium: claude-sonnet-5.*Complex: gpt-5.6-sol.*Reasoning: gpt-5.6-sol/),
-    ).toBeInTheDocument();
+    expectTierModel("Simple", "gpt-5.6-luna");
+    expectTierModel("Medium", "claude-sonnet-5");
+    expectTierModel("Complex", "gpt-5.6-sol");
+    expectTierModel("Reasoning", "gpt-5.6-sol");
+  });
+
+  it("opens Detailed Configuration on the tiers automatic setup just filled in", async () => {
+    const simpleModel = "gpt-5.6-luna";
+    mockFetchAvailableModels.mockResolvedValue([...ALL_FAMILY_MODELS, { model_group: simpleModel, mode: "chat" }]);
+    renderWithProviders(<Harness />);
+
+    expect(screen.queryByText("Complexity Tier Configuration")).not.toBeInTheDocument();
+
+    await userEvent.click(await screen.findByTestId("configure-automatically-button"));
+
+    expect(screen.getByText("Complexity Tier Configuration")).toBeInTheDocument();
+    expectTierModel("Simple", simpleModel);
   });
 
   // Nothing is filled in, so there is nothing to submit. The button reports that itself instead of
diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
index 60780f9cbe0..2d0d6bc2fd8 100644
--- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
@@ -207,10 +207,6 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
   const [isSubmitting, setIsSubmitting] = useState(false);
 
   const [selectedPreset, setSelectedPreset] = useState<string | undefined>(undefined);
-  // Closed by default: a caller opens it deliberately, either by clicking it or by choosing Custom
-  // (which expands it automatically, since there's nothing else to show them their config from). A
-  // preset re-collapses it after prefilling, offering the same "here's what got filled in, expand to
-  // change it" affordance. A caller can always toggle it manually at any point.
   const [detailsExpanded, setDetailsExpanded] = useState<boolean>(false);
 
   const [isRoutingTestVisible, setIsRoutingTestVisible] = useState<boolean>(false);
@@ -335,7 +331,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
     if (automaticRouterConfig === null) return;
     setSelectedPreset(undefined);
     applyPrefill({ ...buildEmptyPrefill(), complexityRouterConfig: automaticRouterConfig });
-    setDetailsExpanded(false);
+    setDetailsExpanded(true);
     toast.success("Automatic setup created", { description: tierConfigSummary(automaticRouterConfig) });
   };
 
@@ -543,14 +539,15 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
                 </FormField>
 
                 {!automaticSetupLoading && automaticRouterConfig && (
-                  <button
-                    type="button"
-                    className="mt-3 rounded-sm text-sm font-medium text-blue-600 hover:text-blue-700 hover:underline focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-blue-500 focus-visible:ring-offset-2"
-                    data-testid="configure-automatically-button"
-                    onClick={handleAutomaticSetup}
-                  >
-                    Configure automatically
-                  </button>
+                  <div className="mt-5 flex flex-wrap items-center justify-between gap-3 rounded-lg border border-border bg-muted px-4 py-3">
+                    <div>
+                      <p className="text-sm font-medium text-foreground">Not sure where to start?</p>
+                      <p className="text-sm text-muted-foreground">Let us pick models for each complexity tier.</p>
+                    </div>
+                    <Button type="button" data-testid="configure-automatically-button" onClick={handleAutomaticSetup}>
+                      Configure automatically
+                    </Button>
+                  </div>
                 )}
 
                 <div className="mt-5">

From f896df1b065e2627e3f018e7f504c44666f5d35c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jean-R=C3=A9mi=20Larcelet-Prost?= <jrlarcelet@gmail.com>
Date: Mon, 7 Sep 2026 23:19:46 +0200
Subject: [PATCH 139/164] docs: fix stale file paths in ARCHITECTURE.md
 (#40157)

Several file references under proxy/management_helpers/ and other
paths no longer exist; the code moved to proxy/common_utils/,
proxy/db/db_transaction_queue/, litellm_enterprise/proxy/common_utils/,
proxy/hooks/litellm_skills/, and litellm_core_utils/.
---
 ARCHITECTURE.md | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md
index 3d2fa3e51c8..b04e004aa1a 100644
--- a/ARCHITECTURE.md
+++ b/ARCHITECTURE.md
@@ -149,7 +149,7 @@ graph TD
 | `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user |
 | `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation |
 | `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation |
-| `litellm_skills` | `proxy/hooks/skills_injection.py` | Skills injection |
+| `litellm_skills` | `proxy/hooks/litellm_skills/main.py` | Skills injection |
 
 To add a new proxy hook, implement `CustomLogger` and register in `PROXY_HOOKS`.
 
@@ -220,20 +220,20 @@ graph LR
 | Job | Interval | Purpose | Key Files |
 |-----|----------|---------|-----------|
 | `update_spend` | 60s | Batch write spend logs to PostgreSQL | `proxy/db/db_spend_update_writer.py` |
-| `reset_budget` | 10-12min | Reset budgets for keys/users/teams | `proxy/management_helpers/budget_reset_job.py` |
+| `reset_budget` | 10-12min | Reset budgets for keys/users/teams | `proxy/common_utils/reset_budget_job.py` |
 | `add_deployment` | 10s | Sync new model deployments from DB | `proxy/proxy_server.py` (`ProxyConfig`) |
-| `cleanup_old_spend_logs` | cron/interval | Delete old spend logs | `proxy/management_helpers/spend_log_cleanup.py` |
-| `check_batch_cost` | 30min | Calculate costs for batch jobs | `proxy/management_helpers/check_batch_cost_job.py` |
-| `check_responses_cost` | 30min | Calculate costs for responses API | `proxy/management_helpers/check_responses_cost_job.py` |
-| `process_rotations` | 1hr | Auto-rotate API keys | `proxy/management_helpers/key_rotation_manager.py` |
+| `cleanup_old_spend_logs` | cron/interval | Delete old spend logs | `proxy/db/db_transaction_queue/spend_log_cleanup.py` |
+| `check_batch_cost` | 30min | Calculate costs for batch jobs | `enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py` |
+| `check_responses_cost` | 30min | Calculate costs for responses API | `enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py` |
+| `process_rotations` | 1hr | Auto-rotate API keys | `proxy/common_utils/key_rotation_manager.py` |
 | `_run_background_health_check` | continuous | Health check model deployments | `proxy/proxy_server.py` |
 | `send_weekly_spend_report` | weekly | Slack spend alerts | `proxy/utils.py` (`SlackAlerting`) |
 | `send_monthly_spend_report` | monthly | Slack spend alerts | `proxy/utils.py` (`SlackAlerting`) |
 
 **Cost Attribution Flow:**
 1. LLM response returns to `utils.py` wrapper after `litellm.acompletion()` completes
-2. `update_response_metadata()` (`llm_response_utils/response_metadata.py`) is called
-3. `logging_obj._response_cost_calculator()` (`litellm_logging.py`) calculates cost via `litellm.completion_cost()` (`cost_calculator.py`)
+2. `update_response_metadata()` (`litellm_core_utils/llm_response_utils/response_metadata.py`) is called
+3. `logging_obj._response_cost_calculator()` (`litellm_core_utils/litellm_logging.py`) calculates cost via `litellm.completion_cost()` (`cost_calculator.py`)
 4. Cost is stored in `response._hidden_params["response_cost"]`
 5. `proxy/common_request_processing.py` extracts cost from `hidden_params` and adds to response headers (`x-litellm-response-cost`)
 6. `logging_obj.async_success_handler()` triggers callbacks including `_ProxyDBLogger.async_log_success_event()`

From 038025ba5e2a6796186a86909de03e1a5eeb915d Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Mon, 7 Sep 2026 14:33:02 -0700
Subject: [PATCH 140/164] fix(guardrails): accept on_violation block and alert
 for mcp_security (#40155)

* fix(guardrails): accept on_violation block and alert for mcp_security

The MCP Security policy template sends on_violation: "block", but the shared
LitellmParams model only allowed the /v1/realtime values "warn" and
"end_session", so POST /guardrails returned 422 before the MCP guardrail was
initialized. Widen the literal to include the MCP actions, map every non-alert
value to MCP's default "block" at init, and regenerate the lazy OpenAPI
snapshot and dashboard API types

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(guardrails): restrict on_violation block/alert to mcp_security and keep legacy MCP mapping

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(guardrails): return 422 when PATCH sets an mcp_security-only on_violation on another guardrail

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm/proxy/_lazy_openapi_snapshot.json     | 12 +++++---
 .../proxy/guardrails/guardrail_endpoints.py   | 10 +++++--
 .../guardrail_hooks/mcp_security/__init__.py  |  7 ++---
 litellm/types/guardrails.py                   | 20 +++++++++++--
 .../guardrail_hooks/test_mcp_security.py      | 25 +++++++++++++++-
 .../guardrails/test_guardrail_endpoints.py    | 16 ++++++++++
 .../test_guardrails_case_normalization.py     | 30 +++++++++++++++++++
 ui/litellm-dashboard/src/lib/http/schema.d.ts |  8 ++---
 8 files changed, 110 insertions(+), 18 deletions(-)

diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json
index 4093f2c5248..c71761a7adb 100644
--- a/litellm/proxy/_lazy_openapi_snapshot.json
+++ b/litellm/proxy/_lazy_openapi_snapshot.json
@@ -9652,7 +9652,9 @@
                 {
                   "enum": [
                     "warn",
-                    "end_session"
+                    "end_session",
+                    "block",
+                    "alert"
                   ],
                   "type": "string"
                 },
@@ -9660,7 +9662,7 @@
                   "type": "null"
                 }
               ],
-              "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.",
+              "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.",
               "title": "On Violation"
             },
             "only_scan_new_messages": {
@@ -11969,7 +11971,9 @@
                 {
                   "enum": [
                     "warn",
-                    "end_session"
+                    "end_session",
+                    "block",
+                    "alert"
                   ],
                   "type": "string"
                 },
@@ -11977,7 +11981,7 @@
                   "type": "null"
                 }
               ],
-              "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.",
+              "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.",
               "title": "On Violation"
             },
             "only_scan_new_messages": {
diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py
index 744b2959c73..afb9997f2e6 100644
--- a/litellm/proxy/guardrails/guardrail_endpoints.py
+++ b/litellm/proxy/guardrails/guardrail_endpoints.py
@@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union,
 from urllib.parse import urlparse
 
 from fastapi import APIRouter, Depends, HTTPException, Request
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError
 
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
@@ -1202,7 +1202,13 @@ async def patch_guardrail(
             litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True)
             litellm_params_dict.update(requested_litellm_params)
             merged_litellm_params: Final = _as_str_object_mapping(litellm_params_dict)
-            litellm_params = LitellmParams(**merged_litellm_params)
+            try:
+                litellm_params = LitellmParams(**merged_litellm_params)
+            except ValidationError as validation_error:
+                raise HTTPException(
+                    status_code=422,
+                    detail=f"Invalid guardrail configuration, update rejected: {validation_error}",
+                ) from validation_error
 
         # Update guardrail_info if provided
         guardrail_info: Final = (
diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py
index d53a4157e0e..1607dfff63e 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Final, Literal, Optional, cast
+from typing import TYPE_CHECKING, Final, Literal, Optional
 
 import litellm
 from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import (
@@ -20,10 +20,7 @@ def initialize_guardrail(
     if not guardrail_name:
         raise ValueError("MCP Security: guardrail_name is required")
 
-    on_violation: Final[Literal["block", "alert"]] = cast(
-        Literal["block", "alert"],
-        getattr(litellm_params, "on_violation", "block"),
-    )
+    on_violation: Final[Literal["block", "alert"]] = "block" if litellm_params.on_violation == "block" else "alert"
 
     mcp_security_guardrail: Final = MCPSecurityGuardrail(
         guardrail_name=guardrail_name,
diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py
index ef28181eba5..02dee40f2a3 100644
--- a/litellm/types/guardrails.py
+++ b/litellm/types/guardrails.py
@@ -778,6 +778,9 @@ class ContentFilterConfigModel(BaseModel):
     )
 
 
+MCP_SECURITY_ON_VIOLATION: Final = frozenset({"block", "alert"})
+
+
 class BaseLitellmParams(ContentFilterConfigModel):  # works for new and patch update guardrails
     api_key: str | None = Field(default=None, description="API key for the guardrail service")
     api_base: str | None = Field(default=None, description="Base URL for the guardrail service API")
@@ -886,9 +889,13 @@ class BaseLitellmParams(ContentFilterConfigModel):  # works for new and patch up
         default=None,
         description="For /v1/realtime sessions: automatically close the session after this many guardrail violations.",
     )
-    on_violation: Literal["warn", "end_session"] | None = Field(
+    on_violation: Literal["warn", "end_session", "block", "alert"] | None = Field(
         default=None,
-        description="For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.",
+        description=(
+            "For /v1/realtime sessions: 'warn' speaks the violation message and continues; "
+            "'end_session' speaks the message and closes the connection. "
+            "For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning."
+        ),
     )
     realtime_violation_message: str | None = Field(
         default=None,
@@ -1093,6 +1100,15 @@ class LitellmParams(  # pyright: ignore[reportIncompatibleVariableOverride]  # o
         except (TypeError, ValueError) as e:
             raise ValueError(f"timeout must be numeric, got {v!r}") from e
 
+    @model_validator(mode="after")
+    def validate_on_violation_for_guardrail(self) -> "LitellmParams":
+        if (
+            self.on_violation in MCP_SECURITY_ON_VIOLATION
+            and self.guardrail != SupportedGuardrailIntegrations.MCP_SECURITY.value
+        ):
+            raise ValueError(f"on_violation={self.on_violation!r} is only supported by guardrail='mcp_security'")
+        return self
+
     def __init__(self, **kwargs) -> None:
         default_on: Final = kwargs.pop("default_on", None)
         if default_on is not None:
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py
index 4444cd693ff..d57a91d45bf 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py
@@ -6,16 +6,19 @@ and allows requests with only registered servers. Covers both /chat/completions
 and /responses API paths (same pre_call_hook logic, different call_type).
 """
 
+from typing import Literal
 from unittest.mock import MagicMock, patch
 
 import pytest
 from fastapi import HTTPException
 
+import litellm
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.guardrails.guardrail_hooks.mcp_security import initialize_guardrail
 from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import (
     MCPSecurityGuardrail,
 )
-from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams
 
 
 @pytest.fixture
@@ -182,3 +185,23 @@ class TestMCPSecurityGuardrailPreCall:
             call_type="acompletion",
         )
         assert result == data
+
+
+class TestInitializeGuardrail:
+    @pytest.mark.parametrize(
+        "configured,expected",
+        [("block", "block"), ("alert", "alert"), (None, "alert"), ("warn", "alert"), ("end_session", "alert")],
+    )
+    def test_on_violation_from_litellm_params(
+        self,
+        configured: Literal["block", "alert", "warn", "end_session"] | None,
+        expected: Literal["block", "alert"],
+    ):
+        litellm_params = LitellmParams(guardrail="mcp_security", mode="pre_call", on_violation=configured)
+        guardrail = Guardrail(guardrail_name="mcp-security-block", litellm_params=litellm_params)
+
+        result = initialize_guardrail(litellm_params=litellm_params, guardrail=guardrail)
+
+        assert isinstance(result, MCPSecurityGuardrail)
+        assert result.on_violation == expected
+        assert result in litellm.callbacks
diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
index fe2cd819717..530f8ffd854 100644
--- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
+++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
@@ -1361,6 +1361,22 @@ async def test_patch_guardrail_endpoint(
             assert "Failed to update" in str(mock_logger.warning.call_args)
 
 
+@pytest.mark.asyncio
+async def test_patch_guardrail_rejects_mcp_only_on_violation_with_422(mocker, mock_guardrail_registry):
+    mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock())  # test-quality-ok: endpoint has no DI seam
+    mocker.patch(  # test-quality-ok: endpoint has no DI seam
+        "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry
+    )
+    request = PatchGuardrailRequest(litellm_params=BaseLitellmParams(on_violation="block"))
+
+    with pytest.raises(HTTPException) as exc_info:
+        await patch_guardrail("test-guardrail-id", request, user_api_key_dict=MOCK_ADMIN_USER)
+
+    assert exc_info.value.status_code == 422
+    assert "only supported by guardrail='mcp_security'" in str(exc_info.value.detail)
+    mock_guardrail_registry.update_guardrail_in_db.assert_not_called()
+
+
 @pytest.mark.parametrize(
     "scenario,expected_result,expected_exception",
     [
diff --git a/tests/test_litellm/types/test_guardrails_case_normalization.py b/tests/test_litellm/types/test_guardrails_case_normalization.py
index 3e7a573ea8e..26c1d395320 100644
--- a/tests/test_litellm/types/test_guardrails_case_normalization.py
+++ b/tests/test_litellm/types/test_guardrails_case_normalization.py
@@ -2,6 +2,8 @@
 Test case normalization in LitellmParams for all guardrail types
 """
 
+from typing import Literal
+
 import pytest
 from pydantic import ValidationError
 
@@ -93,6 +95,34 @@ class TestLitellmParamsCaseNormalization:
             assert params.on_disallowed_action.islower()
 
 
+class TestOnViolationAcceptedValues:
+    """on_violation is shared by /v1/realtime guardrails and the mcp_security guardrail"""
+
+    @pytest.mark.parametrize("action", ["block", "alert"])
+    def test_mcp_security_policy_template_on_violation_is_accepted(self, action: Literal["block", "alert"]):
+        params = LitellmParams(
+            guardrail="mcp_security",
+            mode="pre_call",
+            default_on=True,
+            on_violation=action,
+        )
+        assert params.on_violation == action
+
+    @pytest.mark.parametrize("action", ["warn", "end_session"])
+    def test_realtime_on_violation_still_accepted(self, action: Literal["warn", "end_session"]):
+        params = LitellmParams(guardrail="presidio", mode="pre_call", on_violation=action)
+        assert params.on_violation == action
+
+    @pytest.mark.parametrize("action", ["block", "alert"])
+    def test_mcp_only_on_violation_is_rejected_for_other_guardrails(self, action: Literal["block", "alert"]):
+        with pytest.raises(ValidationError, match="only supported by guardrail='mcp_security'"):
+            LitellmParams(guardrail="presidio", mode="pre_call", on_violation=action)
+
+    def test_unknown_on_violation_is_rejected(self):
+        with pytest.raises(ValidationError):
+            LitellmParams(guardrail="mcp_security", mode="pre_call", on_violation="ignore")
+
+
 class TestSensitiveDataRoutingValidation:
     """on_sensitive_data='route' requires a target model to be set"""
 
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 9de988fde4c..6fb08445aff 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -23759,9 +23759,9 @@ export interface components {
             on_sensitive_data?: ("block" | "route") | null;
             /**
              * On Violation
-             * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.
+             * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.
              */
-            on_violation?: ("warn" | "end_session") | null;
+            on_violation?: ("warn" | "end_session" | "block" | "alert") | null;
             /**
              * Only Scan New Messages
              * @description When True, the guardrail only scans messages that have not already been scanned earlier in the same session (identified by litellm_session_id / session_id). Message content is hashed per session and cached; only the diff (new or edited messages) is sent to the guardrail provider on follow-up calls. Falls back to a full scan when the request has no session id or the cache is unavailable. Intended for blocking/detection guardrails; not applied when mask_request_content is set.
@@ -30801,9 +30801,9 @@ export interface components {
             on_sensitive_data?: ("block" | "route") | null;
             /**
              * On Violation
-             * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.
+             * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.
              */
-            on_violation?: ("warn" | "end_session") | null;
+            on_violation?: ("warn" | "end_session" | "block" | "alert") | null;
             /**
              * Only Scan New Messages
              * @description When True, the guardrail only scans messages that have not already been scanned earlier in the same session (identified by litellm_session_id / session_id). Message content is hashed per session and cached; only the diff (new or edited messages) is sent to the guardrail provider on follow-up calls. Falls back to a full scan when the request has no session id or the cache is unavailable. Intended for blocking/detection guardrails; not applied when mask_request_content is set.

From e238d20fbd3ee28f0a8bbb8621e166b8b4710868 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Mon, 7 Sep 2026 14:34:45 -0700
Subject: [PATCH 141/164] ci: follow the default branch in development tooling

---
 CLAUDE.md                                  |   6 +-
 CONTRIBUTING.md                            |   4 +-
 Makefile                                   |  48 +++--
 ci_cd/run_migration.py                     |  62 ++++--
 litellm-proxy-extras/migration_runbook.md  |  11 +-
 scripts/budget_ratchet_check.py            |  16 +-
 scripts/default_branch.py                  |  61 ++++++
 scripts/pre_commit_lint.sh                 |  20 +-
 scripts/ruff_strict_gate.py                |   9 +-
 scripts/test_quality_gate.py               |  12 +-
 scripts/type_check_gate.py                 |  13 +-
 scripts/type_discipline_gate.py            |   9 +-
 terraform/provider/RELEASING.md            |   2 +-
 tests/test_litellm/test_default_branch.py  | 212 +++++++++++++++++++++
 tests/test_litellm/test_pre_commit_lint.py |  46 ++++-
 15 files changed, 437 insertions(+), 94 deletions(-)
 create mode 100644 scripts/default_branch.py
 create mode 100644 tests/test_litellm/test_default_branch.py

diff --git a/CLAUDE.md b/CLAUDE.md
index 2bc39332817..41678432989 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -29,7 +29,7 @@ Never test structure of code only function of it
 
 End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`
 
-When creating PRs, don't set base to `main`. `litellm_internal_staging` is the default base branch and serves that purpose for both internal and external / OSS contributions
+When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
 
 When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
 
@@ -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
 
-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
+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 the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. 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
 
@@ -70,7 +70,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests,
 
 Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
 
-If you're an internal contributor, when creating a new PR, the typical flow is to branch off litellm_internal_staging and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
+If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
 
 Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
 
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 9ef1d5ae2b8..0443f1bed75 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -315,10 +315,12 @@ Ensure the UI builds successfully before submitting your PR:
 npm run build
 ```
 
+Local lint and budget checks follow origin's current default branch. They refresh it from the remote instead of trusting cached `origin/HEAD`. For an intentional comparison against another branch or commit, use `make check BASE_REF=<ref>` or the standalone gate's `--base <ref>` option. An explicit ref can also be used offline once it has been fetched locally. Without an override, unavailable remote metadata stops the check
+
 ## Submitting Your PR
 
 1. **Push your branch**: `git push origin your-feature-branch`
-2. **Create a PR**: Go to GitHub and open a pull request against [`litellm_internal_staging`](https://github.com/BerriAI/litellm/tree/litellm_internal_staging), which is the default base branch. Do not target `main`.
+2. **Create a PR**: Go to GitHub and open a pull request against the repository's current default branch. Run `python3 scripts/default_branch.py --branch` to check its name
 3. **Fill out the PR template**: Provide clear description of changes
 4. **Wait for review**: Maintainers will review and provide feedback
 5. **Address feedback**: Make requested changes and push updates
diff --git a/Makefile b/Makefile
index e17fdba3c85..50d431a7c98 100644
--- a/Makefile
+++ b/Makefile
@@ -34,7 +34,7 @@ help:
 	@echo "  make lint-basedpyright-budget-update - Ratchet basedpyright limits down by what this branch fixed"
 	@echo "  make lint-format        - Check ruff format formatting (matches CI)"
 	@echo "  make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its limit"
-	@echo "  make lint-gate        - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)"
+	@echo "  make lint-gate        - Strict ruff gate in CI-parity mode (fetches the default branch, simulates the merge)"
 	@echo "  make lint-ruff-budget-update - Ratchet ruff-strict-budget.json limits down by what this branch fixed"
 	@echo "  make lint-test-quality  - Gate the test suite against test-quality-budget.json"
 	@echo "  make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + test quality + basedpyright)"
@@ -60,6 +60,9 @@ help:
 
 UV := uv
 UV_RUN := $(UV) run --no-sync
+BASE_REF ?=
+export BASE_REF
+RESOLVE_BASE = python3 scripts/default_branch.py --base "$(BASE_REF)"
 
 # Machine-wide slot queue for the heavy targets below; python3 + stdlib only, so
 # it runs before any venv exists. See scripts/gate_slot_lock.py.
@@ -133,7 +136,7 @@ format-check: install-dev
 # Single fetch of the PR base so the delta-based gates below share one network round
 # trip instead of each re-fetching when chained from `lint`.
 lint-fetch-base:
-	git fetch origin litellm_internal_staging
+	@$(RESOLVE_BASE)
 
 # Mirror test-linting.yml's lint job environment: the proxy-dev group plus a generated
 # Prisma client, so `basedpyright tests/e2e` resolves the same modules CI does. The
@@ -150,7 +153,9 @@ lint-install:
 # recursively, so 'litellm/*.py' covers nested modules and the top-level files that
 # CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step.
 lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
-	@files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/*.py' | grep -v '^litellm/enterprise/' || true); \
+	@base_ref=$$($(RESOLVE_BASE)) && \
+	changed=$$(git diff --name-only --diff-filter=ACMR "$$base_ref...HEAD" -- 'litellm/*.py') && \
+	files=$$(printf '%s\n' "$$changed" | grep -v '^litellm/enterprise/' || true) || exit $$?; \
 	if [ -z "$$files" ]; then \
 		echo "No changed litellm Python files to format-check."; \
 	else \
@@ -167,7 +172,9 @@ lint-ruff: $(LINT_DEP_INSTALL)
 # https://github.com/astral-sh/ruff/discussions/10977
 # https://github.com/astral-sh/ruff/discussions/4049
 lint-format-changed: install-dev
-	@git diff origin/main --unified=0 --no-color -- '*.py' | \
+	@base_ref=$$($(RESOLVE_BASE)) && \
+	diff=$$(git diff "$$base_ref" --unified=0 --no-color -- '*.py') && \
+	printf '%s\n' "$$diff" | \
 	perl -ne '\
 		if (/^diff --git a\/(.*) b\//) { $$file = $$1; } \
 		if (/^@@ .* \+(\d+)(?:,(\d+))? @@/) { \
@@ -182,20 +189,22 @@ lint-format-changed: install-dev
 		done
 
 lint-ruff-dev: install-dev
-	@tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \
+	@base_ref=$$($(RESOLVE_BASE)) || exit $$?; \
+	tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \
 	cd litellm && \
 	($(UV_RUN) ruff check . --output-format=pylint || true) > "$$tmpfile" && \
-	$(UV_RUN) diff-quality --violations=pylint "$$tmpfile" --compare-branch=origin/main && \
+	$(UV_RUN) diff-quality --violations=pylint "$$tmpfile" --compare-branch="$$base_ref" && \
 	cd .. ; \
 	rm -f "$$tmpfile"
 
 lint-ruff-FULL-dev: install-dev
-	@files=$$(git diff --name-only origin/main -- '*.py'); \
+	@base_ref=$$($(RESOLVE_BASE)) && \
+	files=$$(git diff --name-only "$$base_ref" -- '*.py') || exit $$?; \
 	if [ -n "$$files" ]; then echo "$$files" | xargs $(UV_RUN) ruff check; \
 	else echo "No changed .py files to check."; fi
 
 lint-basedpyright: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
-	$(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging
+	$(UV_RUN) python scripts/type_check_gate.py --base "$(BASE_REF)"
 
 lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
 	$(UV_RUN) basedpyright tests/e2e
@@ -203,37 +212,37 @@ lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
 # Type-discipline budget (mutable collections / casts / type guards / kwargs /
 # unexplained suppressions), the test-linting.yml step `make lint` used to omit.
 lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
-	$(UV_RUN) python scripts/type_discipline_gate.py --base origin/litellm_internal_staging
+	$(UV_RUN) python scripts/type_discipline_gate.py --base "$(BASE_REF)"
 
 # Test-quality budget (zero-assert / mock-echo tests, sys.path.insert, raw env writes,
 # litellm module-global mutation, credential-gated skips, conftest snapshot
 # inventory), counted across tests/ the same delta-vs-base way.
 lint-test-quality: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
-	$(UV_RUN) python scripts/test_quality_gate.py --base origin/litellm_internal_staging
+	$(UV_RUN) python scripts/test_quality_gate.py --base "$(BASE_REF)"
 
 # --update lowers each limit by what this branch fixed since its branch point, so
 # it needs the base ref fetched to resolve the merge-base.
 lint-basedpyright-budget-update: install-dev lint-fetch-base
-	$(UV_RUN) python scripts/type_check_gate.py --update
+	$(UV_RUN) python scripts/type_check_gate.py --update --base "$(BASE_REF)"
 
 lint-format: format-check
 
 lint-ruff-budget: install-dev
-	$(UV_RUN) python scripts/ruff_strict_gate.py
+	$(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)"
 
 # Strict gate, invoked the same way CI does in test-linting.yml so a local pass
 # means the CI check will pass too.
 lint-gate: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
-	$(UV_RUN) python scripts/ruff_strict_gate.py --base origin/litellm_internal_staging
+	$(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)"
 
 lint-ruff-budget-update: install-dev lint-fetch-base
-	$(UV_RUN) python scripts/ruff_strict_gate.py --update
+	$(UV_RUN) python scripts/ruff_strict_gate.py --update --base "$(BASE_REF)"
 
 lint-type-discipline-budget-update: install-dev lint-fetch-base
-	$(UV_RUN) python scripts/type_discipline_gate.py --update
+	$(UV_RUN) python scripts/type_discipline_gate.py --update --base "$(BASE_REF)"
 
 lint-test-quality-budget-update: install-dev lint-fetch-base
-	$(UV_RUN) python scripts/test_quality_gate.py --update
+	$(UV_RUN) python scripts/test_quality_gate.py --update --base "$(BASE_REF)"
 
 # Ratchet all budgets in one shot (ruff strict + type-discipline + test quality + basedpyright)
 lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-test-quality-budget-update lint-basedpyright-budget-update
@@ -249,14 +258,15 @@ check-import-safety: $(LINT_DEP_INSTALL)
 # runs the diff-scoped ruff format check, whole-tree ruff check, the strict-rule /
 # type-discipline / basedpyright budgets as a delta vs the base, then the circular-import
 # and import-safety checks. Steps that compare against the base resolve it the same way CI
-# does (merge-base with origin/litellm_internal_staging). Setup (env sync, Prisma client,
+# does (merge-base with origin's current default branch). Setup (env sync, Prisma client,
 # base fetch) runs once up front; the checks themselves are independent, so a sub-make
 # fans them out with -j and the fast ones finish under basedpyright's shadow.
 lint:
 	@$(GATE_SLOT_LOCK) $(MAKE) lint-inner
 
-lint-inner: lint-install lint-fetch-base
-	$(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
+lint-inner: lint-install
+	@base_ref=$$($(RESOLVE_BASE)) && \
+	$(MAKE) BASE_REF="$$base_ref" -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
 
 lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
 
diff --git a/ci_cd/run_migration.py b/ci_cd/run_migration.py
index feec4046ee1..c737050f3a8 100644
--- a/ci_cd/run_migration.py
+++ b/ci_cd/run_migration.py
@@ -6,12 +6,9 @@ import subprocess
 import sys
 from datetime import datetime
 from pathlib import Path
-
-import testing.postgresql
-
+from typing import Final
 
 DESTRUCTIVE_PATTERN = re.compile(r"\bDROP\s+(COLUMN|TABLE|INDEX)\b", re.IGNORECASE)
-DEFAULT_BASE_BRANCH = "litellm_internal_staging"
 
 
 def _find_destructive_statements(sql: str) -> list:
@@ -94,31 +91,57 @@ def _print_stale_branch_refusal(base_branch: str, behind: int) -> None:
     print(banner, file=out)
 
 
-def _check_branch_freshness(root_dir: Path, base_branch: str) -> None:
+def _default_base_branch(root_dir: Path) -> str:
+    try:
+        result: Final = subprocess.run(
+            [
+                sys.executable,
+                str(Path(__file__).resolve().parents[1] / "scripts" / "default_branch.py"),
+                "--repo-root",
+                str(root_dir),
+                "--branch",
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+            timeout=90,
+        )
+    except (OSError, subprocess.SubprocessError) as exc:
+        _print_freshness_failure(
+            "default branch",
+            "Could not discover origin's default branch. Pass --base-branch <name> to choose one.",
+            exc.stderr if isinstance(exc, subprocess.CalledProcessError) else str(exc),
+        )
+        sys.exit(3)
+    return result.stdout.strip()
+
+
+def _check_branch_freshness(root_dir: Path, base_branch: str | None = None) -> None:
     """Fetch origin/<base_branch> and exit 3 if HEAD is behind it."""
+    resolved_branch: Final = base_branch or _default_base_branch(root_dir)
     cwd = str(root_dir)
     try:
         subprocess.run(
-            ["git", "fetch", "origin", base_branch],
+            ["git", "fetch", "origin", f"+refs/heads/{resolved_branch}:refs/remotes/origin/{resolved_branch}"],
             check=True,
             capture_output=True,
             text=True,
             cwd=cwd,
         )
     except FileNotFoundError:
-        _print_freshness_failure(base_branch, "git executable not found on PATH")
+        _print_freshness_failure(resolved_branch, "git executable not found on PATH")
         sys.exit(3)
     except subprocess.CalledProcessError as e:
         _print_freshness_failure(
-            base_branch,
-            f"`git fetch origin {base_branch}` failed",
+            resolved_branch,
+            f"`git fetch origin {resolved_branch}` failed",
             e.stderr or "",
         )
         sys.exit(3)
 
     try:
         result = subprocess.run(
-            ["git", "rev-list", "--count", f"HEAD..origin/{base_branch}"],
+            ["git", "rev-list", "--count", f"HEAD..origin/{resolved_branch}"],
             check=True,
             capture_output=True,
             text=True,
@@ -127,23 +150,23 @@ def _check_branch_freshness(root_dir: Path, base_branch: str) -> None:
         behind = int(result.stdout.strip())
     except subprocess.CalledProcessError as e:
         _print_freshness_failure(
-            base_branch,
-            f"`git rev-list HEAD..origin/{base_branch}` failed",
+            resolved_branch,
+            f"`git rev-list HEAD..origin/{resolved_branch}` failed",
             e.stderr or "",
         )
         sys.exit(3)
     except ValueError:
         _print_freshness_failure(
-            base_branch,
+            resolved_branch,
             "could not parse commit count from `git rev-list`",
         )
         sys.exit(3)
 
     if behind > 0:
-        _print_stale_branch_refusal(base_branch, behind)
+        _print_stale_branch_refusal(resolved_branch, behind)
         sys.exit(3)
 
-    print(f"Branch freshness OK: up to date with origin/{base_branch}.")
+    print(f"Branch freshness OK: up to date with origin/{resolved_branch}.")
 
 
 def _print_destructive_refusal(destructive_lines: list) -> None:
@@ -198,7 +221,7 @@ def _print_destructive_refusal(destructive_lines: list) -> None:
 def create_migration(
     migration_name: str = None,
     allow_destructive: bool = False,
-    base_branch: str = DEFAULT_BASE_BRANCH,
+    base_branch: str | None = None,
     skip_freshness_check: bool = False,
 ):
     """
@@ -211,7 +234,7 @@ def create_migration(
             DROP COLUMN, DROP TABLE, or DROP INDEX statements. Without this
             flag, the script exits non-zero and prints guidance.
         base_branch (str): Branch to check freshness against
-            (default: "litellm_internal_staging").
+            (default: origin's current default branch).
         skip_freshness_check (bool): Skip the "branch is up to date" check.
             Only for intentional migrations against an older base.
     """
@@ -225,6 +248,8 @@ def create_migration(
     else:
         _check_branch_freshness(root_dir, base_branch)
 
+    import testing.postgresql
+
     try:
         migrations_dir = (
             root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
@@ -342,9 +367,8 @@ if __name__ == "__main__":
     )
     parser.add_argument(
         "--base-branch",
-        default=DEFAULT_BASE_BRANCH,
         help=(
-            f"Branch to check freshness against (default: {DEFAULT_BASE_BRANCH}). "
+            "Branch to check freshness against (default: origin's current default branch). "
             "The script fetches origin/<base-branch> and refuses to run if HEAD "
             "is behind it."
         ),
diff --git a/litellm-proxy-extras/migration_runbook.md b/litellm-proxy-extras/migration_runbook.md
index a277441b164..b1e9236e520 100644
--- a/litellm-proxy-extras/migration_runbook.md
+++ b/litellm-proxy-extras/migration_runbook.md
@@ -48,7 +48,7 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n
 
 ## What It Does
 
-1. **Verifies the current branch is up to date with `origin/litellm_internal_staging`** (see [Branch freshness](#branch-freshness-check))
+1. **Verifies the current branch is up to date with origin's current default branch** (see [Branch freshness](#branch-freshness-check))
 2. Creates temp PostgreSQL DB
 3. Applies existing migrations
 4. Compares with `schema.prisma`
@@ -57,11 +57,11 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n
 
 ## Branch Freshness Check
 
-Before generating anything, `run_migration.py` runs `git fetch origin <base>` and refuses to proceed if `HEAD` is behind `origin/<base>`. Default base is `litellm_internal_staging` (the branch PRs target). A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense.
+Before generating anything, `run_migration.py` runs `git fetch origin <base>` and refuses to proceed if `HEAD` is behind `origin/<base>`. The default base is discovered from origin's advertised HEAD on each run, so an existing clone follows a default-branch change without trusting cached `origin/HEAD`. If discovery or fetching fails, migration generation stops. A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense.
 
 Flags:
 
-- `--base-branch <name>` — check against a different base (e.g. `main`). Default is `litellm_internal_staging`.
+- `--base-branch <name>` — check against a different base (e.g. a release branch). Defaults to origin's current default branch
 - `--skip-freshness-check` — bypass entirely. Only for intentional migrations against an older base.
 
 When the guard fires:
@@ -69,8 +69,9 @@ When the guard fires:
 1. Update your branch:
 
    ```bash
-   git fetch origin && git rebase origin/litellm_internal_staging
-   # or git merge origin/litellm_internal_staging — whichever matches your workflow
+   base_branch=$(python3 scripts/default_branch.py --branch) &&
+     git fetch origin "+refs/heads/$base_branch:refs/remotes/origin/$base_branch" &&
+     git rebase "origin/$base_branch"
    ```
 2. Re-run `run_migration.py`.
 
diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py
index adc4c0664be..485e118efd2 100644
--- a/scripts/budget_ratchet_check.py
+++ b/scripts/budget_ratchet_check.py
@@ -34,7 +34,7 @@ import subprocess
 import sys
 from pathlib import Path
 from types import MappingProxyType
-from typing import NamedTuple
+from typing import Final, NamedTuple
 
 if sys.version_info >= (3, 11):
     import tomllib
@@ -42,7 +42,6 @@ else:
     import tomli as tomllib
 
 REPO_ROOT = Path(__file__).resolve().parent.parent
-DEFAULT_BASE = "origin/litellm_internal_staging"
 DEFAULT_BUDGETS: tuple[str, ...] = (
     "ruff-strict-budget.json",
     "type-discipline-budget.json",
@@ -182,12 +181,15 @@ def regressions_for(
 
 def main() -> int:
     parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--base", default=DEFAULT_BASE)
+    parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)")
     parser.add_argument("budgets", nargs="*", help="budget files to check")
     args = parser.parse_args()
+    from default_branch import resolve_base_ref
+
+    base_ref: Final = resolve_base_ref(args.base, REPO_ROOT)
     budgets = args.budgets or list(DEFAULT_BUDGETS)
 
-    ref = _merge_base(args.base)
+    ref = _merge_base(base_ref)
     if not _ref_is_commit(ref):
         print(
             f"FAIL: base ref {ref!r} does not resolve to a commit, so the ratchet has nothing "
@@ -204,14 +206,14 @@ def main() -> int:
         if base is None and head is None:
             continue
         if base is None:
-            print(f"skip {rel}: new file (no base at {args.base} to ratchet against)")
+            print(f"skip {rel}: new file (no base at {base_ref} to ratchet against)")
             continue
         checked.append(rel)
         regressions.extend(regressions_for(rel, base, head, graduated_selectors(rel)))
 
     if regressions:
         print(
-            f"FAIL: budget limit(s) loosened vs base {args.base} (merge-base {ref[:12]}):"
+            f"FAIL: budget limit(s) loosened vs base {base_ref} (merge-base {ref[:12]}):"
         )
         for reg in regressions:
             print(f"  {reg.budget}  {reg.rule}: {reg.detail}")
@@ -223,7 +225,7 @@ def main() -> int:
         return 1
 
     suffix = f" ({', '.join(checked)})" if checked else ""
-    print(f"OK: no budget limit increased vs base {args.base}{suffix}")
+    print(f"OK: no budget limit increased vs base {base_ref}{suffix}")
     return 0
 
 
diff --git a/scripts/default_branch.py b/scripts/default_branch.py
new file mode 100644
index 00000000000..fb1852fd21c
--- /dev/null
+++ b/scripts/default_branch.py
@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import argparse
+import os
+import subprocess
+from pathlib import Path
+from typing import Final
+
+
+def _git(repo_root: Path, *args: str) -> str:
+    try:
+        result: Final = subprocess.run(
+            ["git", *args],
+            cwd=repo_root,
+            env={**os.environ, "GIT_TERMINAL_PROMPT": "0"},
+            check=True,
+            capture_output=True,
+            text=True,
+            timeout=60,
+        )
+    except (OSError, subprocess.SubprocessError) as exc:
+        raise SystemExit(
+            "Cannot verify the base branch against origin. Check remote access, "
+            "or supply an explicit base ref (--base / BASE_REF). "
+            f"Git operation failed: {exc}"
+        ) from exc
+    return result.stdout.strip()
+
+
+def default_branch(repo_root: Path) -> str:
+    output: Final = _git(repo_root, "ls-remote", "--symref", "origin", "HEAD")
+    branches: Final = tuple(
+        line.removeprefix("ref: refs/heads/").removesuffix("\tHEAD")
+        for line in output.splitlines()
+        if line.startswith("ref: refs/heads/") and line.endswith("\tHEAD")
+    )
+    if len(branches) != 1:
+        raise SystemExit("Origin did not advertise a default branch. Supply an explicit base ref (--base / BASE_REF).")
+    _git(repo_root, "check-ref-format", f"refs/heads/{branches[0]}")
+    return branches[0]
+
+
+def resolve_base_ref(base_ref: str | None, repo_root: Path) -> str:
+    if base_ref:
+        return base_ref
+    branch: Final = default_branch(repo_root)
+    _git(repo_root, "fetch", "--quiet", "origin", f"+refs/heads/{branch}:refs/remotes/origin/{branch}")
+    return f"origin/{branch}"
+
+
+def main() -> None:
+    parser: Final = argparse.ArgumentParser(description="Resolve the live default branch of origin.")
+    parser.add_argument("--base", help="Explicit comparison ref; skips default-branch discovery")
+    parser.add_argument("--repo-root", type=Path, default=Path.cwd())
+    parser.add_argument("--branch", action="store_true", help="Print only the default branch name, without fetching")
+    args: Final = parser.parse_args()
+    print(default_branch(args.repo_root) if args.branch else resolve_base_ref(args.base, args.repo_root))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh
index f245803408c..1abd415d237 100755
--- a/scripts/pre_commit_lint.sh
+++ b/scripts/pre_commit_lint.sh
@@ -7,7 +7,7 @@
 #   - anything staged -> scope is the staged files; changed-but-unstaged files
 #     whose checks were skipped are called out
 #   - nothing staged  -> scope is the working tree's diff against the merge base
-#     with origin/litellm_internal_staging, untracked files included
+#     with origin's current default branch, untracked files included
 # The per-area checks:
 #   - litellm/ Python  -> `make lint` (test-linting.yml's lint job)
 #   - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step)
@@ -33,8 +33,8 @@ set -eu
 # at a time instead of thrashing the machine. The wrapper exports
 # LITELLM_GATE_SLOT_HELD, so this re-exec happens exactly once and everything this
 # script spawns (make lint, the budget gates) skips its own acquisition.
+script_dir=$(python3 -c 'import os, sys; print(os.path.dirname(os.path.realpath(sys.argv[1])))' "$0")
 if [ -z "${LITELLM_GATE_SLOT_HELD:-}" ]; then
-    script_dir=$(python3 -c 'import os, sys; print(os.path.dirname(os.path.realpath(sys.argv[1])))' "$0")
     exec python3 "$script_dir/gate_slot_lock.py" "$0" "$@"
 fi
 
@@ -65,20 +65,24 @@ untracked=$(git ls-files --others --exclude-standard)
 if [ -n "$staged" ]; then
     scope=$staged
 else
-    git fetch --quiet origin litellm_internal_staging 2>/dev/null || true
-    merge_base=$(git merge-base origin/litellm_internal_staging HEAD 2>/dev/null) || {
-        echo "check: cannot resolve the merge base with origin/litellm_internal_staging." >&2
-        echo "  Fix: git fetch origin litellm_internal_staging" >&2
+    base_ref=$(python3 "$script_dir/default_branch.py" --base "${BASE_REF:-}") || {
+        echo "check: FAIL"
+        exit 1
+    }
+    export BASE_REF="$base_ref"
+    merge_base=$(git merge-base "$base_ref" HEAD 2>/dev/null) || {
+        echo "check: cannot resolve the merge base with $base_ref." >&2
+        echo "  Fix: fetch the base ref and provide BASE_REF=<ref>" >&2
         echo "check: FAIL"
         exit 1
     }
     scope=$(printf '%s\n' "$(git diff --name-only --diff-filter=ACMRD "$merge_base")" "$untracked" | sed '/^$/d' | sort -u)
     if [ -z "$scope" ]; then
-        echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs origin/litellm_internal_staging)"
+        echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs $base_ref)"
         echo "check: PASS"
         exit 0
     fi
-    echo "check: nothing staged; scoping to the working tree's diff against the merge base with origin/litellm_internal_staging:"
+    echo "check: nothing staged; scoping to the working tree's diff against the merge base with $base_ref:"
     printf '%s\n' "$scope" | sed 's/^/    /'
 fi
 
diff --git a/scripts/ruff_strict_gate.py b/scripts/ruff_strict_gate.py
index bf070beeb0f..8da10dd76f0 100644
--- a/scripts/ruff_strict_gate.py
+++ b/scripts/ruff_strict_gate.py
@@ -24,7 +24,6 @@ REPO_ROOT = Path(__file__).resolve().parent.parent
 STRICT_CONFIG = REPO_ROOT / "ruff-strict.toml"
 BUDGET_PATH = REPO_ROOT / "ruff-strict-budget.json"
 TARGET = "litellm"
-DEFAULT_BASE = "origin/litellm_internal_staging"
 
 _HUNK = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@")
 
@@ -193,7 +192,7 @@ def ratcheted_budget(budget: dict, current: dict, base: dict) -> dict:
     }
 
 
-def cmd_update(base_ref: str = DEFAULT_BASE) -> None:
+def cmd_update(base_ref: str) -> None:
     """Ratchet each rule's limit down by the violations this branch fixed.
 
     The working-tree count is compared against a ruff pass over a detached
@@ -212,13 +211,15 @@ def cmd_update(base_ref: str = DEFAULT_BASE) -> None:
 
 def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--base", default=DEFAULT_BASE)
+    parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)")
     parser.add_argument("--update", action="store_true")
     args = parser.parse_args()
+    from default_branch import resolve_base_ref
     from gate_slot_lock import held_slot
 
+    base_ref: Final = resolve_base_ref(args.base, REPO_ROOT)
     with held_slot():
-        cmd_update(args.base) if args.update else cmd_check(args.base)
+        cmd_update(base_ref) if args.update else cmd_check(base_ref)
 
 
 if __name__ == "__main__":
diff --git a/scripts/test_quality_gate.py b/scripts/test_quality_gate.py
index 4f4eeb17ec1..e486324c741 100644
--- a/scripts/test_quality_gate.py
+++ b/scripts/test_quality_gate.py
@@ -14,7 +14,7 @@ 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
+against the repository's default branch, 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.
@@ -43,7 +43,6 @@ REPO_ROOT: Final = Path(__file__).resolve().parent.parent
 CHECKER: Final = REPO_ROOT / "scripts" / "check_test_quality.py"
 BUDGET_PATH: Final = REPO_ROOT / "test-quality-budget.json"
 TARGET: Final = "tests"
-DEFAULT_BASE: Final = "origin/litellm_internal_staging"
 TERMINATION_SIGNALS: Final = (signal.SIGTERM, signal.SIGHUP)
 
 _HUNK: Final = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@", re.MULTILINE)
@@ -240,7 +239,7 @@ def ratcheted_budget(
     })
 
 
-def cmd_update(base_ref: str = DEFAULT_BASE) -> None:
+def cmd_update(base_ref: str) -> None:
     """Ratchet each rule's limit down by the violations this branch fixed."""
     budget: Final = json.loads(BUDGET_PATH.read_text())
     base_point: Final = resolve_base_point(base_ref)
@@ -264,19 +263,20 @@ def cmd_seed() -> None:
 
 def main() -> None:
     parser: Final = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--base", default=DEFAULT_BASE)
+    parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)")
     parser.add_argument("--update", action="store_true")
     parser.add_argument("--seed", action="store_true")
     args: Final = parser.parse_args()
+    from default_branch import resolve_base_ref
     from gate_slot_lock import held_slot
 
     with held_slot():
         if args.seed:
             cmd_seed()
         elif args.update:
-            cmd_update(args.base)
+            cmd_update(resolve_base_ref(args.base, REPO_ROOT))
         else:
-            cmd_check(args.base)
+            cmd_check(resolve_base_ref(args.base, REPO_ROOT))
 
 
 if __name__ == "__main__":
diff --git a/scripts/type_check_gate.py b/scripts/type_check_gate.py
index 763835e6d2e..78f74ec65a1 100644
--- a/scripts/type_check_gate.py
+++ b/scripts/type_check_gate.py
@@ -71,7 +71,6 @@ REPO_ROOT = Path(__file__).resolve().parent.parent
 BUDGET_PATH = REPO_ROOT / "basedpyright-code-budget.json"
 PYRIGHT_CONFIG = REPO_ROOT / "pyrightconfig.json"
 UV_LOCK = REPO_ROOT / "uv.lock"
-DEFAULT_BASE = "origin/litellm_internal_staging"
 CACHE_FILE_PREFIX = "basedpyright-base-"
 CACHE_KEEP_ENTRIES = 8
 ARTIFACT_NAME_PREFIX = "basedpyright-counts-"
@@ -578,7 +577,7 @@ def ratcheted_budget(
     }
 
 
-def cmd_update(current: Mapping[str, int], base_ref: str = DEFAULT_BASE) -> None:
+def cmd_update(current: Mapping[str, int], base_ref: str) -> None:
     """Ratchet each rule's limit down by the errors this branch fixed.
 
     `current` is the working-tree count; the reference count comes
@@ -666,12 +665,14 @@ def cmd_check(head: Mapping[str, int], base_ref: str) -> None:
 
 def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--base", default=DEFAULT_BASE)
+    parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)")
     parser.add_argument("--update", action="store_true")
     parser.add_argument("--emit-counts-dir", type=Path)
     args = parser.parse_args()
+    from default_branch import resolve_base_ref
     from gate_slot_lock import held_slot
 
+    base_ref: Final = None if args.emit_counts_dir is not None else resolve_base_ref(args.base, REPO_ROOT)
     with held_slot():
         ensure_typecheck_env()
         head = count_basedpyright(run_basedpyright())
@@ -679,10 +680,8 @@ def main() -> None:
             cmd_emit_counts(
                 head, args.emit_counts_dir, _run(["git", "rev-parse", "HEAD"]).strip()
             )
-        elif args.update:
-            cmd_update(head, args.base)
-        else:
-            cmd_check(head, args.base)
+        elif base_ref is not None:
+            cmd_update(head, base_ref) if args.update else cmd_check(head, base_ref)
 
 
 if __name__ == "__main__":
diff --git a/scripts/type_discipline_gate.py b/scripts/type_discipline_gate.py
index 5f6474f20bc..40e61cf7265 100644
--- a/scripts/type_discipline_gate.py
+++ b/scripts/type_discipline_gate.py
@@ -44,7 +44,6 @@ REPO_ROOT = Path(__file__).resolve().parent.parent
 CHECKER = REPO_ROOT / "scripts" / "check_type_discipline.py"
 BUDGET_PATH = REPO_ROOT / "type-discipline-budget.json"
 TARGET = "litellm"
-DEFAULT_BASE = "origin/litellm_internal_staging"
 
 _HUNK = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@")
 _LINE = re.compile(r"^(?P<file>.+?):(?P<line>\d+): (?P<code>LIT\d+) ")
@@ -239,7 +238,7 @@ def _base_budget_rules(base_point: str) -> frozenset:
     return frozenset(json.loads(proc.stdout))
 
 
-def cmd_update(base_ref: str = DEFAULT_BASE) -> None:
+def cmd_update(base_ref: str) -> None:
     """Ratchet each rule's limit down by the violations this branch fixed.
 
     The working-tree count is compared against a checker pass over a detached
@@ -264,13 +263,15 @@ def cmd_update(base_ref: str = DEFAULT_BASE) -> None:
 
 def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--base", default=DEFAULT_BASE)
+    parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)")
     parser.add_argument("--update", action="store_true")
     args = parser.parse_args()
+    from default_branch import resolve_base_ref
     from gate_slot_lock import held_slot
 
+    base_ref: Final = resolve_base_ref(args.base, REPO_ROOT)
     with held_slot():
-        cmd_update(args.base) if args.update else cmd_check(args.base)
+        cmd_update(base_ref) if args.update else cmd_check(base_ref)
 
 
 if __name__ == "__main__":
diff --git a/terraform/provider/RELEASING.md b/terraform/provider/RELEASING.md
index 59f4c5f066c..5b72621c291 100644
--- a/terraform/provider/RELEASING.md
+++ b/terraform/provider/RELEASING.md
@@ -79,7 +79,7 @@ Before publishing to the Terraform Registry:
 
 ## What a change needs
 
-1. **Land it in `BerriAI/litellm`.** Open a PR against `litellm_internal_staging` with the source change and a `CHANGELOG.md` entry under `[Unreleased]`. CI runs `gofmt`, `go vet`, build, tests and the endpoint-drift audit. A change that breaks existing configurations or state must say so in the changelog: the version number cannot signal it any more
+1. **Land it in `BerriAI/litellm`.** Open a PR against the repository's current default branch with the source change and a `CHANGELOG.md` entry under `[Unreleased]`. CI runs `gofmt`, `go vet`, build, tests and the endpoint-drift audit. A change that breaks existing configurations or state must say so in the changelog: the version number cannot signal it any more
 2. **Wait for the next LiteLLM release.** The nightly dev release carries it within a day; it reaches a stable version on the next stable cut
 3. **Verify** (optional): the version appears at https://registry.terraform.io/providers/BerriAI/litellm and https://github.com/BerriAI/terraform-provider-litellm/releases. If the tag is on the mirror but there is no release, the goreleaser run failed: https://github.com/BerriAI/terraform-provider-litellm/actions
 
diff --git a/tests/test_litellm/test_default_branch.py b/tests/test_litellm/test_default_branch.py
new file mode 100644
index 00000000000..ac673894067
--- /dev/null
+++ b/tests/test_litellm/test_default_branch.py
@@ -0,0 +1,212 @@
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+from typing import Final
+
+import pytest
+
+ROOT: Final = Path(__file__).resolve().parents[2]
+
+
+def _git(repo: Path, *args: str) -> str:
+    return subprocess.run(["git", *args], cwd=repo, check=True, capture_output=True, text=True).stdout.strip()
+
+
+def _commit(repo: Path, message: str) -> None:
+    _git(repo, "add", ".")
+    _git(repo, "-c", "user.name=Test", "-c", "user.email=test@example.com", "commit", "-qm", message)
+
+
+@pytest.fixture
+def remote_and_clone(tmp_path: Path) -> tuple[Path, Path]:
+    seed: Final = tmp_path / "seed"
+    seed.mkdir()
+    _git(seed, "init", "-q", "-b", "litellm_internal_staging")
+    (seed / "scripts").mkdir()
+    for name in (
+        "default_branch.py",
+        "budget_ratchet_check.py",
+        "ruff_strict_gate.py",
+        "type_discipline_gate.py",
+        "test_quality_gate.py",
+        "type_check_gate.py",
+        "gate_slot_lock.py",
+    ):
+        shutil.copyfile(ROOT / "scripts" / name, seed / "scripts" / name)
+    shutil.copyfile(ROOT / "Makefile", seed / "Makefile")
+    (seed / "litellm").mkdir()
+    (seed / "litellm" / "example.py").write_text("value = 0\n")
+    (seed / "ruff-strict-budget.json").write_text('{"C901": {"limit": 1}}\n')
+    _commit(seed, "staging base")
+    _git(seed, "checkout", "-qb", "main")
+    (seed / "litellm" / "example.py").write_text("value = 1\n")
+    (seed / "ruff-strict-budget.json").write_text('{"C901": {"limit": 0}}\n')
+    _commit(seed, "main base")
+    remote: Final = tmp_path / "remote.git"
+    _git(tmp_path, "clone", "-q", "--bare", str(seed), str(remote))
+    _git(remote, "symbolic-ref", "HEAD", "refs/heads/litellm_internal_staging")
+    repo: Final = tmp_path / "clone"
+    _git(tmp_path, "clone", "-q", "--single-branch", str(remote), str(repo))
+    return remote, repo
+
+
+def _resolve(repo: Path, *args: str) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        [sys.executable, str(ROOT / "scripts" / "default_branch.py"), *args],
+        cwd=repo,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+
+
+def _make(repo: Path, target: str, *args: str) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        ["make", target, "LINT_DEP_INSTALL=", "LINT_DEP_BASE=", *args],
+        cwd=repo,
+        capture_output=True,
+        text=True,
+        check=False,
+        env={key: value for key, value in os.environ.items() if key != "BASE_REF"},
+    )
+
+
+def test_existing_single_branch_clone_follows_remote_switch(remote_and_clone: tuple[Path, Path]) -> None:
+    remote, repo = remote_and_clone
+    before: Final = _resolve(repo)
+    assert before.returncode == 0, before.stderr
+    assert before.stdout.strip() == "origin/litellm_internal_staging"
+    _git(remote, "symbolic-ref", "HEAD", "refs/heads/main")
+    after: Final = _resolve(repo)
+    assert after.returncode == 0, after.stderr
+    assert after.stdout.strip() == "origin/main"
+    assert _git(repo, "rev-parse", "origin/main") == _git(remote, "rev-parse", "main")
+    assert _git(repo, "symbolic-ref", "refs/remotes/origin/HEAD").endswith("/litellm_internal_staging")
+
+
+@pytest.mark.parametrize("missing_head", [False, True])
+def test_unverifiable_default_never_uses_cached_head(
+    remote_and_clone: tuple[Path, Path],
+    missing_head: bool,
+) -> None:
+    remote, repo = remote_and_clone
+    if missing_head:
+        _git(remote, "symbolic-ref", "HEAD", "refs/heads/missing")
+    else:
+        _git(repo, "remote", "set-url", "origin", str(remote / "missing"))
+    result: Final = _resolve(repo)
+    assert result.returncode != 0
+    assert not result.stdout
+    assert "explicit base ref" in result.stderr
+    checked: Final = _make(repo, "lint-format-check-changed")
+    assert checked.returncode != 0
+    assert "No changed" not in checked.stdout
+
+
+@pytest.mark.parametrize("base_ref", ["HEAD", "origin/litellm_internal_staging"])
+def test_explicit_base_works_without_remote_access(
+    remote_and_clone: tuple[Path, Path],
+    base_ref: str,
+) -> None:
+    remote, repo = remote_and_clone
+    _git(repo, "remote", "set-url", "origin", str(remote / "missing"))
+    result: Final = _resolve(repo, "--base", base_ref)
+    assert result.returncode == 0, result.stderr
+    assert result.stdout.strip() == base_ref
+    checked: Final = _make(repo, "lint-format-check-changed", f"BASE_REF={base_ref}")
+    assert checked.returncode == 0, checked.stderr
+    assert "No changed litellm Python files" in checked.stdout
+
+
+def test_budget_ratchet_compares_against_new_default(remote_and_clone: tuple[Path, Path]) -> None:
+    remote, repo = remote_and_clone
+    _git(remote, "symbolic-ref", "HEAD", "refs/heads/main")
+    resolved: Final = _resolve(repo)
+    assert resolved.returncode == 0, resolved.stderr
+    _git(repo, "checkout", "-qb", "litellm_feature", "origin/main")
+    (repo / "ruff-strict-budget.json").write_text('{"C901": {"limit": 1}}\n')
+    command: Final = [sys.executable, "scripts/budget_ratchet_check.py"]
+    checked: Final = subprocess.run(command, cwd=repo, capture_output=True, text=True, check=False)
+    assert checked.returncode == 1
+    assert "limit raised 0 -> 1" in checked.stdout
+    assert "base origin/main" in checked.stdout
+    overridden: Final = subprocess.run(
+        [*command, "--base", "origin/litellm_internal_staging"],
+        cwd=repo,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert overridden.returncode == 0, overridden.stdout + overridden.stderr
+
+
+def _freshness(repo: Path, *args: str) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import sys; from pathlib import Path; "
+            "from ci_cd.run_migration import _check_branch_freshness; "
+            "_check_branch_freshness(Path(sys.argv[1]), sys.argv[2] if len(sys.argv) > 2 else None)",
+            str(repo),
+            *args,
+        ],
+        cwd=ROOT,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+
+
+def test_migration_freshness_refuses_stale_branch_after_switch(remote_and_clone: tuple[Path, Path]) -> None:
+    remote, repo = remote_and_clone
+    before: Final = _freshness(repo)
+    assert before.returncode == 0, before.stderr
+    assert "Branch freshness OK" in before.stdout
+    _git(remote, "symbolic-ref", "HEAD", "refs/heads/main")
+    after: Final = _freshness(repo)
+    assert after.returncode == 3
+    assert "1 commit(s) behind origin/main" in after.stderr
+    overridden: Final = _freshness(repo, "litellm_internal_staging")
+    assert overridden.returncode == 0, overridden.stderr
+    _git(repo, "merge", "--ff-only", "origin/main")
+    updated: Final = _freshness(repo)
+    assert updated.returncode == 0, updated.stderr
+    assert "up to date with origin/main" in updated.stdout
+
+
+def test_migration_freshness_refuses_unavailable_remote(remote_and_clone: tuple[Path, Path]) -> None:
+    remote, repo = remote_and_clone
+    _git(repo, "remote", "set-url", "origin", str(remote / "missing"))
+    result: Final = _freshness(repo)
+    assert result.returncode == 3
+    assert "Could not discover origin's default branch" in result.stderr
+    explicit: Final = _freshness(repo, "litellm_internal_staging")
+    assert explicit.returncode == 3
+    assert "git fetch origin litellm_internal_staging" in explicit.stderr
+
+
+@pytest.mark.parametrize(
+    "gate",
+    [
+        "budget_ratchet_check",
+        "ruff_strict_gate",
+        "type_discipline_gate",
+        "test_quality_gate",
+        "type_check_gate",
+    ],
+)
+def test_each_gate_refuses_an_unverifiable_default(remote_and_clone: tuple[Path, Path], gate: str) -> None:
+    remote, repo = remote_and_clone
+    _git(repo, "remote", "set-url", "origin", str(remote / "missing"))
+    result: Final = subprocess.run(
+        [sys.executable, f"scripts/{gate}.py"],
+        cwd=repo,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert result.returncode != 0
+    assert "Cannot verify the base branch against origin" in result.stderr
diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py
index b84cb8aa657..e12da0833dc 100644
--- a/tests/test_litellm/test_pre_commit_lint.py
+++ b/tests/test_litellm/test_pre_commit_lint.py
@@ -159,12 +159,12 @@ def _commit_all(repo: Path, message: str) -> None:
     )
 
 
-def _set_base_ref(repo: Path) -> None:
-    subprocess.run(
-        ["git", "update-ref", "refs/remotes/origin/litellm_internal_staging", "HEAD"],
-        cwd=repo,
-        check=True,
-    )
+def _set_base_ref(repo: Path, branch: str = "litellm_internal_staging") -> None:
+    remote = repo.parent / "remote.git"
+    subprocess.run(["git", "clone", "-q", "--bare", str(repo), str(remote)], check=True)
+    subprocess.run(["git", "update-ref", f"refs/heads/{branch}", "HEAD"], cwd=remote, check=True)
+    subprocess.run(["git", "symbolic-ref", "HEAD", f"refs/heads/{branch}"], cwd=remote, check=True)
+    subprocess.run(["git", "remote", "add", "origin", str(remote)], cwd=repo, check=True)
 
 
 def _stage_file(repo: Path, relative: str, body: str) -> None:
@@ -174,10 +174,11 @@ def _stage_file(repo: Path, relative: str, body: str) -> None:
     subprocess.run(["git", "add", relative], cwd=repo, check=True)
 
 
-def test_nothing_staged_scopes_to_working_tree_diff_and_runs_checks(tmp_path: Path) -> None:
+@pytest.mark.parametrize("branch", ["litellm_internal_staging", "main"])
+def test_nothing_staged_scopes_to_working_tree_diff_and_runs_checks(tmp_path: Path, branch: str) -> None:
     repo, bin_dir = _sandbox(tmp_path)
     _commit_all(repo, "base")
-    _set_base_ref(repo)
+    _set_base_ref(repo, branch)
     (repo / "litellm" / "foo.py").write_text("x = 2\n")
     proc = _run(repo, bin_dir, {})
     assert proc.returncode == 0, proc.stdout + proc.stderr
@@ -261,8 +262,8 @@ def test_nothing_staged_without_a_base_ref_fails_with_a_fetch_hint(tmp_path: Pat
     _commit_all(repo, "base")
     proc = _run(repo, bin_dir, {})
     assert proc.returncode == 1
-    assert "cannot resolve the merge base" in proc.stdout
-    assert "git fetch origin litellm_internal_staging" in proc.stdout
+    assert "Cannot verify the base branch against origin" in proc.stdout
+    assert "explicit base ref" in proc.stdout
     assert "check: FAIL" in proc.stdout
 
 
@@ -622,3 +623,28 @@ def test_failing_run_ends_with_a_fail_verdict(tmp_path: Path) -> None:
     assert proc.returncode == 1
     assert "check: FAIL" in proc.stdout
     assert "check: PASS" not in proc.stdout
+
+
+
+def test_explicit_base_scopes_offline_without_a_remote(tmp_path: Path) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    _commit_all(repo, "base")
+    (repo / "litellm" / "foo.py").write_text("x = 2\n")
+    proc = _run(repo, bin_dir, {"BASE_REF": "HEAD"})
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+    assert "merge base with HEAD" in proc.stdout
+    assert "linting Python" in proc.stdout
+
+
+def test_symlinked_hook_can_resolve_default_branch(tmp_path: Path) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    _commit_all(repo, "base")
+    _set_base_ref(repo, "main")
+    hook = repo / ".git" / "hooks" / "pre-commit"
+    hook.symlink_to(SCRIPT)
+    proc = subprocess.run(
+        [str(hook)], cwd=repo, capture_output=True, text=True,
+        env=_env(repo, bin_dir, {}), timeout=120,
+    )
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+    assert "no branch changes vs origin/main" in proc.stdout

From c5ec2eedc14fa68707d6cf4d77dcfe57b6b33ba5 Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Mon, 7 Sep 2026 14:41:05 -0700
Subject: [PATCH 142/164] fix(spend): price caching savings on the billed
 request basis (#40160)

Resolves LIT-7137

Co-authored-by: Claude Code <noreply@anthropic.com>
---
 .../litellm_core_utils/llm_cost_calc/utils.py |  86 ++++++++---
 litellm/proxy/db/db_spend_update_writer.py    |   2 +
 litellm/proxy/spend_tracking/savings.py       | 105 +++++--------
 .../llm_cost_calc/test_llm_cost_calc_utils.py |  38 +++++
 .../proxy/spend_tracking/test_savings.py      | 143 +++++++++++++++++-
 5 files changed, 288 insertions(+), 86 deletions(-)

diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py
index 9432fefc368..68dc27ec25e 100644
--- a/litellm/litellm_core_utils/llm_cost_calc/utils.py
+++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py
@@ -514,6 +514,7 @@ def _get_token_base_cost(
     current_time: datetime | None = None,
     *,
     threshold_is_inclusive: bool = False,
+    missing_cache_read_uses_input: bool = False,
 ) -> tuple[float, float, float, float, float]:
     """
     Return prompt cost, completion cost, and cache costs for a given model and usage.
@@ -524,6 +525,9 @@ def _get_token_base_cost(
     `threshold_is_inclusive` switches that comparison to >=, for providers such as xAI
     that bill the higher tier once the prompt reaches the threshold.
 
+    `missing_cache_read_uses_input` resolves an absent cache-read rate to the resolved
+    input rate instead of 0.0; an explicit 0.0 rate stays a real price either way.
+
     Returns:
         Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost)
     """
@@ -551,29 +555,16 @@ def _get_token_base_cost(
         float,
         _get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"),
     )
-    cache_read_cost = cast(float, _get_cost_per_unit(model_info, cache_read_cost_key))
+    cache_read_cost = _get_cost_per_unit(model_info, cache_read_cost_key, default_value=None)
 
     ## CHECK IF ABOVE THRESHOLD
     # Optimization: collect threshold keys first to avoid sorting all model_info keys.
-    # Most models don't have threshold pricing, so we can return early.
     # Exclude service_tier-specific variants (e.g. input_cost_per_token_above_200k_tokens_priority)
     # so that the threshold detection loop only processes standard keys.  The
     # service_tier-specific above-threshold key is resolved later via _get_service_tier_cost_key.
     threshold_keys: Final = [
         k for k in model_info if k.startswith("input_cost_per_token_above_") and not k.endswith(_SERVICE_TIER_SUFFIXES)
     ]
-    if not threshold_keys:
-        return _apply_off_peak_to_base_costs(
-            model_info,
-            current_time,
-            (
-                prompt_base_cost,
-                completion_base_cost,
-                cache_creation_cost,
-                cache_creation_cost_above_1hr,
-                cache_read_cost,
-            ),
-        )
 
     # Only sort the threshold keys (typically 1-2 keys instead of 66+)
     threshold: float | None = None
@@ -662,10 +653,7 @@ def _get_token_base_cost(
                         ),
                     )
 
-                    cache_read_cost = cast(
-                        float,
-                        _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost),
-                    )
+                    cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost)
 
                     break
             except (IndexError, ValueError):
@@ -673,6 +661,17 @@ def _get_token_base_cost(
             except Exception:
                 continue
 
+    if cache_read_cost is None:
+        cache_read_cost = (
+            _off_peak_rate(
+                _open_off_peak_block(model_info, current_time) or MappingProxyType({}),
+                "input_cost_per_token",
+                prompt_base_cost,
+            )
+            if missing_cache_read_uses_input
+            else 0.0
+        )
+
     return _apply_off_peak_to_base_costs(
         model_info,
         current_time,
@@ -1416,6 +1415,57 @@ def get_token_type_cost_breakdown(
     )
 
 
+def calculate_prompt_caching_savings(
+    model_info: ModelInfo,
+    usage: Usage,
+    custom_llm_provider: str | None,
+    service_tier: str | None = None,
+    data_residency: str | None = None,
+    vertex_location: str | None = None,
+    billed_at: datetime | None = None,
+) -> float:
+    """Read discount minus write premium, using the biller's rate and TTL resolution.
+
+    Missing reads and unpublished (missing/zero) writes claim no saving or premium;
+    explicit zero reads remain free. An unpublished 1h price uses the ordinary write rate.
+    ``billed_at`` is the request's completion time, so off-peak windows resolve as the
+    biller saw them rather than at the later spend write.
+    """
+    prompt_base_cost, _, cache_creation_cost, cache_creation_cost_above_1hr, cache_read_cost = _get_token_base_cost(
+        model_info=model_info,
+        usage=usage,
+        service_tier=service_tier,
+        current_time=billed_at,
+        threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider),
+        missing_cache_read_uses_input=True,
+    )
+    write_rate: Final = cache_creation_cost or prompt_base_cost
+    write_rate_1h: Final = cache_creation_cost_above_1hr or write_rate
+    prompt_tokens_details: Final = parse_prompt_tokens_details(usage)
+    cache_read_tokens: Final = max(prompt_tokens_details["cache_hit_tokens"], 0)
+    cache_creation_tokens: Final = max(prompt_tokens_details["cache_creation_tokens"], 0)
+    details: Final = prompt_tokens_details["cache_creation_token_details"]
+    cache_creation_details: Final = (
+        CacheCreationTokenDetails(
+            ephemeral_5m_input_tokens=max(details.ephemeral_5m_input_tokens or 0, 0),
+            ephemeral_1h_input_tokens=max(details.ephemeral_1h_input_tokens or 0, 0),
+        )
+        if details is not None
+        else None
+    )
+    read_discount: Final = cache_read_tokens * max(prompt_base_cost - cache_read_cost, 0.0)
+    write_premium: Final = calculate_cache_writing_cost(
+        cache_creation_tokens=cache_creation_tokens,
+        cache_creation_token_details=cache_creation_details,
+        cache_creation_cost_above_1hr=write_rate_1h - prompt_base_cost,
+        cache_creation_cost=write_rate - prompt_base_cost,
+    )
+    uplift: Final = _get_regional_uplift_multiplier(model_info, data_residency) * get_vertex_regional_endpoint_uplift(
+        model_info, vertex_location
+    )
+    return (read_discount - write_premium) * uplift
+
+
 def calculate_image_response_cost_from_usage(
     model: str,
     image_response: ImageResponse,
diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index 9230be8055e..914c961b145 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/litellm/proxy/db/db_spend_update_writer.py
@@ -502,6 +502,7 @@ class DBSpendUpdateWriter:
                 llm_router=get_llm_router,
                 cost_breakdown=metadata.get("cost_breakdown"),
                 recorded_autorouter_savings=metadata.get("autorouter_savings"),
+                billed_at=payload.get("endTime"),
             )
             transaction: Final = build_autorouter_turn_transaction(
                 payload=payload,
@@ -2188,6 +2189,7 @@ class DBSpendUpdateWriter:
                 usage_object=usage_obj,
                 cost_breakdown=_metadata.get("cost_breakdown"),
                 recorded_autorouter_savings=_metadata.get("autorouter_savings"),
+                billed_at=payload.get("endTime"),
             )
 
             daily_transaction: Final = BaseDailySpendTransaction(
diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py
index 1d0eb12da75..c541b9b40e5 100644
--- a/litellm/proxy/spend_tracking/savings.py
+++ b/litellm/proxy/spend_tracking/savings.py
@@ -9,12 +9,17 @@ have been aggregated across models.
 """
 
 from collections.abc import Callable, Mapping
+from datetime import datetime
 from typing import TYPE_CHECKING, Final, NamedTuple
 
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
-from litellm.litellm_core_utils.llm_cost_calc.utils import _get_cost_per_unit, generic_cost_per_token
+from litellm.litellm_core_utils.llm_cost_calc.utils import (
+    _get_cost_per_unit,
+    calculate_prompt_caching_savings,
+    generic_cost_per_token,
+)
 from litellm.types.integrations.anthropic_cache_control_hook import (
     GATEWAY_INJECTED_CACHE_METADATA_KEY,
     GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT,
@@ -32,42 +37,13 @@ class SavingsSpend(NamedTuple):
     gateway_injected_caching: float = 0.0
 
 
-def _input_cache_read_and_write_cost(info: ModelInfo | None) -> tuple[float, float, float]:
-    """
-    Return ``(input_cost, cache_read_cost, cache_write_cost)`` per token.
-
-    ``info`` is whatever pricing the caller resolved -- deployment rates when the
-    request came through a router deployment, public rates otherwise -- so a
-    negotiated price is honoured here rather than silently replaced by the list rate.
-    ``None`` falls open to ``(0.0, 0.0, 0.0)`` so savings degrade to zero rather than
-    raising inside the spend writer.
-
-    Prices are read through ``_get_cost_per_unit``, the same accessor the cost
-    calculator uses, which coerces the string prices a ``config.yaml`` can produce
-    (``"3e-7"``) and resolves service-tier suffixes.
-
-    An absent cache price mirrors the input cost, which yields a zero discount on the
-    read leg and a zero premium on the write leg. Mirroring rather than taking
-    ``_get_cost_per_unit``'s 0.0 default is load-bearing on the write leg: a zero write
-    price would make the premium ``0 - input_cost``, turning a model that simply has no
-    write pricing into a spurious extra saving.
-
-    The two legs then differ on an explicit ``0.0``, and the asymmetry is deliberate. A
-    free cache *write* does not exist -- entries carrying a literal zero (``deepseek-chat``
-    does) mean "no separate price", so a falsy write price also mirrors input. A free
-    cache *read* is real: 15 models charge for input and serve reads for nothing, which
-    is the largest discount available, so the read leg keeps its literal zero.
-    """
-    if info is None:
-        return 0.0, 0.0, 0.0
-    input_cost: Final = _get_cost_per_unit(info, "input_cost_per_token") or 0.0
-    cache_read_cost: Final = _get_cost_per_unit(info, "cache_read_input_token_cost", default_value=None)
-    cache_write_cost: Final = _get_cost_per_unit(info, "cache_creation_input_token_cost", default_value=None)
-    return (
-        input_cost,
-        input_cost if cache_read_cost is None else cache_read_cost,
-        cache_write_cost if cache_write_cost else input_cost,
-    )
+def _coerce_billed_at(value: datetime | str | None) -> datetime | None:
+    if isinstance(value, datetime) or value is None:
+        return value
+    try:
+        return datetime.fromisoformat(value.replace("Z", "+00:00"))
+    except ValueError:
+        return None
 
 
 class _ModelIdentity(NamedTuple):
@@ -586,6 +562,7 @@ def compute_savings_spend(
     llm_router: "Callable[[], Router | None] | None" = None,
     cost_breakdown: Mapping[str, object] | None = None,
     recorded_autorouter_savings: object = None,
+    billed_at: datetime | str | None = None,
 ) -> SavingsSpend:
     """
     Dollar savings for one request, split by optimization driver.
@@ -595,24 +572,10 @@ def compute_savings_spend(
     premium paid to write those entries, both derived here from ``usage_object`` so no
     caller can hand in a count that disagrees with the usage record.
 
-    The net form follows from what the request would have cost with caching off. The
-    provider reports ``prompt_tokens`` as the inclusive total of three disjoint
-    partitions (uncached text, cache reads, cache writes), so an uncached counterfactual
-    bills every one of those tokens at the flat input rate::
-
-        would_have_cost = (text + reads + writes) * input
-        actually_cost   = text * input + reads * read_rate + writes * write_rate
-        savings         = reads * (input - read_rate) - writes * (write_rate - input)
-
-    So the write leg subtracts the write PREMIUM, not the whole write cost: those tokens
-    had to be sent either way, and the counterfactual already pays the input rate for
-    them. The premium stays signed, because a handful of models price writes below their
-    input rate and there the write is a genuine extra saving.
-
-    A request that only writes cache and gets no hits therefore reports negative savings,
-    which is accurate: it really did cost more than the uncached call would have. The
-    daily rollup increments arithmetically, so those rows offset positive ones in the
-    same bucket.
+    The uncached counterfactual pays the ordinary input rate for the same prompt size
+    and tier. Cache writes subtract only the premium over that rate, split by TTL.
+    Savings stay signed: a write-only request can lose money, and daily rollups net
+    those losses against read savings.
 
     Caching is reported twice. ``prompt_caching`` is every net dollar caching saved,
     whoever caused it, which is what a customer means by "what did caching save me".
@@ -638,12 +601,9 @@ def compute_savings_spend(
     calls this and only auto-routed ones need one, so looking it up eagerly at the call
     site would fetch and discard it on the rest.
 
-    ``cost_breakdown`` is what the cost calculator recorded for this request, and it
-    carries both what the request really cost and the tier and region it was priced on.
-    Only the auto-router driver reads it. Compression and prompt caching price a
-    hypothetical token delta off flat rate keys, so they are blind to tiered pricing in
-    the same way; that is pre-existing behaviour on two shipped drivers rather than
-    something introduced here, and moving those numbers is its own change.
+    ``cost_breakdown`` supplies the biller's tier and region to caching and auto-router
+    savings. Caching also uses the logged prompt size and TTL split. Compression retains
+    its flat input-rate estimate; changing that counterfactual is a separate concern.
 
     ``recorded_autorouter_savings`` is the figure the logging path stamped on the spend
     log's metadata, honoured over recomputation so the rollup, the turn table and the
@@ -658,13 +618,24 @@ def compute_savings_spend(
     pricing: Final = _effective_model_info(router_instance, model_id, model or "") or (
         _model_info(identity) if identity else None
     )
-    input_cost, cache_read_cost, cache_write_cost = _input_cache_read_and_write_cost(pricing)
+    input_cost: Final = (_get_cost_per_unit(pricing, "input_cost_per_token") or 0.0) if pricing else 0.0
     compression: Final = max(compression_saved_tokens, 0) * input_cost
-    cache_read_input_tokens: Final = extract_cache_read_tokens(usage_object)
-    cache_creation_input_tokens: Final = extract_cache_creation_tokens(usage_object)
-    read_discount: Final = max(cache_read_input_tokens, 0) * max(input_cost - cache_read_cost, 0.0)
-    write_premium: Final = max(cache_creation_input_tokens, 0) * (cache_write_cost - input_cost)
-    prompt_caching: Final = read_discount - write_premium
+    usage: Final = _usage_from_spend_log(usage_object)
+    basis: Final = _pricing_basis(cost_breakdown)
+    billed_at_datetime: Final = _coerce_billed_at(billed_at)
+    prompt_caching: Final = (
+        calculate_prompt_caching_savings(
+            model_info=pricing,
+            usage=usage,
+            custom_llm_provider=identity.provider if identity else custom_llm_provider,
+            service_tier=basis.service_tier,
+            data_residency=basis.data_residency,
+            vertex_location=basis.vertex_location,
+            billed_at=billed_at_datetime,
+        )
+        if pricing is not None and usage is not None
+        else 0.0
+    )
     gateway_injected_caching: Final = prompt_caching if gateway_injected_cache else 0.0
 
     # The figure the logging path recorded wins, before the usage gate on purpose: a row
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 e1bd20ece6f..59f0938e338 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
@@ -49,6 +49,44 @@ def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
 
 
+@pytest.mark.parametrize("prompt_tokens", [100, 200000, 200001])
+@pytest.mark.parametrize("read_rate", [None, 0.0, 0.25e-6])
+@pytest.mark.parametrize("service_tier", [None, "priority"])
+def test_missing_cache_read_policy_preserves_billing(prompt_tokens, read_rate, service_tier):
+    info = {
+        "input_cost_per_token": 3e-6,
+        "input_cost_per_token_priority": 4e-6,
+        "input_cost_per_token_above_200k_tokens": 6e-6,
+        "input_cost_per_token_above_200k_tokens_priority": 8e-6,
+        "output_cost_per_token": 1e-6,
+        "cache_read_input_token_cost": read_rate,
+    }
+    usage = Usage(prompt_tokens=prompt_tokens, prompt_tokens_details={"cached_tokens": 100})
+    billed = _get_token_base_cost(info, usage, service_tier=service_tier)
+    savings = _get_token_base_cost(info, usage, service_tier=service_tier, missing_cache_read_uses_input=True)
+    prompt_cost, _ = generic_cost_per_token("policy-fixture", usage, "openai", service_tier=service_tier, model_info=info)
+    assert billed[4] == pytest.approx(read_rate or 0.0)
+    assert savings[:4] == billed[:4]
+    assert savings[4] == pytest.approx(billed[0] if read_rate is None else read_rate)
+    assert prompt_cost == pytest.approx((prompt_tokens - 100) * billed[0] + 100 * billed[4])
+
+
+def test_missing_cache_read_uses_off_peak_input_rate():
+    from datetime import datetime, timezone
+
+    info = {
+        "input_cost_per_token": 3e-6,
+        "off_peak_pricing": {"hours_utc": "00:00-23:59", "input_cost_per_token": 5e-6},
+    }
+    when = datetime(2026, 9, 7, 12, tzinfo=timezone.utc)
+    billed = _get_token_base_cost(info, Usage(prompt_tokens=100), current_time=when)
+    savings = _get_token_base_cost(
+        info, Usage(prompt_tokens=100), current_time=when, missing_cache_read_uses_input=True
+    )
+    assert billed[4] == 0.0
+    assert savings[0] == savings[4] == 5e-6
+
+
 def test_reasoning_tokens_no_price_set(_local_model_cost_map):
     # Use o1 - o1-mini was deprecated/renamed; o1 has same reasoning-token semantics
     # (no separate output_cost_per_reasoning_token, so all completion tokens use output_cost_per_token)
diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py
index 3f775d82b7f..cc8fdeb0160 100644
--- a/tests/test_litellm/proxy/spend_tracking/test_savings.py
+++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py
@@ -1,4 +1,4 @@
-
+from typing import Final
 
 import pytest
 
@@ -121,6 +121,147 @@ def _caching_usage(read: int, written: int, text: int = 10, out: int = 100) -> d
     }
 
 
+@pytest.mark.parametrize(
+    "model,provider,prompt,reads,writes_5m,writes_1h,tier,region,location",
+    [
+        pytest.param("claude-sonnet-4-5", "anthropic", 50000, 20000, 20000, 0, None, None, None, id="5m"),
+        pytest.param("claude-sonnet-4-5", "anthropic", 50000, 20000, 0, 20000, None, None, None, id="1h"),
+        pytest.param("claude-sonnet-4-5", "anthropic", 50000, 20000, 12000, 8000, None, None, None, id="mixed-ttl"),
+        pytest.param("claude-sonnet-4-5", "anthropic", 199999, 80000, 20000, 0, None, None, None, id="below-200k"),
+        pytest.param("claude-sonnet-4-5", "anthropic", 200000, 80000, 20000, 0, None, None, None, id="exactly-200k"),
+        pytest.param("claude-sonnet-4-5", "anthropic", 200001, 80000, 20000, 0, None, None, None, id="above-200k"),
+        pytest.param("claude-sonnet-4-5", "anthropic", 250000, 80000, 12000, 8000, None, None, None, id="ttl-and-200k"),
+        pytest.param(
+            "claude-sonnet-4-5", "anthropic", 250000, 80000, 0, 20000, "priority", None, None, id="absent-tier"
+        ),
+        pytest.param("gpt-5.5", "openai", 100000, 80000, 0, 0, "priority", None, None, id="priority"),
+        pytest.param("gpt-5.5", "openai", 100000, 80000, 0, 0, "flex", None, None, id="flex"),
+        pytest.param("gpt-5.5", "openai", 100000, 80000, 0, 0, "batch", None, None, id="batch-resolver-fallback"),
+        pytest.param("gpt-5.5", "openai", 300000, 80000, 0, 0, "flex", None, None, id="flex-and-272k"),
+        pytest.param("gpt-5.5", "openai", 100000, 80000, 0, 0, "priority", "eu", None, id="priority-and-eu"),
+        pytest.param("gemini-2.5-pro", "vertex_ai", 250000, 80000, 20000, 0, None, None, None, id="variant-only-write"),
+        pytest.param("gemini-3.5-flash", "vertex_ai", 100000, 80000, 0, 0, None, None, "us-east5", id="vertex-region"),
+        pytest.param("gpt-5.5", "openai", 300000, 0, 0, 0, "priority", "eu", None, id="no-cache"),
+    ],
+)
+def test_caching_savings_agree_with_biller_on_the_request_pricing_basis(
+    model: str,
+    provider: str,
+    prompt: int,
+    reads: int,
+    writes_5m: int,
+    writes_1h: int,
+    tier: str | None,
+    region: str | None,
+    location: str | None,
+) -> None:
+    pricing: Final = litellm.get_model_info(model=model, custom_llm_provider=provider)
+    usage: Final = Usage(
+        prompt_tokens=prompt,
+        completion_tokens=100,
+        total_tokens=prompt + 100,
+        prompt_tokens_details={
+            "cached_tokens": reads,
+            "cache_creation_tokens": writes_5m + writes_1h,
+            "text_tokens": prompt - reads - writes_5m - writes_1h,
+            "cache_creation_token_details": {
+                "ephemeral_5m_input_tokens": writes_5m,
+                "ephemeral_1h_input_tokens": writes_1h,
+            },
+        },
+    )
+    uncached: Final = Usage(
+        prompt_tokens=prompt,
+        completion_tokens=100,
+        total_tokens=prompt + 100,
+        prompt_tokens_details={"text_tokens": prompt, "cached_tokens": 0, "cache_creation_tokens": 0},
+    )
+    costs: Final = tuple(
+        sum(
+            generic_cost_per_token(
+                model=model,
+                usage=arm,
+                custom_llm_provider=provider,
+                model_info=pricing,
+                service_tier=tier,
+                data_residency=region,
+                vertex_location=location,
+            )
+        )
+        for arm in (uncached, usage)
+    )
+    expected: Final = costs[0] - costs[1]
+    for attributed in (False, True):
+        result: Final = compute_savings_spend(
+            model=model,
+            custom_llm_provider=provider,
+            compression_saved_tokens=4389,
+            gateway_injected_cache=attributed,
+            usage_object=usage.model_dump(),
+            cost_breakdown={"service_tier": tier, "data_residency": region, "vertex_location": location},
+            billed_at="2026-09-07T12:00:00+00:00",
+        )
+        assert result.prompt_caching == pytest.approx(expected)
+        assert result.gateway_injected_caching == pytest.approx(expected if attributed else 0.0)
+        assert result.compression == pytest.approx(4389 * (pricing["input_cost_per_token"] or 0.0))
+        assert result.autorouter == 0.0
+    if reads + writes_5m + writes_1h == 0:
+        assert expected == 0.0
+
+
+def test_negative_ttl_counts_do_not_become_cache_write_credits() -> None:
+    results: Final = tuple(
+        compute_savings_spend(
+            model="claude-sonnet-4-5",
+            custom_llm_provider="anthropic",
+            compression_saved_tokens=0,
+            gateway_injected_cache=True,
+            usage_object={
+                "prompt_tokens": 6000,
+                "completion_tokens": 100,
+                "prompt_tokens_details": {
+                    "text_tokens": 1000,
+                    "cache_creation_tokens": 5000,
+                    "cache_creation_token_details": {
+                        "ephemeral_5m_input_tokens": short_count,
+                        "ephemeral_1h_input_tokens": 5000,
+                    },
+                },
+            },
+        )
+        for short_count in (-5000, 0)
+    )
+    assert results[0] == results[1]
+    assert results[0].prompt_caching < 0
+
+
+def test_unpublished_one_hour_price_uses_the_ordinary_write_price() -> None:
+    model: Final = "claude-4-opus-20250514"
+    pricing: Final = litellm.get_model_info(model=model, custom_llm_provider="anthropic")
+    assert pricing.get("cache_creation_input_token_cost_above_1hr") is None
+    assert pricing["cache_creation_input_token_cost"] > pricing["input_cost_per_token"]
+    results: Final = tuple(
+        compute_savings_spend(
+            model=model,
+            custom_llm_provider="anthropic",
+            compression_saved_tokens=0,
+            gateway_injected_cache=True,
+            usage_object={
+                "prompt_tokens": 6000,
+                "completion_tokens": 100,
+                "prompt_tokens_details": {
+                    "text_tokens": 1000,
+                    "cache_creation_tokens": 5000,
+                    "cache_creation_token_details": ttl,
+                },
+            },
+        )
+        for ttl in (None, {"ephemeral_1h_input_tokens": 5000})
+    )
+    assert results[0] == results[1]
+    assert results[0].prompt_caching < 0
+
+
 def test_prompt_caching_savings_nets_out_the_cache_write_premium():
     """A cache-writing request is only credited the read discount minus the write premium."""
     input_cost, cache_read_cost = _anthropic_costs("claude-sonnet-5")

From 8e5a12057ab5733cb3c71e05aed8c29c5f295740 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Mon, 7 Sep 2026 14:55:43 -0700
Subject: [PATCH 143/164] feat(ui): list the ChatGPT subscription provider in
 the Add Model form

The Add Model provider dropdown is driven entirely by provider_create_fields.json,
and chatgpt had no entry there, so the documented ChatGPT subscription setup was
unreachable from the Admin UI. Add the entry plus the dashboard enum, slug, logo and
placeholder mappings so the provider can be selected and its cost-map models listed.

The entry carries no credential fields on purpose: the chatgpt backend ignores
api_key and api_base and signs in through the device-code auth file on the proxy
host, so any field here would be inert.

Add a parity test that every LlmProviders value is either listed for Add Model or
frozen in an explicit unlisted set, so a new backend provider cannot silently miss
the dropdown again.

Claude-Session: https://claude.ai/code/session_011Tn3657NkV6ojLqewL64Kb
---
 .../provider_create_fields.json               |  7 ++
 .../public_endpoints/test_public_endpoints.py | 86 +++++++++++++++++++
 .../components/provider_info_helpers.test.tsx | 14 +++
 .../src/components/provider_info_helpers.tsx  |  4 +
 4 files changed, 111 insertions(+)

diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json
index 66f8c2ea36f..cd781abee26 100644
--- a/litellm/proxy/public_endpoints/provider_create_fields.json
+++ b/litellm/proxy/public_endpoints/provider_create_fields.json
@@ -688,6 +688,13 @@
     ],
     "default_model_placeholder": "gpt-3.5-turbo"
   },
+  {
+    "provider": "CHATGPT",
+    "provider_display_name": "ChatGPT Subscription",
+    "litellm_provider": "chatgpt",
+    "credential_fields": [],
+    "default_model_placeholder": "chatgpt/gpt-5.4"
+  },
   {
     "provider": "CLARIFAI",
     "provider_display_name": "Clarifai",
diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
index 4a19ad3541c..fade7c9e7ee 100644
--- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
+++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
@@ -1,5 +1,6 @@
 import re
 from datetime import datetime, timezone
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -327,6 +328,91 @@ def test_cognition_provider_fields():
     assert fields_by_key["api_base"]["required"] is False
 
 
+def test_chatgpt_provider_fields():
+    """The ChatGPT subscription provider must be selectable in the Add Model flow (LIT-7127).
+
+    Its backend signs in through the device-code auth file on the proxy host and ignores
+    api_key/api_base, so the entry carries no credential fields: any field here would be inert.
+    """
+    app_instance = FastAPI()
+    app_instance.include_router(router)
+    test_client = TestClient(app_instance)
+
+    response = test_client.get("/public/providers/fields")
+    assert response.status_code == 200
+    providers = response.json()
+
+    chatgpt = next((p for p in providers if p["provider"] == "CHATGPT"), None)
+    assert chatgpt is not None, "ChatGPT provider entry not found"
+
+    assert chatgpt["provider_display_name"] == "ChatGPT Subscription"
+    assert chatgpt["litellm_provider"] == LlmProviders.CHATGPT.value
+    assert chatgpt["default_model_placeholder"].startswith("chatgpt/")
+    assert chatgpt["credential_fields"] == []
+
+
+ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset(
+    {
+        "a2a",
+        "a2a_agent",
+        "amazon_nova",
+        "apertis",
+        "aws_polly",
+        "black_forest_labs",
+        "charity_engine",
+        "chutes",
+        "darkbloom",
+        "gdc",
+        "helicone",
+        "inception",
+        "langflow",
+        "langgraph",
+        "libertai",
+        "litellm_agent",
+        "manus",
+        "meta",
+        "modelscope",
+        "mongodb",
+        "nano-gpt",
+        "neosantara",
+        "parasail",
+        "pinstripes",
+        "poe",
+        "publicai",
+        "ragflow",
+        "reducto",
+        "s3_vectors",
+        "sagemaker_nova",
+        "scaleway",
+        "stability",
+        "synthetic",
+        "tencent",
+        "tensormesh",
+        "text-completion-inception",
+        "valkey",
+        "xiaomi_mimo",
+        "zai",
+    }
+)
+
+
+def test_every_backend_provider_is_listed_in_add_model_or_frozen_as_unlisted():
+    """A provider LiteLLM ships must be reachable from the Add Model dropdown, which is driven
+    entirely by /public/providers/fields (LIT-7127). Providers that predate this check are frozen
+    in ADD_MODEL_UNLISTED_PROVIDERS; a new provider gets a JSON entry rather than a line here.
+    """
+    app_instance = FastAPI()
+    app_instance.include_router(router)
+    test_client = TestClient(app_instance)
+
+    response = test_client.get("/public/providers/fields")
+    assert response.status_code == 200
+    listed = {p["litellm_provider"] for p in response.json()}
+
+    unlisted = {provider.value for provider in LlmProviders} - listed
+    assert unlisted == ADD_MODEL_UNLISTED_PROVIDERS
+
+
 def test_google_ai_studio_provider_fields_expose_api_base():
     """The Google AI Studio (gemini) credential form must let admins set a custom
     api_base so they can point at a Gemini-compatible gateway (e.g. a self-hosted
diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx
index dfc737ddd45..4c68e302267 100644
--- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx
+++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx
@@ -89,6 +89,16 @@ describe("provider_info_helpers", () => {
       expect(result.logo).toBe(providerLogoMap[Providers.BedrockMantle]);
     });
 
+    it("should map the chatgpt slug and CHATGPT enum key to the ChatGPT Subscription name and OpenAI logo", () => {
+      const fromSlug = getProviderLogoAndName("chatgpt");
+      expect(fromSlug.displayName).toBe("ChatGPT Subscription");
+      expect(fromSlug.logo).toContain("openai_small");
+
+      const fromEnumKey = getProviderLogoAndName("CHATGPT");
+      expect(fromEnumKey.displayName).toBe("ChatGPT Subscription");
+      expect(fromEnumKey.logo).toContain("openai_small");
+    });
+
     it("should handle provider values case-insensitively", () => {
       const result = getProviderLogoAndName("OPENAI");
       expect(result.displayName).toBe(Providers.OpenAI);
@@ -272,6 +282,10 @@ describe("provider_info_helpers", () => {
       expect(getPlaceholder(Providers.Cognition)).toBe("cognition/swe-1.7");
     });
 
+    it("should return a chatgpt/ placeholder for the CHATGPT dropdown key", () => {
+      expect(getPlaceholder("CHATGPT")).toBe("chatgpt/gpt-5.4");
+    });
+
     it("should return default gpt-3.5-turbo placeholder for unknown provider", () => {
       expect(getPlaceholder("UnknownProvider" as any)).toBe("gpt-3.5-turbo");
     });
diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx
index d01a6a34cbe..72ec5990557 100644
--- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx
+++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx
@@ -84,6 +84,7 @@ export enum Providers {
   BASETEN = "Baseten",
   BYTEZ = "Bytez",
   Cerebras = "Cerebras",
+  CHATGPT = "ChatGPT Subscription",
   CLARIFAI = "Clarifai",
   CLOUDFLARE = "Cloudflare",
   CODESTRAL = "Codestral",
@@ -198,6 +199,7 @@ export const provider_map: Record<string, string> = {
   BedrockMantle: "bedrock_mantle",
   BYTEZ: "bytez",
   Cerebras: "cerebras",
+  CHATGPT: "chatgpt",
   CLARIFAI: "clarifai",
   CLOUDFLARE: "cloudflare",
   CODESTRAL: "codestral",
@@ -314,6 +316,7 @@ export const providerLogoMap: Partial<Record<Providers, string>> = {
   [Providers.BedrockMantle]: bedrockLogo.src,
   [Providers.SageMaker]: bedrockLogo.src,
   [Providers.Cerebras]: cerebrasLogo.src,
+  [Providers.CHATGPT]: openaiSmallLogo.src,
   [Providers.CLOUDFLARE]: cloudflareLogo.src,
   [Providers.CODESTRAL]: mistralLogo.src,
   [Providers.Cohere]: cohereLogo.src,
@@ -425,6 +428,7 @@ const providerPlaceholderMap: Partial<Record<Providers, string>> = {
   [Providers.Azure]: "my-deployment",
   [Providers.Azure_AI_Studio]: "azure_ai/command-r-plus",
   [Providers.Bedrock]: "claude-3-opus",
+  [Providers.CHATGPT]: "chatgpt/gpt-5.4",
   [Providers.Cognition]: "cognition/swe-1.7",
   [Providers.Cursor]: "cursor/claude-4-sonnet",
   [Providers.DeepInfra]: "deepinfra/<any-model-on-deepinfra>",

From bb2db2d3f8e218c6a781e029223a8af903e9d6dc Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Mon, 7 Sep 2026 15:12:41 -0700
Subject: [PATCH 144/164] test(proxy): drop docstrings from the Add Model
 provider tests

Move the only guidance worth keeping into the parity assertion message so a
failing run tells the contributor to add a catalog entry instead of growing
the frozen unlisted set.

Claude-Session: https://claude.ai/code/session_011Tn3657NkV6ojLqewL64Kb
---
 .../public_endpoints/test_public_endpoints.py      | 14 ++++----------
 1 file changed, 4 insertions(+), 10 deletions(-)

diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
index fade7c9e7ee..0d82ed778f5 100644
--- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
+++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
@@ -329,11 +329,6 @@ def test_cognition_provider_fields():
 
 
 def test_chatgpt_provider_fields():
-    """The ChatGPT subscription provider must be selectable in the Add Model flow (LIT-7127).
-
-    Its backend signs in through the device-code auth file on the proxy host and ignores
-    api_key/api_base, so the entry carries no credential fields: any field here would be inert.
-    """
     app_instance = FastAPI()
     app_instance.include_router(router)
     test_client = TestClient(app_instance)
@@ -397,10 +392,6 @@ ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset(
 
 
 def test_every_backend_provider_is_listed_in_add_model_or_frozen_as_unlisted():
-    """A provider LiteLLM ships must be reachable from the Add Model dropdown, which is driven
-    entirely by /public/providers/fields (LIT-7127). Providers that predate this check are frozen
-    in ADD_MODEL_UNLISTED_PROVIDERS; a new provider gets a JSON entry rather than a line here.
-    """
     app_instance = FastAPI()
     app_instance.include_router(router)
     test_client = TestClient(app_instance)
@@ -410,7 +401,10 @@ def test_every_backend_provider_is_listed_in_add_model_or_frozen_as_unlisted():
     listed = {p["litellm_provider"] for p in response.json()}
 
     unlisted = {provider.value for provider in LlmProviders} - listed
-    assert unlisted == ADD_MODEL_UNLISTED_PROVIDERS
+    assert unlisted == ADD_MODEL_UNLISTED_PROVIDERS, (
+        "Add Model dropdown drift: give the new provider an entry in provider_create_fields.json "
+        "rather than adding it to ADD_MODEL_UNLISTED_PROVIDERS"
+    )
 
 
 def test_google_ai_studio_provider_fields_expose_api_base():

From 15ba07f20d3a241593aadaa02ce6c7a4c177c000 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Mon, 7 Sep 2026 15:28:01 -0700
Subject: [PATCH 145/164] ci: avoid duplicate default branch fetches

---
 Makefile                                  | 12 ++--
 tests/test_litellm/test_default_branch.py | 28 +++++++++
 tests/test_litellm/test_gate_slot_lock.py | 76 ++++++++++++++---------
 3 files changed, 79 insertions(+), 37 deletions(-)

diff --git a/Makefile b/Makefile
index 50d431a7c98..ab11220821f 100644
--- a/Makefile
+++ b/Makefile
@@ -70,7 +70,7 @@ GATE_SLOT_LOCK := python3 scripts/gate_slot_lock.py
 
 LINT_DEP_INSTALL ?= install-dev
 LINT_E2E_DEP_INSTALL ?= lint-install
-LINT_DEP_BASE ?= lint-fetch-base
+LINT_DEP_BASE ?=
 LINT_JOBS := $(shell sysctl -n hw.ncpu 2>/dev/null || nproc 2>/dev/null || echo 4)
 LINT_OUTPUT_SYNC := $(if $(filter output-sync,$(.FEATURES)),--output-sync=target,)
 
@@ -133,8 +133,6 @@ format: install-dev
 format-check: install-dev
 	cd litellm && $(UV_RUN) ruff format --check --exclude '/enterprise/' . && cd ..
 
-# Single fetch of the PR base so the delta-based gates below share one network round
-# trip instead of each re-fetching when chained from `lint`.
 lint-fetch-base:
 	@$(RESOLVE_BASE)
 
@@ -222,7 +220,7 @@ lint-test-quality: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
 
 # --update lowers each limit by what this branch fixed since its branch point, so
 # it needs the base ref fetched to resolve the merge-base.
-lint-basedpyright-budget-update: install-dev lint-fetch-base
+lint-basedpyright-budget-update: install-dev
 	$(UV_RUN) python scripts/type_check_gate.py --update --base "$(BASE_REF)"
 
 lint-format: format-check
@@ -235,13 +233,13 @@ lint-ruff-budget: install-dev
 lint-gate: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
 	$(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)"
 
-lint-ruff-budget-update: install-dev lint-fetch-base
+lint-ruff-budget-update: install-dev
 	$(UV_RUN) python scripts/ruff_strict_gate.py --update --base "$(BASE_REF)"
 
-lint-type-discipline-budget-update: install-dev lint-fetch-base
+lint-type-discipline-budget-update: install-dev
 	$(UV_RUN) python scripts/type_discipline_gate.py --update --base "$(BASE_REF)"
 
-lint-test-quality-budget-update: install-dev lint-fetch-base
+lint-test-quality-budget-update: install-dev
 	$(UV_RUN) python scripts/test_quality_gate.py --update --base "$(BASE_REF)"
 
 # Ratchet all budgets in one shot (ruff strict + type-discipline + test quality + basedpyright)
diff --git a/tests/test_litellm/test_default_branch.py b/tests/test_litellm/test_default_branch.py
index ac673894067..a1b2a8c5c91 100644
--- a/tests/test_litellm/test_default_branch.py
+++ b/tests/test_litellm/test_default_branch.py
@@ -1,3 +1,4 @@
+import json
 import os
 import shutil
 import subprocess
@@ -210,3 +211,30 @@ def test_each_gate_refuses_an_unverifiable_default(remote_and_clone: tuple[Path,
     )
     assert result.returncode != 0
     assert "Cannot verify the base branch against origin" in result.stderr
+
+
+@pytest.mark.parametrize(
+    "target", ["lint-format-check-changed", "lint-test-quality", "lint-test-quality-budget-update"]
+)
+def test_direct_make_target_fetches_default_once(remote_and_clone: tuple[Path, Path], target: str) -> None:
+    _, repo = remote_and_clone
+    trace: Final = repo.parent / "git-trace.jsonl"
+    shutil.copyfile(ROOT / "scripts" / "check_test_quality.py", repo / "scripts" / "check_test_quality.py")
+    shutil.copyfile(ROOT / "test-quality-budget.json", repo / "test-quality-budget.json")
+    (repo / "tests").mkdir()
+    result: Final = subprocess.run(
+        ["make", "-o", "install-dev", target, "LINT_DEP_INSTALL=", "UV_RUN=env"],
+        cwd=repo,
+        capture_output=True,
+        text=True,
+        check=False,
+        env={**{key: value for key, value in os.environ.items() if key != "BASE_REF"}, "GIT_TRACE2_EVENT": str(trace)},
+    )
+    assert result.returncode == 0, result.stdout + result.stderr
+    commands: Final = tuple(
+        event["argv"][1:]
+        for line in trace.read_text().splitlines()
+        if (event := json.loads(line)).get("event") == "start"
+    )
+    assert sum(command[0] == "ls-remote" for command in commands) == 1
+    assert sum(command[0] == "fetch" for command in commands) == 1
diff --git a/tests/test_litellm/test_gate_slot_lock.py b/tests/test_litellm/test_gate_slot_lock.py
index 17fa8547ce7..c80e876700b 100644
--- a/tests/test_litellm/test_gate_slot_lock.py
+++ b/tests/test_litellm/test_gate_slot_lock.py
@@ -1,6 +1,8 @@
 import fcntl
 import importlib.util
+import json
 import os
+import shlex
 import signal
 import subprocess
 import sys
@@ -301,33 +303,47 @@ def test_held_slot_context_manager_releases_on_exit(tmp_path: Path, monkeypatch:
         fcntl.flock(probe, fcntl.LOCK_UN)
 
 
-def _make_rule(target: str) -> tuple[list[str], list[str]]:
-    database = subprocess.run(
-        ["make", "--dry-run", "--print-data-base", "info"],
-        cwd=ROOT,
-        capture_output=True,
-        text=True,
-        check=True,
-    ).stdout
-    lines = database.splitlines()
-    for index, line in enumerate(lines):
-        if line != f"{target}:" and not line.startswith(f"{target}: "):
-            continue
-        recipe: list[str] = []
-        for follower in lines[index + 1 :]:
-            if follower.startswith("#"):
-                continue
-            if not follower.startswith("\t"):
-                break
-            recipe.append(follower.strip())
-        return line.split(":", 1)[1].split(), recipe
-    raise AssertionError(f"target {target} not found in make database")
-
-
-def test_direct_make_lint_takes_a_slot_before_any_setup() -> None:
-    lint_prerequisites, lint_recipe = _make_rule("lint")
-    assert lint_prerequisites == []
-    assert any("$(GATE_SLOT_LOCK)" in line for line in lint_recipe)
-    inner_prerequisites, _ = _make_rule("lint-inner")
-    assert "lint-install" in inner_prerequisites
-    assert "lint-fetch-base" in inner_prerequisites
+def test_direct_make_lint_takes_a_slot_before_any_setup(tmp_path: Path) -> None:
+    lock_dir = tmp_path / "locks"
+    lock_dir.mkdir()
+    events_file = tmp_path / "setup.jsonl"
+    stderr_file = tmp_path / "make.stderr"
+    probe = tmp_path / "probe.py"
+    probe.write_text(
+        "import fcntl, json, os, pathlib, sys\n"
+        "with (pathlib.Path(os.environ['LITELLM_GATE_SLOT_DIR']) / 'slot-0.lock').open('wb') as slot:\n"
+        "    try:\n"
+        "        fcntl.flock(slot, fcntl.LOCK_EX | fcntl.LOCK_NB)\n"
+        "        locked = False\n"
+        "    except BlockingIOError:\n"
+        "        locked = True\n"
+        "with open(os.environ['EVENTS_FILE'], 'a') as events:\n"
+        "    events.write(json.dumps({'phase': sys.argv[1], 'locked': locked}) + '\\n')\n"
+        "if sys.argv[1] == 'base':\n"
+        "    print('HEAD')\n"
+    )
+    (tmp_path / "Makefile").write_text((ROOT / "Makefile").read_text())
+    command = [
+        "make", "-o", "lint-checks", "lint", "MAKE=make -o lint-checks",
+        f"GATE_SLOT_LOCK={shlex.join([sys.executable, str(HELPER)])}",
+        f"UV={shlex.join([sys.executable, str(probe), 'setup'])}",
+        f"UV_RUN={shlex.join([sys.executable, str(probe), 'setup'])}",
+        f"RESOLVE_BASE={shlex.join([sys.executable, str(probe), 'base'])}",
+    ]
+    with (lock_dir / "slot-0.lock").open("wb") as held, stderr_file.open("wb") as stderr:
+        fcntl.flock(held, fcntl.LOCK_EX)
+        process = subprocess.Popen(
+            command, cwd=tmp_path, stdout=subprocess.DEVNULL, stderr=stderr,
+            env={**_env(lock_dir, "1"), "EVENTS_FILE": str(events_file)},
+        )
+        try:
+            assert _wait_until(lambda: "queueing" in stderr_file.read_text(), 10)
+            assert not events_file.exists()
+            fcntl.flock(held, fcntl.LOCK_UN)
+            assert process.wait(timeout=30) == 0, stderr_file.read_text()
+        finally:
+            fcntl.flock(held, fcntl.LOCK_UN)
+            _reap(process)
+    events = tuple(json.loads(line) for line in events_file.read_text().splitlines())
+    assert {event["phase"] for event in events} == {"setup", "base"}
+    assert all(event["locked"] for event in events)

From cd681a573fd9f5b6f15a1355f46178e4e9d374d2 Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Mon, 7 Sep 2026 16:03:06 -0700
Subject: [PATCH 146/164] fix(mcp): encrypt stored static headers and stdio
 environment (#40164)

Encrypt secret maps at the shared persistence boundary, preserve plaintext API/runtime views, and extend rotation and migration scanning to legacy rows.

Co-authored-by: Claude Code <noreply@anthropic.com>
---
 litellm/models/mcp_server.py                  |  13 +-
 litellm/proxy/_experimental/mcp_server/db.py  |  56 +++---
 .../mcp_server/mcp_server_manager.py          |  10 +-
 .../common_utils/encrypt_decrypt_utils.py     |  40 ++++
 .../credential_migration.py                   |  41 +++-
 .../mcp_server/test_db_credentials.py         | 177 +++++++++++++++++-
 .../mcp_server/test_mcp_env_vars.py           |  22 ++-
 .../mcp_server/test_mcp_partial_update.py     |   9 +-
 .../mcp_server/test_mcp_server_manager.py     |  62 ++++++
 .../mcp_server/test_mcp_sigv4_auth.py         |  17 +-
 .../test_credential_migration.py              |  60 ++++++
 11 files changed, 454 insertions(+), 53 deletions(-)

diff --git a/litellm/models/mcp_server.py b/litellm/models/mcp_server.py
index 6bf21a19896..7ccff9434a7 100644
--- a/litellm/models/mcp_server.py
+++ b/litellm/models/mcp_server.py
@@ -6,10 +6,12 @@ Canonical definition for ``litellm_mcpservertable``. Re-exported from
 """
 
 import enum
+from collections.abc import Mapping
 from datetime import datetime
+from types import MappingProxyType
 from typing import Literal
 
-from pydantic import Field
+from pydantic import Field, ValidationInfo, field_validator
 
 from litellm.types.llms.base import LiteLLMPydanticObjectBase
 from litellm.types.mcp import MCPAuthType, MCPCredentials, MCPTransportType
@@ -115,3 +117,12 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
     submitted_at: datetime | None = None
     reviewed_at: datetime | None = None
     review_notes: str | None = None
+
+    @field_validator("static_headers", "env", mode="before")
+    @classmethod
+    def decode_stored_secret_map(cls, value: object, info: ValidationInfo) -> Mapping[str, str] | None:
+        from litellm.proxy.common_utils.encrypt_decrypt_utils import decode_secret_map
+
+        if value is None and info.field_name == "env":
+            return MappingProxyType({})
+        return decode_secret_map(value, key=info.field_name or "secret map")
diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py
index 41d0b78b555..082a90fdcfb 100644
--- a/litellm/proxy/_experimental/mcp_server/db.py
+++ b/litellm/proxy/_experimental/mcp_server/db.py
@@ -23,8 +23,11 @@ from litellm.proxy._types import (
     UserAPIKeyAuth,
 )
 from litellm.proxy.common_utils.encrypt_decrypt_utils import (
+    SecretMapDecodeError,
     _get_salt_key,
+    decode_secret_map,
     decrypt_value_helper,
+    encrypt_secret_map,
     encrypt_value_helper,
 )
 from litellm.proxy.utils import PrismaClient
@@ -360,7 +363,7 @@ def _prepare_mcp_server_data(
     # exclude_unset filter is respected. Reading back from ``data`` would
     # reintroduce defaults (e.g. ``env={}``) for fields the caller never set.
     if data_dict.get("static_headers") is not None:
-        data_dict["static_headers"] = safe_dumps(data_dict["static_headers"])
+        data_dict["static_headers"] = encrypt_secret_map(data_dict["static_headers"])
 
     # env_vars is read from ``data_dict`` (not ``data``) like every other JSON
     # column so the exclude_unset filter is respected: a partial update that
@@ -376,7 +379,7 @@ def _prepare_mcp_server_data(
         data_dict["mcp_info"] = safe_dumps(data_dict["mcp_info"])
 
     if data_dict.get("env") is not None:
-        data_dict["env"] = safe_dumps(data_dict["env"])
+        data_dict["env"] = encrypt_secret_map(data_dict["env"])
 
     if "tool_name_to_display_name" in data_dict:
         data_dict["tool_name_to_display_name"] = safe_dumps(data_dict["tool_name_to_display_name"] or {})
@@ -589,6 +592,19 @@ def decrypt_credentials(
     return credentials
 
 
+def _readable_mcp_servers(
+    rows: Iterable["prisma_db_models.LiteLLM_MCPServerTable"],
+) -> Iterable[LiteLLM_MCPServerTable]:
+    for row in rows:
+        try:
+            table = LiteLLM_MCPServerTable.model_validate(row.model_dump())
+        except SecretMapDecodeError:
+            verbose_proxy_logger.warning("Skipping MCP server %s: cannot decrypt secret map", row.server_id)
+            continue
+        decrypt_global_env_var_values(table.env_vars)
+        yield table
+
+
 async def get_all_mcp_servers(
     prisma_client: PrismaClient,
     approval_status: str | None = None,
@@ -609,10 +625,7 @@ async def get_all_mcp_servers(
     )
     mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where)
 
-    tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers]
-    for table in tables:
-        decrypt_global_env_var_values(table.env_vars)
-    return tables
+    return list(_readable_mcp_servers(mcp_servers))
 
 
 async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> LiteLLM_MCPServerTable | None:
@@ -638,13 +651,7 @@ async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]
             "server_id": {"in": server_ids},
         }
     )
-    final_mcp_servers: Final[list[LiteLLM_MCPServerTable]] = []
-    for _mcp_server in _mcp_servers:
-        table = LiteLLM_MCPServerTable.model_validate(_mcp_server.model_dump())
-        decrypt_global_env_var_values(table.env_vars)
-        final_mcp_servers.append(table)
-
-    return final_mcp_servers
+    return list(_readable_mcp_servers(_mcp_servers))
 
 
 async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, token: str) -> list[str]:
@@ -852,12 +859,10 @@ async def create_mcp_server(
     data_dict["created_by"] = touched_by
     data_dict["updated_by"] = touched_by
 
-    new_mcp_server: Final[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.create(
-        data=data_dict,  # pyright: ignore[reportAssignmentType]  # prisma row, not domain LiteLLM_MCPServerTable
-    )
+    new_mcp_server: Final = await MCPServerRepository(prisma_client).table.create(data=data_dict)
 
     _decrypt_env_vars_on_returned_row(new_mcp_server)
-    return new_mcp_server
+    return LiteLLM_MCPServerTable.model_validate(new_mcp_server.model_dump())
 
 
 async def create_draft_mcp_server(
@@ -1066,13 +1071,13 @@ async def update_mcp_server(
 
         data_dict["credentials"] = Json(None)
 
-    updated_mcp_server: Final[LiteLLM_MCPServerTable | None] = await MCPServerRepository(prisma_client).table.update(
+    updated_mcp_server: Final = await MCPServerRepository(prisma_client).table.update(
         where={"server_id": data.server_id},
-        data=data_dict,  # pyright: ignore[reportAssignmentType]  # prisma row, not domain LiteLLM_MCPServerTable
+        data=data_dict,
     )
 
     _decrypt_env_vars_on_returned_row(updated_mcp_server)
-    return updated_mcp_server
+    return LiteLLM_MCPServerTable.model_validate(updated_mcp_server.model_dump()) if updated_mcp_server else None
 
 
 async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, server_id: str) -> object | None:
@@ -1144,6 +1149,13 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient,
         if rotated_env_vars is not None:
             update_data["env_vars"] = safe_dumps(rotated_env_vars)
 
+        for field in ("static_headers", "env"):
+            try:
+                if secret_map := decode_secret_map(getattr(mcp_server, field, None), key=field):
+                    update_data[field] = encrypt_secret_map(secret_map, new_encryption_key=new_master_key)
+            except SecretMapDecodeError:
+                verbose_proxy_logger.warning("Cannot rotate MCP %s for server %s", field, mcp_server.server_id)
+
         if not update_data:
             continue
 
@@ -1894,9 +1906,7 @@ async def get_mcp_submissions(
         order={"submitted_at": "desc"},
         take=500,  # safety cap; paginate if needed in a future iteration
     )
-    items: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in rows]
-    for item in items:
-        decrypt_global_env_var_values(item.env_vars)
+    items: Final = list(_readable_mcp_servers(rows))
 
     pending: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.pending_review)
     active: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.active)
diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
index e7fd650a324..d7f238f142c 100644
--- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
@@ -6272,8 +6272,7 @@ class MCPServerManager:
                 ]
             }
         )
-        db_mcp_servers: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in raw_rows]
-        verbose_logger.info("Found %s MCP servers in database", len(db_mcp_servers))
+        verbose_logger.info("Found %s MCP servers in database", len(raw_rows))
 
         previous_registry: Final = self.registry
         new_registry: Final[dict[str, MCPServer]] = {}
@@ -6281,8 +6280,9 @@ class MCPServerManager:
         # Stage one: build every server.  Stage two assigns short prefixes
         # against the *full* set so dedup is deterministic regardless of
         # iteration order.
-        for server in db_mcp_servers:
+        for row in raw_rows:
             try:
+                server = LiteLLM_MCPServerTable.model_validate(row.model_dump())
                 existing_server = previous_registry.get(server.server_id)
 
                 if (
@@ -6320,8 +6320,8 @@ class MCPServerManager:
             except Exception as e:
                 verbose_logger.exception(
                     "Skipping MCP server %s (%s) during DB reload: %s",
-                    server.server_id,
-                    getattr(server, "alias", None),
+                    getattr(row, "server_id", None),
+                    getattr(row, "alias", None),
                     e,
                 )
 
diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py
index 836a4a778bb..fd9b3beee46 100644
--- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py
+++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py
@@ -1,7 +1,10 @@
 import base64
 import os
+from collections.abc import Mapping
 from typing import Final, Literal, cast
 
+from pydantic import TypeAdapter, ValidationError
+
 from litellm._logging import verbose_proxy_logger
 
 # Versioned ciphertext marker for AES-256-GCM values.
@@ -203,3 +206,40 @@ def decrypt_value(value: bytes, signing_key: str) -> str:
         return plaintext
     except Exception as e:
         raise e
+
+
+class SecretMapDecodeError(RuntimeError):
+    pass
+
+
+_SECRET_MAP: Final = TypeAdapter(Mapping[str, str])
+_STORED_SECRET_MAP: Final = TypeAdapter(Mapping[str, str] | str)
+_SECRET_STRING: Final = TypeAdapter(str)
+
+
+def encrypt_secret_map(value: Mapping[str, str], new_encryption_key: str | None = None) -> str:
+    if not value:
+        return "{}"
+    ciphertext: Final = _SECRET_STRING.validate_python(
+        encrypt_value_helper(_SECRET_MAP.dump_json(value).decode(), new_encryption_key=new_encryption_key), strict=True
+    )
+    return _SECRET_STRING.dump_json(ciphertext).decode()
+
+
+def decode_secret_map(value: object, *, key: str) -> Mapping[str, str] | None:
+    if value is None:
+        return None
+    try:
+        stored: Final = (
+            _STORED_SECRET_MAP.validate_json(value, strict=True)
+            if isinstance(value, str) and value.lstrip().startswith(("{", '"'))
+            else _STORED_SECRET_MAP.validate_python(value, strict=True)
+        )
+        if not isinstance(stored, str):
+            return stored
+        decrypted: Final = decrypt_value_helper(
+            value=stored, key=key, exception_type="debug", return_original_value=False
+        )
+        return _SECRET_MAP.validate_json(decrypted, strict=True)
+    except ValidationError:
+        raise SecretMapDecodeError(f"Cannot decode encrypted MCP {key}; check LITELLM_SALT_KEY") from None
diff --git a/litellm/proxy/management_endpoints/credential_migration.py b/litellm/proxy/management_endpoints/credential_migration.py
index e0725119576..915cce87dbd 100644
--- a/litellm/proxy/management_endpoints/credential_migration.py
+++ b/litellm/proxy/management_endpoints/credential_migration.py
@@ -43,7 +43,9 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
     _ALGO_AES_GCM,
     _ENCRYPTION_ALGORITHM_SETTING,
     _V2_GCM_PREFIX,
+    SecretMapDecodeError,
     _get_salt_key,
+    decode_secret_map,
     decrypt_value_helper,
     encrypt_value_helper,
 )
@@ -65,6 +67,20 @@ class LocationReport:
     # Used by --check (read-only classification):
     legacy: int = 0  # nacl ciphertext still awaiting migration
 
+    def count(self, classification: ValueClass | None) -> None:
+        if classification is None:
+            return
+        self.scanned += 1
+        match classification:
+            case "migrated":
+                self.already_v2 += 1
+            case "legacy":
+                self.legacy += 1
+            case "undecryptable":
+                self.undecryptable += 1
+            case _:
+                self.plaintext += 1
+
     def as_dict(self) -> dict[str, int]:
         return {
             "scanned": self.scanned,
@@ -441,7 +457,7 @@ def _classify_callback_value(value: object) -> ValueClass:
 _COVERED_TABLE_SPECS: Final = [
     ("model_table", "litellm_proxymodeltable", ("litellm_params",), ()),
     ("credentials", "litellm_credentialstable", ("credential_values",), ()),
-    ("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars"), ()),
+    ("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars", "static_headers", "env"), ()),
     ("mcp_user_credentials", "litellm_mcpusercredentials", (), ("credential_b64",)),
     ("mcp_user_env_vars", "litellm_mcpuserenvvars", (), ("values_b64",)),
 ]
@@ -472,14 +488,18 @@ def _classify_into_report(report: LocationReport, value: str) -> None:
     names, base URLs, …) do not decrypt and fall through to ``plaintext``, so
     over-scanning a column is harmless to the residual count.
     """
-    report.scanned += 1
-    cls: Final = classify_value(value, key="scan")
-    if cls == "migrated":
-        report.already_v2 += 1
-    elif cls == "legacy":
-        report.legacy += 1
-    else:  # plaintext / not-a-string
-        report.plaintext += 1
+    report.count(classify_value(value, key="scan"))
+
+
+def _classify_secret_map(value: object, key: str) -> ValueClass | None:
+    try:
+        decoded: Final = decode_secret_map(value, key=key)
+    except SecretMapDecodeError:
+        return "undecryptable"
+    if not decoded:
+        return None
+    ciphertext: Final = json.loads(value) if isinstance(value, str) and value.lstrip().startswith('"') else value
+    return "migrated" if is_migrated(ciphertext) else "legacy"
 
 
 async def _scan_one_table(
@@ -503,6 +523,9 @@ async def _scan_one_table(
             raw = getattr(row, col, None)
             if raw is None:
                 continue
+            if db_attr == "litellm_mcpservertable" and col in ("static_headers", "env"):
+                report.count(_classify_secret_map(raw, col))
+                continue
             if isinstance(raw, str):
                 try:
                     raw = json.loads(raw)
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py
index e20f6646310..355b3bfd30e 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py
@@ -12,28 +12,39 @@ import base64
 import json
 from datetime import datetime, timedelta, timezone
 from types import SimpleNamespace
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
+from prisma.models import LiteLLM_MCPServerTable as PrismaMCPServer
 
 from litellm.proxy._experimental.mcp_server.db import (
     _decode_user_credential,
     _prepare_mcp_server_data,
+    create_mcp_server,
     decrypt_credentials,
     encrypt_credentials,
+    get_all_mcp_servers,
+    get_mcp_servers,
+    get_mcp_submissions,
     get_user_credential,
     get_user_oauth_credential,
     is_oauth_credential_expired,
     list_user_oauth_credentials,
     resolve_valid_user_oauth_token,
+    rotate_mcp_server_credentials_master_key,
     rotate_mcp_user_credentials_master_key,
     rotate_mcp_user_env_vars_master_key,
     store_user_credential,
     store_user_oauth_credential,
+    update_mcp_server,
 )
-from litellm.proxy._types import NewMCPServerRequest, UpdateMCPServerRequest
+from litellm.proxy._types import LiteLLM_MCPServerTable, NewMCPServerRequest, UpdateMCPServerRequest
 from litellm.proxy.common_utils.encrypt_decrypt_utils import (
+    SecretMapDecodeError,
+    decode_secret_map,
     decrypt_value_helper,
+    encrypt_secret_map,
     encrypt_value_helper,
 )
 from litellm.types.mcp import MCPAuth, MCPTransport
@@ -44,6 +55,7 @@ SALT_KEY = "test-salt-key-for-byok-credential-tests-1234"
 @pytest.fixture(autouse=True)
 def _set_salt_key(monkeypatch):
     monkeypatch.setenv("LITELLM_SALT_KEY", SALT_KEY)
+    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"encryption_algorithm": "xsalsa20-poly1305"})
 
 
 def _make_prisma_with_existing(row):
@@ -368,6 +380,169 @@ def test_client_private_key_encrypted_at_rest():
     assert decrypted["client_secret"] == "shh"
 
 
+@pytest.fixture(params=["xsalsa20-poly1305", "aes-256-gcm"])
+def map_algorithm(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> str:
+    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"encryption_algorithm": request.param})
+    return request.param
+
+
+def _prisma_map_row(data: dict[str, object], quoted: bool = False) -> PrismaMCPServer:
+    return PrismaMCPServer.model_validate({
+        "transport": "http", "mcp_access_groups": [], "allowed_tools": [], "extra_headers": [], "args": [],
+        "allow_all_keys": False, "available_on_public_internet": True, "delegate_auth_to_upstream": False,
+        "oauth_passthrough": False, "per_server_oauth_discovery": False, "is_byok": False, "byok_description": [],
+        **data,
+        **{field: json.dumps(data[field]) for field in ("static_headers", "env") if quoted and data.get(field)},
+    })
+
+
+class _MapTable:
+    def __init__(self, *rows: dict[str, object], quoted: bool = False) -> None:
+        self.rows = {row["server_id"]: row for row in rows}
+        self.quoted = quoted
+
+    async def create(self, *, data: dict[str, object]) -> PrismaMCPServer:
+        self.rows = {**self.rows, data["server_id"]: dict(data)}
+        return _prisma_map_row(data, self.quoted)
+
+    async def update(self, *, where: dict[str, str], data: dict[str, object]) -> PrismaMCPServer:
+        return await self.create(data={**self.rows[where["server_id"]], **data})
+
+    async def find_many(self, where: object = None) -> list[PrismaMCPServer]:
+        return [_prisma_map_row(row, self.quoted) for row in self.rows.values()]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("field", ["static_headers", "env"])
+@pytest.mark.parametrize("quoted", [False, True])
+async def test_secret_maps_create_update_round_trip(map_algorithm: str, field: str, quoted: bool) -> None:
+    table: Final = _MapTable(quoted=quoted)
+    prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table))
+    original: Final = {"TOKEN": "  sensitive-secret\n", "PREFIX": "v2:gcm:literal", "TEMPLATE": "Bearer ${TOKEN}"}
+    create: Final = NewMCPServerRequest.model_validate({
+        "server_id": "srv-map", "transport": "http", "url": "https://up.example.com/mcp", field: original,
+    })
+    created: Final = await create_mcp_server(prisma, create, touched_by="test")
+    first: Final = table.rows["srv-map"][field]
+    assert isinstance(first, str) and isinstance(json.loads(first), str)
+    assert json.loads(first).startswith("v2:gcm:") is (map_algorithm == "aes-256-gcm")
+    assert "sensitive-secret" not in first and "TEMPLATE" not in first
+    assert getattr(created, field) == original == getattr(create, field)
+    assert decode_secret_map(first, key=field) == original
+    replacement: Final = {**original, "TOKEN": "updated-sensitive-secret"}
+    update: Final = UpdateMCPServerRequest.model_validate({"server_id": "srv-map", field: replacement})
+    updated: Final = await update_mcp_server(prisma, update, touched_by="test")
+    second: Final = table.rows["srv-map"][field]
+    assert second != first and "updated-sensitive-secret" not in second
+    assert decode_secret_map(second, key=field) == replacement
+    assert getattr(updated, field) == replacement == getattr(update, field)
+    assert original["TOKEN"] == "  sensitive-secret\n"
+    omitted: Final = await update_mcp_server(prisma, UpdateMCPServerRequest(server_id="srv-map"), touched_by="test")
+    assert table.rows["srv-map"][field] == second and getattr(omitted, field) == replacement
+    cleared: Final = await update_mcp_server(
+        prisma, UpdateMCPServerRequest.model_validate({"server_id": "srv-map", field: {}}), touched_by="test"
+    )
+    assert table.rows["srv-map"][field] == "{}" and getattr(cleared, field) == {}
+
+
+@pytest.mark.parametrize("field", ["static_headers", "env"])
+@pytest.mark.parametrize("as_json", [False, True])
+def test_secret_map_legacy_model_read_preserves_exact_values(field: str, as_json: bool) -> None:
+    original: Final = {"PREFIX": "v2:gcm:literal", "SPACE": "  secret\n", "TEMPLATE": "${TOKEN}", "B64": "YWJjZA=="}
+    incoming: Final = {
+        "server_id": "srv-map", "transport": "http", field: json.dumps(original) if as_json else original,
+    }
+    snapshot: Final = json.dumps(incoming)
+    parsed: Final = LiteLLM_MCPServerTable.model_validate(incoming)
+    assert getattr(parsed, field) == original
+    assert json.dumps(incoming) == snapshot
+    assert LiteLLM_MCPServerTable.model_validate(parsed.model_dump()).model_dump() == parsed.model_dump()
+    empty: Final = LiteLLM_MCPServerTable.model_validate({"server_id": "srv-map", "transport": "http", field: None})
+    assert getattr(empty, field) == ({} if field == "env" else None)
+
+
+@pytest.mark.parametrize("field", ["static_headers", "env"])
+@pytest.mark.parametrize("failure", ["wrong-key", "corrupt", "invalid-values", "invalid-shape", "invalid-json"])
+def test_secret_map_model_read_fails_closed(map_algorithm: str, field: str, failure: str) -> None:
+    plaintext: Final = {"invalid-values": '{"TOKEN": ["sensitive-secret"]}', "invalid-shape": '["sensitive-secret"]',
+                       "invalid-json": "sensitive-secret"}.get(failure, '{"TOKEN": "sensitive-secret"}')
+    ciphertext: Final = encrypt_value_helper(
+        plaintext, new_encryption_key="wrong-map-key" if failure == "wrong-key" else None
+    )
+    stored: Final = json.dumps(ciphertext[:-8] if failure == "corrupt" else ciphertext)
+    with pytest.raises(SecretMapDecodeError) as exc:
+        LiteLLM_MCPServerTable.model_validate({"server_id": "srv-map", "transport": "http", field: stored})
+    assert field in str(exc.value) and "LITELLM_SALT_KEY" in str(exc.value)
+    assert all(secret not in str(exc.value) for secret in (plaintext, ciphertext, "sensitive-secret"))
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("field,other", [("static_headers", "env"), ("env", "static_headers")])
+async def test_secret_map_rotation_migrates_rekeys_and_preserves_corrupt(
+    map_algorithm: str, field: str, other: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    values: Final = {"TOKEN": "rotation-sensitive-secret", "TEMPLATE": "Bearer ${TOKEN}"}
+    old: Final = encrypt_secret_map(values)
+    corrupt: Final = json.dumps(json.loads(old)[:-8])
+    table: Final = _MapTable(
+        {"server_id": "broken", field: corrupt, other: old},
+        {"server_id": "legacy", field: json.dumps(values), other: "{}"},
+        {"server_id": "encrypted", field: old, other: None},
+    )
+    prisma: Final = SimpleNamespace(db=SimpleNamespace(
+        litellm_mcpservertable=table, litellm_mcpserveroauthclient=SimpleNamespace(find_many=AsyncMock(return_value=[]))
+    ))
+    await rotate_mcp_server_credentials_master_key(prisma, touched_by="test", new_master_key="rotated-map-key")
+    assert table.rows["broken"][field] == corrupt
+    assert table.rows["legacy"][other] == "{}" and table.rows["encrypted"][other] is None
+    for server_id, map_field in (("broken", other), ("legacy", field), ("encrypted", field)):
+        stored: Final = table.rows[server_id][map_field]
+        assert isinstance(json.loads(stored), str) and stored != old and "rotation-sensitive-secret" not in stored
+        with pytest.raises(SecretMapDecodeError):
+            decode_secret_map(stored, key=map_field)
+    monkeypatch.setenv("LITELLM_SALT_KEY", "rotated-map-key")
+    for server_id, map_field in (("broken", other), ("legacy", field), ("encrypted", field)):
+        assert decode_secret_map(table.rows[server_id][map_field], key=map_field) == values
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("reader", [get_all_mcp_servers, get_mcp_servers, get_mcp_submissions])
+@pytest.mark.parametrize("field", ["static_headers", "env"])
+async def test_bulk_reads_isolate_corrupt_secret_maps(reader, field, map_algorithm, caplog):
+    secret = {"TOKEN": "bulk-sensitive-secret"}
+    encrypted = encrypt_secret_map(secret)
+    corrupt = encrypt_secret_map(secret, new_encryption_key="wrong-bulk-key")
+    rows = [
+        _prisma_map_row({"server_id": "broken", field: corrupt, "approval_status": "pending_review"}),
+        _prisma_map_row({"server_id": "healthy", field: encrypted, "approval_status": "active"}),
+    ]
+    snapshot = [row.model_dump() for row in rows]
+    table = SimpleNamespace(find_many=AsyncMock(return_value=rows))
+    prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table))
+    result = await reader(prisma, ["broken", "healthy"]) if reader is get_mcp_servers else await reader(prisma)
+    items = result.items if reader is get_mcp_submissions else result
+    assert [row.server_id for row in items] == ["healthy"]
+    assert getattr(items[0], field) == secret
+    assert [row.model_dump() for row in rows] == snapshot
+    assert "broken" in caplog.text
+    assert all(value not in caplog.text for value in ("bulk-sensitive-secret", corrupt, encrypted))
+    if reader is get_mcp_submissions:
+        assert (result.total, result.pending_review, result.active, result.rejected) == (1, 0, 1, 0)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("reader", [get_all_mcp_servers, get_mcp_servers, get_mcp_submissions])
+async def test_bulk_reads_do_not_swallow_unrelated_validation_errors(reader):
+    from pydantic import ValidationError
+
+    row = _prisma_map_row({"server_id": "invalid", "transport": "unsupported"})
+    table = SimpleNamespace(find_many=AsyncMock(return_value=[row]))
+    prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table))
+    request = reader(prisma, ["invalid"]) if reader is get_mcp_servers else reader(prisma)
+    with pytest.raises(ValidationError, match="transport"):
+        await request
+
+
 # ── BYOK round-trip ───────────────────────────────────────────────────────────
 
 
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py
index 76cc235f7eb..36b545ad031 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py
@@ -861,6 +861,7 @@ _SALT_KEY = "test-salt-key-for-env-vars-tests-1234"
 @pytest.fixture
 def env_vars_salt_key(monkeypatch):
     monkeypatch.setenv("LITELLM_SALT_KEY", _SALT_KEY)
+    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"encryption_algorithm": "xsalsa20-poly1305"})
 
 
 def _mock_env_vars_prisma(row=None):
@@ -1518,9 +1519,16 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri
     assert "s3cr3t-p@ss" not in encrypted_env_vars_str
 
     def _prisma_row_with_json_string_env_vars():
-        row = MagicMock()
-        row.env_vars = encrypted_env_vars_str
-        return row
+        import json
+
+        from prisma.models import LiteLLM_MCPServerTable
+
+        return LiteLLM_MCPServerTable.model_validate({
+            "server_id": "srv-returned", "transport": "http", "mcp_access_groups": [], "allowed_tools": [],
+            "extra_headers": [], "args": [], "allow_all_keys": False, "available_on_public_internet": True,
+            "delegate_auth_to_upstream": False, "oauth_passthrough": False, "per_server_oauth_discovery": False,
+            "is_byok": False, "byok_description": [], "env_vars": json.dumps(encrypted_env_vars_str),
+        })
 
     mock_prisma = MagicMock()
     mock_prisma.db.litellm_mcpservertable.create = AsyncMock(
@@ -1537,7 +1545,9 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri
         touched_by="test-user",
     )
     assert isinstance(created.env_vars, list)
-    assert created.env_vars[0]["value"] == "s3cr3t-p@ss"
+    assert created.env_vars[0].value == "s3cr3t-p@ss"
+    assert created.env_vars[0].name == "DB_PASSWORD"
+    assert created.env == {}
 
     mock_prisma_upd = MagicMock()
     mock_prisma_upd.db.litellm_mcpservertable.update = AsyncMock(
@@ -1549,7 +1559,9 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri
         touched_by="test-user",
     )
     assert isinstance(updated.env_vars, list)
-    assert updated.env_vars[0]["value"] == "s3cr3t-p@ss"
+    assert updated.env_vars[0].value == "s3cr3t-p@ss"
+    assert updated.env_vars[0].name == "DB_PASSWORD"
+    assert updated.env == {}
 
 
 def test_reencrypt_global_env_var_values_handles_json_string(env_vars_salt_key):
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py
index f6bd79c5d2d..c0d055edb7f 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py
@@ -11,7 +11,7 @@ import json
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
-from prisma import Json
+from prisma import Json, models
 
 from litellm.proxy._experimental.mcp_server.db import (
     create_mcp_server,
@@ -28,8 +28,11 @@ def _credentials_cleared(value) -> bool:
 def _mock_prisma():
     mock_prisma = MagicMock()
     mock_prisma.db.litellm_mcpservertable = AsyncMock()
-    mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
-    mock_prisma.db.litellm_mcpservertable.create = AsyncMock(return_value=MagicMock())
+    row = models.LiteLLM_MCPServerTable.model_construct(
+        server_id="test-server", transport="http", env={}, env_vars=[]
+    )
+    mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=row)
+    mock_prisma.db.litellm_mcpservertable.create = AsyncMock(return_value=row)
     return mock_prisma
 
 
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 d19363d3b5f..46fef83092d 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
@@ -905,6 +905,68 @@ class TestMCPServerManager:
         assert retry_slot is not None
         assert retry_slot.generation > old_generation
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("corrupt_column", ("static_headers", "env"))
+    async def test_database_reload_drops_cached_server_whose_secret_map_stops_decoding(
+        self, monkeypatch, caplog, corrupt_column
+    ):
+        from types import SimpleNamespace
+
+        from litellm.proxy import proxy_server
+        from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_secret_map
+
+        monkeypatch.setenv("LITELLM_SALT_KEY", "sk-reload-secret-map-salt")
+        monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": "aes-256-gcm"})
+        headers = {"Authorization": "Bearer dummy-header-secret-4f1c"}
+        env = {"UPSTREAM_TOKEN": "dummy-env-secret-9a2b"}
+        stamp = datetime.now()
+        cached = MCPServer(
+            server_id="cached-server",
+            name="cached_server",
+            url="https://up.example.com/mcp",
+            transport=MCPTransport.http,
+            static_headers=dict(headers),
+            env=dict(env),
+            updated_at=stamp,
+        )
+        manager = MCPServerManager()
+        manager.registry[cached.server_id] = cached
+        stored = {"static_headers": encrypt_secret_map(headers), "env": encrypt_secret_map(env)}
+        corrupted = {**stored, corrupt_column: stored[corrupt_column][:-6] + 'AAAAA"'}
+
+        def _row(server_id, maps):
+            row = MagicMock()
+            row.server_id = server_id
+            row.alias = server_id
+            row.model_dump.return_value = {
+                "server_id": server_id,
+                "alias": server_id,
+                "server_name": server_id,
+                "url": "https://up.example.com/mcp",
+                "transport": MCPTransport.http,
+                "updated_at": stamp,
+                **maps,
+            }
+            return row
+
+        table = SimpleNamespace(
+            find_many=AsyncMock(return_value=[_row(cached.server_id, corrupted), _row("healthy-sibling", stored)])
+        )
+        prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table))
+        monkeypatch.setattr(proxy_server, "prisma_client", prisma)
+
+        with caplog.at_level(logging.DEBUG, logger="LiteLLM"):
+            await manager.reload_servers_from_database()
+
+        assert set(manager.registry) == {"healthy-sibling"}
+        sibling = manager.registry["healthy-sibling"]
+        assert dict(sibling.static_headers) == headers
+        assert dict(sibling.env) == env
+        logged = "\n".join(caplog.messages)
+        assert cached.server_id in logged
+        for secret in (*headers.values(), *env.values(), *stored.values(), corrupted[corrupt_column]):
+            assert secret not in logged
+
     @pytest.mark.asyncio
     async def test_lazy_oauth_discovery_preserves_manual_authorization_url_gate(self):
         with patch.dict(os.environ, {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": "false"}):
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py
index f6b61c1d9f7..e814425c9a2 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py
@@ -15,6 +15,11 @@ import httpx
 
 from litellm.experimental_mcp_client.client import MCPSigV4Auth, MCPClient
 from litellm.types.mcp import MCPAuth, MCPTransport
+from prisma import models
+
+
+def _updated_row() -> models.LiteLLM_MCPServerTable:
+    return models.LiteLLM_MCPServerTable.model_construct(server_id="test-server", transport="http", env={}, env_vars=[])
 
 
 class TestMCPSigV4Auth:
@@ -600,7 +605,7 @@ class TestCredentialMergeOnUpdate:
 
         mock_prisma = MagicMock()
         mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record)
-        mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
+        mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row())
 
         data = UpdateMCPServerRequest(
             server_id="test-server",
@@ -639,7 +644,7 @@ class TestCredentialMergeOnUpdate:
         from litellm.proxy._types import UpdateMCPServerRequest
 
         mock_prisma = MagicMock()
-        mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
+        mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row())
 
         data = UpdateMCPServerRequest(
             server_id="test-server",
@@ -667,7 +672,7 @@ class TestCredentialMergeOnUpdate:
 
         mock_prisma = MagicMock()
         mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record)
-        mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
+        mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row())
 
         data = UpdateMCPServerRequest(
             server_id="test-server",
@@ -709,7 +714,7 @@ class TestCredentialMergeOnUpdate:
 
         mock_prisma = MagicMock()
         mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record)
-        mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
+        mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row())
 
         data = UpdateMCPServerRequest(
             server_id="test-server",
@@ -752,7 +757,7 @@ class TestCredentialMergeOnUpdate:
 
         mock_prisma = MagicMock()
         mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record)
-        mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
+        mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row())
 
         data = UpdateMCPServerRequest(
             server_id="test-server",
@@ -1083,7 +1088,7 @@ class TestAuthTypeSwitchClearsCredentials:
 
         mock_prisma = MagicMock()
         mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record)
-        mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
+        mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row())
 
         data = UpdateMCPServerRequest(
             server_id="test-server",
diff --git a/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py b/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py
index 81226981089..0ecc4f8d7cb 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py
@@ -8,6 +8,7 @@ proof-of-fix (real proxy + DB) is performed separately on the repro server.
 
 import json
 from types import SimpleNamespace
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -457,6 +458,65 @@ async def test_scan_covered_tables_classifies_legacy_and_v2(salt_key, monkeypatc
     assert by_loc["credentials"].legacy == 0
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("column", ("static_headers", "env"))
+@pytest.mark.parametrize("algorithm", ("xsalsa20-poly1305", "aes-256-gcm"))
+@pytest.mark.parametrize("as_json", (False, True))
+@pytest.mark.parametrize(
+    "case", ("legacy", "encrypted", "wrong-key", "corrupt", "invalid-shape", "invalid-scalar", "empty", "null")
+)
+async def test_check_classifies_mcp_secret_maps(
+    salt_key: str,
+    monkeypatch: pytest.MonkeyPatch,
+    column: str,
+    algorithm: str,
+    as_json: bool,
+    case: str,
+) -> None:
+    monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": algorithm})
+    plaintext: Final = {"Authorization": "v2:gcm:operator-text", "CUSTOM": "litellm_enc::literal\n café "}
+    ciphertext: Final = encrypt_value_helper(json.dumps(plaintext))
+    cases: Final[dict[str, object]] = {
+        "legacy": plaintext,
+        "encrypted": ciphertext,
+        "wrong-key": encrypt_value_helper(json.dumps(plaintext), new_encryption_key="different-map-salt"),
+        "corrupt": ciphertext[:-4] + "AAAA",
+        "invalid-shape": encrypt_value_helper(json.dumps({"Authorization": 42})),
+        "invalid-scalar": "null",
+        "empty": {},
+        "null": None,
+    }
+    value: Final = json.dumps(cases[case]) if as_json and case != "null" else cases[case]
+    row: Final = SimpleNamespace(**{column: value})
+    client: Final = MagicMock()
+    _empty_covered_tables(client)
+    client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row])
+    client.db.litellm_mcpservertable.update = AsyncMock()
+    client.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
+    client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
+    client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None)
+    client.db.litellm_config.find_unique = AsyncMock(return_value=None)
+    monkeypatch.setattr(proxy_server, "general_settings", {})
+
+    report: Final = await cm.check_encryption(client)
+    expected_legacy: Final = int(case == "legacy" or (case == "encrypted" and algorithm == "xsalsa20-poly1305"))
+    expected_v2: Final = int(case == "encrypted" and algorithm == "aes-256-gcm")
+    expected_invalid: Final = int(case in ("wrong-key", "corrupt", "invalid-shape", "invalid-scalar"))
+
+    assert report.as_dict()["locations"]["mcp_server"] == {
+        "scanned": int(case not in ("empty", "null")),
+        "migrated": 0,
+        "already_v2": expected_v2,
+        "plaintext": 0,
+        "undecryptable": expected_invalid,
+        "legacy": expected_legacy,
+    }
+    assert report.residual_legacy == expected_legacy
+    assert report.total_undecryptable == expected_invalid
+    assert getattr(row, column) == value
+    client.db.litellm_mcpservertable.update.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_check_counts_covered_table_residual(salt_key, monkeypatch):
     """check_encryption now scans the rotation-covered tables (model table here),

From 9d0c9b938283dbadaeb03c0bffc1c3012dab250c Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Mon, 7 Sep 2026 16:29:19 -0700
Subject: [PATCH 147/164] feat(ui): itemize auto-router classification spend
 (#40168)

Resolves LIT-7141

Co-authored-by: Claude Code <noreply@anthropic.com>
---
 .../migration.sql                             |  3 +
 .../litellm_proxy_extras/schema.prisma        |  2 +
 litellm/proxy/db/autorouter_session_rollup.py | 10 +-
 .../auto_router_endpoints.py                  |  6 ++
 litellm/proxy/schema.prisma                   |  2 +
 .../auto_router_endpoints.py                  |  4 +
 schema.prisma                                 |  2 +
 .../spend/test_autorouter_session_rollup.py   | 98 +++++++++++++++++--
 .../db/test_autorouter_session_rollup.py      | 52 ++++++----
 .../test_auto_router_endpoints.py             | 50 ++++++++--
 .../AutoRouterBenchmarksTab.test.tsx          | 44 ++++++++-
 .../_components/AutoRouterBenchmarksTab.tsx   | 38 +++++--
 .../_components/autoRouterBenchmarks.test.ts  |  1 +
 ...KeyAutoRouterUsageTab.integration.test.tsx |  6 ++
 ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 ++
 15 files changed, 282 insertions(+), 46 deletions(-)
 create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql

diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql
new file mode 100644
index 00000000000..5503167ce09
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql
@@ -0,0 +1,3 @@
+ALTER TABLE "LiteLLM_AutoRouterSession"
+ADD COLUMN IF NOT EXISTS "classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0,
+ADD COLUMN IF NOT EXISTS "classifier_cost_recorded_turns" INTEGER NOT NULL DEFAULT 0;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
index 06ac177cca4..ccbab0fef10 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
@@ -1509,6 +1509,8 @@ model LiteLLM_AutoRouterSession {
   total_tokens          BigInt   @default(0)
   spend                 Float    @default(0)
   saved_spend           Float    @default(0)
+  classifier_cost       Float    @default(0)
+  classifier_cost_recorded_turns Int @default(0)
   tier_turns            Json     @default("{}")
 
   @@id([api_key, session_id, router_name])
diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py
index b33bffbaaa1..b866ecc741f 100644
--- a/litellm/proxy/db/autorouter_session_rollup.py
+++ b/litellm/proxy/db/autorouter_session_rollup.py
@@ -75,6 +75,8 @@ SELECT
     COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens,
     COALESCE(SUM(spend), 0)::float8 AS spend,
     COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend,
+    COALESCE(SUM(classifier_cost), 0)::float8 AS classifier_cost,
+    COALESCE(SUM(classifier_cost_recorded_turns), 0)::int AS classifier_cost_recorded_turns,
     COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds
 FROM windowed
 GROUP BY router_name, router_type
@@ -95,6 +97,7 @@ class AutoRouterTurnTransaction:
     total_tokens: int
     spend: float
     saved_spend: float
+    classifier_cost: float
     covered: bool
     cache_hit: bool
     cache_ttl_seconds: int | None
@@ -225,6 +228,7 @@ def build_autorouter_turn_transaction(
         total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0),
         spend=float(payload.get("spend") or 0.0) + (classifier_cost or 0.0),
         saved_spend=saved_spend,
+        classifier_cost=classifier_cost or 0.0,
         covered=cache.covered,
         cache_hit=cache.read_tokens > 0,
         cache_ttl_seconds=cache.write_ttl_seconds,
@@ -266,7 +270,7 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t (
     last_model, models, turns, unordered_turns, covered_turns, cache_hits,
     same_model_turns, same_model_hits, first_visit_turns, first_visit_hits,
     return_turns, return_hits, return_expired_misses, return_within_ttl_misses,
-    ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, tier_turns
+    ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns
 )
 VALUES (
     {_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp,
@@ -277,13 +281,15 @@ VALUES (
     (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END),
     (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END),
     {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8,
-    {_TIER_DELTA}
+    {_p("classifier_cost")}::float8, 1, {_TIER_DELTA}
 )
 ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
     turns = t.turns + 1,
     total_tokens = t.total_tokens + EXCLUDED.total_tokens,
     spend = t.spend + EXCLUDED.spend,
     saved_spend = t.saved_spend + EXCLUDED.saved_spend,
+    classifier_cost = t.classifier_cost + EXCLUDED.classifier_cost,
+    classifier_cost_recorded_turns = t.classifier_cost_recorded_turns + 1,
     covered_turns = t.covered_turns + EXCLUDED.covered_turns,
     cache_hits = t.cache_hits + EXCLUDED.cache_hits,
     ttl_5m_turns = t.ttl_5m_turns + EXCLUDED.ttl_5m_turns,
diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py
index 0f0323b45f8..bbc914a772a 100644
--- a/litellm/proxy/management_endpoints/auto_router_endpoints.py
+++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py
@@ -484,6 +484,8 @@ class _SessionAggRow(BaseModel):
     total_tokens: int
     spend: float
     saved_spend: float
+    classifier_cost: float
+    classifier_cost_recorded_turns: int
     session_seconds: float
 
 
@@ -520,6 +522,7 @@ def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals:
         avg_tokens_per_session=row.total_tokens / sessions if sessions else 0.0,
         spend=row.spend,
         saved_spend=row.saved_spend,
+        classifier_cost=row.classifier_cost if row.classifier_cost_recorded_turns == row.turns else None,
         baseline_spend=baseline_spend,
         saved_pct=_pct(row.saved_spend, baseline_spend),
         saved_per_session=row.saved_spend / sessions if sessions else 0.0,
@@ -552,6 +555,7 @@ def _benchmark_group(row: _SessionAggRow) -> AutoRouterBenchmarkGroup:
         avg_tokens_per_session=totals.avg_tokens_per_session,
         spend=totals.spend,
         saved_spend=totals.saved_spend,
+        classifier_cost=totals.classifier_cost,
         baseline_spend=totals.baseline_spend,
         saved_pct=totals.saved_pct,
         saved_per_session=totals.saved_per_session,
@@ -582,6 +586,8 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow:
         total_tokens=sum(row.total_tokens for row in rows),
         spend=sum(row.spend for row in rows),
         saved_spend=sum(row.saved_spend for row in rows),
+        classifier_cost=sum(row.classifier_cost for row in rows),
+        classifier_cost_recorded_turns=sum(row.classifier_cost_recorded_turns for row in rows),
         session_seconds=sum(row.session_seconds for row in rows),
     )
 
diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index 06ac177cca4..ccbab0fef10 100644
--- a/litellm/proxy/schema.prisma
+++ b/litellm/proxy/schema.prisma
@@ -1509,6 +1509,8 @@ model LiteLLM_AutoRouterSession {
   total_tokens          BigInt   @default(0)
   spend                 Float    @default(0)
   saved_spend           Float    @default(0)
+  classifier_cost       Float    @default(0)
+  classifier_cost_recorded_turns Int @default(0)
   tier_turns            Json     @default("{}")
 
   @@id([api_key, session_id, router_name])
diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py
index 50c3515cf01..6306658ad0b 100644
--- a/litellm/types/management_endpoints/auto_router_endpoints.py
+++ b/litellm/types/management_endpoints/auto_router_endpoints.py
@@ -195,6 +195,10 @@ class AutoRouterBenchmarkTotals(BaseModel):
     avg_session_seconds: float
     avg_tokens_per_session: float
     spend: float = Field(description="What the routed traffic actually cost")
+    classifier_cost: float | None = Field(
+        description="Recorded LLM classifier cost already included in spend; null when any session turns predate "
+        "subtotal recording, and zero for an empty window"
+    )
     saved_spend: float = Field(
         description="Signed dollars saved versus each router's savings baseline (derived from its hardest "
         "tier, or the configured override), from the same per-request savings record the usage tab reads"
diff --git a/schema.prisma b/schema.prisma
index 06ac177cca4..ccbab0fef10 100644
--- a/schema.prisma
+++ b/schema.prisma
@@ -1509,6 +1509,8 @@ model LiteLLM_AutoRouterSession {
   total_tokens          BigInt   @default(0)
   spend                 Float    @default(0)
   saved_spend           Float    @default(0)
+  classifier_cost       Float    @default(0)
+  classifier_cost_recorded_turns Int @default(0)
   tier_turns            Json     @default("{}")
 
   @@id([api_key, session_id, router_name])
diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py
index c2272f3d20d..9ac6476a03c 100644
--- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py
+++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py
@@ -40,12 +40,26 @@ async def _turn(
     tokens: int = 100,
     spend: float = 0.01,
     saved: float = 0.02,
+    classifier_cost: float = 0.0,
     tier: "str | None" = None,
 ) -> None:
     touched: Final = 1 if (hit or ttl is not None or not covered) else 0
     await db.execute_raw(
         UPSERT_AUTOROUTER_SESSION_SQL,
-        key, session_id, router, router_type, model, at.isoformat(), tokens, spend, saved, covered, hit, ttl, touched,
+        key,
+        session_id,
+        router,
+        router_type,
+        model,
+        at.isoformat(),
+        tokens,
+        spend,
+        saved,
+        classifier_cost,
+        covered,
+        hit,
+        ttl,
+        touched,
         tier,
     )
 
@@ -53,7 +67,9 @@ async def _turn(
 async def _row(db, key: str, session_id: str = "s1", router: str = "auto-1") -> dict:
     rows = await db.query_raw(
         'SELECT * FROM "LiteLLM_AutoRouterSession" WHERE api_key = $1 AND session_id = $2 AND router_name = $3',
-        key, session_id, router,
+        key,
+        session_id,
+        router,
     )
     assert len(rows) == 1
     return rows[0]
@@ -143,9 +159,9 @@ async def test_out_of_order_turns_do_not_rewind_the_session(db):
 
 async def test_concurrent_writers_compose_without_losing_turns(db):
     key = f"k-{uuid.uuid4()}"
-    await _turn(db, key, "A", T0)
+    await _turn(db, key, "A", T0, classifier_cost=0.001)
     await asyncio.gather(
-        *(_turn(db, key, "A", T0 + timedelta(seconds=1 + offset), hit=1) for offset in range(30))
+        *(_turn(db, key, "A", T0 + timedelta(seconds=1 + offset), hit=1, classifier_cost=0.002) for offset in range(30))
     )
     row = await _row(db, key)
     assert row["turns"] == 31
@@ -154,6 +170,51 @@ async def test_concurrent_writers_compose_without_losing_turns(db):
         == row["turns"]
     )
     assert row["spend"] == pytest.approx(0.31)
+    assert row["saved_spend"] == pytest.approx(0.62)
+    assert row["classifier_cost"] == pytest.approx(0.061)
+    assert row["classifier_cost_recorded_turns"] == 31
+
+
+async def _legacy_turn(db, key: str, at: datetime, session_id: str = "s1", router: str = "auto-1") -> None:
+    await db.execute_raw(
+        """INSERT INTO "LiteLLM_AutoRouterSession" AS t (
+            api_key, session_id, router_name, router_type, first_turn_at, last_turn_at, last_model, turns, spend, saved_spend
+        ) VALUES ($1, $2, $3, 'complexity', $4::timestamp, $4::timestamp, 'A', 1, 0.01, 0.02)
+        ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
+            turns = t.turns + 1, spend = t.spend + EXCLUDED.spend, saved_spend = t.saved_spend + EXCLUDED.saved_spend,
+            last_turn_at = EXCLUDED.last_turn_at""",
+        key,
+        session_id,
+        router,
+        at.isoformat(),
+    )
+
+
+@pytest.mark.parametrize("writers", [(False,), (True,), (False, True), (True, False)])
+async def test_subtotal_coverage_survives_legacy_and_rolling_writers(db, writers: tuple[bool, ...]):
+    key: Final = f"k-{uuid.uuid4()}"
+    for offset, records_cost in enumerate(writers):
+        at: Final = T0 + timedelta(seconds=offset)
+        if records_cost:
+            await _turn(db, key, "A", at, classifier_cost=0.004)
+        else:
+            await _legacy_turn(db, key, at)
+
+    row: Final = await _row(db, key)
+    assert row["turns"] == len(writers)
+    assert row["spend"] == pytest.approx(0.01 * len(writers))
+    assert row["saved_spend"] == pytest.approx(0.02 * len(writers))
+    assert row["classifier_cost"] == pytest.approx(0.004 * sum(writers))
+    assert row["classifier_cost_recorded_turns"] == sum(writers)
+    groups: Final = await db.query_raw(
+        AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key
+    )
+    assert len(groups) == 1
+    assert groups[0]["classifier_cost"] == row["classifier_cost"]
+    assert groups[0]["classifier_cost_recorded_turns"] == sum(writers)
+    assert groups[0]["turns"] == len(writers)
+    assert groups[0]["spend"] == row["spend"]
+    assert groups[0]["saved_spend"] == row["saved_spend"]
 
 
 async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db):
@@ -161,9 +222,19 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db):
     router = f"r-{uuid.uuid4()}"
     in_window = f"s-{uuid.uuid4()}"
     out_of_window = f"s-{uuid.uuid4()}"
-    await _turn(db, key, "A", T0, session_id=in_window, router=router, saved=0.5, spend=0.25)
-    await _turn(db, key, "B", T0 + timedelta(seconds=60), session_id=in_window, router=router, saved=0.5, spend=0.25)
-    await _turn(db, key, "A", T0 - timedelta(days=40), session_id=out_of_window, router=router)
+    await _turn(
+        db,
+        key,
+        "B",
+        T0 + timedelta(seconds=60),
+        session_id=in_window,
+        router=router,
+        saved=0.5,
+        spend=0.25,
+        classifier_cost=0.01,
+    )
+    await _turn(db, key, "A", T0, session_id=in_window, router=router, saved=0.5, spend=0.25, classifier_cost=0.02)
+    await _turn(db, key, "A", T0 - timedelta(days=40), session_id=out_of_window, router=router, classifier_cost=9.0)
 
     rows = await db.query_raw(
         AUTOROUTER_BENCHMARKS_SQL,
@@ -179,6 +250,9 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db):
     assert grouped["turns"] == 2
     assert grouped["spend"] == pytest.approx(0.5)
     assert grouped["saved_spend"] == pytest.approx(1.0)
+    assert grouped["classifier_cost"] == pytest.approx(0.03)
+    assert grouped["classifier_cost_recorded_turns"] == 2
+    assert grouped["unordered_turns"] == 1
     assert grouped["session_seconds"] == pytest.approx(60.0)
 
 
@@ -186,8 +260,8 @@ async def test_the_benchmarks_aggregate_can_filter_to_one_key(db):
     router = f"r-{uuid.uuid4()}"
     first_key = f"k-{uuid.uuid4()}"
     second_key = f"k-{uuid.uuid4()}"
-    await _turn(db, first_key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, saved=0.5)
-    await _turn(db, second_key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, saved=9.0)
+    await _turn(db, first_key, "A", T0, router=router, saved=0.5, classifier_cost=0.01)
+    await _turn(db, second_key, "A", T0, router=router, saved=9.0, classifier_cost=0.09)
 
     rows = await db.query_raw(
         AUTOROUTER_BENCHMARKS_SQL,
@@ -199,6 +273,8 @@ async def test_the_benchmarks_aggregate_can_filter_to_one_key(db):
     assert len(matching) == 1
     assert matching[0]["sessions"] == 1
     assert matching[0]["saved_spend"] == pytest.approx(0.5)
+    assert matching[0]["classifier_cost"] == pytest.approx(0.01)
+    assert matching[0]["classifier_cost_recorded_turns"] == 1
 
     unknown_key_rows = await db.query_raw(
         AUTOROUTER_BENCHMARKS_SQL,
@@ -213,7 +289,9 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db
     key = f"k-{uuid.uuid4()}"
     router = f"r-{uuid.uuid4()}"
     await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, router_type="complexity")
-    await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, router_type="quality")
+    await _turn(
+        db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, router_type="quality"
+    )
 
     rows = await db.query_raw(
         AUTOROUTER_BENCHMARKS_SQL,
diff --git a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py
index 2ed4f843711..4507892bd0f 100644
--- a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py
+++ b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py
@@ -9,13 +9,15 @@ request-time transaction builder and the flush contract with an injected fake cl
 import asyncio
 import json
 from datetime import datetime
+from types import SimpleNamespace
+from typing import Final
 
 import httpx
 import pytest
 
 from litellm.proxy.db.autorouter_session_rollup import (
-    AutoRouterTurnTransaction,
     UPSERT_AUTOROUTER_SESSION_SQL,
+    AutoRouterTurnTransaction,
     build_autorouter_turn_transaction,
     flush_autorouter_turn_transactions,
 )
@@ -70,6 +72,7 @@ class TestBuildTransaction:
             total_tokens=100,
             spend=0.01,
             saved_spend=0.02,
+            classifier_cost=0.0,
             covered=True,
             cache_hit=True,
             cache_ttl_seconds=300,
@@ -111,6 +114,9 @@ class TestBuildTransaction:
         folded once into the turn that paid for it (GH #38816)."""
         transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "classifier_cost": 0.005}))
         assert transaction is not None and transaction.spend == pytest.approx(0.015)
+        assert transaction.classifier_cost == 0.005
+        assert transaction.spend - transaction.classifier_cost == pytest.approx(0.01)
+        assert transaction.saved_spend == 0.02
 
     @pytest.mark.parametrize(
         "decision_extra", [{}, {"classifier_cost": 0.0}, {"classifier_cost": "bogus"}, {"classifier_cost": True}]
@@ -118,6 +124,8 @@ class TestBuildTransaction:
     def test_an_unpriced_classifier_leaves_the_spend_alone(self, decision_extra: dict):
         transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, **decision_extra}))
         assert transaction is not None and transaction.spend == pytest.approx(0.01)
+        assert transaction.classifier_cost == 0.0
+        assert transaction.saved_spend == 0.02
 
     def test_every_turn_carries_its_own_classifier_charge(self):
         first = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "classifier_cost": 0.005}))
@@ -127,6 +135,7 @@ class TestBuildTransaction:
         )
         assert first is not None and first.spend == pytest.approx(0.015)
         assert second is not None and second.spend == pytest.approx(0.027)
+        assert (first.classifier_cost, second.classifier_cost) == (0.005, 0.007)
 
     def test_router_name_falls_back_to_the_payload_model_group(self):
         transaction = _build(metadata=_metadata(routing_decision={"router_type": "complexity"}))
@@ -217,6 +226,7 @@ def _transaction(
         total_tokens=100,
         spend=0.01,
         saved_spend=0.02,
+        classifier_cost=0.005,
         covered=True,
         cache_hit=False,
         cache_ttl_seconds=None,
@@ -249,6 +259,7 @@ class TestFlush:
             100,
             0.01,
             0.02,
+            0.005,
             1,
             0,
             None,
@@ -279,26 +290,31 @@ class TestFlush:
 
 class TestEnqueueSeam:
     @pytest.mark.asyncio
-    async def test_update_database_seam_enqueues_only_auto_routed_success(self, monkeypatch: pytest.MonkeyPatch):
+    @pytest.mark.parametrize("classifier_cost", [0.005, 0.0, None])
+    async def test_update_database_seam_enqueues_only_auto_routed_success(self, classifier_cost: float | None):
         from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
-        from litellm.proxy.utils import PrismaClient
 
-        monkeypatch.setattr(PrismaClient, "autorouter_turn_transactions", [])
-        writer = DBSpendUpdateWriter()
-        fake_prisma = type("P", (), {})()
-        fake_prisma._autorouter_turn_transactions_lock = asyncio.Lock()
-        fake_prisma.autorouter_turn_transactions = []
+        writer: Final = DBSpendUpdateWriter()
+        fake_prisma: Final = SimpleNamespace(
+            _autorouter_turn_transactions_lock=asyncio.Lock(), autorouter_turn_transactions=[]
+        )
+        metadata: Final = _metadata(
+            routing_decision={**ROUTING_DECISION, "classifier_cost": classifier_cost}, autorouter_savings=-0.003
+        )
+        for payload in (
+            _payload(metadata=json.dumps(metadata)),
+            _payload(metadata=json.dumps({"usage_object": {"prompt_tokens": 9}})),
+            _payload(status="failure", metadata=json.dumps(metadata)),
+            _payload(metadata=json.dumps({**metadata, "internal_call_origin": "autorouter_classifier"})),
+        ):
+            await writer._enqueue_autorouter_turn_transaction(payload=payload, prisma_client=fake_prisma)
 
-        routed = _payload()
-        routed["metadata"] = json.dumps(_metadata())
-        await writer._enqueue_autorouter_turn_transaction(payload=routed, prisma_client=fake_prisma)
-
-        plain = _payload()
-        plain["metadata"] = json.dumps({"usage_object": {"prompt_tokens": 9}})
-        await writer._enqueue_autorouter_turn_transaction(payload=plain, prisma_client=fake_prisma)
-
-        assert [t.router_name for t in fake_prisma.autorouter_turn_transactions] == ["live-auto"]
-        assert fake_prisma.autorouter_turn_transactions[0].saved_spend == 0.0
+        assert len(fake_prisma.autorouter_turn_transactions) == 1
+        transaction: Final = fake_prisma.autorouter_turn_transactions[0]
+        assert transaction.router_name == "live-auto"
+        assert transaction.spend == pytest.approx(0.01 + (classifier_cost or 0.0))
+        assert transaction.classifier_cost == (classifier_cost or 0.0)
+        assert transaction.saved_spend == -0.003
 
 
 def test_every_drain_trigger_reads_the_one_queue_census_owner():
diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py
index 35e76c96c14..dc18e0f7d4a 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py
@@ -10,7 +10,6 @@ import pytest
 from fastapi import HTTPException
 from pydantic import ValidationError
 
-
 from litellm.proxy._types import (
     LitellmUserRoles,
     ProxyErrorTypes,
@@ -21,11 +20,11 @@ from litellm.proxy.management_endpoints.auto_router_endpoints import (
     preview_auto_router_routing,
 )
 from litellm.router import Router
-from litellm.types.utils import Choices, Message, ModelResponse
 from litellm.types.management_endpoints.auto_router_endpoints import (
     AutoRouterBenchmarksResponse,
     AutoRouterRoutingTestRequest,
 )
+from litellm.types.utils import Choices, Message, ModelResponse
 
 ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test", user_id="admin")
 
@@ -529,6 +528,8 @@ class TestAutoRouterBenchmarks:
         total_tokens=4000,
         spend=10.0,
         saved_spend=30.0,
+        classifier_cost=0.4,
+        classifier_cost_recorded_turns=40,
         session_seconds=400.0,
     )
 
@@ -567,6 +568,7 @@ class TestAutoRouterBenchmarks:
         totals = _benchmark_totals(losing)
         assert totals.baseline_spend == 5.0
         assert totals.saved_pct == -100.0
+        assert totals.classifier_cost == 0.4
 
     def test_an_empty_window_folds_to_zeros(self):
         from litellm.proxy.management_endpoints.auto_router_endpoints import (
@@ -579,6 +581,7 @@ class TestAutoRouterBenchmarks:
         assert totals.turns == 0
         assert totals.saved_pct == 0.0
         assert totals.cache.hit_rate_pct == 0.0
+        assert totals.classifier_cost == 0.0
 
     def test_totals_sum_counters_across_groups_before_deriving_ratios(self):
         from litellm.proxy.management_endpoints.auto_router_endpoints import (
@@ -660,6 +663,38 @@ class TestAutoRouterBenchmarks:
         assert response.routers_in_scope == 1
         assert response.groups[0].router_name == "live-auto"
         assert response.groups[0].saved_pct == response.totals.saved_pct == 75.0
+        assert response.groups[0].classifier_cost == response.totals.classifier_cost == 0.4
+        assert response.totals.spend - response.totals.classifier_cost == pytest.approx(9.6)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("recorded_turns", [0, 3, 10])
+    async def test_classifier_subtotals_require_every_included_turn_to_be_recorded(
+        self, recorded_turns: int, monkeypatch: pytest.MonkeyPatch
+    ):
+        other: Final = self.ROW.model_copy(
+            update={
+                "router_name": "other-auto",
+                "sessions": 1,
+                "turns": 10,
+                "spend": 2.0,
+                "saved_spend": -0.5,
+                "classifier_cost": recorded_turns * 0.02,
+                "classifier_cost_recorded_turns": recorded_turns,
+            }
+        )
+        response: Final = await self._benchmarks(
+            monkeypatch, rows=[self.ROW.model_dump(), other.model_dump()], model_list=[]
+        )
+        wire: Final = response.model_dump()
+        assert wire["groups"][0]["classifier_cost"] == 0.4
+        assert wire["groups"][1]["classifier_cost"] == (pytest.approx(0.2) if recorded_turns == 10 else None)
+        assert wire["totals"]["classifier_cost"] == (pytest.approx(0.6) if recorded_turns == 10 else None)
+        assert response.totals.turns == 50
+        assert response.totals.spend == 12.0
+        assert response.totals.saved_spend == 29.5
+        assert response.totals.baseline_spend == 41.5
+        assert response.totals.saved_pct == 71.1
+        assert response.totals.saved_per_session == 5.9
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -723,6 +758,7 @@ class TestAutoRouterBenchmarks:
             assert (idle.cache.hit_rate_pct, idle.cache.coverage_pct) == (0.0, 0.0)
             assert idle.cache.same_model.turns == idle.cache.return_to_tier.hits == 0
             assert idle.tier_turns == {}
+            assert idle.classifier_cost == 0.0
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -799,7 +835,6 @@ class TestAutoRouterBenchmarks:
 from datetime import datetime, timedelta, timezone
 from unittest.mock import AsyncMock, MagicMock
 
-
 from litellm.proxy.management_endpoints.auto_router_endpoints import (
     get_shadow_eval_job,
     list_shadow_eval_jobs,
@@ -1134,7 +1169,9 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp
     assert (
         len(
             {
-                frozenset((k, tuple(v) if isinstance(v, list) else v) for k, v in row.items() if k not in ("target_id", "id"))
+                frozenset(
+                    (k, tuple(v) if isinstance(v, list) else v) for k, v in row.items() if k not in ("target_id", "id")
+                )
                 for row in rows
             }
         )
@@ -1261,8 +1298,8 @@ async def test_start_shadow_eval_accepts_an_sdk_judge_when_anthropic_secret_look
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     import litellm
-    from litellm.integrations.custom_secret_manager import CustomSecretManager
     import litellm.proxy.proxy_server as proxy_server
+    from litellm.integrations.custom_secret_manager import CustomSecretManager
     from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
 
     class AnthropicSecretManager(CustomSecretManager):
@@ -1824,9 +1861,10 @@ async def test_list_shadow_eval_jobs_rejects_a_lone_filter_half(monkeypatch: pyt
 
 @pytest.mark.asyncio
 async def test_start_shadow_eval_concurrent_unique_violation_is_a_409(monkeypatch: pytest.MonkeyPatch):
-    import litellm.proxy.proxy_server as proxy_server
     from prisma.errors import UniqueViolationError
 
+    import litellm.proxy.proxy_server as proxy_server
+
     _configure_anthropic_sdk_judge(monkeypatch)
     prisma = _shadow_prisma()
     prisma.db.litellm_shadowevaljob.create_many = AsyncMock(
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx
index e7c6adc478c..006da4f2725 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx
@@ -68,6 +68,7 @@ const totals = (overrides: Partial<Totals> = {}): Totals => ({
   avg_session_seconds: 7560,
   avg_tokens_per_session: 5_300_000,
   spend: 359.86,
+  classifier_cost: 6.146,
   saved_spend: 2174.59,
   baseline_spend: 2534.45,
   saved_pct: 85.8,
@@ -99,6 +100,7 @@ const zeroTotals: Totals = {
   avg_session_seconds: 0,
   avg_tokens_per_session: 0,
   spend: 0,
+  classifier_cost: 0,
   saved_spend: 0,
   baseline_spend: 0,
   saved_pct: 0,
@@ -184,6 +186,37 @@ describe("AutoRouterBenchmarksTab", () => {
     expect(screen.getByText("5.3M")).toBeInTheDocument();
   });
 
+  it.each([
+    { spend: 20665.28, classifier_cost: 342.18, turns: 140815, llm: "$20,323.10", cost: "$342.18" },
+    { spend: 0, classifier_cost: 0, turns: 0, llm: "$0.00", cost: "$0.00" },
+    { spend: 0.002, classifier_cost: 0.0004, turns: 100, llm: "$0.0016", cost: "$0.0004" },
+  ])("shows total classification cost across $turns turns without a per-turn rate", ({ llm, cost, ...values }) => {
+    const stats = totals({ ...values, saved_spend: 10126.28, baseline_spend: values.spend + 10126.28 });
+    mockHook({ data: response([group(stats)], stats) });
+    renderTab();
+
+    expect(
+      screen
+        .getAllByRole("definition")
+        .map((node) => node.textContent)
+        .slice(1, 3),
+    ).toEqual([llm, cost]);
+    expect(screen.queryByText(/1K turns/)).not.toBeInTheDocument();
+    expect(screen.getAllByText("$10,126.28").length).toBeGreaterThan(0);
+  });
+
+  it.each([null, undefined])("keeps totals when the classification breakdown is %s", (classifier_cost) => {
+    const stats = totals({ classifier_cost });
+    mockHook({ data: response([group(stats)], stats) });
+    renderTab();
+
+    expect(screen.getAllByText("Unavailable")).toHaveLength(2);
+    expect(screen.queryByText(/\/ 1K turns/)).not.toBeInTheDocument();
+    expect(screen.getByText("$359.86")).toBeInTheDocument();
+    expect(screen.getByText("$2,174.59")).toBeInTheDocument();
+    expect(screen.getByText(/some usage predates classification-cost tracking/)).toBeInTheDocument();
+  });
+
   it("pairs the savings with the session count it was earned over, in its own tile", () => {
     mockHook({ data: response([group(), group({ router_name: "gpt-auto" })]) });
     renderTab();
@@ -201,8 +234,13 @@ describe("AutoRouterBenchmarksTab", () => {
 
     const terms = screen.getAllByRole("term").map((node) => node.textContent);
     const values = screen.getAllByRole("definition").map((node) => node.textContent);
-    expect(terms).toEqual(["Actual auto-router spend", "Estimated spend at highest-tier model"]);
-    expect(values).toEqual(["$359.86", "$2,534.45"]);
+    expect(terms).toEqual([
+      "Actual auto-router spend",
+      "LLM spend",
+      "Classification cost",
+      "Estimated spend at highest-tier model",
+    ]);
+    expect(values).toEqual(["$359.86", "$353.71", "$6.15", "$2,534.45"]);
   });
 
   it("lets both hero columns shrink below their content so a large total cannot clip", () => {
@@ -339,7 +377,7 @@ describe("AutoRouterBenchmarksTab", () => {
     renderTab();
 
     expect(screen.getByText("Total estimated savings")).toBeInTheDocument();
-    expect(screen.getAllByText("$0.00")).toHaveLength(4);
+    expect(screen.getAllByText("$0.00")).toHaveLength(6);
     expect(screen.getByText("· 0 sessions")).toBeInTheDocument();
     expect(screen.getByText("0s")).toBeInTheDocument();
     expect(screen.getByText(/turns measured/)).toBeInTheDocument();
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx
index 39e6b0fd390..33ad1bfe555 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx
@@ -52,10 +52,14 @@ const Metric: React.FC<{ label: string; value: string; hint?: string }> = ({ lab
   </Card>
 );
 
-const SpendRow: React.FC<{ label: string; value: string }> = ({ label, value }) => (
-  <dl className="flex items-baseline justify-between gap-6 py-3">
-    <dt className="text-sm text-muted-foreground">{label}</dt>
-    <dd className="text-base font-semibold tabular-nums text-foreground">{value}</dd>
+const SpendRow: React.FC<{ label: string; value: string; subdued?: boolean }> = ({ label, value, subdued }) => (
+  <dl className="flex flex-wrap items-baseline justify-between gap-x-6 gap-y-1 py-2">
+    <dt className="min-w-0 text-sm text-muted-foreground">{label}</dt>
+    <dd
+      className={`min-w-0 break-all tabular-nums ${subdued ? "text-sm font-normal text-muted-foreground" : "text-base font-semibold text-foreground"}`}
+    >
+      {value}
+    </dd>
   </dl>
 );
 
@@ -70,7 +74,9 @@ const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => {
             Total estimated savings
           </p>
           <div className="flex flex-wrap items-center justify-center gap-3">
-            <p className="text-6xl font-semibold tracking-tight text-foreground">{usd(stats.saved_spend)}</p>
+            <p className="min-w-0 break-all text-center text-4xl font-semibold tracking-tight text-foreground xl:text-6xl">
+              {usd(stats.saved_spend)}
+            </p>
             <Badge
               variant="secondary"
               className={`h-6 px-2.5 text-sm ${cheaper ? "bg-success/10 text-success" : "bg-destructive/10 text-destructive"}`}
@@ -83,6 +89,23 @@ const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => {
 
         <div className="flex flex-col justify-center border-t p-6 md:border-t-0 md:border-l">
           <SpendRow label="Actual auto-router spend" value={usd(stats.spend)} />
+          <div className="mb-3 border-l-2 pl-4">
+            <SpendRow
+              subdued
+              label="LLM spend"
+              value={stats.classifier_cost == null ? "Unavailable" : usd(stats.spend - stats.classifier_cost)}
+            />
+            <SpendRow
+              subdued
+              label="Classification cost"
+              value={stats.classifier_cost == null ? "Unavailable" : usd(stats.classifier_cost)}
+            />
+          </div>
+          {stats.classifier_cost == null && (
+            <p className="mb-3 text-xs text-muted-foreground">
+              Breakdown unavailable because some usage predates classification-cost tracking.
+            </p>
+          )}
           <Separator />
           <SpendRow label="Estimated spend at highest-tier model" value={usd(stats.baseline_spend)} />
         </div>
@@ -254,8 +277,9 @@ const BenchmarksBody: React.FC<BenchmarksBodyProps> = ({ isPending, error, data,
       <p className="text-xs text-muted-foreground">
         Compares your actual routed spend with the estimated cost of using only the most expensive model configured in
         the auto-router. It accounts for both the cache savings from staying on one model and the added cache costs from
-        switching models. The range counts whole sessions that overlap it, so totals can differ slightly from the
-        Overall tab, which buckets savings by UTC day.
+        switching models. Savings are net of recorded LLM classification cost, which is included in actual spend. The
+        range counts whole sessions that overlap it, so totals can differ slightly from the Overall tab, which buckets
+        savings by UTC day.
       </p>
 
       <div className="space-y-4">
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.test.ts
index 9a70fd9289a..22d6336e86f 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.test.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.test.ts
@@ -37,6 +37,7 @@ const totals = (overrides: Partial<AutoRouterBenchmarkGroup> = {}) => ({
   avg_session_seconds: 7560,
   avg_tokens_per_session: 5_300_000,
   spend: 359.86,
+  classifier_cost: 6.146,
   saved_spend: 2174.59,
   baseline_spend: 2534.45,
   saved_pct: 85.8,
diff --git a/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx
index 8cfd9d1941e..be95c0e600d 100644
--- a/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx
+++ b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx
@@ -32,6 +32,7 @@ const stats = {
   avg_session_seconds: 30,
   avg_tokens_per_session: 100,
   spend: 1.25,
+  classifier_cost: 0.25,
   saved_spend: 8.75,
   baseline_spend: 10,
   saved_pct: 87.5,
@@ -84,6 +85,11 @@ describe("KeyAutoRouterUsageTab", () => {
     expect(await screen.findByText("$8.75")).toBeInTheDocument();
     expect(screen.getByText("Actual auto-router spend")).toBeInTheDocument();
     expect(screen.getByText("$1.25")).toBeInTheDocument();
+    expect(screen.getByText("LLM spend")).toBeInTheDocument();
+    expect(screen.getByText("$1.00")).toBeInTheDocument();
+    expect(screen.getByText("Classification cost")).toBeInTheDocument();
+    expect(screen.getByText("$0.2500")).toBeInTheDocument();
+    expect(screen.queryByText(/1K turns/)).not.toBeInTheDocument();
     expect(screen.getByText("Estimated spend at highest-tier model")).toBeInTheDocument();
     expect(screen.getByText("$10.00")).toBeInTheDocument();
     expect(screen.getByText("Auto-router prompt caching")).toBeInTheDocument();
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 6fb08445aff..424ee773896 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -23303,6 +23303,11 @@ export interface components {
              */
             baseline_spend: number;
             cache: components["schemas"]["AutoRouterCacheStats"];
+            /**
+             * Classifier Cost
+             * @description Recorded LLM classifier cost already included in spend; null when any session turns predate subtotal recording, and zero for an empty window
+             */
+            classifier_cost: number | null;
             /**
              * Router Name
              * @description The auto-router alias requests were sent to
@@ -23359,6 +23364,11 @@ export interface components {
              */
             baseline_spend: number;
             cache: components["schemas"]["AutoRouterCacheStats"];
+            /**
+             * Classifier Cost
+             * @description Recorded LLM classifier cost already included in spend; null when any session turns predate subtotal recording, and zero for an empty window
+             */
+            classifier_cost: number | null;
             /**
              * Saved Pct
              * @description saved_spend over baseline_spend, as a percentage

From 1009976c497592207804593ea7ffbd639b1f68d6 Mon Sep 17 00:00:00 2001
From: yucheng-berri <yucheng@berri.ai>
Date: Mon, 7 Sep 2026 17:16:47 -0700
Subject: [PATCH 148/164] fix(bedrock): keep x-amzn-RequestId on chat error
 responses (#40089)

* fix(bedrock): keep x-amzn-RequestId on chat error responses

Bedrock chat error paths built BedrockError from only a status code and a
message, so the provider response headers were gone before exception mapping
ran and the proxy had nothing to forward. AWS support needs x-amzn-RequestId
to investigate a server-side error.

- converse and invoke chat handlers pass the real headers and response when
  they turn an httpx.HTTPStatusError into a BedrockError, and read the body
  through error_response_text so a streamed body nobody read does not throw
- every bedrock chat get_error_class honors the headers it is already handed:
  invoke, moonshot, bedrock-hosted openai, agentcore and the invoke agent
- BedrockError carries those headers into the response it synthesizes when a
  caller has headers but no response, skipping values httpx cannot carry
- the bedrock 500 mapping forwards the provider response like its 4xx and 503
  siblings instead of fabricating a blank one

The proxy now returns llm_provider-x-amzn-requestid on Bedrock chat errors.

* fix(bedrock): keep request-id on text-classified errors

The context-window and image branches of _map_bedrock_exception built their
litellm exception without the provider response, so a Bedrock 400 classified
by its body text lost x-amzn-RequestId while the sibling branches kept it.

Also narrows the new BedrockError types and trims its docstrings.

* chore(bedrock): drop the docstrings on the new error helpers

* fix(bedrock): keep request-id on every error path that has one

The ticket's root cause is that every BedrockError raise site under
litellm/llms/bedrock/ was built from status and message alone. The first
commits covered the chat and invoke handlers; this covers the rest.

Embeddings, rerank, image generation, image edit, count tokens, search and
the transformation layers now hand on the provider response or its headers,
and both bedrock_mantle configs return a BedrockError instead of the
OpenAI error that drops them.

Two blockers surfaced while verifying the streaming path. The trailing
`except Exception` in make_call and make_sync_call swallowed the BedrockError
raised a few lines above, relabelling a provider status as a 500, and the
non-200 branch read an unread streamed body, which throws.

The raise sites left alone have no provider response to carry: timeouts,
credential and config errors, and mid-stream event frames.

* fix(bedrock): forward provider headers from the count tokens route

The count tokens route converts BedrockError into an HTTPException, and dropped
the headers the handler had just kept, so that route still lost the request id.

get_response_headers now takes a Mapping so an httpx.Headers can be handed to it
without a copy.

* fix(bedrock): classify every bedrock surface through BedrockError

Eleven bedrock configs still inherited a provider-agnostic get_error_class
that builds a blank response, so the request id was gone before the proxy
read it. Claude platform, bedrock anthropic-messages, both image edit
configs, passthrough, realtime, vector stores and agentcore search now
return BedrockError, and a parametrized audit drives all 36 configs.

* fix(proxy): keep provider headers on the httpx status error branch

_handle_llm_api_exception forwards safe_headers on every branch except the
httpx.HTTPStatusError one, which the bedrock passthrough route reaches, so
the request id was dropped before the client saw the response.

* fix(bedrock): keep the request id on the timeout mappings

Timeout takes no response argument, so the three bedrock timeout branches
dropped the provider headers even when the upstream answered 408 or 504
with an x-amzn-RequestId. They now ride on the exception, already
llm_provider-prefixed, which is the form the proxy emits.

* fix(bedrock): keep the provider response on mapped timeouts

The previous round attached llm_provider-prefixed headers directly to the
Timeout. That shadowed the raw upstream headers for _get_response_headers,
so router cooldown and fallback cooldown stopped honouring retry-after on
bedrock 408/504 replies.

Give Timeout an optional response instead, the way every other mapped
bedrock exception already carries one. Retry logic reads the raw
retry-after off the response, and the proxy prefixes those headers on the
way out, so clients still see llm_provider-x-amzn-requestid.

* chore(bedrock): drop the explanatory comment on Timeout.response
---
 litellm/exceptions.py                         |   3 +
 .../exception_mapping_utils.py                |  11 +-
 .../llm_response_utils/get_headers.py         |   5 +-
 .../bedrock/chat/agentcore/transformation.py  |  19 +-
 litellm/llms/bedrock/chat/converse_handler.py |  23 +-
 .../bedrock/chat/converse_transformation.py   |   1 +
 .../chat/invoke_agent/transformation.py       |   3 +-
 litellm/llms/bedrock/chat/invoke_handler.py   |  33 ++-
 .../amazon_moonshot_transformation.py         |   2 +-
 .../amazon_openai_transformation.py           |   2 +-
 ...mazon_twelvelabs_pegasus_transformation.py |   2 +
 .../base_invoke_transformation.py             |  10 +-
 .../bedrock/claude_platform/common_utils.py   |  11 +
 litellm/llms/bedrock/common_utils.py          |  47 +++-
 litellm/llms/bedrock/count_tokens/handler.py  |   4 +
 litellm/llms/bedrock/embed/embedding.py       |  14 +-
 ...n_nova_canvas_image_edit_transformation.py |   9 +
 litellm/llms/bedrock/image_edit/handler.py    |  14 +-
 .../image_edit/stability_transformation.py    |   9 +
 .../bedrock/image_generation/image_handler.py |  14 +-
 .../anthropic_claude3_transformation.py       |   9 +
 .../bedrock/passthrough/transformation.py     |  11 +-
 .../llms/bedrock/realtime/transformation.py   |  10 +
 litellm/llms/bedrock/rerank/handler.py        |  14 +-
 litellm/llms/bedrock/search/transformation.py |   6 +-
 .../bedrock/vector_stores/transformation.py   |   9 +
 .../bedrock_mantle/chat/transformation.py     |   9 +
 .../responses/transformation.py               |   8 +
 litellm/proxy/common_request_processing.py    |   4 +
 .../llm_passthrough_endpoints.py              |   9 +-
 .../test_exception_mapping_utils.py           | 155 +++++++++++
 .../test_base_invoke_transformation.py        |  13 +
 .../chat/test_converse_transformation.py      |   4 +
 .../llms/bedrock/chat/test_invoke_handler.py  | 168 ++++++++++++
 .../llms/bedrock/test_bedrock_common_utils.py | 257 ++++++++++++++++++
 .../llms/chat/test_converse_handler.py        |  61 +++++
 .../test_llm_pass_through_endpoints.py        |  34 +++
 .../proxy/test_common_request_processing.py   |  35 +++
 .../test_exception_header_preservation.py     |  83 ++++++
 39 files changed, 1100 insertions(+), 35 deletions(-)

diff --git a/litellm/exceptions.py b/litellm/exceptions.py
index 16202321709..f9215267bf3 100644
--- a/litellm/exceptions.py
+++ b/litellm/exceptions.py
@@ -338,6 +338,7 @@ class Timeout(openai.APITimeoutError):
         num_retries: int | None = None,
         headers: dict | None = None,
         exception_status_code: int | None = None,
+        response: httpx.Response | None = None,
     ):
         request: Final = httpx.Request(
             method="POST",
@@ -352,6 +353,8 @@ class Timeout(openai.APITimeoutError):
         self.max_retries = max_retries
         self.num_retries = num_retries
         self.headers = headers
+        if response is not None:
+            self.response = response
 
     # custom function to convert to str
     def __str__(self):
diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py
index 8f8c955d971..82708d412c9 100644
--- a/litellm/litellm_core_utils/exception_mapping_utils.py
+++ b/litellm/litellm_core_utils/exception_mapping_utils.py
@@ -860,6 +860,7 @@ def _map_bedrock_exception(
                 message=mantle_context_window_message,
                 model=model,
                 llm_provider=custom_llm_provider,
+                response=getattr(original_exception, "response", None),
             )
     if (
         "too many tokens" in error_str
@@ -873,6 +874,7 @@ def _map_bedrock_exception(
             message=f"BedrockException: Context Window Error - {error_str}",
             model=model,
             llm_provider="bedrock",
+            response=getattr(original_exception, "response", None),
         )
     elif "Conversation blocks and tool result blocks cannot be provided in the same turn." in error_str:
         raise BadRequestError(
@@ -924,12 +926,14 @@ def _map_bedrock_exception(
             message=f"BedrockException: Timeout Error - {error_str}",
             model=model,
             llm_provider="bedrock",
+            response=getattr(original_exception, "response", None),
         )
     elif "Could not process image" in error_str:
         raise litellm.InternalServerError(
             message=f"BedrockException - {error_str}",
             model=model,
             llm_provider="bedrock",
+            response=getattr(original_exception, "response", None),
         )
     elif hasattr(original_exception, "status_code"):
         if original_exception.status_code == 500:
@@ -937,10 +941,7 @@ def _map_bedrock_exception(
                 message=f"BedrockException - {original_exception.message}",
                 llm_provider="bedrock",
                 model=model,
-                response=httpx.Response(
-                    status_code=500,
-                    request=httpx.Request(method="POST", url="https://api.openai.com/v1/"),
-                ),
+                response=getattr(original_exception, "response", None),
             )
         elif original_exception.status_code == 401:
             raise AuthenticationError(
@@ -969,6 +970,7 @@ def _map_bedrock_exception(
                 model=model,
                 llm_provider=custom_llm_provider,
                 litellm_debug_info=extra_information,
+                response=getattr(original_exception, "response", None),
             )
         elif original_exception.status_code == 422:
             raise BadRequestError(
@@ -1001,6 +1003,7 @@ def _map_bedrock_exception(
                 llm_provider=custom_llm_provider,
                 litellm_debug_info=extra_information,
                 exception_status_code=original_exception.status_code,
+                response=getattr(original_exception, "response", None),
             )
 
 
diff --git a/litellm/litellm_core_utils/llm_response_utils/get_headers.py b/litellm/litellm_core_utils/llm_response_utils/get_headers.py
index f1ae6492e4e..d04abcb6e7b 100644
--- a/litellm/litellm_core_utils/llm_response_utils/get_headers.py
+++ b/litellm/litellm_core_utils/llm_response_utils/get_headers.py
@@ -1,7 +1,8 @@
+from collections.abc import Mapping
 from typing import Final
 
 
-def get_response_headers(_response_headers: dict | None = None) -> dict:
+def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict:
     """
 
     Sets the Appropriate OpenAI headers for the response and forward all headers as llm_provider-{header}
@@ -31,7 +32,7 @@ def get_response_headers(_response_headers: dict | None = None) -> dict:
     return {**llm_provider_headers, **openai_headers}
 
 
-def _get_llm_provider_headers(response_headers: dict) -> dict:
+def _get_llm_provider_headers(response_headers: Mapping[str, str]) -> dict:
     """
     Adds a llm_provider-{header} to all headers that are not already prefixed with llm_provider
 
diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py
index 690040dd93b..6aa17372258 100644
--- a/litellm/llms/bedrock/chat/agentcore/transformation.py
+++ b/litellm/llms/bedrock/chat/agentcore/transformation.py
@@ -667,7 +667,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
         )
 
         if response.status_code != 200:
-            raise BedrockError(status_code=response.status_code, message=str(response.read()))
+            raise BedrockError(
+                status_code=response.status_code,
+                message=str(response.read()),
+                headers=response.headers,
+                response=response,
+            )
 
         # LOGGING
         logging_obj.post_call(
@@ -690,6 +695,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
                 raise BedrockError(
                     status_code=response.status_code,
                     message=f"AgentCore: Failed to read/parse JSON response body: {e}",
+                    headers=response.headers,
                 )
             parsed: Final = self._parse_json_response(response_json)
 
@@ -880,7 +886,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
         )
 
         if response.status_code != 200:
-            raise BedrockError(status_code=response.status_code, message=str(await response.aread()))
+            raise BedrockError(
+                status_code=response.status_code,
+                message=str(await response.aread()),
+                headers=response.headers,
+                response=response,
+            )
 
         # LOGGING
         logging_obj.post_call(
@@ -903,6 +914,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
                 raise BedrockError(
                     status_code=response.status_code,
                     message=f"AgentCore: Failed to read/parse JSON response body: {e}",
+                    headers=response.headers,
                 )
             parsed: Final = self._parse_json_response(response_json)
 
@@ -1031,6 +1043,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
             raise BedrockError(
                 message=f"Error processing response: {e}",
                 status_code=raw_response.status_code,
+                headers=raw_response.headers,
             )
 
     def validate_environment(
@@ -1046,7 +1059,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
         return headers
 
     def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
-        return BedrockError(status_code=status_code, message=error_message)
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
 
     def should_fake_stream(
         self,
diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py
index a75124325ae..984ba371898 100644
--- a/litellm/llms/bedrock/chat/converse_handler.py
+++ b/litellm/llms/bedrock/chat/converse_handler.py
@@ -22,7 +22,7 @@ from litellm.types.utils import ModelResponse
 from litellm.utils import CustomStreamWrapper
 
 from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
-from ..common_utils import BedrockError, _get_all_bedrock_regions
+from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text
 from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
 
 
@@ -66,7 +66,12 @@ def make_sync_call(
     )
 
     if response.status_code != 200:
-        raise BedrockError(status_code=response.status_code, message=str(response.read()))
+        raise BedrockError(
+            status_code=response.status_code,
+            message=str(response.read()),
+            headers=response.headers,
+            response=response,
+        )
 
     if fake_stream:
         model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
@@ -247,7 +252,12 @@ class BedrockConverseLLM(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=error_response_text(err.response),
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
@@ -594,7 +604,12 @@ class BedrockConverseLLM(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=error_response_text(err.response),
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py
index e097805f54a..fa24f8be893 100644
--- a/litellm/llms/bedrock/chat/converse_transformation.py
+++ b/litellm/llms/bedrock/chat/converse_transformation.py
@@ -2255,6 +2255,7 @@ class AmazonConverseConfig(BaseConfig):
             raise BedrockError(
                 message=f"Error converting to valid response block={e}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues",
                 status_code=422,
+                headers=response.headers,
             )
 
         """
diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py
index e30ec731d8c..d489e47c3b5 100644
--- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py
+++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py
@@ -470,6 +470,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM):
             raise BedrockError(
                 message=f"Error processing response: {e}",
                 status_code=raw_response.status_code,
+                headers=raw_response.headers,
             )
 
     def validate_environment(
@@ -485,7 +486,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM):
         return headers
 
     def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
-        return BedrockError(status_code=status_code, message=error_message)
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
 
     def should_fake_stream(
         self,
diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py
index c39c88240c5..5f8a5544d65 100644
--- a/litellm/llms/bedrock/chat/invoke_handler.py
+++ b/litellm/llms/bedrock/chat/invoke_handler.py
@@ -42,6 +42,7 @@ from litellm.types.utils import GenericStreamingChunk as GChunk
 from ..common_utils import (
     BedrockError,
     build_bedrock_stream_error,
+    error_response_text,
     get_bedrock_response_stream_shape,
     get_bedrock_tool_name,
 )
@@ -184,7 +185,12 @@ async def make_call(
         )
 
         if response.status_code != 200:
-            raise BedrockError(status_code=response.status_code, message=response.text)
+            raise BedrockError(
+                status_code=response.status_code,
+                message=error_response_text(response),
+                headers=response.headers,
+                response=response,
+            )
 
         if fake_stream:
             model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
@@ -228,9 +234,16 @@ async def make_call(
         )
 
         return completion_stream, response.headers
+    except BedrockError:
+        raise
     except httpx.HTTPStatusError as err:
         error_code: Final = err.response.status_code
-        raise BedrockError(status_code=error_code, message=err.response.text)
+        raise BedrockError(
+            status_code=error_code,
+            message=error_response_text(err.response),
+            headers=err.response.headers,
+            response=err.response,
+        )
     except httpx.TimeoutException:
         raise BedrockError(status_code=408, message="Timeout error occurred.")
     except Exception as e:
@@ -270,7 +283,12 @@ def make_sync_call(
         )
 
         if response.status_code != 200:
-            raise BedrockError(status_code=response.status_code, message=response.text)
+            raise BedrockError(
+                status_code=response.status_code,
+                message=error_response_text(response),
+                headers=response.headers,
+                response=response,
+            )
 
         if fake_stream:
             model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
@@ -314,9 +332,16 @@ def make_sync_call(
         )
 
         return completion_stream, response.headers
+    except BedrockError:
+        raise
     except httpx.HTTPStatusError as err:
         error_code: Final = err.response.status_code
-        raise BedrockError(status_code=error_code, message=err.response.text)
+        raise BedrockError(
+            status_code=error_code,
+            message=error_response_text(err.response),
+            headers=err.response.headers,
+            response=err.response,
+        )
     except httpx.TimeoutException:
         raise BedrockError(status_code=408, message="Timeout error occurred.")
     except Exception as e:
diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py
index 04c6ec86a13..5d39b68d9d5 100644
--- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py
+++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py
@@ -247,4 +247,4 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig):
 
     def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError:
         """Return the appropriate error class for Bedrock."""
-        return BedrockError(status_code=status_code, message=error_message)
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py
index 1671585be2d..4bf1a1cba73 100644
--- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py
+++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py
@@ -182,4 +182,4 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM):
 
     def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError:
         """Return the appropriate error class for Bedrock."""
-        return BedrockError(status_code=status_code, message=error_message)
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py
index cd8066cda4d..d12c8aee48c 100644
--- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py
+++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py
@@ -212,6 +212,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
             raise BedrockError(
                 message=f"Error parsing response: {raw_response.text}, error: {e}",
                 status_code=raw_response.status_code,
+                headers=raw_response.headers,
             )
 
         verbose_logger.debug(
@@ -241,6 +242,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
             raise BedrockError(
                 message=f"Error setting response content: {e}. Response: {completion_response}",
                 status_code=raw_response.status_code,
+                headers=raw_response.headers,
             )
 
         # Calculate usage from headers
diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
index 37121d2ece7..a0e32c8aa22 100644
--- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
+++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
@@ -295,7 +295,11 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
         try:
             completion_response: Final = raw_response.json()
         except Exception:
-            raise BedrockError(message=raw_response.text, status_code=raw_response.status_code)
+            raise BedrockError(
+                message=raw_response.text,
+                status_code=raw_response.status_code,
+                headers=raw_response.headers,
+            )
         verbose_logger.debug(
             "bedrock invoke response % s",
             json.dumps(completion_response, indent=4, default=str),
@@ -363,6 +367,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
             raise BedrockError(
                 message=f"Error processing={raw_response.text}, Received error={e}",
                 status_code=422,
+                headers=raw_response.headers,
             )
 
         try:
@@ -384,6 +389,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
             raise BedrockError(
                 message=f"Error parsing received text={outputText}.\nError-{e}",
                 status_code=raw_response.status_code,
+                headers=raw_response.headers,
             )
 
         ## CALCULATING USAGE - bedrock returns usage in the headers
@@ -431,7 +437,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
         return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names)
 
     def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
-        return BedrockError(status_code=status_code, message=error_message)
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
 
     @track_llm_api_timing()
     async def get_async_custom_stream_wrapper(
diff --git a/litellm/llms/bedrock/claude_platform/common_utils.py b/litellm/llms/bedrock/claude_platform/common_utils.py
index 311f3a56b84..fb7f2185ec5 100644
--- a/litellm/llms/bedrock/claude_platform/common_utils.py
+++ b/litellm/llms/bedrock/claude_platform/common_utils.py
@@ -1,7 +1,10 @@
 from typing import Final
 
+import httpx
+
 import litellm
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.secret_managers.main import get_secret_str
 
 CLAUDE_PLATFORM_SERVICE_NAME: Final = "aws-external-anthropic"
@@ -15,6 +18,14 @@ def strip_claude_platform_route(model: str) -> str:
 
 
 class BedrockClaudePlatformMixin(BaseAWSLLM):
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     @staticmethod
     def _get_workspace_id(optional_params: dict, litellm_params: dict) -> str | None:
         workspace_id = (
diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py
index fe675a30a00..be4f0f32689 100644
--- a/litellm/llms/bedrock/common_utils.py
+++ b/litellm/llms/bedrock/common_utils.py
@@ -33,8 +33,53 @@ if TYPE_CHECKING:
     from litellm.types.llms.openai import AllMessageValues
 
 
+_ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs"
+
+
+def error_response_text(response: httpx.Response) -> str:
+    try:
+        return response.text
+    except httpx.ResponseNotRead:
+        return response.reason_phrase
+
+
+def _synthesize_error_response(
+    *, status_code: int, headers: dict[str, object] | httpx.Headers, request: httpx.Request | None
+) -> tuple[httpx.Request, httpx.Response]:
+    error_request: Final = request or httpx.Request(method="POST", url=_ERROR_REQUEST_URL)
+    safe_headers: Final = (
+        headers
+        if isinstance(headers, httpx.Headers)
+        else tuple((key, value) for key, value in headers.items() if isinstance(value, (str, bytes)))
+    )
+    return error_request, httpx.Response(status_code=status_code, headers=safe_headers, request=error_request)
+
+
 class BedrockError(BaseLLMException):
-    pass
+    def __init__(
+        self,
+        status_code: int,
+        message: str,
+        headers: dict[str, object] | httpx.Headers | None = None,
+        request: httpx.Request | None = None,
+        response: httpx.Response | None = None,
+        body: dict[str, object] | None = None,
+        status_code_is_synthesized: bool = False,
+    ) -> None:
+        error_request, error_response = (
+            _synthesize_error_response(status_code=status_code, headers=headers, request=request)
+            if response is None and headers
+            else (request, response)
+        )
+        super().__init__(
+            status_code=status_code,
+            message=message,
+            headers=headers,
+            request=error_request,
+            response=error_response,
+            body=body,
+            status_code_is_synthesized=status_code_is_synthesized,
+        )
 
 
 _BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = (
diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py
index 2383350b3a3..1fb53f6ff0a 100644
--- a/litellm/llms/bedrock/count_tokens/handler.py
+++ b/litellm/llms/bedrock/count_tokens/handler.py
@@ -102,6 +102,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
                 raise BedrockError(
                     status_code=response.status_code,
                     message=error_text,
+                    headers=response.headers,
+                    response=response,
                 )
 
             bedrock_response: Final = response.json()
@@ -124,6 +126,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
             raise BedrockError(
                 status_code=e.response.status_code,
                 message=e.response.text,
+                headers=e.response.headers,
+                response=e.response,
             )
         except Exception as e:
             verbose_logger.error("Error in CountTokens handler: %s", e)
diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py
index 5fb86d476f4..d3725434498 100644
--- a/litellm/llms/bedrock/embed/embedding.py
+++ b/litellm/llms/bedrock/embed/embedding.py
@@ -132,7 +132,12 @@ class BedrockEmbedding(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
@@ -161,7 +166,12 @@ class BedrockEmbedding(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
diff --git a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py
index 18d47301ee5..acb0cc8dcb7 100644
--- a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py
+++ b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py
@@ -20,6 +20,7 @@ import httpx
 
 from litellm._logging import verbose_logger
 from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.types.images.main import ImageEditOptionalRequestParams
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import FileTypes, ImageObject, ImageResponse
@@ -228,6 +229,14 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig):
         """
         return _supports_nova_canvas_image_edit_from_model_cost(model or "")
 
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def get_supported_openai_params(self, model: str) -> list:
         return [
             "n",
diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py
index 5c517f2049c..be6489f20ae 100644
--- a/litellm/llms/bedrock/image_edit/handler.py
+++ b/litellm/llms/bedrock/image_edit/handler.py
@@ -114,7 +114,12 @@ class BedrockImageEdit(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
@@ -156,7 +161,12 @@ class BedrockImageEdit(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py
index 24e7ba73075..bc9a64f587a 100644
--- a/litellm/llms/bedrock/image_edit/stability_transformation.py
+++ b/litellm/llms/bedrock/image_edit/stability_transformation.py
@@ -27,6 +27,7 @@ from typing import TYPE_CHECKING, Any, Final
 import httpx
 
 from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.types.images.main import ImageEditOptionalRequestParams
 from litellm.types.llms.stability import (
     OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO,
@@ -84,6 +85,14 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
                 return True
         return False
 
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def get_supported_openai_params(self, model: str) -> list:
         """
         Return list of OpenAI params supported by Bedrock Stability.
diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py
index c78e3c147cb..87762b648e0 100644
--- a/litellm/llms/bedrock/image_generation/image_handler.py
+++ b/litellm/llms/bedrock/image_generation/image_handler.py
@@ -119,7 +119,12 @@ class BedrockImageGeneration(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
         ### FORMAT RESPONSE TO OPENAI FORMAT ###
@@ -162,7 +167,12 @@ class BedrockImageGeneration(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
index 6ff9f0155f9..a715d150b4c 100644
--- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
+++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
@@ -29,6 +29,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation
     AmazonInvokeConfig,
 )
 from litellm.llms.bedrock.common_utils import (
+    BedrockError,
     apply_bedrock_invoke_structured_output,
     ensure_bedrock_anthropic_messages_tool_names,
     get_anthropic_beta_from_headers,
@@ -79,6 +80,14 @@ class AmazonAnthropicClaudeMessagesConfig(
 
     BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(BedrockInvokeAnthropicMessagesRequest.__annotations__.keys())
 
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def __init__(self, **kwargs):
         BaseAnthropicMessagesConfig.__init__(self, **kwargs)
         AmazonInvokeConfig.__init__(self, **kwargs)
diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py
index d0a3c37ffb3..fb8bc4f191f 100644
--- a/litellm/llms/bedrock/passthrough/transformation.py
+++ b/litellm/llms/bedrock/passthrough/transformation.py
@@ -2,13 +2,14 @@ import json
 from collections.abc import Mapping
 from typing import TYPE_CHECKING, Final, Optional, cast
 
+import httpx
 from httpx import Response
 
 from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
 
 from ..base_aws_llm import BaseAWSLLM
-from ..common_utils import BedrockEventStreamDecoderBase, BedrockModelInfo
+from ..common_utils import BedrockError, BedrockEventStreamDecoderBase, BedrockModelInfo
 
 if TYPE_CHECKING:
     from httpx import URL
@@ -18,6 +19,14 @@ if TYPE_CHECKING:
 
 
 class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig):
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
         return "stream" in endpoint
 
diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py
index 1f4c81d6491..3b972961940 100644
--- a/litellm/llms/bedrock/realtime/transformation.py
+++ b/litellm/llms/bedrock/realtime/transformation.py
@@ -9,12 +9,14 @@ import json
 import uuid as uuid_lib
 from typing import Final, cast
 
+import httpx
 from pydantic import BaseModel
 
 from litellm._logging import verbose_logger
 from litellm._uuid import uuid
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm
 from litellm.types.llms.openai import (
     OpenAIRealtimeContentPartDone,
@@ -121,6 +123,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
         self._cumulative_usage = BedrockUsageEvent()
         self._reported_usage = BedrockUsageEvent()
 
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict:
         """Validate environment - no special validation needed for Bedrock."""
         return headers
diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py
index 4860c99268e..8847381cbc9 100644
--- a/litellm/llms/bedrock/rerank/handler.py
+++ b/litellm/llms/bedrock/rerank/handler.py
@@ -46,7 +46,12 @@ class BedrockRerankHandler(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
@@ -117,7 +122,12 @@ class BedrockRerankHandler(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py
index 920e566c9dd..e7d706c3731 100644
--- a/litellm/llms/bedrock/search/transformation.py
+++ b/litellm/llms/bedrock/search/transformation.py
@@ -39,7 +39,6 @@ from typing import Final
 import httpx
 
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
-from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.base_llm.search.transformation import (
     BaseSearchConfig,
     SearchResponse,
@@ -380,6 +379,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
             raise BedrockError(
                 status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
                 message=f"AgentCore gateway MCP error: {error}",
+                headers=raw_response.headers,
             )
 
         # A failed tools/call is reported in-band, as HTTP 200 with result.isError
@@ -389,6 +389,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
             raise BedrockError(
                 status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
                 message=f"AgentCore web search tool error: {self._tool_error_message(response_json)}",
+                headers=raw_response.headers,
             )
 
         text_items: Final = tuple(
@@ -440,6 +441,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
         raise BedrockError(
             status_code=502,
             message=f"AgentCore gateway returned SSE without a JSON data frame: {text[:200]}",
+            headers=raw_response.headers,
         )
 
     def get_error_class(
@@ -448,7 +450,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
         status_code: int,
         headers: dict,  # mutable-ok: BaseSearchConfig.get_error_class takes the response headers as a dict
     ) -> Exception:
-        return BaseLLMException(
+        return BedrockError(
             status_code=status_code,
             message=error_message,
             headers=headers,
diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py
index 6940077391f..27c90c9d71e 100644
--- a/litellm/llms/bedrock/vector_stores/transformation.py
+++ b/litellm/llms/bedrock/vector_stores/transformation.py
@@ -8,6 +8,7 @@ from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.url_utils import encode_url_path_segment
 from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.types.integrations.rag.bedrock_knowledgebase import (
     BedrockKBContent,
     BedrockKBResponse,
@@ -38,6 +39,14 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
         BaseVectorStoreConfig.__init__(self)
         BaseAWSLLM.__init__(self)
 
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials:
         return {}
 
diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py
index 64d7ef2bed6..d91157c3d10 100644
--- a/litellm/llms/bedrock_mantle/chat/transformation.py
+++ b/litellm/llms/bedrock_mantle/chat/transformation.py
@@ -13,6 +13,8 @@ Auth: Bearer token (litellm_params.api_key, BEDROCK_MANTLE_API_KEY, or the
 from collections.abc import AsyncIterator, Iterator
 from typing import Any, Final
 
+import httpx
+
 import litellm
 from litellm._logging import verbose_logger
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
@@ -24,6 +26,8 @@ from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.router import GenericLiteLLMParams
 
+from ...base_llm.chat.transformation import BaseLLMException
+from ...bedrock.common_utils import BedrockError
 from ...openai_like.chat.transformation import OpenAILikeChatConfig
 from ..common_utils import mantle_base_segment
 
@@ -45,6 +49,11 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
     def get_config(cls):
         return super().get_config()
 
+    def get_error_class(
+        self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers
+    ) -> BaseLLMException:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def _get_openai_compatible_provider_info(
         self,
         api_base: str | None,
diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py
index 5179c966584..bbbda4d14b6 100644
--- a/litellm/llms/bedrock_mantle/responses/transformation.py
+++ b/litellm/llms/bedrock_mantle/responses/transformation.py
@@ -19,11 +19,14 @@ import json
 from collections.abc import Mapping
 from typing import Any, Final
 
+import httpx
 from typing_extensions import ReadOnly, TypedDict
 
 import litellm
 from litellm._logging import verbose_logger
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.llms.bedrock_mantle.common_utils import (
     MANTLE_HOST_RE,
     BedrockMantleAuthMixin,
@@ -98,6 +101,11 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
     def custom_llm_provider(self) -> LlmProviders:
         return LlmProviders.BEDROCK_MANTLE
 
+    def get_error_class(
+        self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers
+    ) -> BaseLLMException:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def get_complete_url(
         self,
         api_base: str | None,
diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py
index 9720e4b1cf8..d0e3914e9fb 100644
--- a/litellm/proxy/common_request_processing.py
+++ b/litellm/proxy/common_request_processing.py
@@ -3508,9 +3508,13 @@ class ProxyBaseLLMRequestProcessing:
             error_body: Final = await http_status_error.response.aread()
             error_text: Final = error_body.decode("utf-8")
 
+            error_headers: Final = {  # mutable-ok: HTTPException takes a plain header dict
+                k: v if isinstance(v, str) else str(v) for k, v in safe_headers.items()
+            }
             raise HTTPException(
                 status_code=http_status_error.response.status_code,
                 detail={"error": error_text},
+                headers=error_headers,
             )
         error_msg: Final = f"{e}"
         # Check for AttributeError in the exception chain.
diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
index b95547e2b54..2ea46b740a8 100644
--- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
@@ -946,7 +946,14 @@ async def handle_bedrock_count_tokens(
     except BedrockError as e:
         # Convert BedrockError to HTTPException for FastAPI
         verbose_proxy_logger.error("BedrockError in handle_bedrock_count_tokens: %s", e)
-        raise HTTPException(status_code=e.status_code, detail={"error": e.message})
+        from litellm.litellm_core_utils.llm_response_utils.get_headers import get_response_headers
+
+        provider_headers: Final = getattr(getattr(e, "response", None), "headers", None)
+        raise HTTPException(
+            status_code=e.status_code,
+            detail={"error": e.message},
+            headers=get_response_headers(provider_headers) if provider_headers else None,
+        )
     except HTTPException:
         # Re-raise HTTP exceptions as-is
         raise
diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py
index 1778eca25ef..42d3df76902 100644
--- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py
+++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py
@@ -9,9 +9,11 @@ import litellm
 from litellm.litellm_core_utils.exception_mapping_utils import (
     ExceptionCheckers,
     _get_body_error_code,
+    _get_response_headers,
     exception_type,
     extract_and_raise_litellm_exception,
 )
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.llms.openai.common_utils import OpenAIError
 from litellm.types.utils import LlmProviders
 
@@ -1254,3 +1256,156 @@ def test_handle_error_marks_only_a_status_code_it_never_received():
         raise handler._handle_error(e=upstream, provider_config=None)
     assert received.value.status_code == 500
     assert received.value.status_code_is_synthesized is False
+
+
+def test_bedrock_500_preserves_provider_response_headers():
+    """A Bedrock 5xx must keep x-amzn-RequestId so AWS support can trace it (LIT-5428)."""
+    provider_response = httpx.Response(
+        status_code=500,
+        headers={"x-amzn-RequestId": "req-map-500"},
+        text='{"message":"Amazon Bedrock is unable to process your request."}',
+        request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
+    )
+    original_exception = BedrockError(
+        status_code=500,
+        message=provider_response.text,
+        headers=provider_response.headers,
+        response=provider_response,
+    )
+
+    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
+        exception_type(
+            model="anthropic.claude-haiku-4-5-20251001-v1:0",
+            original_exception=original_exception,
+            custom_llm_provider="bedrock",
+            completion_kwargs={},
+            extra_kwargs={},
+        )
+
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-map-500"
+
+
+@pytest.mark.parametrize(
+    "custom_llm_provider, status_code, provider_message, expected_exception",
+    [
+        (
+            "bedrock_mantle",
+            400,
+            (
+                '{"error":{"code":"validation_error",'
+                '"message":"prompt tokens (1055489) exceed model maximum (1050000) for openai.gpt-5.6-sol",'
+                '"param":null,"type":"invalid_request_error"}}'
+            ),
+            litellm.ContextWindowExceededError,
+        ),
+        (
+            "bedrock",
+            400,
+            '{"message":"Input is too long for requested model."}',
+            litellm.ContextWindowExceededError,
+        ),
+        (
+            "bedrock",
+            400,
+            '{"message":"Could not process image"}',
+            litellm.InternalServerError,
+        ),
+    ],
+)
+def test_bedrock_classified_errors_preserve_provider_response_headers(
+    custom_llm_provider, status_code, provider_message, expected_exception
+):
+    """Branches that classify a Bedrock error by its text must keep x-amzn-RequestId (LIT-5428)."""
+    provider_response = httpx.Response(
+        status_code=status_code,
+        headers={"x-amzn-RequestId": "req-classified"},
+        text=provider_message,
+        request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
+    )
+    original_exception = BedrockError(
+        status_code=status_code,
+        message=provider_message,
+        headers=provider_response.headers,
+        response=provider_response,
+    )
+
+    with pytest.raises(expected_exception) as exc_info:
+        exception_type(
+            model="anthropic.claude-haiku-4-5-20251001-v1:0",
+            original_exception=original_exception,
+            custom_llm_provider=custom_llm_provider,
+            completion_kwargs={},
+            extra_kwargs={},
+        )
+
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-classified"
+
+
+@pytest.mark.parametrize(
+    "status_code, provider_message",
+    [
+        (504, '{"message":"Gateway timeout"}'),
+        (408, '{"message":"Bedrock did not answer in time"}'),
+        (408, '{"message":"Connect timeout on endpoint URL"}'),
+    ],
+)
+def test_bedrock_timeout_mapping_preserves_provider_headers(status_code, provider_message):
+    """A mapped bedrock timeout keeps the upstream response, like every other mapped bedrock error.
+
+    The proxy prefixes those headers on the way out, while retry and cooldown
+    logic still reads the raw retry-after off the response.
+    """
+    provider_response = httpx.Response(
+        status_code=status_code,
+        headers={"x-amzn-RequestId": "req-timeout", "set-cookie": "session=attacker"},
+        text=provider_message,
+        request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
+    )
+    original_exception = BedrockError(
+        status_code=status_code,
+        message=provider_message,
+        headers=provider_response.headers,
+        response=provider_response,
+    )
+
+    with pytest.raises(litellm.Timeout) as exc_info:
+        exception_type(
+            model="anthropic.claude-haiku-4-5-20251001-v1:0",
+            original_exception=original_exception,
+            custom_llm_provider="bedrock",
+            completion_kwargs={},
+            extra_kwargs={},
+        )
+
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-timeout"
+    assert exc_info.value.headers is None
+
+
+@pytest.mark.parametrize("status_code", [504, 408])
+def test_bedrock_timeout_mapping_keeps_retry_after_readable(status_code):
+    """Cooldown and retry timing read retry-after through _get_response_headers."""
+    provider_response = httpx.Response(
+        status_code=status_code,
+        headers={"x-amzn-RequestId": "req-retry-after", "retry-after": "7"},
+        text='{"message":"Bedrock did not answer in time"}',
+        request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
+    )
+    original_exception = BedrockError(
+        status_code=status_code,
+        message='{"message":"Bedrock did not answer in time"}',
+        headers=provider_response.headers,
+        response=provider_response,
+    )
+
+    with pytest.raises(litellm.Timeout) as exc_info:
+        exception_type(
+            model="anthropic.claude-haiku-4-5-20251001-v1:0",
+            original_exception=original_exception,
+            custom_llm_provider="bedrock",
+            completion_kwargs={},
+            extra_kwargs={},
+        )
+
+    exception_headers = _get_response_headers(original_exception=exc_info.value)
+    assert exception_headers is not None
+    assert litellm.utils._get_retry_after_from_exception_header(response_headers=exception_headers) == 7
diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py
index aba51689094..c2c448cd7e2 100644
--- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py
+++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py
@@ -177,3 +177,16 @@ def test_guardrail_config_flows_to_headers_not_request_body(model):
     assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "ff6ujrregl1q"
     assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "DRAFT"
     assert headers["X-Amzn-Bedrock-Trace"] == "DISABLED"
+
+
+def test_get_error_class_preserves_provider_headers():
+    """The invoke handler path hands real provider headers to get_error_class (LIT-5428)."""
+    error = AmazonInvokeConfig().get_error_class(
+        error_message="Amazon Bedrock is unable to process your request.",
+        status_code=500,
+        headers={"x-amzn-RequestId": "req-invoke-500"},
+    )
+
+    assert isinstance(error, BedrockError)
+    assert error.headers == {"x-amzn-RequestId": "req-invoke-500"}
+    assert error.response.headers["x-amzn-requestid"] == "req-invoke-500"
diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
index cb05cdb9451..f0e361ceb88 100644
--- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
+++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
@@ -2,6 +2,7 @@ import asyncio
 import json
 import os
 
+import httpx
 import pytest
 from fastapi.testclient import TestClient
 
@@ -6039,6 +6040,8 @@ def test_transform_response_does_not_leak_body_on_parse_failure():
     leaky_body = {"output": {"message": {"content": [{"text": "secret content"}]}}}
 
     class MockResponse:
+        headers = httpx.Headers({"x-amzn-RequestId": "req-parse-failure"})
+
         def json(self):
             return leaky_body
 
@@ -6067,6 +6070,7 @@ def test_transform_response_does_not_leak_body_on_parse_failure():
     msg = str(exc_info.value)
     assert "secret content" not in msg
     assert "Error converting to valid response block" in msg
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-parse-failure"
 
 
 def test_converse_drops_sampling_params_for_models_that_removed_them():
diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py
index 4bef59842f1..d0adabe7b4e 100644
--- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py
+++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py
@@ -496,3 +496,171 @@ async def test_async_invoke_streaming_forwards_bedrock_response_headers():
 
     assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-987"
 
+
+def _bedrock_stream_error_response(status_code: int, request_id: str) -> httpx.Response:
+    return httpx.Response(
+        status_code=status_code,
+        headers={
+            "x-amzn-RequestId": request_id,
+            "x-amzn-ErrorType": "InternalServerException",
+        },
+        text='{"message":"Amazon Bedrock is unable to process your request."}',
+        request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
+    )
+
+
+def test_invoke_streaming_error_forwards_bedrock_response_headers():
+    error_response = _bedrock_stream_error_response(500, "req-stream-err-1")
+    client = HTTPHandler()
+    client.post = MagicMock(
+        side_effect=httpx.HTTPStatusError(
+            "server error",
+            request=error_response.request,
+            response=error_response,
+        )
+    )
+
+    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
+        litellm.completion(
+            model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
+            messages=[{"role": "user", "content": "hi"}],
+            stream=True,
+            client=client,
+            aws_access_key_id="fake",
+            aws_secret_access_key="fake",
+            aws_region_name="us-east-1",
+        )
+
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-stream-err-1"
+
+
+@pytest.mark.asyncio
+async def test_async_invoke_streaming_error_forwards_bedrock_response_headers():
+    error_response = _bedrock_stream_error_response(500, "req-stream-err-2")
+    client = AsyncHTTPHandler()
+    client.post = AsyncMock(
+        side_effect=httpx.HTTPStatusError(
+            "server error",
+            request=error_response.request,
+            response=error_response,
+        )
+    )
+
+    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
+        await litellm.acompletion(
+            model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
+            messages=[{"role": "user", "content": "hi"}],
+            stream=True,
+            client=client,
+            aws_access_key_id="fake",
+            aws_secret_access_key="fake",
+            aws_region_name="us-east-1",
+        )
+
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-stream-err-2"
+
+
+def _unread_bedrock_stream_error_response(status_code: int, request_id: str) -> httpx.Response:
+    return httpx.Response(
+        status_code=status_code,
+        headers={
+            "x-amzn-RequestId": request_id,
+            "x-amzn-ErrorType": "InternalServerException",
+        },
+        stream=httpx.ByteStream(b'{"message":"Amazon Bedrock is unable to process your request."}'),
+        request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
+    )
+
+
+def test_invoke_streaming_error_forwards_headers_when_body_was_never_read():
+    """A retried streamed request raises HTTPStatusError over a body nobody read, so
+    reading it for the error message throws and loses the request id (LIT-5428)."""
+    error_response = _unread_bedrock_stream_error_response(500, "req-unread-sync")
+    client = HTTPHandler()
+    client.post = MagicMock(
+        side_effect=httpx.HTTPStatusError(
+            "server error",
+            request=error_response.request,
+            response=error_response,
+        )
+    )
+
+    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
+        litellm.completion(
+            model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
+            messages=[{"role": "user", "content": "hi"}],
+            stream=True,
+            client=client,
+            aws_access_key_id="fake",
+            aws_secret_access_key="fake",
+            aws_region_name="us-east-1",
+        )
+
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-unread-sync"
+
+
+@pytest.mark.asyncio
+async def test_async_invoke_streaming_error_forwards_headers_when_body_was_never_read():
+    error_response = _unread_bedrock_stream_error_response(500, "req-unread-async")
+    client = AsyncHTTPHandler()
+    client.post = AsyncMock(
+        side_effect=httpx.HTTPStatusError(
+            "server error",
+            request=error_response.request,
+            response=error_response,
+        )
+    )
+
+    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
+        await litellm.acompletion(
+            model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
+            messages=[{"role": "user", "content": "hi"}],
+            stream=True,
+            client=client,
+            aws_access_key_id="fake",
+            aws_secret_access_key="fake",
+            aws_region_name="us-east-1",
+        )
+
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-unread-async"
+
+
+def test_invoke_streaming_non_200_forwards_bedrock_response_headers():
+    """A caller-supplied client that returns a failure instead of raising still reaches the
+    provider's headers, and reading the streamed body for the message must not throw (LIT-5428)."""
+    error_response = _unread_bedrock_stream_error_response(500, "req-non200-sync")
+    client = HTTPHandler()
+    client.post = MagicMock(return_value=error_response)
+
+    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
+        litellm.completion(
+            model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
+            messages=[{"role": "user", "content": "hi"}],
+            stream=True,
+            client=client,
+            aws_access_key_id="fake",
+            aws_secret_access_key="fake",
+            aws_region_name="us-east-1",
+        )
+
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-non200-sync"
+
+
+@pytest.mark.asyncio
+async def test_async_invoke_streaming_non_200_forwards_bedrock_response_headers():
+    error_response = _unread_bedrock_stream_error_response(500, "req-non200-async")
+    client = AsyncHTTPHandler()
+    client.post = AsyncMock(return_value=error_response)
+
+    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
+        await litellm.acompletion(
+            model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
+            messages=[{"role": "user", "content": "hi"}],
+            stream=True,
+            client=client,
+            aws_access_key_id="fake",
+            aws_secret_access_key="fake",
+            aws_region_name="us-east-1",
+        )
+
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-non200-async"
diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py
index 9302dc01abe..3f03305423a 100644
--- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py
+++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py
@@ -614,3 +614,260 @@ def test_sign_aws_request_assumes_role_with_external_id(monkeypatch):
     authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"]
     assert "ASIABATCHSIGNROLE" in authorization
     assert signed_data == b'{"jobName": "litellm-batch-job"}'
+
+
+# --------------------------------------------------------------------------- #
+# Provider error headers (LIT-5428)                                            #
+# --------------------------------------------------------------------------- #
+
+
+def _bedrock_chat_error_configs():
+    from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
+    from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
+    from litellm.llms.bedrock.chat.invoke_agent.transformation import AmazonInvokeAgentConfig
+    from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import (
+        AmazonMoonshotConfig,
+    )
+    from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import (
+        AmazonBedrockOpenAIConfig,
+    )
+    from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
+        AmazonInvokeConfig,
+    )
+
+    return [
+        AmazonInvokeConfig,
+        AmazonConverseConfig,
+        AmazonMoonshotConfig,
+        AmazonBedrockOpenAIConfig,
+        AmazonAgentCoreConfig,
+        AmazonInvokeAgentConfig,
+    ]
+
+
+@pytest.mark.parametrize("config", _bedrock_chat_error_configs())
+def test_bedrock_chat_get_error_class_keeps_provider_headers(config):
+    """Every Bedrock chat route must carry x-amzn-RequestId out to the caller (LIT-5428).
+
+    A config that drops the headers it is handed shadows the fix for its own models.
+    """
+    error = config().get_error_class(
+        error_message="Amazon Bedrock is unable to process your request.",
+        status_code=500,
+        headers={"x-amzn-RequestId": "req-chat-500"},
+    )
+
+    assert error.response.headers["x-amzn-requestid"] == "req-chat-500"
+
+
+def test_error_response_text_reads_a_read_response():
+    import httpx
+
+    from litellm.llms.bedrock.common_utils import error_response_text
+
+    response = httpx.Response(status_code=500, text="Amazon Bedrock is unable to process your request.")
+
+    assert error_response_text(response) == "Amazon Bedrock is unable to process your request."
+
+
+def test_error_response_text_falls_back_when_a_streamed_response_was_never_read():
+    """A retried streamed request raises HTTPStatusError over an unread body; reading it
+    throws ResponseNotRead and would lose the status and headers this fix preserves."""
+    import httpx
+
+    from litellm.llms.bedrock.common_utils import error_response_text
+
+    request = httpx.Request(method="POST", url="https://bedrock-runtime.amazonaws.com")
+    response = httpx.Response(
+        status_code=500,
+        headers={"x-amzn-RequestId": "req-unread-500"},
+        stream=httpx.ByteStream(b"never read"),
+        request=request,
+    )
+
+    with pytest.raises(httpx.ResponseNotRead):
+        _ = response.text
+
+    assert error_response_text(response) == "Internal Server Error"
+
+
+def test_bedrock_error_skips_header_values_httpx_cannot_carry():
+    """The shared HTTP handler copies an arbitrary exception's header values in verbatim,
+    so a non-str value must not take down the whole error (LIT-5428)."""
+    import httpx
+
+    from litellm.llms.bedrock.common_utils import BedrockError
+
+    error = BedrockError(
+        status_code=500,
+        message="boom",
+        headers={"x-amzn-RequestId": "req-mixed-500", "x-retry-count": 3, "x-nothing": None},
+    )
+
+    assert error.response.headers["x-amzn-requestid"] == "req-mixed-500"
+    assert "x-retry-count" not in error.response.headers
+    assert isinstance(error.response, httpx.Response)
+
+
+def test_bedrock_error_keeps_duplicate_httpx_header_values():
+    import httpx
+
+    from litellm.llms.bedrock.common_utils import BedrockError
+
+    error = BedrockError(
+        status_code=500,
+        message="boom",
+        headers=httpx.Headers([("x-amzn-RequestId", "req-dup-500"), ("set-cookie", "a=1"), ("set-cookie", "b=2")]),
+    )
+
+    assert error.response.headers.get_list("set-cookie") == ["a=1", "b=2"]
+
+
+def _bedrock_httpx_status_error_sites():
+    """Every `except httpx.HTTPStatusError as err` that raises a BedrockError, across bedrock."""
+    import ast
+    import pathlib
+
+    sites = []
+    for path in sorted(pathlib.Path("litellm/llms/bedrock").rglob("*.py")):
+        tree = ast.parse(path.read_text())
+        for handler in (n for n in ast.walk(tree) if isinstance(n, ast.ExceptHandler)):
+            caught = ast.unparse(handler.type) if handler.type is not None else ""
+            if "HTTPStatusError" not in caught or handler.name is None:
+                continue
+            for call in (
+                n
+                for n in ast.walk(handler)
+                if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "BedrockError"
+            ):
+                sites.append((str(path), call.lineno, handler.name, {k.arg for k in call.keywords}))
+    return sites
+
+
+def test_every_bedrock_httpx_status_error_site_keeps_provider_headers():
+    """A raise site holding the provider's failed response must hand its headers on (LIT-5428).
+
+    These sites are the only place x-amzn-RequestId still exists; a site that drops it
+    silently shadows the fix for that whole surface.
+    """
+    sites = _bedrock_httpx_status_error_sites()
+
+    assert len(sites) >= 12
+    dropped = [f"{path}:{lineno}" for path, lineno, _, kwargs in sites if "headers" not in kwargs]
+    assert dropped == []
+
+
+@pytest.mark.parametrize("is_async", [False, True])
+@pytest.mark.asyncio
+async def test_bedrock_embedding_call_keeps_provider_headers(is_async):
+    """The embeddings surface raises from the same shape as chat and lost the same header."""
+    import httpx
+
+    from litellm.llms.bedrock.common_utils import BedrockError
+    from litellm.llms.bedrock.embed.embedding import BedrockEmbedding
+    from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
+
+    failure = httpx.Response(
+        status_code=500,
+        headers={"x-amzn-RequestId": "req-embed-500"},
+        text='{"message":"Amazon Bedrock is unable to process your request."}',
+        request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
+    )
+
+    class _SyncUpstream(HTTPHandler):
+        def post(self, *args, **kwargs):
+            return failure
+
+    class _AsyncUpstream(AsyncHTTPHandler):
+        async def post(self, *args, **kwargs):
+            return failure
+
+    async def _drive():
+        embedding = BedrockEmbedding()
+        kwargs = dict(
+            timeout=None,
+            api_base="https://bedrock-runtime.us-east-1.amazonaws.com/",
+            headers={},
+            data={},
+        )
+        if is_async:
+            return await embedding._make_async_call(client=_AsyncUpstream(), **kwargs)
+        return embedding._make_sync_call(client=_SyncUpstream(), **kwargs)
+
+    with pytest.raises(BedrockError) as exc_info:
+        await _drive()
+
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-embed-500"
+
+
+def _bedrock_mantle_error_configs():
+    from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig
+    from litellm.llms.bedrock_mantle.responses.transformation import BedrockMantleResponsesAPIConfig
+
+    return [BedrockMantleChatConfig, BedrockMantleResponsesAPIConfig]
+
+
+@pytest.mark.parametrize("config", _bedrock_mantle_error_configs())
+def test_bedrock_mantle_get_error_class_keeps_provider_headers(config):
+    """bedrock_mantle rides the OpenAI-compatible surfaces, whose errors drop the headers.
+
+    A chat request for a responses-API model is bridged onto the responses config, so
+    fixing only the chat one leaves the model the customer actually calls uncovered.
+    """
+    error = config().get_error_class(
+        error_message="prompt tokens exceed model maximum",
+        status_code=400,
+        headers={"x-amzn-RequestId": "req-mantle-400"},
+    )
+
+    assert error.response.headers["x-amzn-requestid"] == "req-mantle-400"
+
+
+def _bedrock_configs_with_get_error_class():
+    import importlib
+    import inspect
+    import pathlib
+
+    import litellm
+
+    llms_root = pathlib.Path(inspect.getfile(litellm)).parent / "llms"
+    configs = []
+    for package in ("bedrock", "bedrock_mantle"):
+        for path in sorted((llms_root / package).rglob("*.py")):
+            module_name = "litellm.llms." + ".".join(path.relative_to(llms_root).with_suffix("").parts)
+            module = importlib.import_module(module_name)
+            for name, obj in vars(module).items():
+                if not inspect.isclass(obj) or obj.__module__ != module_name:
+                    continue
+                if getattr(obj, "get_error_class", None) is None:
+                    continue
+                configs.append(pytest.param(obj, id=f"{module_name}.{name}"))
+    return configs
+
+
+@pytest.mark.parametrize("config", _bedrock_configs_with_get_error_class())
+def test_every_bedrock_config_get_error_class_keeps_provider_headers(config):
+    """Every bedrock surface must classify errors through BedrockError, not a header-dropping base.
+
+    A config that inherits get_error_class from a provider-agnostic base builds a blank
+    response, so the request id is gone before the proxy ever reads it.
+    """
+    try:
+        instance = config()
+    except Exception:
+        instance = config.__new__(config)
+
+    try:
+        error = instance.get_error_class(
+            error_message="boom",
+            status_code=500,
+            headers={"x-amzn-RequestId": "req-audit-500"},
+        )
+    except Exception as raised:  # some bases raise the exception instead of returning it
+        error = raised
+
+    assert error.response.headers["x-amzn-requestid"] == "req-audit-500"
+
+
+def test_bedrock_get_error_class_audit_covers_every_surface():
+    assert len(_bedrock_configs_with_get_error_class()) >= 30
diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/test_litellm/llms/chat/test_converse_handler.py
index ca79c8d7025..12b5f03aedc 100644
--- a/tests/test_litellm/llms/chat/test_converse_handler.py
+++ b/tests/test_litellm/llms/chat/test_converse_handler.py
@@ -308,3 +308,64 @@ def test_completion_plumbs_stream_chunk_size_through_converse():
         stream_chunk_size=2048,
     )
     iter_bytes_spy.assert_called_once_with(chunk_size=2048)
+
+
+def _bedrock_error_response(status_code: int, request_id: str) -> httpx.Response:
+    return httpx.Response(
+        status_code=status_code,
+        headers={
+            "x-amzn-RequestId": request_id,
+            "x-amzn-ErrorType": "InternalServerException",
+        },
+        text=json.dumps({"message": "Amazon Bedrock is unable to process your request."}),
+        request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
+    )
+
+
+def test_converse_completion_error_forwards_bedrock_response_headers():
+    error_response = _bedrock_error_response(500, "req-err-123")
+    client = HTTPHandler()
+    client.post = MagicMock(
+        side_effect=httpx.HTTPStatusError(
+            "server error",
+            request=error_response.request,
+            response=error_response,
+        )
+    )
+
+    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
+        litellm.completion(
+            model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
+            messages=[{"role": "user", "content": "hi"}],
+            client=client,
+            aws_access_key_id="fake",
+            aws_secret_access_key="fake",
+            aws_region_name="us-east-1",
+        )
+
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-err-123"
+
+
+@pytest.mark.asyncio
+async def test_async_converse_completion_error_forwards_bedrock_response_headers():
+    error_response = _bedrock_error_response(500, "req-err-456")
+    client = AsyncHTTPHandler()
+    client.post = AsyncMock(
+        side_effect=httpx.HTTPStatusError(
+            "server error",
+            request=error_response.request,
+            response=error_response,
+        )
+    )
+
+    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
+        await litellm.acompletion(
+            model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
+            messages=[{"role": "user", "content": "hi"}],
+            client=client,
+            aws_access_key_id="fake",
+            aws_secret_access_key="fake",
+            aws_region_name="us-east-1",
+        )
+
+    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-err-456"
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
index acb45038df0..d9969dd1dc9 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
@@ -5206,3 +5206,37 @@ class TestAzureRouterModelStreamingKeepalive:
 
         assert result.headers["x-upstream"] == "kept"
         assert chunks == [b"data: hello\n\n"]
+
+
+@pytest.mark.asyncio
+async def test_bedrock_count_tokens_error_forwards_provider_headers():
+    """The count tokens route converts BedrockError into an HTTPException, and dropping the
+    headers there loses x-amzn-RequestId after the handler went to the trouble of keeping it."""
+    from fastapi import HTTPException
+
+    from litellm.llms.bedrock.common_utils import BedrockError
+    from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
+        handle_bedrock_count_tokens,
+    )
+
+    failure = BedrockError(
+        status_code=500,
+        message="Amazon Bedrock is unable to process your request.",
+        headers={"x-amzn-RequestId": "req-count-tokens-500"},
+    )
+
+    with patch(  # test-quality-ok: the route's BedrockError branch is only reachable when the handler raises
+        "litellm.llms.bedrock.count_tokens.handler.BedrockCountTokensHandler.handle_count_tokens_request",
+        new=AsyncMock(side_effect=failure),
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await handle_bedrock_count_tokens(
+                endpoint="v1/messages/count_tokens",
+                request=MagicMock(),
+                fastapi_response=MagicMock(),
+                user_api_key_dict=MagicMock(),
+                request_body={"model": "anthropic.claude-haiku-4-5-20251001-v1:0"},
+            )
+
+    assert exc_info.value.status_code == 500
+    assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-count-tokens-500"
diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py
index 6acd9d7258e..bfae42f64f1 100644
--- a/tests/test_litellm/proxy/test_common_request_processing.py
+++ b/tests/test_litellm/proxy/test_common_request_processing.py
@@ -8259,3 +8259,38 @@ class TestPassthroughHeadersAcceptImmutableMappings:
         assert merged["content-type"] == "text/event-stream"
         # the excluded hop-by-hop header is still dropped
         assert "transfer-encoding" not in merged
+
+
+@pytest.mark.asyncio
+async def test_handle_llm_api_exception_forwards_provider_headers_on_http_status_error():
+    """The httpx.HTTPStatusError branch dropped the headers its sibling branches forward.
+
+    A Bedrock passthrough failure reaches this branch, so the request id was gone
+    before the client saw the response.
+    """
+    import httpx
+
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    request = httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/model/m/converse")
+    response = httpx.Response(
+        status_code=500,
+        headers={"x-amzn-RequestId": "req-passthrough-500"},
+        content=b'{"message": "Amazon Bedrock is unable to process your request."}',
+        request=request,
+    )
+
+    processor = ProxyBaseLLMRequestProcessing(data={})
+    proxy_logging_obj = MagicMock()
+    proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
+    proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
+
+    with pytest.raises(HTTPException) as exc_info:
+        await processor._handle_llm_api_exception(
+            e=httpx.HTTPStatusError("boom", request=request, response=response),
+            user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
+            proxy_logging_obj=proxy_logging_obj,
+        )
+
+    assert exc_info.value.headers is not None
+    assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-passthrough-500"
diff --git a/tests/test_litellm/test_exception_header_preservation.py b/tests/test_litellm/test_exception_header_preservation.py
index 6ea478c633b..dd142d9d40b 100644
--- a/tests/test_litellm/test_exception_header_preservation.py
+++ b/tests/test_litellm/test_exception_header_preservation.py
@@ -18,6 +18,7 @@ from litellm.exceptions import (
     ImageFetchError,
     MidStreamFallbackError,
     RateLimitError,
+    ServiceUnavailableError,
 )
 
 
@@ -312,3 +313,85 @@ class TestProxyHeaderExtraction:
         # Verify headers are extracted and prefixed correctly
         assert headers.get("llm_provider-x-request-id") == "req-abc123"
         assert headers.get("llm_provider-x-ms-region") == "eastus"
+
+
+class TestBedrockErrorHeaders:
+    """A BedrockError built with headers but no response still exposes them (LIT-5428)."""
+
+    def test_synthesized_response_carries_headers(self):
+        from litellm.llms.bedrock.common_utils import BedrockError
+
+        error = BedrockError(
+            status_code=500,
+            message="Amazon Bedrock is unable to process your request.",
+            headers={"x-amzn-RequestId": "req-base-500"},
+        )
+
+        assert error.response.headers["x-amzn-requestid"] == "req-base-500"
+        assert str(error.request.url) == str(BedrockError(status_code=500, message="boom").request.url)
+        assert str(error.response.request.url) == str(error.request.url)
+
+    def test_synthesized_response_without_headers_stays_empty(self):
+        from litellm.llms.bedrock.common_utils import BedrockError
+
+        error = BedrockError(status_code=500, message="boom")
+
+        assert dict(error.response.headers) == {}
+
+    def test_explicit_response_is_kept(self):
+        from litellm.llms.bedrock.common_utils import BedrockError
+
+        provider_response = httpx.Response(
+            status_code=500,
+            headers={"x-amzn-RequestId": "from-response"},
+            request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
+        )
+        error = BedrockError(
+            status_code=500,
+            message="boom",
+            headers={"x-amzn-RequestId": "from-headers"},
+            response=provider_response,
+        )
+
+        assert error.response is provider_response
+
+    def test_proxy_extraction_surfaces_bedrock_request_id(self):
+        """End-to-end shape the proxy error handler returns to the caller."""
+        from litellm.litellm_core_utils.exception_mapping_utils import exception_type
+        from litellm.litellm_core_utils.llm_response_utils.get_headers import (
+            get_response_headers,
+        )
+        from litellm.llms.bedrock.common_utils import BedrockError
+
+        provider_response = httpx.Response(
+            status_code=500,
+            headers={"x-amzn-RequestId": "req-proxy-500"},
+            text='{"message":"Amazon Bedrock is unable to process your request."}',
+            request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
+        )
+
+        with pytest.raises(ServiceUnavailableError) as exc_info:
+            exception_type(
+                model="anthropic.claude-haiku-4-5-20251001-v1:0",
+                original_exception=BedrockError(
+                    status_code=500,
+                    message=provider_response.text,
+                    headers=provider_response.headers,
+                    response=provider_response,
+                ),
+                custom_llm_provider="bedrock",
+                completion_kwargs={},
+                extra_kwargs={},
+            )
+
+        # Mirrors ProxyBaseLLMRequestProcessing._handle_llm_api_exception
+        error = exc_info.value
+        headers = getattr(error, "headers", None) or {}
+        if not headers:
+            _response = getattr(error, "response", None)
+            if _response is not None:
+                _response_headers = getattr(_response, "headers", None)
+                if _response_headers:
+                    headers = get_response_headers(dict(_response_headers))
+
+        assert headers.get("llm_provider-x-amzn-requestid") == "req-proxy-500"

From 7da6fe54b5201dc42c2c55baac943fc46ec5e097 Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Mon, 7 Sep 2026 18:03:43 -0700
Subject: [PATCH 149/164] fix: skip one-shot Claude Code cache injection
 (#40175)

---
 .../anthropic_cache_control_hook.py           |  82 +++++-
 litellm/llms/anthropic/common_utils.py        |  90 +++++++
 litellm/proxy/litellm_pre_call_utils.py       |   8 +-
 .../prompt_caching_deployment_check.py        |   1 +
 .../test_anthropic_cache_control_hook.py      | 237 ++++++++++++++++++
 .../anthropic/test_anthropic_common_utils.py  |  24 ++
 .../test_prompt_caching_deployment_check.py   |  61 +++++
 7 files changed, 492 insertions(+), 11 deletions(-)

diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py
index 3519240dda9..4f9b18713d0 100644
--- a/litellm/integrations/anthropic_cache_control_hook.py
+++ b/litellm/integrations/anthropic_cache_control_hook.py
@@ -16,6 +16,8 @@ from collections.abc import Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, cast
 from urllib.parse import urlparse
 
+from pydantic import TypeAdapter, ValidationError
+
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.integrations.custom_prompt_management import CustomPromptManagement
@@ -23,6 +25,7 @@ from litellm.integrations.prompt_management_base import PromptManagementClient
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     with_prompt_cache_breakpoint,
 )
+from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request
 from litellm.types.integrations.anthropic_cache_control_hook import (
     GATEWAY_INJECTED_CACHE_METADATA_KEY,
     GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT,
@@ -62,10 +65,26 @@ OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset(
 )
 OPENAI_API_HOST: Final = "api.openai.com"
 OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE")
+_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
+_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
 
 AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues
 
 
+def _validated_object_mapping(value: object) -> dict[object, object] | None:
+    try:
+        return _OBJECT_MAPPING_ADAPTER.validate_python(value)
+    except ValidationError:
+        return None
+
+
+def _validated_object_list(value: object) -> list[object] | None:
+    try:
+        return _OBJECT_LIST_ADAPTER.validate_python(value)
+    except ValidationError:
+        return None
+
+
 def supports_openai_prompt_cache_breakpoint(model: str) -> bool:
     model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model)
     if model_map_flag is not None:
@@ -114,6 +133,36 @@ CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_
 
 
 class AnthropicCacheControlHook(CustomPromptManagement):
+    @staticmethod
+    def _request_value(request_kwargs: object, key: str) -> object:
+        request_mapping: Final = _validated_object_mapping(request_kwargs)
+        if request_mapping is None:
+            return None
+        return request_mapping.get(key)
+
+    @staticmethod
+    def _request_user_agent(request_kwargs: object) -> str | None:
+        proxy_server_request: Final = AnthropicCacheControlHook._request_value(request_kwargs, "proxy_server_request")
+        proxy_server_request_mapping: Final = _validated_object_mapping(proxy_server_request)
+        if proxy_server_request_mapping is None:
+            return None
+        headers: Final = proxy_server_request_mapping.get("headers")
+        headers_mapping: Final = _validated_object_mapping(headers)
+        if headers_mapping is None:
+            return None
+        user_agent: Final = next(
+            (value for key, value in headers_mapping.items() if isinstance(key, str) and key.lower() == "user-agent"),
+            None,
+        )
+        return user_agent if isinstance(user_agent, str) else None
+
+    @staticmethod
+    def _request_system(request_kwargs: object) -> str | list[object] | None:
+        system: Final = AnthropicCacheControlHook._request_value(request_kwargs, "system")
+        if isinstance(system, str):
+            return system
+        return _validated_object_list(system)
+
     def get_chat_completion_prompt(
         self,
         model: str,
@@ -520,12 +569,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         points: Sequence[CacheControlInjectionPoint],
         messages: list[AllMessageValues],
         tools: list[object] | None,
+        cache_control: object,
         model: str,
         custom_llm_provider: str | None,
         api_base: object,
         prompt_cache_options: object,
     ) -> Sequence[Mapping[str, object]] | None:
-        if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools):
+        if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control):
             return None
         return AnthropicCacheControlHook._stamped_with_dialect(
             points, model, custom_llm_provider, api_base, prompt_cache_options
@@ -561,6 +611,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         messages: list[AllMessageValues],
         system: str | list | None,
         tools: list | None,
+        cache_control: object = None,
     ) -> bool:
         """Whether configured injection points must yield to client-set cache_control.
 
@@ -573,13 +624,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         """
         if all(point.get("_litellm_judged") for point in points):
             return False
-        return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools)
+        return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control)
 
     @staticmethod
     def _request_has_cache_control(
         messages: list[AllMessageValues],
         system: str | list | None,
         tools: list | None = None,
+        cache_control: object = None,
     ) -> bool:
         """Return True if the request already carries any client-supplied cache_control.
 
@@ -591,6 +643,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         carry the mark either at the top level (Anthropic shape) or nested under
         ``function`` (OpenAI shape); the Anthropic chat transform accepts both.
         """
+        if cache_control is not None:
+            return True
         if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0:
             return True
         if tools is not None:
@@ -612,6 +666,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         custom_llm_provider: str | None,
         tools: list | None = None,
         enable_prompt_caching: bool | None = None,
+        cache_control: object = None,
+        request_kwargs: object = None,
     ) -> list[CacheControlInjectionPoint]:
         """Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on.
 
@@ -649,7 +705,12 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         if not supports_prompt_caching(model=model, custom_llm_provider=provider):
             return []
 
-        if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools):
+        if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control):
+            return []
+
+        if is_claude_code_one_shot_subagent_request(
+            messages, system, tools, AnthropicCacheControlHook._request_user_agent(request_kwargs)
+        ):
             return []
 
         control: Final = AnthropicCacheControlHook._default_control()
@@ -665,6 +726,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         models: Iterable[str],
         tools: list[AllToolParamValues] | None = None,
         enable_prompt_caching: bool | None = None,
+        request_kwargs: object = None,
     ) -> list[AllMessageValues]:
         """Return the messages auto prompt caching will send, default breakpoints included.
 
@@ -681,11 +743,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
                 for candidate in (
                     AnthropicCacheControlHook.get_default_injection_points(
                         messages=messages,
-                        system=None,
                         model=model,
                         custom_llm_provider=None,
                         tools=tools,
                         enable_prompt_caching=enable_prompt_caching,
+                        system=AnthropicCacheControlHook._request_system(request_kwargs),
+                        cache_control=AnthropicCacheControlHook._request_value(request_kwargs, "cache_control"),
+                        request_kwargs=request_kwargs,
                     )
                     for model in models
                 )
@@ -730,6 +794,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
                 non_default_params["cache_control_injection_points"],
                 messages,
                 tools,
+                non_default_params.get("cache_control"),
                 model,
                 custom_llm_provider,
                 api_base,
@@ -747,6 +812,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
             custom_llm_provider=custom_llm_provider,
             tools=tools,
             enable_prompt_caching=enable_prompt_caching,
+            cache_control=non_default_params.get("cache_control"),
+            request_kwargs=non_default_params,
         )
         if points:
             non_default_params["cache_control_injection_points"] = points
@@ -853,10 +920,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         enable_prompt_caching: Final = cast(  # cast-ok: kwargs is untyped; key stamped as bool by the proxy
             bool | None, kwargs.pop("enable_prompt_caching", None)
         )
+        cache_control: Final = kwargs.get("cache_control")
         configured: Final = cast(  # cast-ok: kwargs is untyped; this key only holds the documented injection-point list
             list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None)
         )
-        if configured and AnthropicCacheControlHook._should_stand_down(configured, typed_messages, system, tools):
+        if configured and AnthropicCacheControlHook._should_stand_down(
+            configured, typed_messages, system, tools, cache_control
+        ):
             return messages, system
         injection_points: list[CacheControlInjectionPoint] = configured or []
         if not injection_points and model is not None:
@@ -867,6 +937,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
                 model=model,
                 custom_llm_provider=custom_llm_provider,
                 enable_prompt_caching=enable_prompt_caching,
+                cache_control=cache_control,
+                request_kwargs=kwargs,
             )
         if not injection_points:
             return messages, system
diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py
index d9424d6a243..2b57883cc13 100644
--- a/litellm/llms/anthropic/common_utils.py
+++ b/litellm/llms/anthropic/common_utils.py
@@ -67,6 +67,96 @@ _BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$")
 _INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$")
 _DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$")
 _DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)")
+_CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:"
+_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
+_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
+
+
+def is_claude_code_user_agent(user_agent: str) -> bool:
+    return user_agent.startswith("claude-cli/")
+
+
+def _validated_claude_code_mapping(value: object) -> dict[object, object] | None:
+    try:
+        return _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER.validate_python(value)
+    except ValidationError:
+        return None
+
+
+def _validated_claude_code_list(value: object) -> list[object] | None:
+    try:
+        return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value)
+    except ValidationError:
+        return None
+
+
+def _claude_code_billing_fields(text: str) -> tuple[tuple[str, str], ...] | None:
+    stripped: Final = text.strip()
+    if "\n" in stripped or "\r" in stripped or not stripped.startswith(_CLAUDE_CODE_BILLING_HEADER_PREFIX):
+        return None
+    fields: Final = tuple(
+        field
+        for raw_field in stripped.removeprefix(_CLAUDE_CODE_BILLING_HEADER_PREFIX).split(";")
+        if (field := raw_field.strip())
+    )
+    if not fields or any("=" not in field for field in fields):
+        return None
+    parsed_fields: Final = tuple(
+        (parts[0].strip(), parts[1].strip()) for field in fields for parts in (field.split("=", 1),)
+    )
+    if any(not key or not value for key, value in parsed_fields):
+        return None
+    return parsed_fields
+
+
+def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None:
+    if isinstance(system, str):
+        return (system,)
+    blocks: Final = _validated_claude_code_list(system)
+    if blocks is None:
+        return None
+    block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks)
+    if any(block is None for block in block_mappings):
+        return None
+    text_values: Final = tuple(
+        block.get("text") for block in block_mappings if block is not None and block.get("type") == "text"
+    )
+    if len(text_values) != len(blocks) or any(not isinstance(text, str) for text in text_values):
+        return None
+    meaningful_text: Final = tuple(text for text in text_values if isinstance(text, str) and text.strip())
+    return meaningful_text or None
+
+
+def _is_claude_code_subagent_billing_system(system: object) -> bool:
+    billing_texts: Final = _claude_code_billing_texts(system)
+    if billing_texts is None:
+        return False
+    billing_fields: Final = tuple(
+        fields for text in billing_texts if (fields := _claude_code_billing_fields(text)) is not None
+    )
+    if len(billing_fields) != len(billing_texts):
+        return False
+    subagent_values: Final = tuple(
+        value for fields in billing_fields for key, value in fields if key == "cc_is_subagent"
+    )
+    return subagent_values == ("true",)
+
+
+def is_claude_code_one_shot_subagent_request(
+    messages: list[AllMessageValues],
+    system: object,
+    tools: object,
+    user_agent: str | None,
+) -> bool:
+    only_message: Final = _validated_claude_code_mapping(messages[0]) if len(messages) == 1 else None
+    return (
+        user_agent is not None
+        and is_claude_code_user_agent(user_agent)
+        and not tools
+        and only_message is not None
+        and only_message.get("role") == "user"
+        and _is_claude_code_subagent_billing_system(system)
+    )
 
 
 def _strip_bedrock_id_suffixes(model: str) -> str:
diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py
index 56512570448..3c186e19829 100644
--- a/litellm/proxy/litellm_pre_call_utils.py
+++ b/litellm/proxy/litellm_pre_call_utils.py
@@ -789,12 +789,6 @@ def apply_missing_session_id_policy(
             )
 
 
-def is_claude_code_user_agent(user_agent: str) -> bool:
-    """Claude Code identifies itself as ``claude-cli/<version> ...``; the IDE
-    extensions and the Agent SDK run through the same CLI and share that prefix."""
-    return user_agent.startswith("claude-cli/")
-
-
 def is_codex_user_agent(user_agent: str) -> bool:
     """Codex builds its user agent as ``<originator>/<version> ...`` and ships
     several first-party originators: ``codex-tui``, ``codex_cli_rs``,
@@ -811,6 +805,8 @@ def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_c
     requests routed to providers that reject them. An explicit drop_params
     from the caller or in the operator's ``litellm_settings`` always wins
     over this default."""
+    from litellm.llms.anthropic.common_utils import is_claude_code_user_agent
+
     if not (is_claude_code_user_agent(user_agent) or is_codex_user_agent(user_agent)):
         return False
     if "drop_params" in data:
diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py
index 0788c8db710..70362e60495 100644
--- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py
+++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py
@@ -90,6 +90,7 @@ class PromptCachingDeploymentCheck(CustomLogger):
                 enable_prompt_caching=(
                     request_kwargs.get("enable_prompt_caching") is True if request_kwargs is not None else None
                 ),
+                request_kwargs=request_kwargs,
             )
 
             model_id_dict: Final = await prompt_cache.async_get_model_id(
diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py
index de8b654987b..6b7780acd20 100644
--- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py
+++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py
@@ -1777,6 +1777,224 @@ class TestEnableAnthropicPromptCaching:
         assert messages == before
 
 
+class TestClaudeCodeOneShotAutoCaching:
+    BILLING_TEXT = "x-anthropic-billing-header: cc_version=2.1.263; cc_entrypoint=cli; cc_is_subagent=true;"
+    BILLING_SYSTEM = [{"type": "text", "text": BILLING_TEXT}]
+    MESSAGES = [{"role": "user", "content": [{"type": "text", "text": "unique fetched document"}]}]
+
+    @staticmethod
+    def _kwargs(configured=None):
+        kwargs = {
+            "litellm_metadata": {},
+            "proxy_server_request": {
+                "headers": {
+                    "user-agent": "claude-cli/2.1.263 (external, cli)",
+                    "x-app": "cli-bg",
+                }
+            },
+        }
+        if configured is not None:
+            kwargs["cache_control_injection_points"] = configured
+        return kwargs
+
+    @pytest.mark.parametrize(
+        "system",
+        [
+            BILLING_TEXT,
+            BILLING_SYSTEM,
+            [*BILLING_SYSTEM, {"type": "text", "text": "  "}],
+            [
+                *BILLING_SYSTEM,
+                {"type": "text", "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_entrypoint=cli;"},
+            ],
+        ],
+        ids=["string", "text_block", "whitespace_block", "multiple_billing_blocks"],
+    )
+    @pytest.mark.parametrize("tools", [None, []], ids=["absent_tools", "empty_tools"])
+    def test_skips_defaults_and_attribution_for_one_shot_subagent(self, monkeypatch, system, tools):
+        monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
+        messages = copy.deepcopy(self.MESSAGES)
+        kwargs = self._kwargs()
+
+        result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control(
+            messages,
+            copy.deepcopy(system),
+            kwargs,
+            model="claude-sonnet-4-5",
+            custom_llm_provider="anthropic",
+            tools=tools,
+        )
+
+        assert result_messages == self.MESSAGES
+        assert result_system == system
+        assert "litellm_gateway_injected_cache" not in kwargs["litellm_metadata"]
+
+    def test_user_agent_header_lookup_is_case_insensitive(self, monkeypatch):
+        monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
+        kwargs = self._kwargs()
+        user_agent = kwargs["proxy_server_request"]["headers"].pop("user-agent")
+        kwargs["proxy_server_request"]["headers"]["User-Agent"] = user_agent
+
+        result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control(
+            copy.deepcopy(self.MESSAGES),
+            copy.deepcopy(self.BILLING_SYSTEM),
+            kwargs,
+            model="claude-sonnet-4-5",
+            custom_llm_provider="anthropic",
+        )
+
+        assert result_messages == self.MESSAGES
+        assert result_system == self.BILLING_SYSTEM
+
+    def test_router_affinity_skips_string_billing_system(self, monkeypatch):
+        monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
+        messages = copy.deepcopy(self.MESSAGES)
+        kwargs = self._kwargs()
+        kwargs["system"] = self.BILLING_TEXT
+
+        result = AnthropicCacheControlHook.messages_with_default_injections(
+            messages=messages,
+            models=("claude-sonnet-4-5",),
+            request_kwargs=kwargs,
+        )
+
+        assert result == messages
+
+    @pytest.mark.parametrize(
+        "headers,system",
+        [
+            ("not-a-mapping", BILLING_SYSTEM),
+            (
+                {"user-agent": "claude-cli/2.1.263 (external, cli)"},
+                [{"type": "text", "text": "x-anthropic-billing-header: malformed"}],
+            ),
+            ({"user-agent": "claude-cli/2.1.263 (external, cli)"}, None),
+            ({"user-agent": "claude-cli/2.1.263 (external, cli)"}, ["not-a-mapping"]),
+            (
+                {"user-agent": "claude-cli/2.1.263 (external, cli)"},
+                [{"type": "image", "text": BILLING_TEXT}],
+            ),
+        ],
+        ids=["malformed_headers", "malformed_billing", "missing_system", "malformed_block", "non_text_block"],
+    )
+    def test_malformed_untrusted_context_keeps_defaults(self, monkeypatch, headers, system):
+        monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
+
+        points = AnthropicCacheControlHook.get_default_injection_points(
+            messages=copy.deepcopy(self.MESSAGES),
+            system=copy.deepcopy(system),
+            model="claude-sonnet-4-5",
+            custom_llm_provider="anthropic",
+            request_kwargs={"proxy_server_request": {"headers": headers}},
+        )
+
+        assert len(points) == 2
+
+    def test_message_without_role_keeps_defaults(self, monkeypatch):
+        monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
+
+        points = AnthropicCacheControlHook.get_default_injection_points(
+            messages=[{"content": "missing role"}],
+            system=copy.deepcopy(self.BILLING_SYSTEM),
+            model="claude-sonnet-4-5",
+            custom_llm_provider="anthropic",
+            request_kwargs=self._kwargs(),
+        )
+
+        assert len(points) == 2
+
+    @pytest.mark.parametrize(
+        "messages,system,tools",
+        [
+            (
+                MESSAGES,
+                BILLING_SYSTEM,
+                [{"name": "WebFetch", "description": "fetch", "input_schema": {"type": "object"}}],
+            ),
+            (MESSAGES, [*BILLING_SYSTEM, {"type": "text", "text": "Explore the repository"}], None),
+            (
+                [
+                    {"role": "user", "content": "first turn"},
+                    {"role": "assistant", "content": "reply"},
+                    *MESSAGES,
+                ],
+                BILLING_SYSTEM,
+                None,
+            ),
+        ],
+        ids=["tools", "real_system", "history"],
+    )
+    def test_keeps_defaults_for_reusable_subagents(self, monkeypatch, messages, system, tools):
+        monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
+        kwargs = self._kwargs()
+
+        result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control(
+            copy.deepcopy(messages),
+            copy.deepcopy(system),
+            kwargs,
+            model="claude-sonnet-4-5",
+            custom_llm_provider="anthropic",
+            tools=copy.deepcopy(tools),
+        )
+
+        assert AnthropicCacheControlHook.count_request_cache_breakpoints(result_messages, result_system) == 2
+        assert kwargs["litellm_metadata"]["litellm_gateway_injected_cache"] == ""
+
+    @pytest.mark.parametrize(
+        "user_agent,system",
+        [
+            ("anthropic-sdk-python/0.75.0", BILLING_SYSTEM),
+            (
+                "claude-cli/2.1.263 (external, cli)",
+                [
+                    {
+                        "type": "text",
+                        "text": f"{BILLING_TEXT}\nadditional system instructions",
+                    }
+                ],
+            ),
+            (
+                "claude-cli/2.1.263 (external, cli)",
+                [
+                    {
+                        "type": "text",
+                        "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_is_subagent=false;",
+                    }
+                ],
+            ),
+        ],
+        ids=["different_client", "appended_instructions", "not_a_subagent"],
+    )
+    def test_ambiguous_or_unmatched_signals_fail_open(self, monkeypatch, user_agent, system):
+        monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
+        kwargs = self._kwargs()
+        kwargs["proxy_server_request"]["headers"]["user-agent"] = user_agent
+
+        result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control(
+            copy.deepcopy(self.MESSAGES),
+            copy.deepcopy(system),
+            kwargs,
+            model="claude-sonnet-4-5",
+            custom_llm_provider="anthropic",
+        )
+
+        assert AnthropicCacheControlHook.count_request_cache_breakpoints(result_messages, result_system) == 2
+
+    def test_explicit_injection_points_remain_authoritative(self, monkeypatch):
+        monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
+        kwargs = self._kwargs([{"location": "message", "role": "user"}])
+
+        result_messages, _ = AnthropicCacheControlHook.maybe_inject_cache_control(
+            copy.deepcopy(self.MESSAGES),
+            copy.deepcopy(self.BILLING_SYSTEM),
+            kwargs,
+            model="claude-sonnet-4-5",
+            custom_llm_provider="anthropic",
+        )
+
+        assert result_messages[0]["content"][-1]["cache_control"] == {"type": "ephemeral"}
+
+
 class TestPerKeyEnablePromptCaching:
     """Per-request enable_prompt_caching override (stamped from key metadata) with the global flag off."""
 
@@ -1977,6 +2195,25 @@ class TestConfiguredInjectionPointsStandDown:
         _, result_sys = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs)
         assert result_sys == [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}]
 
+    @pytest.mark.parametrize(
+        "configured",
+        [None, CONFIGURED],
+        ids=["automatic_defaults", "configured_points"],
+    )
+    def test_v1_messages_stands_down_for_root_cache_control(self, monkeypatch, configured):
+        monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
+        root_cache_control = {"type": "ephemeral"}
+        kwargs = {"cache_control": root_cache_control, "litellm_metadata": {}}
+        if configured is not None:
+            kwargs["cache_control_injection_points"] = copy.deepcopy(configured)
+
+        result_messages, result_system = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs)
+
+        assert result_messages == self.V1_MESSAGES
+        assert result_system == "sys"
+        assert kwargs["cache_control"] is root_cache_control
+        assert "litellm_gateway_injected_cache" not in kwargs["litellm_metadata"]
+
     def test_v1_messages_reentry_flow_preserves_tool_config_remainder(self):
         """The advisor interceptor re-enters anthropic_messages() with the outer
         request's kwargs and post-injection messages. The first pass applies the
diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py
index 794613942a1..ae620fdd6dc 100644
--- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py
+++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py
@@ -26,6 +26,30 @@ FAKE_REGULAR_KEY = "sk-ant-api03-regular-key-for-testing-123456789"
 FAKE_AUTH_TOKEN = "sk-ant-aut01-fake-auth-token-for-testing-123456789"
 
 
+@pytest.mark.parametrize(
+    "messages,system,expected",
+    [
+        ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: cc_is_subagent=true;", True),
+        ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: =junk; cc_is_subagent=true;", False),
+        ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: cc_version=; cc_is_subagent=true;", False),
+        ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: malformed", False),
+        ([{"content": "missing role"}], "x-anthropic-billing-header: cc_is_subagent=true;", False),
+        (["not-a-mapping"], "x-anthropic-billing-header: cc_is_subagent=true;", False),
+        ([{"role": "user", "content": "hi"}], ["not-a-mapping"], False),
+        ([{"role": "user", "content": "hi"}], None, False),
+    ],
+)
+def test_is_claude_code_one_shot_subagent_request(messages, system, expected):
+    from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request
+
+    assert is_claude_code_one_shot_subagent_request(
+        messages=messages,
+        system=system,
+        tools=None,
+        user_agent="claude-cli/2.1.263 (external, cli)",
+    ) is expected
+
+
 class TestOptionallyHandleAnthropicOAuth:
     """Tests for optionally_handle_anthropic_oauth function."""
 
diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
index 79ae00e155c..030bdfe03e9 100644
--- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
+++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
@@ -332,6 +332,67 @@ async def test_per_request_enable_prompt_caching_reaches_the_affinity_key(monkey
     assert filtered == [deployments[1]]
 
 
+@pytest.mark.asyncio
+async def test_claude_code_one_shot_subagent_does_not_reuse_an_auto_injected_affinity_key(monkeypatch):
+    monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
+    cache = DualCache()
+    check = PromptCachingDeploymentCheck(cache=cache)
+    deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)
+    messages = cast(List[AllMessageValues], [{"role": "user", "content": "unique " * 3000}])
+    request_kwargs = {
+        "system": [
+            {
+                "type": "text",
+                "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_is_subagent=true;",
+            }
+        ],
+        "proxy_server_request": {"headers": {"user-agent": "claude-cli/2.1.263 (external, cli)"}},
+    }
+    auto_injected_messages = AnthropicCacheControlHook.messages_with_default_injections(
+        messages=messages,
+        models=(AUTO_CACHING_MODEL,),
+    )
+    assert auto_injected_messages != messages
+    await PromptCachingCache(cache=cache).async_add_model_id(
+        model_id="dep-2", messages=auto_injected_messages, tools=None
+    )
+
+    filtered = await check.async_filter_deployments(
+        model=MODEL_GROUP_ALIAS,
+        healthy_deployments=deployments,
+        messages=messages,
+        request_kwargs=request_kwargs,
+    )
+
+    assert filtered == deployments
+
+
+@pytest.mark.asyncio
+async def test_root_cache_control_does_not_reuse_an_auto_injected_affinity_key(monkeypatch):
+    monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
+    cache = DualCache()
+    check = PromptCachingDeploymentCheck(cache=cache)
+    deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)
+    messages = _auto_caching_messages()
+    auto_injected_messages = AnthropicCacheControlHook.messages_with_default_injections(
+        messages=messages,
+        models=(AUTO_CACHING_MODEL,),
+    )
+    assert auto_injected_messages != messages
+    await PromptCachingCache(cache=cache).async_add_model_id(
+        model_id="dep-2", messages=auto_injected_messages, tools=None
+    )
+
+    filtered = await check.async_filter_deployments(
+        model=MODEL_GROUP_ALIAS,
+        healthy_deployments=deployments,
+        messages=messages,
+        request_kwargs={"cache_control": {"type": "ephemeral"}},
+    )
+
+    assert filtered == deployments
+
+
 @pytest.mark.asyncio
 async def test_tool_marked_cache_control_keeps_routing_off_another_requests_prefix(monkeypatch, local_model_cost_map):
     """

From 1761fe236f1db25c5ca8c76a1e3b43f303d621dd Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Mon, 7 Sep 2026 18:17:30 -0700
Subject: [PATCH 150/164] feat(complexity_router): add declarative custom
 dimensions to the heuristic scorer (#40156)

Co-authored-by: Claude Code <noreply@anthropic.com>
---
 .../complexity_router/README.md               |  24 +++
 .../complexity_router/complexity_router.py    |  25 ++-
 .../complexity_router/config.py               | 157 +++++++++++++-
 .../auto_router_tuning_baseline.py            |   1 +
 .../router_strategy/test_complexity_router.py | 203 +++++++++++++++++-
 .../test_auto_router_tuning_baseline.py       |  24 +++
 ui/litellm-dashboard/src/lib/http/schema.d.ts |  23 ++
 7 files changed, 450 insertions(+), 7 deletions(-)

diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md
index 93dddfb3d20..88ed374dd3f 100644
--- a/litellm/router_strategy/complexity_router/README.md
+++ b/litellm/router_strategy/complexity_router/README.md
@@ -195,6 +195,30 @@ model_list:
         session_affinity_ttl_seconds: 300
 ```
 
+## Custom dimensions
+
+Add `custom_dimensions` under `complexity_router_config` to give domain keywords or regex patterns their own weighted signal
+
+```yaml
+custom_dimensions:
+  - name: internalFrameworks
+    weight: 0.9
+    keywords: [orbitmesh, fluxgate]
+  - name: sqlMigration
+    weight: 0.7
+    patterns: ['\b(create|alter|drop)\s{1,4}table\b']
+```
+
+Each dimension contributes its weight once when any matcher hits the current ask. Repeated matches do not increase it. The built-in score and tier boundaries are unchanged, and the total score is not renormalized. Keywords use the existing case-insensitive word-boundary and CJK rules. Regexes search the first 2048 characters case-insensitively and compile during configuration validation and router initialization, never per request
+
+Only `heuristic`, `heuristic_first` and `hybrid` accept custom dimensions. Each name must be a unique ASCII identifier starting with a letter, at most 64 characters, and cannot reuse a built-in dimension name or a key in `dimension_weights`. Set its weight inline, greater than zero and at most one
+
+Patterns are checked at configuration time against a grammar whose worst case stays a few milliseconds on 2048 characters. Every quantifier needs an explicit upper bound of at most 64 and must repeat a single character or character class, so `\s{1,4}` is accepted while `\s+`, `(a|aa){0,12}` and `(?:ab){0,64}` are refused. Backreferences, lookarounds, atomic groups and possessive quantifiers are refused as well. Each pattern is then costed: alternation branches and repeat lengths multiply the ways the engine can retry, and every later piece of the pattern is charged once per path that can reach it, so `a?a?a?a?a?a?a?a?` followed by a long fixed tail is refused even though each quantifier is small. The budget is 2048 work units per pattern and 8192 across the router. An invalid or over-budget pattern fails the write with a message naming the pattern and the rule it broke
+
+Limits are 16 dimensions, 32 combined keywords/patterns per dimension, 256 characters per matcher and 4096 matcher characters per dimension. Matching runs inline on the request path with no timeout and no worker thread, because the grammar is what bounds the cost. These are routing hints, not security enforcement rules
+
+The existing heuristic-v1 tuning quota covers custom dimensions and their weights: one changed router without an auto-router license, unlimited with the entitlement. Omitting `custom_dimensions` preserves existing scoring. Routing decisions and spend logs include signals such as `custom (sqlMigration)` without recording the configured pattern or matched text. The field is configured through YAML or the model API; this change adds no dashboard editor
+
 ## Usage
 
 Once configured, use the model name like any other:
diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py
index 7d4497fb6f7..c8644f52c57 100644
--- a/litellm/router_strategy/complexity_router/complexity_router.py
+++ b/litellm/router_strategy/complexity_router/complexity_router.py
@@ -67,6 +67,7 @@ from litellm.types.utils import (
 from .classification_rubrics import BUSINESS_TIER_CRITERIA, calibration_examples_section
 from .config import (
     CALIBRATION_EXAMPLES_HEADING,
+    CUSTOM_PATTERN_SCAN_CHARS,
     DEFAULT_CLASSIFICATION_RUBRIC,
     DEFAULT_CODE_KEYWORDS,
     DEFAULT_ESCALATION_KEYWORDS,
@@ -1119,6 +1120,10 @@ class ComplexityRouter(CustomLogger):
             self.config.custom_technical_keywords,
         )
         self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS
+        self._custom_dimensions = tuple(
+            (dimension, tuple(re.compile(pattern, re.IGNORECASE) for pattern in dimension.patterns))
+            for dimension in self.config.custom_dimensions
+        )
         if self.config.has_custom_tiers:
             self.escalation_keywords: tuple[str, ...] = ()
         elif self.config.escalation_keywords is not None:
@@ -1320,6 +1325,17 @@ class ComplexityRouter(CustomLogger):
         score: Final = score_high if match_count >= high_threshold else score_low
         return DimensionScore(name, score, f"{signal_label} ({detail})"), match_count
 
+    def _score_custom_dimensions(self, prompt: str, user_text: str) -> tuple[tuple[DimensionScore, float], ...]:
+        if not self._custom_dimensions:
+            return ()
+        scanned: Final = prompt[:CUSTOM_PATTERN_SCAN_CHARS]
+        return tuple(
+            (DimensionScore(dimension.name, 1.0, f"custom ({dimension.name})"), dimension.weight)
+            for dimension, patterns in self._custom_dimensions
+            if any(self._keyword_matches(user_text, keyword) for keyword in dimension.keywords)
+            or any(pattern.search(scanned) is not None for pattern in patterns)
+        )
+
     def _score_multi_step(self, text: str) -> DimensionScore:
         """Score based on multi-step patterns."""
         hits: Final = sum(1 for p in self._multi_step_patterns if p.search(text))
@@ -1415,12 +1431,13 @@ class ComplexityRouter(CustomLogger):
             self._score_question_complexity(prompt),
         ]
 
-        # Collect signals
-        signals: Final = [d.signal for d in dimensions if d.signal is not None]
+        custom_dimensions: Final = self._score_custom_dimensions(prompt, user_text)
+        signals: Final = [d.signal for d in (*dimensions, *(d for d, _ in custom_dimensions)) if d.signal is not None]
 
-        # Compute weighted score
         weights: Final = self.config.dimension_weights
-        weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions)
+        weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions) + sum(
+            dimension.score * weight for dimension, weight in custom_dimensions
+        )
 
         boundaries: Final = self._effective_tier_boundaries()
         clears_override_floor: Final = weighted_score >= self._effective_reasoning_override_min_score()
diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py
index c483a0b7073..3ec9f9b5394 100644
--- a/litellm/router_strategy/complexity_router/config.py
+++ b/litellm/router_strategy/complexity_router/config.py
@@ -5,13 +5,21 @@ Contains default keyword lists, weights, tier boundaries, and configuration clas
 All values are configurable via proxy config.yaml.
 """
 
-from collections.abc import Mapping
+import math
+import re
+import warnings
+from collections.abc import Iterable, Mapping
 from enum import Enum
 from types import MappingProxyType
-from typing import Annotated, Final, Literal
+from typing import Annotated, Final, Literal, NamedTuple
 
 from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_serializer, field_validator, model_validator
 
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore", DeprecationWarning)
+    import sre_constants
+    import sre_parse
+
 from litellm.types.llms.openai import REASONING_EFFORT
 from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin
 
@@ -569,6 +577,117 @@ class ClassifierLLMConfig(BaseModel):
         return self
 
 
+MAX_CUSTOM_PATTERN_REPEAT: Final[int] = 64
+MAX_CUSTOM_PATTERN_WORK: Final[int] = 2048
+MAX_CUSTOM_DIMENSIONS_WORK: Final[int] = 8192
+MAX_CUSTOM_PATTERN_DEPTH: Final[int] = 16
+CUSTOM_PATTERN_SCAN_CHARS: Final[int] = 2048
+
+_ATOM_OPCODES: Final = frozenset(
+    {sre_constants.LITERAL, sre_constants.NOT_LITERAL, sre_constants.ANY, sre_constants.IN, sre_constants.CATEGORY}
+)
+_REPEAT_OPCODES: Final = frozenset({sre_constants.MAX_REPEAT, sre_constants.MIN_REPEAT})
+
+
+class _PatternCost(NamedTuple):
+    paths: int
+    steps: int
+
+
+def _atom_steps(node: object) -> int:
+    if isinstance(node, tuple) and len(node) == 2 and node[0] is sre_constants.IN:
+        return 1 + len(node[1])
+    return 1
+
+
+def _repeat_cost(argument: object) -> _PatternCost | str:
+    if not isinstance(argument, tuple) or len(argument) != 3:
+        return "unsupported repeat structure"
+    low, high, body = argument
+    if high > MAX_CUSTOM_PATTERN_REPEAT or len(body) != 1 or body[0][0] not in _ATOM_OPCODES:
+        return "requires a single character or class repeated at most 64 times; use {n,m} instead of *, + or {n,}"
+    choices: Final = high - low + 1
+    return _PatternCost(choices, 1 + high * _atom_steps(body[0]) + choices)
+
+
+def _node_cost(node: object, depth: int) -> _PatternCost | str:
+    if not isinstance(node, tuple) or len(node) != 2:
+        return "unsupported regex structure"
+    opcode, argument = node
+    if opcode in _ATOM_OPCODES or opcode is sre_constants.AT:
+        return _PatternCost(1, _atom_steps(node))
+    if opcode is sre_constants.SUBPATTERN:
+        return _sequence_cost(argument[-1], depth + 1)
+    if opcode is sre_constants.BRANCH:
+        costs: Final = tuple(_sequence_cost(branch, depth + 1) for branch in argument[1])
+        refused: Final = next((cost for cost in costs if isinstance(cost, str)), None)
+        if refused is not None:
+            return refused
+        return _PatternCost(
+            sum(cost.paths for cost in costs if isinstance(cost, _PatternCost)),
+            len(costs) + sum(cost.steps for cost in costs if isinstance(cost, _PatternCost)),
+        )
+    if opcode in _REPEAT_OPCODES:
+        return _repeat_cost(argument)
+    return "contains an unsupported regex construct"
+
+
+def _sequence_cost(nodes: Iterable[object], depth: int) -> _PatternCost | str:
+    if depth > MAX_CUSTOM_PATTERN_DEPTH:
+        return "nests deeper than 16 levels"
+    costs: Final = tuple(_node_cost(node, depth) for node in nodes)
+    refused: Final = next((cost for cost in costs if isinstance(cost, str)), None)
+    if refused is not None:
+        return refused
+    valid: Final = tuple(cost for cost in costs if isinstance(cost, _PatternCost))
+    # Choices multiply across a sequence; every continuation can execute once per preceding path.
+    total: Final = _PatternCost(
+        math.prod(cost.paths for cost in valid),
+        1 + sum(cost.steps * math.prod(prior.paths for prior in valid[:index]) for index, cost in enumerate(valid)),
+    )
+    if total.steps > MAX_CUSTOM_PATTERN_WORK:
+        return "exceeds the per-pattern regex work budget"
+    return total
+
+
+def custom_pattern_work(pattern: str) -> int | str:
+    try:
+        re.compile(pattern, re.IGNORECASE)
+        parsed: Final = sre_parse.parse(pattern, re.IGNORECASE)
+    except (re.error, RecursionError, OverflowError):
+        return "is not a valid regex"
+    cost: Final = _sequence_cost(tuple(parsed), 0)
+    return cost if isinstance(cost, str) else cost.steps
+
+
+class CustomDimension(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    name: str = Field(min_length=1, max_length=64, pattern=r"^[A-Za-z][A-Za-z0-9_]*$")
+    weight: float = Field(gt=0, le=1, allow_inf_nan=False)
+    keywords: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32)
+    patterns: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32)
+
+    @model_validator(mode="after")
+    def _validate_matchers(self) -> "CustomDimension":
+        matchers: Final = (*self.keywords, *self.patterns)
+        if not matchers or any(not matcher.strip() for matcher in matchers):
+            raise ValueError("custom dimensions require nonblank keywords and/or patterns")
+        if len(matchers) > 32 or sum(map(len, matchers)) > 4096:
+            raise ValueError("custom dimensions allow at most 32 matchers and 4096 matcher characters each")
+        costs: Final = tuple((pattern, custom_pattern_work(pattern)) for pattern in self.patterns)
+        rejected: Final = tuple(f"pattern {pattern!r} {work}" for pattern, work in costs if isinstance(work, str))
+        if rejected:
+            raise ValueError("custom dimension " + "; ".join(rejected))
+        return self
+
+    def pattern_work(self) -> int:
+        """Combined work estimate of the validated patterns."""
+        return sum(
+            work for work in (custom_pattern_work(pattern) for pattern in self.patterns) if isinstance(work, int)
+        )
+
+
 class ComplexityRouterConfig(BaseModel):
     """Configuration for the ComplexityRouter."""
 
@@ -671,6 +790,19 @@ class ComplexityRouterConfig(BaseModel):
         description="Weights for each scoring dimension",
     )
 
+    custom_dimensions: tuple[CustomDimension, ...] = Field(
+        default=(),
+        max_length=16,
+        description=(
+            "Named binary dimensions added to the heuristic-v1 score. Each contributes its inline weight once "
+            "when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters. "
+            "Regex quantifiers repeat one character or class at most 64 times. Unbounded quantifiers, repeated groups, "
+            "backreferences and lookarounds are rejected. Conservative work limits include alternation paths, "
+            "repeat lengths and subsequent matching: 2048 units per pattern, 8192 across the router. "
+            "Only heuristic, heuristic_first and hybrid accept this field. Uses the existing heuristic tuning quota."
+        ),
+    )
+
     # Keyword lists (overridable)
     code_keywords: list[str] | None = Field(
         default=None,
@@ -1245,6 +1377,27 @@ class ComplexityRouterConfig(BaseModel):
             )
         return self
 
+    @model_validator(mode="after")
+    def _validate_custom_dimensions(self) -> "ComplexityRouterConfig":
+        if not self.custom_dimensions:
+            return self
+        if self.classifier_type not in ("heuristic", "heuristic_first", "hybrid"):
+            raise ValueError("custom_dimensions requires classifier_type heuristic, heuristic_first or hybrid")
+        names: Final = tuple(dimension.name.casefold() for dimension in self.custom_dimensions)
+        reserved: Final = frozenset(name.casefold() for name in DEFAULT_DIMENSION_WEIGHTS)
+        weighted: Final = frozenset(name.casefold() for name in self.dimension_weights)
+        if len(frozenset(names)) != len(names) or frozenset(names) & reserved:
+            raise ValueError("custom dimension names must be unique and must not shadow built-in dimensions")
+        if frozenset(names) & weighted:
+            raise ValueError("custom dimension weights must be inline, not in dimension_weights")
+        work: Final = sum(dimension.pattern_work() for dimension in self.custom_dimensions)
+        if work > MAX_CUSTOM_DIMENSIONS_WORK:
+            raise ValueError(
+                f"custom_dimensions regex work estimate is {work}; the limit across the router is "
+                f"{MAX_CUSTOM_DIMENSIONS_WORK}"
+            )
+        return self
+
     @field_validator("heuristic_first_max_tier", mode="before")
     @classmethod
     def _coerce_heuristic_first_max_tier(cls, value: object) -> object:
diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_utils/auto_router_tuning_baseline.py
index b82269d5824..74f7b82389a 100644
--- a/litellm/router_utils/auto_router_tuning_baseline.py
+++ b/litellm/router_utils/auto_router_tuning_baseline.py
@@ -20,6 +20,7 @@ HEURISTIC_V1_TUNING_FIELDS: Final = (
     "reasoning_override_min_score",
     "token_thresholds",
     "dimension_weights",
+    "custom_dimensions",
     "code_keywords",
     "reasoning_keywords",
     "technical_keywords",
diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py
index ebfb631f93b..5b1d8562abd 100644
--- a/tests/test_litellm/router_strategy/test_complexity_router.py
+++ b/tests/test_litellm/router_strategy/test_complexity_router.py
@@ -7,7 +7,8 @@ Tests the rule-based complexity scoring and tier assignment logic.
 import asyncio
 import logging
 import sys
-from typing import Dict, List
+import time
+from typing import Dict, Final, List
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -47,6 +48,7 @@ from litellm.router_strategy.complexity_router.config import (
     ClassifierLLMConfig,
     ComplexityRouterConfig,
     ComplexityTier,
+    custom_pattern_work,
 )
 from litellm.router_strategy.complexity_router.tier_predictor import (
     TierGlobalStatistic,
@@ -756,6 +758,205 @@ class TestCustomTechnicalKeywords:
         assert custom_score > baseline_score
 
 
+class TestCustomDimensions:
+    @pytest.mark.parametrize(
+        "matchers,prompt",
+        [
+            pytest.param(
+                {"keywords": ["orbitmesh", "fluxgate"]},
+                "Connect ORBITMESH and fluxgate for the requested change",
+                id="keywords",
+            ),
+            pytest.param(
+                {"patterns": [r"\bCREATE\s{1,4}TABLE\b", r"\bALTER\s{1,4}TABLE\b"]},
+                "create table widgets (id integer); ALTER TABLE widgets ADD label text;",
+                id="regex",
+            ),
+        ],
+    )
+    def test_custom_dimension_changes_only_matching_requests(
+        self, mock_router_instance: MagicMock, matchers: dict[str, object], prompt: str
+    ) -> None:
+        baseline: Final = ComplexityRouter("test-router", mock_router_instance)
+        configured: Final = ComplexityRouter(
+            "test-router",
+            mock_router_instance,
+            {"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, **matchers}]},
+        )
+        baseline_tier, baseline_score, baseline_signals = baseline.classify(prompt)
+        tier, score, signals = configured.classify(prompt)
+        assert baseline_tier == ComplexityTier.SIMPLE
+        assert tier != ComplexityTier.SIMPLE
+        assert score == pytest.approx(baseline_score + 0.7)
+        assert signals == [*baseline_signals, "custom (internalFrameworks)"]
+        plain: Final = "Hello!"
+        assert configured.classify(plain) == baseline.classify(plain)
+        assert configured.classify(plain)[0] == ComplexityTier.SIMPLE
+
+    @pytest.mark.parametrize(
+        "dimension_overrides,config_overrides",
+        [
+            pytest.param({"keywords": []}, {}, id="missing-matchers"),
+            pytest.param({"keywords": ["  "]}, {}, id="blank-keyword"),
+            pytest.param({"patterns": ["\t"]}, {}, id="blank-pattern"),
+            pytest.param({"patterns": ["("]}, {}, id="invalid-regex"),
+            pytest.param({"patterns": [r"a*b"]}, {}, id="unbounded-star"),
+            pytest.param({"patterns": [r"a{2,}b"]}, {}, id="unbounded-brace"),
+            pytest.param({"patterns": [r"a{0,65}b"]}, {}, id="repeat-over-64"),
+            pytest.param({"patterns": [r"(a{0,8}){0,8}b"]}, {}, id="nested-repeat"),
+            pytest.param({"patterns": [r"(a|aa){0,12}b"]}, {}, id="alternation-in-repeat"),
+            pytest.param({"patterns": [r"(?:ab){0,64}c"]}, {}, id="group-repeat"),
+            pytest.param({"patterns": ["a?" * 9 + "b"]}, {}, id="pattern-work-over-budget"),
+            pytest.param({"patterns": ["(?:a|aa)" * 9 + "z"]}, {}, id="ambiguous-alternation-chain"),
+            pytest.param({"patterns": ["a?" * 8 + "a{64}" * 10 + "z"]}, {}, id="cheap-prefix-expensive-tail"),
+            pytest.param({"patterns": [r"(a)\1"]}, {}, id="backreference"),
+            pytest.param({"patterns": [r"(?=x)y"]}, {}, id="lookahead"),
+            pytest.param({"patterns": [r"(?>ab)"]}, {}, id="atomic-group"),
+            pytest.param({"patterns": [r"a*+b"]}, {}, id="possessive"),
+            pytest.param({"name": "CODEPRESENCE"}, {"dimension_weights": {"tokenCount": 0.1}}, id="reserved-name"),
+            pytest.param({}, {"dimension_weights": {"INTERNALFRAMEWORKS": 0.7}}, id="weight-in-map"),
+            pytest.param({"weight": 0}, {}, id="zero-weight"),
+            pytest.param({"weight": 1.1}, {}, id="excess-weight"),
+            pytest.param({"weight": float("nan")}, {}, id="nan-weight"),
+            pytest.param({"weight": float("inf")}, {}, id="infinite-weight"),
+            pytest.param({"name": "bad-name"}, {}, id="invalid-name"),
+            pytest.param({"name": "x" * 65}, {}, id="long-name"),
+            pytest.param({"keywords": [""]}, {}, id="empty-matcher"),
+            pytest.param({"keywords": ["x" * 257]}, {}, id="long-matcher"),
+            pytest.param({"keywords": ["x"] * 32, "patterns": ["y"]}, {}, id="combined-matcher-count"),
+            pytest.param({"keywords": ["x" * 256] * 17}, {}, id="matcher-character-budget"),
+            pytest.param({"unknown": True}, {}, id="extra-field"),
+        ],
+    )
+    def test_custom_dimension_invalid_configuration_rejected(
+        self, dimension_overrides: dict[str, object], config_overrides: dict[str, object]
+    ) -> None:
+        with pytest.raises(ValidationError, match=r"custom_dimensions|custom dimension"):
+            ComplexityRouterConfig.model_validate(
+                {
+                    "custom_dimensions": [
+                        {
+                            "name": "internalFrameworks",
+                            "weight": 0.7,
+                            "keywords": ["orbitmesh"],
+                            **dimension_overrides,
+                        }
+                    ],
+                    **config_overrides,
+                }
+            )
+
+    @pytest.mark.parametrize(
+        "names",
+        [
+            pytest.param(("internalFrameworks", "INTERNALFRAMEWORKS"), id="duplicate-casefolded-name"),
+            pytest.param(tuple(f"dimension{i}" for i in range(17)), id="dimension-count"),
+        ],
+    )
+    def test_custom_dimension_names_and_count_are_bounded(self, names: tuple[str, ...]) -> None:
+        with pytest.raises(ValidationError, match=r"custom_dimensions|custom dimension"):
+            ComplexityRouterConfig.model_validate(
+                {"custom_dimensions": [{"name": name, "weight": 0.7, "keywords": ["orbitmesh"]} for name in names]}
+            )
+
+    @pytest.mark.parametrize("classifier_type", ("heuristic_v2", "llm", "custom"))
+    def test_custom_dimensions_reject_classifiers_outside_the_tuning_gate(self, classifier_type: str) -> None:
+        classifier_config: Final = (
+            {"classifier_plugin": _FixedTierClassifier("SIMPLE")}
+            if classifier_type == "custom"
+            else {"classifier_llm_config": {"model": "judge"}}
+            if classifier_type == "llm"
+            else {}
+        )
+        with pytest.raises(ValidationError, match="custom_dimensions requires classifier_type"):
+            ComplexityRouterConfig.model_validate(
+                {
+                    "classifier_type": classifier_type,
+                    "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}],
+                    **classifier_config,
+                }
+            )
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("current_ask", ("Hello!", "orbitmesh"))
+    async def test_custom_dimensions_public_hook_scores_only_current_ask(
+        self, mock_router_instance: MagicMock, current_ask: str
+    ) -> None:
+        router: Final = ComplexityRouter(
+            "test-router",
+            mock_router_instance,
+            {
+                "tiers": {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"},
+                "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}],
+            },
+        )
+        result: Final = await router.async_pre_routing_hook(
+            model="test-router",
+            request_kwargs={},
+            messages=[
+                {"role": "system", "content": "orbitmesh"},
+                {"role": "user", "content": "orbitmesh"},
+                {"role": "assistant", "content": "orbitmesh is ready"},
+                {"role": "user", "content": current_ask},
+                {"role": "tool", "tool_call_id": "previous", "content": "orbitmesh"},
+            ],
+        )
+        assert result is not None
+        assert result.routing_decision is not None
+        assert ("custom (internalFrameworks)" in result.routing_decision["signals"]) is (current_ask == "orbitmesh")
+        assert result.model == ("top" if current_ask == "orbitmesh" else "cheap")
+        assert "orbitmesh" not in " ".join(result.routing_decision["signals"])
+
+    def test_custom_patterns_scan_only_the_first_2048_characters(self, mock_router_instance: MagicMock) -> None:
+        router: Final = ComplexityRouter(
+            "test-router",
+            mock_router_instance,
+            {"custom_dimensions": [{"name": "late", "weight": 0.7, "patterns": [r"zzz{1,3}"]}]},
+        )
+        assert "custom (late)" in router.classify("a" * 2040 + " zzz")[2]
+        assert "custom (late)" not in router.classify("a" * 2048 + " zzz")[2]
+
+    def test_custom_dimensions_router_wide_regex_work_is_capped(self) -> None:
+        heavy: Final = {"weight": 0.5, "patterns": ["a?" * 8 + "z"]}
+        ComplexityRouterConfig.model_validate({"custom_dimensions": [{"name": f"d{i}", **heavy} for i in range(6)]})
+        with pytest.raises(ValidationError, match="regex work estimate is 8939"):
+            ComplexityRouterConfig.model_validate({"custom_dimensions": [{"name": f"d{i}", **heavy} for i in range(7)]})
+
+    @pytest.mark.parametrize(
+        "pattern,work",
+        [
+            pytest.param(r"\b(create|alter|drop)\s{1,4}table\b", 135, id="sql-ddl"),
+            pytest.param("a?" * 8 + "z", 1277, id="optional-chain-near-cap"),
+            pytest.param(r"a{0,15}a{0,15}z", 801, id="adjacent-bounded-near-cap"),
+            pytest.param(r"[a-z0-9_]{3,63}\.(com|net|io)", 1291, id="class-repeat-plus-alternation"),
+            pytest.param("(?:a|aa)" * 8 + "z", 1787, id="ambiguous-alternation-near-cap"),
+            pytest.param("a{64}" * 10 + "z", 662, id="long-deterministic-tail"),
+        ],
+    )
+    def test_custom_pattern_work_stays_cheap_on_adversarial_text(
+        self, mock_router_instance: MagicMock, pattern: str, work: int
+    ) -> None:
+        assert custom_pattern_work(pattern) == work
+        router: Final = ComplexityRouter(
+            "test-router",
+            mock_router_instance,
+            {
+                "custom_dimensions": [
+                    {"name": "bounded", "weight": 0.7, "patterns": [pattern]},
+                    {"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]},
+                ]
+            },
+        )
+        adversarial: Final = "orbitmesh " + "a" * 4000
+        started: Final = time.perf_counter()
+        tier, score, signals = router.classify(adversarial)
+        elapsed: Final = time.perf_counter() - started
+        assert signals == ["long (1002 tokens)", "custom (internalFrameworks)"]
+        assert score == pytest.approx(0.8)
+        assert tier == ComplexityTier.REASONING
+        assert elapsed < 0.1
+
+
 class TestAsyncPreRoutingHookEdgeCases:
     """Test edge cases for async_pre_routing_hook method."""
 
diff --git a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py
index fa7a96adb20..f686a62db76 100644
--- a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py
+++ b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import Mapping
+from typing import Final
 
 import pytest
 
@@ -58,6 +59,7 @@ class TestTuningFingerprint:
             "reasoning_override_min_score": 0.05,
             "token_thresholds": {"simple": 20, "complex": 500},
             "dimension_weights": {"codePresence": 0.9},
+            "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}],
             "code_keywords": ["orionflow"],
             "reasoning_keywords": ["deduce"],
             "technical_keywords": ["ledgerkit"],
@@ -216,6 +218,28 @@ class TestQuota:
             is None
         )
 
+    def test_custom_dimension_add_edit_and_revert_share_one_quota_slot(self) -> None:
+        baselines: Final = snapshot_tuning_baselines(())
+        original: Final = _router("a", {})
+        config: Final = {
+            "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}]
+        }
+        edited_config: Final = {
+            "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.9, "keywords": ["orbitmesh"]}]
+        }
+        added: Final = _router("a", config)
+        edited: Final = _router("a", edited_config)
+        second: Final = _router("b", config)
+
+        assert tuning_fingerprint(config) != tuning_fingerprint(edited_config)
+        assert mutable_tuned_identities((added,), baselines) == {router_identity(original)}
+        assert tuning_quota_violation(candidate=added, others=(original,), baselines=baselines, limit=1) is None
+        assert tuning_quota_violation(candidate=edited, others=(added,), baselines=baselines, limit=1) is None
+        assert tuning_quota_violation(candidate=second, others=(edited,), baselines=baselines, limit=1) is not None
+        assert tuning_quota_violation(candidate=original, others=(edited,), baselines=baselines, limit=1) is None
+        assert mutable_tuned_identities((original,), baselines) == frozenset()
+        assert tuning_quota_violation(candidate=second, others=(original,), baselines=baselines, limit=1) is None
+
     def test_violation_message_names_the_limit_and_remedy(self) -> None:
         message = tuning_limit_violation(held=2, limit=1)
         assert message is not None
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 3540d2f6aea..7121124e64f 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -26560,6 +26560,23 @@ export interface components {
                 [key: string]: unknown;
             };
         };
+        /** CustomDimension */
+        CustomDimension: {
+            /**
+             * Keywords
+             * @default []
+             */
+            keywords: string[];
+            /** Name */
+            name: string;
+            /**
+             * Patterns
+             * @default []
+             */
+            patterns: string[];
+            /** Weight */
+            weight: number;
+        };
         /**
          * CustomerResponse
          * @description Customer object returned by the /customer read+write endpoints.
@@ -34873,6 +34890,12 @@ export interface components {
              * @default 0.95
              */
             context_window_escalation_buffer: number;
+            /**
+             * Custom Dimensions
+             * @description Named binary dimensions added to the heuristic-v1 score. Each contributes its inline weight once when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters. Regex quantifiers repeat one character or class at most 64 times. Unbounded quantifiers, repeated groups, backreferences and lookarounds are rejected. Conservative work limits include alternation paths, repeat lengths and subsequent matching: 2048 units per pattern, 8192 across the router. Only heuristic, heuristic_first and hybrid accept this field. Uses the existing heuristic tuning quota.
+             * @default []
+             */
+            custom_dimensions: components["schemas"]["CustomDimension"][];
             /**
              * Custom Technical Keywords
              * @description Domain-specific technical keywords appended to the effective base list (technical_keywords if set, otherwise DEFAULT_TECHNICAL_KEYWORDS). Order is preserved; duplicates are removed case-insensitively against the base list and within this list.

From 9bc91041026ee8d2a6444d4fd71394cb94a7b7df Mon Sep 17 00:00:00 2001
From: yucheng-berri <yucheng@berri.ai>
Date: Mon, 7 Sep 2026 18:18:28 -0700
Subject: [PATCH 151/164] fix(proxy): log budget reservation notice once at
 config load (#40167)

* fix(proxy): log disable_budget_reservation notice once at config load

The disabled-budget-reservation reminder fired as a WARNING inside request
authentication, so every authenticated request on a proxy that deliberately
set the flag produced one warning line. The notice now runs once per worker
when general_settings loads, at INFO, and the request path only skips the
reservation. Reservation skipping and read-time budget checks are unchanged

* fix(proxy): keep budget notice sentinel with constants

* fix(proxy): expose shared budget notice state
---
 litellm/constants.py                          |  1 +
 litellm/proxy/_types.py                       |  2 +-
 litellm/proxy/auth/auth_utils.py              | 20 +++++++++-
 litellm/proxy/auth/user_api_key_auth.py       |  8 ----
 litellm/proxy/proxy_server.py                 |  5 +++
 .../proxy/auth/test_auth_utils.py             | 37 +++++++++++++++++++
 .../proxy/auth/test_user_api_key_auth.py      | 30 +++++++++++++++
 .../proxy/proxy_server/test_proxy_config.py   | 27 ++++++++++++++
 ui/litellm-dashboard/src/lib/http/schema.d.ts |  2 +-
 9 files changed, 121 insertions(+), 11 deletions(-)

diff --git a/litellm/constants.py b/litellm/constants.py
index d53686e5e5b..defc9337e9b 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -55,6 +55,7 @@ S3_PREFIX_DIGEST_CHARS: Final = 16
 MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
 DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
+budget_reservation_disabled_info_emitted = False
 DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
 DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512))
 SQS_SEND_MESSAGE_ACTION: Final = "SendMessage"
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index abce10690e5..4dbae6394f6 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -2833,7 +2833,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
             "Enable only if your deployment is experiencing phantom "
             "BudgetExceededError responses caused by leaked reservations "
             "(see GitHub issue #27639). "
-            "A proxy-level WARNING is logged on every request while this flag "
+            "An INFO notice is logged once per worker at config load while this flag "
             "is active as a reminder that hard enforcement is relaxed."
         ),
     )
diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py
index 1e4836654a1..f78c4221f5a 100644
--- a/litellm/proxy/auth/auth_utils.py
+++ b/litellm/proxy/auth/auth_utils.py
@@ -11,7 +11,7 @@ from fastapi import HTTPException, Request, status
 from pydantic import PositiveInt, TypeAdapter, ValidationError
 
 import litellm
-from litellm import Router, provider_list
+from litellm import Router, constants, provider_list
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import (
     BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY,
@@ -1390,6 +1390,24 @@ def warn_once_if_custom_auth_skips_common_checks(
     _custom_auth_common_checks_warning_emitted = True
 
 
+def log_once_if_budget_reservation_disabled(
+    *,
+    disabled: bool,
+    logger: Logger = verbose_proxy_logger,
+) -> None:
+    if constants.budget_reservation_disabled_info_emitted or not disabled:
+        return
+    logger.info(
+        "disable_budget_reservation is enabled: skipping optimistic budget "
+        "reservation. Budget enforcement is read-time only. Concurrent "
+        "requests can each pass the spend check before their cost is recorded, "
+        "so a configured budget may be briefly exceeded under high concurrency. "
+        "Set disable_budget_reservation to False or remove it to restore "
+        "hard per-request budget enforcement."
+    )
+    constants.budget_reservation_disabled_info_emitted = True  # rebind-ok: process-wide one-shot sentinel
+
+
 def is_pass_through_provider_route(route: str) -> bool:
     PROVIDER_SPECIFIC_PASS_THROUGH_ROUTES: Final = [
         "vertex-ai",
diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py
index 93293db24c6..b39b1f330b3 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -2706,14 +2706,6 @@ async def _reserve_budget_after_common_checks(
     if skip_budget_checks:
         return
     if general_settings.get("disable_budget_reservation") is True:
-        verbose_proxy_logger.warning(
-            "disable_budget_reservation is enabled: skipping optimistic budget "
-            "reservation. Budget enforcement is read-time only — concurrent "
-            "requests can each pass the spend check before their cost is recorded, "
-            "so a configured budget may be briefly exceeded under high concurrency. "
-            "Set disable_budget_reservation to False or remove it to restore "
-            "hard per-request budget enforcement."
-        )
         return
 
     from litellm.proxy.spend_tracking.budget_reservation import (
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 0915b8dd1b9..32b6b841af7 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -306,6 +306,7 @@ from litellm.proxy.auth.auth_checks import (
 from litellm.proxy.auth.auth_utils import (
     check_response_size_is_safe,
     is_request_body_safe,
+    log_once_if_budget_reservation_disabled,
     warn_once_if_custom_auth_skips_common_checks,
 )
 from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
@@ -5653,6 +5654,10 @@ class ProxyConfig:
                 run_common_checks=bool(general_settings.get("custom_auth_run_common_checks", False)),
             )
 
+            log_once_if_budget_reservation_disabled(
+                disabled=general_settings.get("disable_budget_reservation") is True,
+            )
+
             custom_key_generate: Final = general_settings.get("custom_key_generate", None)
             if custom_key_generate is not None:
                 user_custom_key_generate = get_instance_fn(value=custom_key_generate, config_file_path=config_file_path)
diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py
index a996de4d40c..aaf630ad29b 100644
--- a/tests/test_litellm/proxy/auth/test_auth_utils.py
+++ b/tests/test_litellm/proxy/auth/test_auth_utils.py
@@ -3,6 +3,7 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext
 """
 
 import base64
+import logging
 from typing import Optional
 from unittest.mock import MagicMock, patch
 
@@ -15,6 +16,7 @@ from litellm.proxy.auth.auth_utils import (
     abbreviate_api_key,
     check_complete_credentials,
     custom_auth_common_checks_warning,
+    log_once_if_budget_reservation_disabled,
     warn_once_if_custom_auth_skips_common_checks,
     get_end_user_id_from_request_body,
     get_key_mcp_rpm_limit,
@@ -101,6 +103,41 @@ class TestWarnOnceIfCustomAuthSkipsCommonChecks:
         assert logger.warning.call_count == 0
 
 
+class TestLogOnceIfBudgetReservationDisabled:
+    @pytest.fixture(autouse=True)
+    def _reset_sentinel(self, monkeypatch):
+        monkeypatch.setattr(
+            "litellm.constants.budget_reservation_disabled_info_emitted",
+            False,
+        )
+
+    def test_logs_info_only_once_when_enabled(self, caplog):
+        with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
+            log_once_if_budget_reservation_disabled(disabled=False)
+            assert not any(
+                "disable_budget_reservation is enabled" in record.message
+                for record in caplog.records
+            )
+            for _ in range(3):
+                log_once_if_budget_reservation_disabled(disabled=True)
+
+        records = [
+            record
+            for record in caplog.records
+            if "disable_budget_reservation is enabled" in record.message
+        ]
+        assert len(records) == 1
+        assert records[0].levelno == logging.INFO
+
+    def test_logs_to_injected_logger_only_once(self):
+        logger = MagicMock()
+        log_once_if_budget_reservation_disabled(disabled=False, logger=logger)
+        for _ in range(3):
+            log_once_if_budget_reservation_disabled(disabled=True, logger=logger)
+        assert logger.info.call_count == 1
+        assert "disable_budget_reservation is enabled" in logger.info.call_args[0][0]
+
+
 class TestGetKeyModelRpmLimit:
     """Tests for get_key_model_rpm_limit function."""
 
diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py
index d44f96d95bf..541aeabcbcd 100644
--- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py
+++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py
@@ -1,5 +1,6 @@
 import asyncio
 import json
+import logging
 from contextlib import contextmanager
 from datetime import datetime, timedelta
 from types import SimpleNamespace
@@ -146,6 +147,35 @@ async def test_disable_budget_reservation_skips_reservation():
     assert user_api_key_auth_obj.budget_reservation is None
 
 
+@pytest.mark.asyncio
+async def test_disable_budget_reservation_does_not_log_per_request(caplog):
+    user_api_key_auth_obj = UserAPIKeyAuth(token="test_token")
+
+    with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
+        for _ in range(3):
+            await _reserve_budget_after_common_checks(
+                user_api_key_auth_obj=user_api_key_auth_obj,
+                request_data={"model": "gpt-4o"},
+                route="/v1/chat/completions",
+                llm_router=None,
+                team_object=None,
+                user_object=None,
+                prisma_client=None,
+                user_api_key_cache=MagicMock(),
+                proxy_logging_obj=MagicMock(),
+                skip_budget_checks=False,
+                general_settings={"disable_budget_reservation": True},
+            )
+
+    records = [
+        record
+        for record in caplog.records
+        if "disable_budget_reservation is enabled" in record.message
+    ]
+    assert records == []
+    assert user_api_key_auth_obj.budget_reservation is None
+
+
 @pytest.mark.asyncio
 async def test_budget_reservation_runs_when_not_disabled():
     """Control for #27639: with the flag absent, the reservation still runs and is stored."""
diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
index 2babfe432f3..770cec1834e 100644
--- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
+++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
@@ -9,6 +9,7 @@ Pins covered:
 from __future__ import annotations
 
 import json
+import logging
 import os
 import re
 from types import SimpleNamespace
@@ -1633,6 +1634,32 @@ async def test_ProxyConfig_load_config_minimal_yaml(tmp_path, monkeypatch):
     }
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("setting", ["true", "false", "null", "'true'", None])
+async def test_load_config_logs_disabled_budget_reservation_once(tmp_path, monkeypatch, caplog, setting):
+    config_file = tmp_path / "budget.yaml"
+    flag = f"  disable_budget_reservation: {setting}\n" if setting is not None else ""
+    config_file.write_text(
+        "model_list: []\nlitellm_settings: {}\ngeneral_settings:\n"
+        "  master_key: null\n" + flag
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
+    monkeypatch.setattr("litellm.constants.budget_reservation_disabled_info_emitted", False)
+    monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
+    config = ProxyConfig()
+
+    with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
+        for _ in range(3):
+            await config.load_config(router=None, config_file_path=str(config_file))
+
+    records = [
+        record for record in caplog.records
+        if "disable_budget_reservation is enabled" in record.message
+    ]
+    assert [record.levelno for record in records] == ([logging.INFO] if setting == "true" else [])
+
+
 @pytest.mark.asyncio
 async def test_ProxyConfig_load_config_resolves_router_settings_plugins(tmp_path, monkeypatch):
     """Regression: router_settings.plugins dotted-path strings must be resolved to
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 7121124e64f..6c2311a0ed9 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -25798,7 +25798,7 @@ export interface components {
             disable_auto_add_proxy_admin_to_teams?: boolean | null;
             /**
              * Disable Budget Reservation
-             * @description If True, disables the optimistic per-request budget reservation introduced in v1.84.0. WARNING: This weakens hard budget enforcement. Without the reservation, a burst of concurrent requests from a single key can each pass the read-time spend check before any of them is charged, allowing a configured budget to be exceeded under high concurrency. Budgets are still evaluated on every request at read time, so an already-exhausted budget is still rejected. Enable only if your deployment is experiencing phantom BudgetExceededError responses caused by leaked reservations (see GitHub issue #27639). A proxy-level WARNING is logged on every request while this flag is active as a reminder that hard enforcement is relaxed.
+             * @description If True, disables the optimistic per-request budget reservation introduced in v1.84.0. WARNING: This weakens hard budget enforcement. Without the reservation, a burst of concurrent requests from a single key can each pass the read-time spend check before any of them is charged, allowing a configured budget to be exceeded under high concurrency. Budgets are still evaluated on every request at read time, so an already-exhausted budget is still rejected. Enable only if your deployment is experiencing phantom BudgetExceededError responses caused by leaked reservations (see GitHub issue #27639). An INFO notice is logged once per worker at config load while this flag is active as a reminder that hard enforcement is relaxed.
              */
             disable_budget_reservation?: boolean | null;
             /**

From 13df85cceb85c85f990ac2a25214f43f03fdfb4f Mon Sep 17 00:00:00 2001
From: yujonglee <yujong@berri.ai>
Date: Mon, 7 Sep 2026 18:46:29 -0700
Subject: [PATCH 152/164] test: add Rust extension pytest contract (#40181)

* test: add Rust extension pytest contract

* test: prove native OCR execution

* test: isolate Rust extension pytest collection

* ci: register Rust extension test coverage

* test: prove native OCR at wire boundary
---
 .github/workflows/test-rust.yml     |  3 ++
 Makefile                            | 13 ++++++
 pyproject.toml                      |  1 +
 tests/test_litellm_rust/conftest.py | 24 ++++++++++
 tests/test_litellm_rust/test_ocr.py | 72 +++++++++++++++++++++++++++++
 5 files changed, 113 insertions(+)
 create mode 100644 tests/test_litellm_rust/conftest.py
 create mode 100644 tests/test_litellm_rust/test_ocr.py

diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml
index 4f56e78ddee..c6901411167 100644
--- a/.github/workflows/test-rust.yml
+++ b/.github/workflows/test-rust.yml
@@ -117,6 +117,9 @@ jobs:
 
       - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
 
+      - name: Run pytest tests/test_litellm_rust with the compiled extension
+        run: make test-rust-extension
+
       - run: >-
           uv build --wheel --out-dir panic-dist
           --config-setting "maturin.build-args=--features panic-test,extension-module"
diff --git a/Makefile b/Makefile
index ab11220821f..91835e19e3c 100644
--- a/Makefile
+++ b/Makefile
@@ -4,6 +4,7 @@
 .PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \
 	test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
 	test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
+	test-rust-extension \
 	info lint lint-inner lint-dev lint-checks format \
 	lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
 	lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
@@ -54,6 +55,7 @@ help:
 	@echo "  make test-proxy-unit-b  - Run proxy_unit_tests (p-z, ~28 files)"
 	@echo "  make test-integration   - Run integration tests"
 	@echo "  make test-unit-helm     - Run helm unit tests"
+	@echo "  make test-rust-extension - Build the Rust extension and run its public Python tests"
 	@echo ""
 	@echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide"
 	@echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine."
@@ -289,6 +291,17 @@ pre-commit:
 	@$(MAKE) check
 
 # Testing targets
+test-rust-extension:
+	@temporary=$$(mktemp -d) && \
+	trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \
+	$(UV) build --python 3.12 --wheel --out-dir "$$temporary/wheels" && \
+	set -- "$$temporary"/wheels/*.whl && \
+	[ "$$#" -eq 1 ] && \
+	UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \
+	$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \
+	LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
+	"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
+
 test: install-test-deps
 	$(UV_RUN) pytest tests/
 
diff --git a/pyproject.toml b/pyproject.toml
index f4f238dd4b9..af35c77d259 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -340,6 +340,7 @@ markers = [
     "asyncio: mark test as an asyncio test",
     "limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')",
     "no_parallel: mark test to run sequentially (not in parallel) - typically for memory measurement tests",
+    "requires_rust_extension: public Python contract requiring an enabled, compiled Rust extension",
 ]
 filterwarnings = [
     # Suppress Pydantic serializer warnings from mock server responses (non-critical for memory tests)
diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py
new file mode 100644
index 00000000000..02d274bb405
--- /dev/null
+++ b/tests/test_litellm_rust/conftest.py
@@ -0,0 +1,24 @@
+import os
+
+import pytest
+
+
+def pytest_collection_modifyitems(items):
+    rust_enabled = os.environ.get("LITELLM_RUST", "").strip().lower() in {
+        "1",
+        "true",
+        "yes",
+        "on",
+    }
+    if not rust_enabled:
+        skip = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension")
+        for item in items:
+            item.add_marker(skip)
+        return
+
+    try:
+        from litellm.rust_bridge import _native  # noqa: F401  # validates the installed extension
+    except ImportError as error:
+        raise pytest.UsageError(
+            "LITELLM_RUST=1 requires a compiled litellm.rust_bridge._native extension"
+        ) from error
diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py
new file mode 100644
index 00000000000..d5b1fce1139
--- /dev/null
+++ b/tests/test_litellm_rust/test_ocr.py
@@ -0,0 +1,72 @@
+import json
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+import pytest
+
+import litellm
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.fixture
+def ocr_server():
+    requests = []
+
+    class Handler(BaseHTTPRequestHandler):
+        def do_POST(self):
+            requests.append(
+                {
+                    "headers": {name.lower(): value for name, value in self.headers.items()},
+                    "body": json.loads(self.rfile.read(int(self.headers["Content-Length"]))),
+                }
+            )
+            if self.headers.get("User-Agent", "").startswith("python-httpx"):
+                self.send_response(418)
+                self.end_headers()
+                return
+            response = json.dumps(
+                {
+                    "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}],
+                    "model": "mistral-ocr-latest",
+                    "usage_info": {"pages_processed": 1, "doc_size_bytes": 3},
+                }
+            ).encode()
+            self.send_response(200)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(response)))
+            self.end_headers()
+            self.wfile.write(response)
+
+        def log_message(self, format, *args):
+            pass
+
+    server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+    thread = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True)
+    thread.start()
+    try:
+        yield server, requests
+    finally:
+        server.shutdown()
+        server.server_close()
+        thread.join()
+
+
+def test_ocr_with_rust_extension(ocr_server):
+    server, requests = ocr_server
+    host, port = server.server_address
+
+    response = litellm.ocr(
+        model="mistral/mistral-ocr-latest",
+        document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
+        api_key="test-key",
+        api_base=f"http://{host}:{port}",
+    )
+
+    assert response.pages[0].markdown == "native OCR response"
+    assert len(requests) == 1
+    assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx")
+    assert requests[0]["body"] == {
+            "model": "mistral-ocr-latest",
+            "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
+    }

From 3023497590a6f7124e9dbcfabd35b6c42b4de9d9 Mon Sep 17 00:00:00 2001
From: mateo <mateo@berri.ai>
Date: Tue, 8 Sep 2026 02:19:32 +0000
Subject: [PATCH 153/164] test: drop static cost-map value assertions

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../llm_translation/test_bedrock_embedding.py |  18 +-
 .../test_bedrock_embedding_pricing.py         |  34 ---
 .../llm_translation/test_bedrock_govcloud.py  | 158 ----------
 tests/llm_translation/test_crusoe.py          |  36 ---
 tests/llm_translation/test_hyperbolic.py      |  30 --
 tests/llm_translation/test_lambda_ai.py       |  41 ---
 tests/llm_translation/test_morph.py           |  16 --
 tests/llm_translation/test_openai_o1.py       |  15 +-
 tests/llm_translation/test_v0.py              |  31 --
 tests/local_testing/test_get_model_info.py    |  26 +-
 .../test_xai_oauth_routing.py                 |   7 -
 .../test_mai_image_generation.py              |  19 --
 .../test_azure_ai_fw_models_metadata.py       | 139 ---------
 .../test_azure_ai_kimi_k26_metadata.py        |  23 --
 ..._cross_region_inference_profile_mapping.py |  70 -----
 ...bedrock_mantle_responses_transformation.py | 101 -------
 .../test_bedrock_mantle_transformation.py     |  70 -----
 .../test_fireworks_ai_chat_transformation.py  |  18 +-
 .../test_fireworks_ai_kimi_model_metadata.py  |   9 -
 .../test_gemini_realtime_transformation.py    |  27 +-
 .../test_inception_chat_transformation.py     |  18 --
 ...est_inception_completion_transformation.py |  16 --
 .../test_moonshot_chat_transformation.py      |  30 --
 .../openai_like/test_cognition_provider.py    |  16 --
 .../llms/openai_like/test_json_providers.py   |  21 +-
 .../openai_like/test_libertai_provider.py     |  29 --
 .../llms/openai_like/test_meta_provider.py    |  13 -
 ...est_perplexity_embedding_transformation.py |  23 --
 .../test_perplexity_cost_calculator.py        |  45 ---
 .../test_vertex_video_transformation.py       |  12 -
 .../xai/test_xai_redirected_slug_pricing.py   |  14 -
 .../llms/zai/test_zai_provider.py             |  38 ---
 .../test_bedrock_extended_beta_models.py      |  54 ----
 .../test_bedrock_nemotron_super.py            |  51 ----
 .../test_bedrock_usgov_haiku_1hr_cache.py     |  47 ---
 .../test_bedrock_usgov_pricing.py             | 271 ------------------
 .../test_claude_fable_5_config.py             | 165 -----------
 .../test_claude_haiku_4_5_config.py           |  82 ------
 .../test_claude_opus_4_6_config.py            | 115 --------
 .../test_claude_opus_4_8_config.py            | 122 --------
 .../test_litellm/test_claude_opus_5_config.py |  91 ------
 .../test_claude_sonnet_5_config.py            |  96 -------
 ...st_cloudflare_workers_ai_model_metadata.py |  45 ---
 .../test_daybreak_model_metadata.py           |  14 -
 .../test_deepseek_model_metadata.py           |  34 ---
 .../test_fireworks_serverless_model_costs.py  |  26 --
 .../test_gpt_5_5_model_metadata.py            |  44 ---
 tests/test_litellm/test_gpt_realtime_mode.py  |  27 --
 .../test_mistral_medium_3_5_model_metadata.py |  45 ---
 .../test_mistral_small_4_0_model_metadata.py  |  21 --
 .../test_muse_spark_1_2_model_metadata.py     |  37 ---
 .../test_muse_spark_1_3_model_metadata.py     |  37 ---
 .../test_replicate_model_key_format.py        |   6 -
 .../test_together_ai_model_metadata.py        |  80 ------
 54 files changed, 9 insertions(+), 2664 deletions(-)
 delete mode 100644 tests/llm_translation/test_bedrock_embedding_pricing.py
 delete mode 100644 tests/test_litellm/test_bedrock_nemotron_super.py
 delete mode 100644 tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py

diff --git a/tests/llm_translation/test_bedrock_embedding.py b/tests/llm_translation/test_bedrock_embedding.py
index 56baed141da..1fc05b43b23 100644
--- a/tests/llm_translation/test_bedrock_embedding.py
+++ b/tests/llm_translation/test_bedrock_embedding.py
@@ -1,14 +1,12 @@
 import json
 import os
-from datetime import datetime
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import Mock, patch
 import pytest
 import base64
-import httpx
 
 
 import litellm
-from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
 
 titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}
 
@@ -394,8 +392,6 @@ def test_bedrock_embedding_uses_correct_region_when_specified():
             os.environ["AWS_REGION_NAME"] = original_region_name
         else:
             os.environ.pop("AWS_REGION_NAME", None)
-
-
 def test_bedrock_embedding_region_bug_reproduction():
     """
     Reproduces the bug where aws_region_name is ignored when passed explicitly.
@@ -458,13 +454,3 @@ def test_bedrock_embedding_region_bug_reproduction():
             os.environ["AWS_REGION_NAME"] = original_region_name
         else:
             os.environ.pop("AWS_REGION_NAME", None)
-
-
-def test_bedrock_titan_g1_text_02_model_info():
-    """Test that amazon.titan-embed-g1-text-02 has correct pricing metadata"""
-    model_info = litellm.get_model_info("amazon.titan-embed-g1-text-02")
-    assert model_info is not None, "Model info should not be None"
-    assert model_info["litellm_provider"] == "bedrock"
-    assert model_info["mode"] == "embedding"
-    assert model_info["input_cost_per_token"] == 1e-07
-    assert model_info["max_input_tokens"] == 8192
diff --git a/tests/llm_translation/test_bedrock_embedding_pricing.py b/tests/llm_translation/test_bedrock_embedding_pricing.py
deleted file mode 100644
index 099d73fed87..00000000000
--- a/tests/llm_translation/test_bedrock_embedding_pricing.py
+++ /dev/null
@@ -1,34 +0,0 @@
-"""
-Tests for AWS Bedrock embedding model pricing in the model cost map.
-
-Regression test for the Amazon Titan Text Embeddings V2 commercial price,
-which was previously set 10x too high (2e-07 instead of 2e-08).
-AWS lists Titan Text Embeddings V2 at $0.02 per 1M input tokens
-(= $0.00002 per 1K tokens = 2e-08 per token).
-"""
-
-import importlib
-
-
-class TestBedrockEmbeddingPricing:
-    """Test suite for Bedrock embedding model pricing in the cost map."""
-
-    def test_titan_embed_v2_commercial_input_cost(self, monkeypatch):
-        """Titan Text Embeddings V2 should be priced at $0.02 / 1M tokens (2e-08)."""
-        # Scope the local-cost-map flag to this test only, so it does not leak
-        # into sibling tests. monkeypatch restores the environment on teardown.
-        monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-
-        import litellm.litellm_core_utils.get_model_cost_map
-        import litellm
-
-        # Reload so the cost map is re-read from the local file with the flag set.
-        importlib.reload(litellm.litellm_core_utils.get_model_cost_map)
-        importlib.reload(litellm)
-
-        model = litellm.model_cost["amazon.titan-embed-text-v2:0"]
-
-        assert model["input_cost_per_token"] == 2e-08
-        assert model["output_cost_per_token"] == 0.0
-        assert model["litellm_provider"] == "bedrock"
-        assert model["mode"] == "embedding"
diff --git a/tests/llm_translation/test_bedrock_govcloud.py b/tests/llm_translation/test_bedrock_govcloud.py
index e69a95c714d..a69b786fd45 100644
--- a/tests/llm_translation/test_bedrock_govcloud.py
+++ b/tests/llm_translation/test_bedrock_govcloud.py
@@ -40,37 +40,6 @@ class TestBedrockGovCloudSupport:
         assert "us-gov-east-1" in all_regions
         assert "us-gov-west-1" in all_regions
 
-    def test_govcloud_models_in_model_cost(self):
-        """Test that GovCloud models are present in model cost configuration"""
-        from litellm import model_cost
-
-        # Test Claude models in GovCloud
-        assert (
-            "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
-            in model_cost
-        )
-        assert (
-            "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0"
-            in model_cost
-        )
-        assert (
-            "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost
-        )
-        assert (
-            "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost
-        )
-        assert "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0" in model_cost
-        assert "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0" in model_cost
-
-        # Test Llama models in GovCloud
-        assert "bedrock/us-gov-east-1/meta.llama3-8b-instruct-v1:0" in model_cost
-        assert "bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0" in model_cost
-        assert "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0" in model_cost
-        assert "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0" in model_cost
-
-        # Test Titan models in GovCloud
-        assert "bedrock/us-gov-east-1/amazon.titan-text-lite-v1" in model_cost
-        assert "bedrock/us-gov-west-1/amazon.titan-text-lite-v1" in model_cost
 
     def test_govcloud_model_routing(self):
         """Test that GovCloud models are routed correctly"""
@@ -148,134 +117,7 @@ class TestBedrockGovCloudSupport:
         assert not any("us-gov-east-1" in model for model in litellm.bedrock_models)
         assert not any("us-gov-west-1" in model for model in litellm.bedrock_models)
 
-    def test_govcloud_model_cost_properties(self):
-        """Test that GovCloud models have proper cost configuration"""
-        from litellm import model_cost
 
-        # Check a specific GovCloud model has all required properties
-        govcloud_model = model_cost[
-            "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
-        ]
-
-        assert "max_tokens" in govcloud_model
-        assert "max_input_tokens" in govcloud_model
-        assert "max_output_tokens" in govcloud_model
-        assert "input_cost_per_token" in govcloud_model
-        assert "output_cost_per_token" in govcloud_model
-        assert govcloud_model["litellm_provider"] == "bedrock"
-        assert govcloud_model["mode"] == "chat"
-
-    def test_govcloud_model_pricing_verification(self):
-        """Test that GovCloud models have correct pricing that differs from base models"""
-        from litellm import model_cost
-
-        # Claude Haiku 4.5 commercial list pricing is under the us.* inference profile id
-        base_model = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
-        gov_east_model = (
-            "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
-        )
-        gov_west_model = (
-            "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0"
-        )
-
-        # Verify base model pricing (us.* inference profile: $1.10/$5.50 per MTok)
-        base_pricing = model_cost[base_model]
-        assert base_pricing["input_cost_per_token"] == 1.1e-06
-        assert base_pricing["output_cost_per_token"] == 5.5e-06
-
-        # Verify GovCloud models have different (higher) pricing
-        gov_east_pricing = model_cost[gov_east_model]
-        gov_west_pricing = model_cost[gov_west_model]
-
-        # GovCloud models should have ~20% higher pricing than base models
-        assert gov_east_pricing["input_cost_per_token"] == 1.2e-06
-        assert gov_east_pricing["output_cost_per_token"] == 6e-06
-        assert gov_west_pricing["input_cost_per_token"] == 1.2e-06
-        assert gov_west_pricing["output_cost_per_token"] == 6e-06
-
-        # Verify the pricing difference is approximately 20%
-        assert (
-            abs(
-                gov_east_pricing["input_cost_per_token"]
-                / base_pricing["input_cost_per_token"]
-                - 1.2
-            )
-            < 0.15
-        )
-        assert (
-            abs(
-                gov_east_pricing["output_cost_per_token"]
-                / base_pricing["output_cost_per_token"]
-                - 1.2
-            )
-            < 0.15
-        )
-        assert (
-            abs(
-                gov_west_pricing["input_cost_per_token"]
-                / base_pricing["input_cost_per_token"]
-                - 1.2
-            )
-            < 0.15
-        )
-        assert (
-            abs(
-                gov_west_pricing["output_cost_per_token"]
-                / base_pricing["output_cost_per_token"]
-                - 1.2
-            )
-            < 0.15
-        )
-
-        # Test Claude 3 Haiku pricing
-        base_haiku_model = "anthropic.claude-3-haiku-20240307-v1:0"
-        gov_east_haiku_model = (
-            "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0"
-        )
-        gov_west_haiku_model = (
-            "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0"
-        )
-
-        # Verify base Haiku model pricing
-        base_haiku_pricing = model_cost[base_haiku_model]
-        assert base_haiku_pricing["input_cost_per_token"] == 2.5e-07  # 0.00000025
-        assert base_haiku_pricing["output_cost_per_token"] == 1.25e-06  # 0.00000125
-
-        # Verify GovCloud Haiku models have different (higher) pricing
-        gov_east_haiku_pricing = model_cost[gov_east_haiku_model]
-        gov_west_haiku_pricing = model_cost[gov_west_haiku_model]
-
-        # GovCloud Haiku models should have 20% higher pricing than base models
-        assert (
-            gov_east_haiku_pricing["input_cost_per_token"] == 3e-07
-        )  # 0.0000003 (20% higher)
-        assert (
-            gov_east_haiku_pricing["output_cost_per_token"] == 1.5e-06
-        )  # 0.0000015 (20% higher)
-        assert (
-            gov_west_haiku_pricing["input_cost_per_token"] == 3e-07
-        )  # 0.0000003 (20% higher)
-        assert (
-            gov_west_haiku_pricing["output_cost_per_token"] == 1.5e-06
-        )  # 0.0000015 (20% higher)
-
-        # Verify the pricing difference is exactly 20%
-        assert (
-            gov_east_haiku_pricing["input_cost_per_token"]
-            == base_haiku_pricing["input_cost_per_token"] * 1.2
-        )
-        assert (
-            gov_east_haiku_pricing["output_cost_per_token"]
-            == base_haiku_pricing["output_cost_per_token"] * 1.2
-        )
-        assert (
-            gov_west_haiku_pricing["input_cost_per_token"]
-            == base_haiku_pricing["input_cost_per_token"] * 1.2
-        )
-        assert (
-            gov_west_haiku_pricing["output_cost_per_token"]
-            == base_haiku_pricing["output_cost_per_token"] * 1.2
-        )
 
     @patch("litellm.completion")
     def test_govcloud_completion_cost_calculation(self, mock_completion):
diff --git a/tests/llm_translation/test_crusoe.py b/tests/llm_translation/test_crusoe.py
index 56aa4e4cd42..576428684fc 100644
--- a/tests/llm_translation/test_crusoe.py
+++ b/tests/llm_translation/test_crusoe.py
@@ -4,7 +4,6 @@ Tests for Crusoe provider integration
 import os
 from unittest import mock
 
-import litellm
 
 CRUSOE_API_BASE = "https://managed-inference-api-proxy.crusoecloud.com/v1"
 
@@ -71,38 +70,3 @@ def test_get_llm_provider_crusoe():
     )
     assert model == "meta-llama/Llama-3.3-70B-Instruct"
     assert provider == "crusoe"
-
-
-def test_crusoe_models_configuration():
-    """Test that Crusoe models are configured correctly"""
-    from litellm import get_model_info
-
-    original_model_cost = litellm.model_cost
-    original_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP")
-    try:
-        os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-        litellm.model_cost = litellm.get_model_cost_map(url="")
-
-        crusoe_models = [
-            "crusoe/meta-llama/Llama-3.3-70B-Instruct",
-            "crusoe/deepseek-ai/DeepSeek-R1-0528",
-            "crusoe/deepseek-ai/DeepSeek-V3-0324",
-            "crusoe/Qwen/Qwen3-235B-A22B-Instruct-2507",
-            "crusoe/moonshotai/Kimi-K2-Thinking",
-            "crusoe/openai/gpt-oss-120b",
-            "crusoe/google/gemma-3-12b-it",
-        ]
-
-        for model in crusoe_models:
-            model_info = get_model_info(model)
-            assert model_info is not None, f"Model info not found for {model}"
-            assert model_info.get("litellm_provider") == "crusoe", (
-                f"{model} should have crusoe as provider"
-            )
-            assert model_info.get("mode") == "chat", f"{model} should be in chat mode"
-    finally:
-        litellm.model_cost = original_model_cost
-        if original_env is None:
-            os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None)
-        else:
-            os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env
diff --git a/tests/llm_translation/test_hyperbolic.py b/tests/llm_translation/test_hyperbolic.py
index 78817fbd902..0dd1c4924c0 100644
--- a/tests/llm_translation/test_hyperbolic.py
+++ b/tests/llm_translation/test_hyperbolic.py
@@ -1,8 +1,4 @@
-import os
-from datetime import datetime
-from unittest.mock import MagicMock
 
-import pytest
 
 
 import litellm
@@ -69,32 +65,6 @@ def test_hyperbolic_in_provider_lists():
     assert "https://api.hyperbolic.xyz/v1" in openai_compatible_endpoints
 
 
-def test_hyperbolic_models_configuration():
-    """Test that Hyperbolic models are properly configured"""
-    import json
-
-    # Load model configuration directly from the JSON file
-    json_path = os.path.join(
-        os.path.dirname(__file__), "../../model_prices_and_context_window.json"
-    )
-    with open(json_path, "r") as f:
-        model_data = json.load(f)
-
-    # Test a few key models
-    test_models = [
-        "hyperbolic/deepseek-ai/DeepSeek-V3",
-        "hyperbolic/Qwen/Qwen2.5-Coder-32B-Instruct",
-        "hyperbolic/deepseek-ai/DeepSeek-R1",
-    ]
-
-    for model in test_models:
-        assert model in model_data
-        model_info = model_data[model]
-        assert model_info["litellm_provider"] == "hyperbolic"
-        assert model_info["mode"] == "chat"
-        assert "max_tokens" in model_info
-        assert "input_cost_per_token" in model_info
-        assert "output_cost_per_token" in model_info
 
 
 def test_hyperbolic_supported_params():
diff --git a/tests/llm_translation/test_lambda_ai.py b/tests/llm_translation/test_lambda_ai.py
index 7ae18828d3f..b2fb72f8412 100644
--- a/tests/llm_translation/test_lambda_ai.py
+++ b/tests/llm_translation/test_lambda_ai.py
@@ -8,7 +8,6 @@ from unittest import mock
 import pytest
 
 import litellm
-from litellm import completion
 from litellm.llms.lambda_ai.chat.transformation import LambdaAIChatConfig
 
 
@@ -103,46 +102,6 @@ async def test_lambda_ai_completion_call():
             raise
 
 
-def test_lambda_ai_models_configuration():
-    """Test that Lambda AI models are configured correctly"""
-    from litellm import get_model_info
-
-    # Reload model cost map to pick up local changes
-    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-
-    # Clear and repopulate lambda_ai_models list after reloading model_cost
-    litellm.lambda_ai_models = set()
-    litellm.add_known_models()
-
-    # Some Lambda AI models to test
-    lambda_ai_models = [
-        "lambda_ai/deepseek-llama3.3-70b",
-        "lambda_ai/hermes3-8b",
-        "lambda_ai/llama3.1-8b-instruct",
-        "lambda_ai/llama3.2-11b-vision-instruct",
-        "lambda_ai/qwen25-coder-32b-instruct",
-    ]
-
-    for model in lambda_ai_models:
-        model_info = get_model_info(model)
-        assert model_info is not None, f"Model info not found for {model}"
-        assert (
-            model_info.get("litellm_provider") == "lambda_ai"
-        ), f"{model} should have lambda_ai as provider"
-        assert model_info.get("mode") == "chat", f"{model} should be in chat mode"
-        assert (
-            model_info.get("supports_function_calling") is True
-        ), f"{model} should support function calling"
-        assert (
-            model_info.get("supports_system_messages") is True
-        ), f"{model} should support system messages"
-
-        # Check vision support for vision models
-        if "vision" in model:
-            assert (
-                model_info.get("supports_vision") is True
-            ), f"{model} should support vision"
 
 
 def test_lambda_ai_model_list_populated():
diff --git a/tests/llm_translation/test_morph.py b/tests/llm_translation/test_morph.py
index b91d1810d38..47ad3a1749b 100644
--- a/tests/llm_translation/test_morph.py
+++ b/tests/llm_translation/test_morph.py
@@ -68,22 +68,6 @@ def test_morph_in_provider_lists():
     )
 
 
-def test_morph_model_info():
-    """Test that morph models have correct configuration."""
-    import litellm
-
-    model_info = litellm.get_model_info("morph/morph-v3-large")
-
-    assert model_info["litellm_provider"] == "morph"
-    assert model_info["mode"] == "chat"
-    assert model_info["max_tokens"] == 16000
-    assert model_info["max_input_tokens"] == 16000
-    assert model_info["max_output_tokens"] == 16000
-    assert model_info["input_cost_per_token"] == 9e-07  # $0.9/1M tokens
-    assert model_info["output_cost_per_token"] == 1.9e-06  # $1.9/1M tokens
-    assert model_info["supports_function_calling"] is False
-    assert model_info["supports_vision"] is False
-    assert model_info["supports_system_messages"] is True
 
 
 def test_morph_supported_params():
diff --git a/tests/llm_translation/test_openai_o1.py b/tests/llm_translation/test_openai_o1.py
index e188a3af647..9de5d5d9431 100644
--- a/tests/llm_translation/test_openai_o1.py
+++ b/tests/llm_translation/test_openai_o1.py
@@ -1,15 +1,12 @@
-import json
 import os
-from datetime import datetime
-from unittest.mock import AsyncMock, patch, MagicMock
+from unittest.mock import patch
 
 
 
-import httpx
 import pytest
 
 import litellm
-from litellm import Choices, Message, ModelResponse
+from litellm import ModelResponse
 from base_llm_unit_tests import BaseLLMChatTest, BaseOSeriesModelsTest
 
 
@@ -74,7 +71,6 @@ async def test_o1_handle_tool_calling_optional_params(
     - max_tokens is translated to 'max_completion_tokens'
     - role 'system' is translated to 'user'
     """
-    from openai import AsyncOpenAI
     from litellm.utils import ProviderConfigManager
     from litellm.types.utils import LlmProviders
 
@@ -186,13 +182,6 @@ class TestOpenAIO3(BaseOSeriesModelsTest, BaseLLMChatTest):
         pass
 
 
-def test_o1_supports_vision():
-    """Test that o1 supports vision"""
-    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-    for k, v in litellm.model_cost.items():
-        if k.startswith("o1") and v.get("litellm_provider") == "openai":
-            assert v.get("supports_vision") is True, f"{k} does not support vision"
 
 
 def test_o3_reasoning_effort():
diff --git a/tests/llm_translation/test_v0.py b/tests/llm_translation/test_v0.py
index 95708dd855a..e96022e1e22 100644
--- a/tests/llm_translation/test_v0.py
+++ b/tests/llm_translation/test_v0.py
@@ -8,7 +8,6 @@ from unittest import mock
 import pytest
 
 import litellm
-from litellm import completion
 from litellm.llms.v0.chat.transformation import V0ChatConfig
 
 
@@ -111,33 +110,3 @@ def test_v0_supported_params():
     ]
 
     assert set(supported_params) == set(expected_params)
-
-
-def test_v0_models_configuration():
-    """Test that v0 models are configured correctly"""
-    from litellm import get_model_info
-
-    # Reload model cost map to pick up local changes
-    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-
-    # All v0 models
-    v0_models = ["v0/v0-1.0-md", "v0/v0-1.5-md", "v0/v0-1.5-lg"]
-
-    for model in v0_models:
-        model_info = get_model_info(model)
-        assert model_info is not None, f"Model info not found for {model}"
-        # All v0 models support vision (multimodal)
-        assert (
-            model_info.get("supports_vision") is True
-        ), f"{model} should support vision"
-        assert (
-            model_info.get("litellm_provider") == "v0"
-        ), f"{model} should have v0 as provider"
-        assert model_info.get("mode") == "chat", f"{model} should be in chat mode"
-        assert (
-            model_info.get("supports_function_calling") is True
-        ), f"{model} should support function calling"
-        assert (
-            model_info.get("supports_system_messages") is True
-        ), f"{model} should support system messages"
diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py
index 2de83778f1c..562ed240b9c 100644
--- a/tests/local_testing/test_get_model_info.py
+++ b/tests/local_testing/test_get_model_info.py
@@ -1,8 +1,6 @@
 # What is this?
 ## Unit testing for the 'get_model_info()' function
 import os
-import traceback
-import json
 
 
 from typing import List, Dict, Any
@@ -11,7 +9,7 @@ import pytest
 
 import litellm
 from litellm import get_model_info
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import MagicMock, patch
 
 
 def test_get_model_info_simple_model_name():
@@ -49,32 +47,12 @@ def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch):
     assert model_info["input_cost_per_token"] == 0.0
 
 
-def test_get_model_info_shows_correct_supports_vision():
-    info = litellm.get_model_info("gemini/gemini-2.0-flash")
-    print("info", info)
-    assert info["supports_vision"] is True
 
 
-def test_get_model_info_shows_assistant_prefill():
-    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-    info = litellm.get_model_info("deepseek/deepseek-chat")
-    print("info", info)
-    assert info.get("supports_assistant_prefill") is True
 
 
-def test_get_model_info_shows_supports_prompt_caching():
-    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-    info = litellm.get_model_info("deepseek/deepseek-chat")
-    print("info", info)
-    assert info.get("supports_prompt_caching") is True
 
 
-def test_get_model_info_finetuned_models():
-    info = litellm.get_model_info("ft:gpt-3.5-turbo:my-org:custom_suffix:id")
-    print("info", info)
-    assert info["input_cost_per_token"] == 0.000003
 
 
 def test_get_model_info_gemini_pro():
@@ -219,7 +197,7 @@ def test_model_info_bedrock_converse_enforcement(monkeypatch):
 def test_get_model_info_custom_provider():
     # Custom provider example copied from https://docs.litellm.ai/docs/providers/custom_llm_server:
     import litellm
-    from litellm import CustomLLM, completion, get_llm_provider
+    from litellm import CustomLLM, completion
 
     class MyCustomLLM(CustomLLM):
         def completion(self, *args, **kwargs) -> litellm.ModelResponse:
diff --git a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py
index ca25ee80c23..83ede898e49 100644
--- a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py
+++ b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py
@@ -1,6 +1,5 @@
 
 
-import litellm
 from litellm import LlmProviders
 from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
 from litellm.litellm_core_utils.get_llm_provider_logic import (
@@ -46,12 +45,6 @@ def test_xai_openai_compatible_provider_info():
     assert dynamic_api_key == "api-key"
 
 
-def test_xai_get_model_info_uses_xai_pricing_metadata():
-    model_info = litellm.get_model_info("xai/grok-3-mini")
-
-    assert model_info["litellm_provider"] == "xai"
-    assert model_info["key"] == "xai/grok-3-mini"
-    assert model_info["mode"] == "chat"
 
 
 def test_xai_validate_environment_reads_api_key(monkeypatch):
diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
index 2a44e77ce09..669c566f96b 100644
--- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
+++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
@@ -1,4 +1,3 @@
-import os
 from unittest.mock import MagicMock
 
 import httpx
@@ -38,24 +37,6 @@ class TestAzureMAIImageGeneration:
         assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro")
         assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-DS-R1")
 
-    def test_mai_flash_and_2e_model_pricing_in_cost_map(self, monkeypatch):
-        monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-        litellm.model_cost = litellm.get_model_cost_map(url="")
-
-        flash_info = litellm.get_model_info(
-            model="azure_ai/MAI-Image-2.5-Flash",
-            custom_llm_provider="azure_ai",
-        )
-        assert flash_info["input_cost_per_token"] == 1.75e-06
-        assert flash_info["input_cost_per_image_token"] == 1.75e-06
-        assert flash_info["output_cost_per_image_token"] == 3.3e-05
-
-        image_2e_info = litellm.get_model_info(
-            model="azure_ai/MAI-Image-2e",
-            custom_llm_provider="azure_ai",
-        )
-        assert image_2e_info["input_cost_per_token"] == 5e-06
-        assert image_2e_info["output_cost_per_image_token"] == 1.95e-05
 
     def test_get_mai_image_generation_url(self):
         url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url(
diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py
index f3618572622..d9b948e212a 100644
--- a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py
+++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py
@@ -12,111 +12,6 @@ from importlib.resources import files
 
 import pytest
 
-FW_MODELS = {
-    "azure_ai/FW-Kimi-K2.5": {
-        "input_cost_per_token": 6.6e-07,
-        "output_cost_per_token": 3.3e-06,
-        "cache_read_input_token_cost": 1.1e-07,
-        "max_input_tokens": 262144,
-        "max_output_tokens": 262144,
-        "supports_vision": True,
-    },
-    "azure_ai/FW-Kimi-K2.6": {
-        "input_cost_per_token": 1.045e-06,
-        "output_cost_per_token": 4.4e-06,
-        "cache_read_input_token_cost": 1.76e-07,
-        "max_input_tokens": 262144,
-        "max_output_tokens": 262144,
-        "supports_vision": True,
-    },
-    "azure_ai/FW-Kimi-K2.7-Code": {
-        "input_cost_per_token": 1.05e-06,
-        "output_cost_per_token": 4.4e-06,
-        "cache_read_input_token_cost": 2.1e-07,
-        "max_input_tokens": 262144,
-        "max_output_tokens": 262144,
-        "supports_vision": True,
-    },
-    "azure_ai/FW-Kimi-K3": {
-        "input_cost_per_token": 3.3e-06,
-        "output_cost_per_token": 1.65e-05,
-        "cache_read_input_token_cost": 3.3e-07,
-        "max_input_tokens": 1048576,
-        "max_output_tokens": 131072,
-        "supports_vision": True,
-    },
-    "azure_ai/FW-Inkling": {
-        "input_cost_per_token": 1e-06,
-        "output_cost_per_token": 4.05e-06,
-        "cache_read_input_token_cost": 1.7e-07,
-        "max_input_tokens": 1048576,
-        "max_output_tokens": 1048576,
-    },
-    "azure_ai/FW-DeepSeek-V3.2": {
-        "input_cost_per_token": 6.2e-07,
-        "output_cost_per_token": 1.85e-06,
-        "cache_read_input_token_cost": 3.1e-07,
-        "max_input_tokens": 163840,
-        "max_output_tokens": 163840,
-    },
-    "azure_ai/FW-DeepSeek-V4-Pro": {
-        "input_cost_per_token": 1.925e-06,
-        "output_cost_per_token": 3.828e-06,
-        "cache_read_input_token_cost": 1.65e-07,
-        "max_input_tokens": 1000000,
-        "max_output_tokens": 384000,
-    },
-    "azure_ai/FW-MiniMax-M3": {
-        "input_cost_per_token": 3.3e-07,
-        "output_cost_per_token": 1.32e-06,
-        "cache_read_input_token_cost": 6.6e-08,
-        "max_input_tokens": 512000,
-        "max_output_tokens": 512000,
-        "supports_vision": True,
-    },
-    "azure_ai/FW-MiniMax-M2.5": {
-        "input_cost_per_token": 3.3e-07,
-        "output_cost_per_token": 1.32e-06,
-        "cache_read_input_token_cost": 3.3e-08,
-        "max_input_tokens": 1000000,
-        "max_output_tokens": 1000000,
-    },
-    "azure_ai/FW-Nemotron-3-Ultra-NVFP4": {
-        "input_cost_per_token": 6e-07,
-        "output_cost_per_token": 2.4e-06,
-        "cache_read_input_token_cost": 1.19e-07,
-        "max_input_tokens": 262144,
-        "max_output_tokens": 262144,
-    },
-    "azure_ai/FW-GLM-5.2-Fast": {
-        "input_cost_per_token": 2.1e-06,
-        "output_cost_per_token": 6.6e-06,
-        "cache_read_input_token_cost": 2.1e-07,
-        "max_input_tokens": 1048576,
-        "max_output_tokens": 131072,
-    },
-    "azure_ai/FW-GLM-5.2": {
-        "input_cost_per_token": 1.54e-06,
-        "output_cost_per_token": 4.84e-06,
-        "cache_read_input_token_cost": 1.5e-07,
-        "max_input_tokens": 1048576,
-        "max_output_tokens": 131072,
-    },
-    "azure_ai/FW-GLM-5.1": {
-        "input_cost_per_token": 1.54e-06,
-        "output_cost_per_token": 4.84e-06,
-        "cache_read_input_token_cost": 2.86e-07,
-        "max_input_tokens": 202800,
-        "max_output_tokens": 131072,
-    },
-    "azure_ai/FW-GLM-5": {
-        "input_cost_per_token": 1.1e-06,
-        "output_cost_per_token": 3.52e-06,
-        "cache_read_input_token_cost": 2.2e-07,
-        "max_input_tokens": 200000,
-        "max_output_tokens": 128000,
-    },
-}
 
 
 @pytest.fixture(scope="module")
@@ -144,26 +39,6 @@ def use_local_model_cost_map():
         monkeypatch.undo()
 
 
-@pytest.mark.parametrize("model_key,expected", list(FW_MODELS.items()))
-def test_azure_ai_fw_model_info(use_local_model_cost_map, model_key, expected):
-    model_info = use_local_model_cost_map.get_model_info(model=model_key)
-
-    assert model_info["litellm_provider"] == "azure_ai"
-    assert model_info["mode"] == "chat"
-    assert model_info["input_cost_per_token"] == pytest.approx(expected["input_cost_per_token"])
-    assert model_info["output_cost_per_token"] == pytest.approx(expected["output_cost_per_token"])
-    assert model_info["cache_read_input_token_cost"] == pytest.approx(
-        expected["cache_read_input_token_cost"]
-    )
-    assert model_info["max_input_tokens"] == expected["max_input_tokens"]
-    assert model_info["max_output_tokens"] == expected["max_output_tokens"]
-    assert model_info["max_tokens"] == expected["max_output_tokens"]
-    assert model_info["supports_function_calling"] is True
-    assert model_info["supports_reasoning"] is True
-    assert model_info["supports_tool_choice"] is True
-    assert model_info["supports_prompt_caching"] is True
-    if expected.get("supports_vision"):
-        assert model_info["supports_vision"] is True
 
 
 @pytest.mark.parametrize(
@@ -197,20 +72,6 @@ def test_azure_ai_fw_cost_per_token(
     assert completion_cost == pytest.approx(expected_completion)
 
 
-def test_azure_ai_fw_nemotron_lightning_model_info(use_local_model_cost_map):
-    model_info = use_local_model_cost_map.get_model_info(model="azure_ai/FW-Nemotron-Lightning-3.5-30B-A3B")
-
-    assert model_info["litellm_provider"] == "azure_ai"
-    assert model_info["mode"] == "chat"
-    assert model_info["input_cost_per_token"] == pytest.approx(6e-08)
-    assert model_info["output_cost_per_token"] == pytest.approx(2.2e-07)
-    assert model_info["cache_read_input_token_cost"] == pytest.approx(1e-08)
-    assert model_info["max_input_tokens"] == 262144
-    assert model_info["supports_function_calling"] is True
-    assert model_info["supports_reasoning"] is True
-    assert model_info["supports_tool_choice"] is True
-    assert model_info["supports_prompt_caching"] is True
-    assert model_info["supports_vision"] is False
 
 
 def test_azure_ai_fw_nemotron_lightning_supports_tool_choice(use_local_model_cost_map):
diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py
index 812b9288ca8..18bdf60e9a0 100644
--- a/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py
+++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py
@@ -33,31 +33,8 @@ def use_local_model_cost_map():
         monkeypatch.undo()
 
 
-def test_azure_ai_kimi_k26_model_info(use_local_model_cost_map):
-    model_info = use_local_model_cost_map.get_model_info(model="azure_ai/kimi-k2.6")
-
-    assert model_info["litellm_provider"] == "azure_ai"
-    assert model_info["mode"] == "chat"
-    assert model_info["max_input_tokens"] == 262144
-    assert model_info["max_output_tokens"] == 262144
-    assert model_info["max_tokens"] == 262144
-    assert model_info["input_cost_per_token"] == pytest.approx(9.5e-07)
-    assert model_info["output_cost_per_token"] == pytest.approx(4e-06)
-    assert model_info["supports_function_calling"] is True
-    assert model_info["supports_reasoning"] is True
-    assert model_info["supports_tool_choice"] is True
-    assert model_info["supports_vision"] is True
 
 
-def test_azure_ai_kimi_k26_raw_model_cost_entry(use_local_model_cost_map):
-    model_info = use_local_model_cost_map.model_cost["azure_ai/kimi-k2.6"]
-
-    assert model_info["supported_modalities"] == ["text", "image"]
-    assert model_info["supported_output_modalities"] == ["text"]
-    assert model_info["supports_function_calling"] is True
-    assert model_info["supports_reasoning"] is True
-    assert model_info["supports_tool_choice"] is True
-    assert model_info["supports_vision"] is True
 
 
 def test_azure_ai_kimi_k26_cost_per_token(use_local_model_cost_map):
diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py
index 5f12ae8566c..c697bcb24b0 100644
--- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py
+++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py
@@ -1,8 +1,5 @@
 """Test Bedrock cross-region inference profile model mapping"""
 
-import json
-from functools import lru_cache
-from pathlib import Path
 from typing import NamedTuple
 
 import pytest
@@ -102,11 +99,6 @@ GPT_5_6_PROFILES = [
 ]
 
 
-@lru_cache(maxsize=1)
-def _packaged_cost_map():
-    """The map litellm actually resolves against, for fields ModelInfoBase drops."""
-    path = Path(litellm.__file__).parent / "model_prices_and_context_window_backup.json"
-    return json.loads(path.read_text())
 
 
 def _bedrock_response(model, usage):
@@ -126,15 +118,6 @@ def _bedrock_response(model, usage):
     )
 
 
-def test_bedrock_cross_region_inference_profile_mapping():
-    """Test that bedrock cross-region inference profile model is mapped"""
-    model = "bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0"
-
-    model_info = _get_model_info_helper(model=model, custom_llm_provider="bedrock")
-
-    assert model_info is not None
-    assert model_info["litellm_provider"] == "bedrock"
-    assert model_info["input_cost_per_token"] == 8e-07
 
 
 def test_proxy_cost_calculation_scenario():
@@ -176,36 +159,6 @@ def test_bedrock_gpt_5_6_profiles_route_to_converse(profile, local_model_cost_ma
     assert BedrockModelInfo.get_bedrock_route(f"bedrock/{profile.model_id}") == "converse"
 
 
-@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id)
-def test_bedrock_gpt_5_6_published_rates(profile, local_model_cost_map):
-    """Geo and Global profiles carry their own published rates, per context tier."""
-    model_info = _get_model_info_helper(
-        model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock"
-    )
-
-    assert model_info["litellm_provider"] == "bedrock_converse"
-    assert model_info["mode"] == "chat"
-    assert model_info["max_input_tokens"] == 1000000
-    assert model_info["input_cost_per_token"] == profile.input_cost
-    assert (
-        model_info["input_cost_per_token_above_272k_tokens"]
-        == profile.input_cost_above_272k
-    )
-    assert model_info["output_cost_per_token"] == profile.output_cost
-    assert (
-        model_info["output_cost_per_token_above_272k_tokens"]
-        == profile.output_cost_above_272k
-    )
-    assert model_info["cache_creation_input_token_cost"] == profile.cache_write
-    assert (
-        model_info["cache_creation_input_token_cost_above_272k_tokens"]
-        == profile.cache_write_above_272k
-    )
-    assert model_info["cache_read_input_token_cost"] == profile.cache_read
-    assert (
-        model_info["cache_read_input_token_cost_above_272k_tokens"]
-        == profile.cache_read_above_272k
-    )
 
 
 def test_bedrock_gpt_5_6_above_272k_tier_applies_to_cost(local_model_cost_map):
@@ -267,29 +220,6 @@ def test_bedrock_gpt_5_6_bills_cache_write_tokens(local_model_cost_map):
     assert cost == pytest.approx(expected, rel=1e-9)
 
 
-@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id)
-def test_bedrock_gpt_5_6_advertises_only_converse_supported_features(
-    profile, local_model_cost_map
-):
-    model_info = _get_model_info_helper(
-        model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock"
-    )
-
-    assert model_info["supports_function_calling"] is True
-    assert model_info["supports_tool_choice"] is True
-    assert model_info["supports_vision"] is True
-
-    # Bedrock rejects an explicit cachePoint block for these models, so the flag that
-    # offers caller-driven caching stays off even though the cache rates are declared.
-    assert not model_info.get("supports_prompt_caching")
-
-    # ModelInfoBase drops these two, so they are read from the map litellm resolves.
-    raw = _packaged_cost_map()[profile.model_id]
-    assert raw["supported_modalities"] == ["text", "image"]
-    assert raw["supported_output_modalities"] == ["text"]
-    # No bedrock_converse entry declares supported_endpoints; these models are reachable
-    # on chat completions and on the Responses API without it.
-    assert "supported_endpoints" not in raw
 
 
 @pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id)
diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
index 1f23d39c631..5994de28ba8 100644
--- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
+++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
@@ -8,9 +8,7 @@ gate, the URL construction for both paths, and the shared Bearer auth.
 """
 
 import copy
-import json
 import logging
-from pathlib import Path
 
 import pytest
 from botocore.exceptions import (
@@ -1777,53 +1775,9 @@ class TestBedrockMantleResponsesSigV4:
 
 
 class TestBedrockMantleResponsesPricing:
-    def test_gpt_5_5_pricing_and_mode(self, local_cost_map):
-        info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5")
-        assert info["mode"] == "responses"
-        assert info["input_cost_per_token"] == pytest.approx(5.5e-06)
-        assert info["output_cost_per_token"] == pytest.approx(3.3e-05)
-        assert info["cache_read_input_token_cost"] == pytest.approx(5.5e-07)
-        assert info["max_input_tokens"] == 1050000
 
-    def test_gpt_5_4_pricing_and_mode(self, local_cost_map):
-        info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.4")
-        assert info["mode"] == "responses"
-        assert info["input_cost_per_token"] == pytest.approx(2.75e-06)
-        assert info["output_cost_per_token"] == pytest.approx(1.65e-05)
-        assert info["cache_read_input_token_cost"] == pytest.approx(2.75e-07)
-        assert info["max_input_tokens"] == 1050000
 
-    def test_gpt_5_6_cyber_pricing_and_mode(self, local_cost_map):
-        info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.6-cyber")
-        assert info["mode"] == "responses"
-        assert info["input_cost_per_token"] == pytest.approx(1.375e-05)
-        assert info["cache_creation_input_token_cost"] == pytest.approx(1.71875e-05)
-        assert info["cache_read_input_token_cost"] == pytest.approx(1.375e-06)
-        assert info["output_cost_per_token"] == pytest.approx(8.25e-05)
-        assert info["max_input_tokens"] == 272000
 
-    @pytest.mark.parametrize(
-        "model, input_cost, cache_creation_cost, cache_read_cost, output_cost",
-        [
-            ("openai.gpt-5.6-sol", 5.5e-06, 6.875e-06, 5.5e-07, 3.3e-05),
-            ("openai.gpt-5.6-terra", 2.2e-06, 2.75e-06, 2.2e-07, 1.32e-05),
-            ("openai.gpt-5.6-luna", 2.2e-07, 2.75e-07, 2.2e-08, 1.32e-06),
-        ],
-    )
-    def test_gpt_5_6_pricing_and_mode(
-        self, local_cost_map, model, input_cost, cache_creation_cost, cache_read_cost, output_cost
-    ):
-        info = litellm.get_model_info(f"bedrock_mantle/{model}")
-        assert info["mode"] == "responses"
-        assert info["input_cost_per_token"] == pytest.approx(input_cost)
-        assert info["cache_creation_input_token_cost"] == pytest.approx(cache_creation_cost)
-        assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost)
-        assert info["output_cost_per_token"] == pytest.approx(output_cost)
-        assert info["max_input_tokens"] == 1050000
-        assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2)
-        assert info["cache_creation_input_token_cost_above_272k_tokens"] == pytest.approx(cache_creation_cost * 2)
-        assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2)
-        assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(output_cost * 1.5)
 
     @pytest.mark.parametrize(
         "model, input_cost, output_cost",
@@ -1861,58 +1815,3 @@ class TestBedrockMantleResponsesPricing:
     def test_models_registered(self, local_cost_map):
         assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models
         assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models
-
-
-def _repo_cost_map(map_name: str) -> dict[str, dict[str, object]]:
-    repo_root = Path(__file__).resolve().parents[4]
-    paths = {
-        "root": repo_root / "model_prices_and_context_window.json",
-        "bundled_backup": repo_root / "litellm" / "model_prices_and_context_window_backup.json",
-    }
-    return json.loads(paths[map_name].read_text())
-
-
-class TestMantleGptRegistryEntries:
-    """Locks the OpenAI GPT entries to Bedrock Mantle's live behavior.
-
-    Mantle enforces a 1,050,000-token prompt maximum for gpt-5.6 sol/terra/luna
-    and for gpt-5.5 and gpt-5.4 (oversize requests 400 with "prompt tokens (N)
-    exceed model maximum (1050000)", and a 1,030,590-token request completes
-    on every one of them), while the AWS model cards still quote 272K for
-    gpt-5.5 and gpt-5.4. mode must stay "responses": Mantle's native
-    /v1/chat/completions rejects function tools unless reasoning_effort is
-    "none", so chat traffic has to keep bridging to the Responses API
-    (see the responses_api_bridge tests above).
-    """
-
-    @pytest.mark.parametrize("map_name", ("root", "bundled_backup"))
-    @pytest.mark.parametrize(
-        "key",
-        (
-            "bedrock_mantle/openai.gpt-5.6-sol",
-            "bedrock_mantle/openai.gpt-5.6-terra",
-            "bedrock_mantle/openai.gpt-5.6-luna",
-        ),
-    )
-    def test_entry_matches_mantle_enforced_limits(self, map_name, key):
-        entry = _repo_cost_map(map_name)[key]
-        assert entry["max_input_tokens"] == 1050000
-        assert entry["max_output_tokens"] == 128000
-        assert entry["mode"] == "responses"
-        assert entry["use_openai_responses_path"] is True
-        assert entry["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"]
-
-    @pytest.mark.parametrize("map_name", ("root", "bundled_backup"))
-    @pytest.mark.parametrize(
-        "key",
-        (
-            "bedrock_mantle/openai.gpt-5.5",
-            "bedrock_mantle/openai.gpt-5.4",
-        ),
-    )
-    def test_gpt_55_and_54_entries_match_mantle_enforced_limits(self, map_name, key):
-        entry = _repo_cost_map(map_name)[key]
-        assert entry["max_input_tokens"] == 1050000
-        assert entry["max_output_tokens"] == 128000
-        assert entry["mode"] == "responses"
-        assert entry["use_openai_responses_path"] is True
diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py
index b88c27e64b9..e370cb22ce7 100644
--- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py
+++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py
@@ -684,39 +684,8 @@ class TestBedrockMantleProviderResolution:
 class TestBedrockMantlePricing:
     """Tests that verify Bedrock Mantle uses correct AWS Bedrock pricing, not OpenAI pricing."""
 
-    def test_gpt_oss_120b_pricing(self, monkeypatch):
-        monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
-        litellm.add_known_models()
-        info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b")
-        # Bedrock pricing: $0.15/M input, $0.60/M output
-        assert info["input_cost_per_token"] == pytest.approx(1.5e-7)
-        assert info["output_cost_per_token"] == pytest.approx(6e-7)
 
-    def test_gpt_oss_20b_pricing(self, monkeypatch):
-        monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
-        litellm.add_known_models()
-        info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-20b")
-        # Bedrock pricing: $0.075/M input, $0.30/M output
-        assert info["input_cost_per_token"] == pytest.approx(7.5e-8)
-        assert info["output_cost_per_token"] == pytest.approx(3e-7)
 
-    def test_pricing_significantly_cheaper_than_openai_native(self, monkeypatch):
-        """
-        Verify Bedrock Mantle pricing is cheaper than OpenAI's direct API pricing.
-        This is the core issue the provider addition fixes — previously users were being
-        billed at OpenAI rates instead of the cheaper Bedrock rates.
-        """
-        monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
-        litellm.add_known_models()
-        bedrock_info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b")
-        # OpenAI direct pricing for gpt-oss-120b is ~$0.039/M input, $0.190/M output
-        # Bedrock should be cheaper at $0.15/M input and $0.60/M output... wait
-        # Actually, Bedrock ADDS value not reduces cost vs OpenAI direct for these models.
-        # The key fix is that we now use Bedrock-specific prices instead of mapping to
-        # some unrelated OpenAI model (like gpt-4) pricing.
-        # Just validate the pricing is as expected from AWS docs.
-        assert bedrock_info["input_cost_per_token"] == pytest.approx(1.5e-7)
-        assert bedrock_info["output_cost_per_token"] == pytest.approx(6e-7)
 
     def test_safeguard_models_have_larger_output_tokens(self, monkeypatch):
         monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
@@ -727,48 +696,9 @@ class TestBedrockMantlePricing:
         )
         assert info_safeguard["max_output_tokens"] > info_120b["max_output_tokens"]
 
-    def test_reasoning_support(self, monkeypatch):
-        monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
-        litellm.add_known_models()
-        info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b")
-        assert info.get("supports_reasoning") is True
-
-    def test_context_window(self, monkeypatch):
-        monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
-        litellm.add_known_models()
-        info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b")
-        assert info["max_input_tokens"] == 131072
 
 
-@pytest.mark.parametrize(
-    "model_id,input_cost,output_cost,max_tokens",
-    [
-        ("google.gemma-4-31b", 1.4e-07, 4e-07, 256000),
-        ("google.gemma-4-26b-a4b", 1.3e-07, 4e-07, 256000),
-        ("google.gemma-4-e2b", 4e-08, 8e-08, 128000),
-    ],
-)
-def test_gemma_4_bedrock_mantle_model_metadata(
-    local_cost_map, model_id, input_cost, output_cost, max_tokens
-):
-    full_model_name = f"bedrock_mantle/{model_id}"
-    info = litellm.get_model_info(full_model_name)
 
-    assert info["mode"] == "chat"
-    assert info["input_cost_per_token"] == pytest.approx(input_cost)
-    assert info["output_cost_per_token"] == pytest.approx(output_cost)
-    assert info["max_input_tokens"] == max_tokens
-    assert info["max_output_tokens"] == max_tokens
-    assert info["supports_function_calling"] is True
-    assert info["supports_reasoning"] is True
-    assert info["supports_tool_choice"] is True
-    assert info["supports_vision"] is True
-    assert (
-        litellm.supports_parallel_function_calling(
-            model=full_model_name, custom_llm_provider="bedrock_mantle"
-        )
-        is False
-    )
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
index e6fe01be4ba..a79baef5ee5 100644
--- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
+++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
@@ -4,7 +4,7 @@ from unittest.mock import MagicMock, patch
 import pytest
 
 import litellm
-from litellm import get_model_info, supports_reasoning, supports_vision
+from litellm import supports_reasoning, supports_vision
 from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
 from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig
 from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
@@ -16,15 +16,6 @@ from litellm.types.utils import (
 )
 
 
-@pytest.fixture(autouse=True)
-def force_local_model_cost(monkeypatch):
-    """Force local model cost map usage for all tests in this file."""
-    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-    # Refresh model_cost from local map
-    import litellm
-    from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
-
-    litellm.model_cost = get_model_cost_map(url=litellm.model_cost_map_url)
 
 
 def test_validate_environment_sets_session_affinity_from_litellm_session_id():
@@ -404,13 +395,6 @@ def test_get_supported_openai_params_parallel_tool_calls_without_tool_choice(
     assert "tool_choice" not in supported_params
 
 
-def test_get_model_info_respects_explicit_fireworks_capabilities():
-    """Test that get_model_info preserves explicit capability flags from the model map."""
-    model_info = get_model_info("fireworks_ai/accounts/fireworks/models/glm-5p1")
-
-    assert model_info["supports_function_calling"] is True
-    assert model_info["supports_reasoning"] is True
-    assert model_info["supports_tool_choice"] is True
 
 
 def test_get_provider_info_omits_false_supports_reasoning(monkeypatch):
diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py
index 5641439aa54..ba40f02ddc1 100644
--- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py
+++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py
@@ -56,15 +56,6 @@ def use_local_model_cost_map():
         monkeypatch.undo()
 
 
-@pytest.mark.parametrize("alias", KIMI_ALIASES)
-def test_fireworks_kimi_raw_cost_entry_limits(use_local_model_cost_map, alias):
-    entry = use_local_model_cost_map.model_cost[alias]
-
-    assert entry["litellm_provider"] == "fireworks_ai"
-    assert entry["max_input_tokens"] == CONTEXT_WINDOW
-    assert entry["max_output_tokens"] == OUTPUT_LIMIT
-    assert entry["max_tokens"] == OUTPUT_LIMIT
-    assert entry["max_output_tokens"] < entry["max_input_tokens"]
 
 
 @pytest.mark.parametrize("alias", KIMI_ALIASES)
diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py
index 3d8200bc474..8295cf72524 100644
--- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py
+++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py
@@ -1,13 +1,11 @@
 import json
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import MagicMock
 
-import httpx
 import pytest
 
 
 import litellm
 from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
-from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents
 
 
 def test_gemini_realtime_transformation_session_created():
@@ -308,18 +306,6 @@ def test_gemini_realtime_transformation_generation_complete():
     assert contains_audio_done_event, "Expected audio done event"
 
 
-def test_gemini_3_1_flash_live_preview_model_cost_map_entry():
-    for key in (
-        "gemini-3.1-flash-live-preview",
-        "gemini/gemini-3.1-flash-live-preview",
-    ):
-        assert key in litellm.model_cost
-        info = litellm.model_cost[key]
-        assert "/v1/realtime" in info.get("supported_endpoints", [])
-        assert info.get("max_input_tokens") == 131072
-        assert info.get("max_output_tokens") == 65536
-        assert "video" in info.get("supported_modalities", [])
-        assert info.get("supports_function_calling") is True
 
 
 def test_gemini_realtime_tool_call_transformation():
@@ -1845,17 +1831,6 @@ def test_is_audio_only_live_model_uses_cost_map(model, expected, patch_gemini_au
     assert GeminiRealtimeConfig._is_audio_only_live_model(model) == expected
 
 
-def test_gemini_live_native_audio_entry_is_vertex_only():
-    import json
-    from pathlib import Path
-    from typing import Final
-
-    catalog_path: Final = Path(__file__).parents[5] / "model_prices_and_context_window.json"
-    catalog: Final = json.loads(catalog_path.read_text())
-    vertex_key: Final = "gemini-live-2.5-flash-native-audio"
-    assert catalog[vertex_key]["litellm_provider"] == "vertex_ai-language-models"
-    assert catalog[vertex_key].get("gemini_native_audio") is True
-    assert "gemini/gemini-live-2.5-flash-native-audio" not in catalog, "the Gemini API does not serve this model"
 
 
 def test_is_setup_message_and_is_content_message():
diff --git a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py
index cff3c6be940..fff352a2f6c 100644
--- a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py
+++ b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py
@@ -231,24 +231,6 @@ def test_inception_in_provider_lists():
     assert "https://api.inceptionlabs.ai/v1" in litellm.openai_compatible_endpoints
 
 
-def test_inception_model_configuration(monkeypatch):
-    from litellm import get_model_info
-
-    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-    litellm.inception_models = set()
-    litellm.add_known_models()
-
-    info = get_model_info("inception/mercury-2")
-    assert info.get("litellm_provider") == "inception"
-    assert info.get("mode") == "chat"
-    assert info.get("max_input_tokens") == 128000
-    assert info.get("input_cost_per_token") == 2.5e-07
-    assert info.get("output_cost_per_token") == 7.5e-07
-    assert info.get("cache_read_input_token_cost") == 2.5e-08
-    assert info.get("supports_function_calling") is True
-    assert info.get("supports_tool_choice") is True
-    assert info.get("supports_response_schema") is True
 
 
 def test_inception_model_list_populated(monkeypatch):
diff --git a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py
index 62688a13c35..347cfe4cfc5 100644
--- a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py
+++ b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py
@@ -143,22 +143,6 @@ async def test_inception_fim_async():
     assert r.choices[0].text == "a + b"
 
 
-def test_inception_fim_model_configuration(monkeypatch):
-    from litellm import get_model_info
-
-    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-    litellm.text_completion_inception_models = set()
-    litellm.add_known_models()
-
-    assert (
-        "text-completion-inception/mercury-edit-2"
-        in litellm.text_completion_inception_models
-    )
-    info = get_model_info("text-completion-inception/mercury-edit-2")
-    assert info.get("litellm_provider") == "text-completion-inception"
-    assert info.get("mode") == "completion"
-    assert info.get("max_input_tokens") == 32000
 
 
 def test_inception_fim_targets_fim_endpoint():
diff --git a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py
index 8c8bea00dea..2d6751fca63 100644
--- a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py
+++ b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py
@@ -708,37 +708,10 @@ class TestKimiK26ModelRegistry:
         """Load directly from the bundled backup so tests don't depend on remote fetch."""
         return GetModelCostMap.load_local_model_cost_map()
 
-    def test_kimi_k26_in_model_cost_map(self, model_cost_map):
-        """kimi-k2.6 should be present in the model cost map."""
-        assert "moonshot/kimi-k2.6" in model_cost_map, "moonshot/kimi-k2.6 not found in model_cost"
 
-    def test_kimi_k26_pricing(self, model_cost_map):
-        """kimi-k2.6 pricing should match official Kimi API rates."""
-        model_info = model_cost_map["moonshot/kimi-k2.6"]
-        assert model_info["input_cost_per_token"] == pytest.approx(9.5e-07)
-        assert model_info["output_cost_per_token"] == pytest.approx(4e-06)
-        assert model_info["cache_read_input_token_cost"] == pytest.approx(1.6e-07)
 
-    def test_kimi_k26_context_window(self, model_cost_map):
-        """kimi-k2.6 should have a 256K (262144 token) context window."""
-        model_info = model_cost_map["moonshot/kimi-k2.6"]
-        assert model_info["max_input_tokens"] == 262144
-        assert model_info["max_output_tokens"] == 262144
-        assert model_info["max_tokens"] == 262144
 
-    def test_kimi_k26_capabilities(self, model_cost_map):
-        """kimi-k2.6 should support function calling, vision, video input, tool choice, and reasoning."""
-        model_info = model_cost_map["moonshot/kimi-k2.6"]
-        assert model_info.get("supports_function_calling") is True
-        assert model_info.get("supports_tool_choice") is True
-        assert model_info.get("supports_vision") is True
-        assert model_info.get("supports_video_input") is True
-        assert model_info.get("supports_reasoning") is True
 
-    def test_kimi_k26_provider(self, model_cost_map):
-        """kimi-k2.6 should be assigned to the moonshot provider."""
-        model_info = model_cost_map["moonshot/kimi-k2.6"]
-        assert model_info["litellm_provider"] == "moonshot"
 
 
 class TestMoonshotResponseSchemaSupport:
@@ -762,9 +735,6 @@ class TestMoonshotResponseSchemaSupport:
     def model_cost_map(self):
         return GetModelCostMap.load_local_model_cost_map()
 
-    @pytest.mark.parametrize("model", LIVE_MODELS)
-    def test_live_model_supports_response_schema(self, model, model_cost_map):
-        assert model_cost_map[model].get("supports_response_schema") is True
 
     def test_supports_response_schema_utility_reports_true(self, model_cost_map, monkeypatch):
         monkeypatch.setattr(litellm, "model_cost", model_cost_map)
diff --git a/tests/test_litellm/llms/openai_like/test_cognition_provider.py b/tests/test_litellm/llms/openai_like/test_cognition_provider.py
index 5c71b60e08a..d392abc6cc5 100644
--- a/tests/test_litellm/llms/openai_like/test_cognition_provider.py
+++ b/tests/test_litellm/llms/openai_like/test_cognition_provider.py
@@ -110,22 +110,6 @@ class TestCognitionProviderIdentity:
 
 
 class TestCognitionCostTracking:
-    @pytest.mark.parametrize(
-        "model, input_cost, output_cost, cache_read_cost",
-        [
-            ("cognition/swe-1.6", 5e-07, 2.5e-06, 2e-07),
-            ("cognition/swe-1.7", 5e-07, 2.5e-06, 2e-07),
-            ("cognition/swe-1.7-lightning", 2.5e-06, 1.25e-05, 1e-06),
-        ],
-    )
-    def test_cost_map_entries(self, model: str, input_cost: float, output_cost: float, cache_read_cost: float):
-        info = litellm.get_model_info(model=model)
-
-        assert info["litellm_provider"] == "cognition"
-        assert info["mode"] == "chat"
-        assert info["input_cost_per_token"] == input_cost
-        assert info["output_cost_per_token"] == output_cost
-        assert info["cache_read_input_token_cost"] == cache_read_cost
 
     @pytest.mark.parametrize(
         "model, expected_prompt_cost, expected_completion_cost",
diff --git a/tests/test_litellm/llms/openai_like/test_json_providers.py b/tests/test_litellm/llms/openai_like/test_json_providers.py
index c8743e1809d..fb5d28b8d3b 100644
--- a/tests/test_litellm/llms/openai_like/test_json_providers.py
+++ b/tests/test_litellm/llms/openai_like/test_json_providers.py
@@ -2,10 +2,9 @@
 Tests for JSON-based provider configuration system.
 """
 
-import json
 import os
 import sys
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 try:
     import pytest
@@ -318,24 +317,6 @@ class TestDarkbloom:
         assert config is not None
         assert config.custom_llm_provider == "darkbloom"
 
-    def test_darkbloom_model_cost_map(self):
-        with open(
-            os.path.join(workspace_path, "model_prices_and_context_window.json")
-        ) as f:
-            model_cost = json.load(f)
-
-        expected_models = {
-            "darkbloom/gemma-4-26b": (3e-08, 1.65e-07),
-            "darkbloom/gpt-oss-20b": (1.45e-08, 7e-08),
-        }
-        for model, (input_cost, output_cost) in expected_models.items():
-            assert model in model_cost
-            assert model_cost[model]["litellm_provider"] == "darkbloom"
-            assert model_cost[model]["max_output_tokens"] == 32768
-            assert model_cost[model]["supports_function_calling"] is True
-            assert model_cost[model]["supports_tool_choice"] is True
-            assert model_cost[model]["input_cost_per_token"] == input_cost
-            assert model_cost[model]["output_cost_per_token"] == output_cost
 
 
 class TestPublicAIIntegration:
diff --git a/tests/test_litellm/llms/openai_like/test_libertai_provider.py b/tests/test_litellm/llms/openai_like/test_libertai_provider.py
index fdbe3046e9b..dc7d5d18f36 100644
--- a/tests/test_litellm/llms/openai_like/test_libertai_provider.py
+++ b/tests/test_litellm/llms/openai_like/test_libertai_provider.py
@@ -59,22 +59,6 @@ class TestLibertAIProviderConfig:
         assert api_base == "https://custom.example.com/v1"
         assert api_key == "sk-test"
 
-    def test_libertai_model_cost_map(self):
-        """Test that libertai models are present in the model cost map"""
-        model_cost = litellm.model_cost
-
-        assert "libertai/qwen3.6-27b" in model_cost
-        info = model_cost["libertai/qwen3.6-27b"]
-        assert info["litellm_provider"] == "libertai"
-        assert info["mode"] == "chat"
-        assert info["max_input_tokens"] == 262144
-        assert info["max_output_tokens"] == 262144
-
-        # thinking variants are marked as reasoning models
-        assert (
-            model_cost["libertai/qwen3.6-27b-thinking"].get("supports_reasoning")
-            is True
-        )
 
     def test_libertai_router_config(self):
         """Test that libertai can be used in Router configuration"""
@@ -95,19 +79,6 @@ class TestLibertAIProviderConfig:
         assert len(router.model_list) == 1
         assert router.model_list[0]["model_name"] == "libertai-chat"
 
-    def test_libertai_model_modes(self):
-        """Chat models carry mode 'chat'; the embedding model carries mode 'embedding'."""
-        model_cost = litellm.model_cost
-
-        # chat model
-        assert model_cost["libertai/qwen3.6-27b"]["mode"] == "chat"
-
-        # embedding model (bge-m3) must be normalized to mode 'embedding' so
-        # /embeddings routing and the supported-endpoints matrix stay consistent
-        assert "libertai/bge-m3" in model_cost
-        bge = model_cost["libertai/bge-m3"]
-        assert bge["litellm_provider"] == "libertai"
-        assert bge["mode"] == "embedding"
 
     def test_libertai_supported_endpoints_matrix(self):
         """The runtime-served backup matrix (GET /public/supported_endpoints) lists libertai."""
diff --git a/tests/test_litellm/llms/openai_like/test_meta_provider.py b/tests/test_litellm/llms/openai_like/test_meta_provider.py
index 11b78828da6..c79e4b77cc5 100644
--- a/tests/test_litellm/llms/openai_like/test_meta_provider.py
+++ b/tests/test_litellm/llms/openai_like/test_meta_provider.py
@@ -193,19 +193,6 @@ class TestMetaAnthropicMessages:
 
 
 class TestMuseSparkModelInfo:
-    def test_muse_spark_pricing_and_capabilities(self):
-        info = litellm.get_model_info("meta/muse-spark-1.1")
-
-        assert info["litellm_provider"] == "meta"
-        assert info["input_cost_per_token"] == 1.25e-06
-        assert info["output_cost_per_token"] == 4.25e-06
-        assert info["cache_read_input_token_cost"] == 1.5e-07
-        assert info["max_input_tokens"] == 1048576
-        assert info["supports_reasoning"] is True
-        assert info["supports_web_search"] is True
-        assert info["supports_vision"] is True
-        assert info["supports_function_calling"] is True
-        assert info["supports_prompt_caching"] is True
 
     def test_muse_spark_cost_calculation(self):
         from litellm import completion_cost
diff --git a/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py b/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py
index 6a6271e95e2..6ca7072e7ab 100644
--- a/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py
+++ b/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py
@@ -3,7 +3,6 @@ Unit tests for Perplexity embedding transformation logic.
 """
 
 import base64
-import json
 import struct
 from unittest.mock import MagicMock
 
@@ -298,25 +297,3 @@ class TestPerplexityEmbeddingProviderConfig:
         )
         assert config is not None
         assert isinstance(config, PerplexityEmbeddingConfig)
-
-
-class TestPerplexityEmbeddingModelInfo:
-    """Test that Perplexity embedding models are in model_prices_and_context_window."""
-
-    def test_model_info_available(self):
-        import litellm
-
-        info = litellm.get_model_info("perplexity/pplx-embed-v1-0.6b")
-        assert info is not None
-        assert info["mode"] == "embedding"
-        assert info["max_input_tokens"] == 32768
-        assert info["output_vector_size"] == 1024
-
-    def test_model_info_4b_available(self):
-        import litellm
-
-        info = litellm.get_model_info("perplexity/pplx-embed-v1-4b")
-        assert info is not None
-        assert info["mode"] == "embedding"
-        assert info["max_input_tokens"] == 32768
-        assert info["output_vector_size"] == 2560
diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py
index 6630039e92e..921022ce562 100644
--- a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py
+++ b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py
@@ -26,7 +26,6 @@ from litellm.types.utils import (
     Usage,
     PromptTokensDetailsWrapper,
 )
-from litellm.utils import get_model_info
 
 
 class TestPerplexityCostCalculator:
@@ -317,20 +316,6 @@ class TestPerplexityCostCalculator:
 
         assert math.isclose(total_cost, expected_total, rel_tol=1e-6)
 
-    def test_model_info_access(self):
-        """Test that model info correctly returns the new cost fields."""
-        model_info = get_model_info(
-            model="sonar-deep-research", custom_llm_provider="perplexity"
-        )
-
-        # Check that the new fields are accessible
-        assert "citation_cost_per_token" in model_info
-        assert model_info["citation_cost_per_token"] == 2e-6
-        assert model_info["search_context_cost_per_query"] == {
-            "search_context_size_low": 0.005,
-            "search_context_size_medium": 0.005,
-            "search_context_size_high": 0.005,
-        }
 
     @pytest.mark.parametrize("citation_tokens", [0, 10, 25, 100])
     @pytest.mark.parametrize("search_queries", [0, 1, 5, 10])
@@ -477,36 +462,6 @@ class TestPerplexityCostCalculator:
         assert math.isclose(prompt_cost, expected_prompt, rel_tol=1e-9)
         assert math.isclose(completion_cost, expected_completion, rel_tol=1e-9)
 
-    @pytest.mark.parametrize(
-        "model_id, usd_per_1m_input, usd_per_1m_output, usd_per_1m_cache_read",
-        [
-            ("deepseek-v4-flash-0731", 0.13, 0.26, 0.028),
-            ("glm-5.2", 1.4, 4.4, 0.14),
-            ("kimi-k3", 3.0, 15.0, 0.3),
-            ("kimi-k2.7-code", 0.95, 4.0, 0.19),
-        ],
-    )
-    def test_agent_api_entries_carry_perplexity_published_rates(
-        self, model_id, usd_per_1m_input, usd_per_1m_output, usd_per_1m_cache_read
-    ):
-        """The Agent API third-party models are priced from Perplexity's own catalog
-        (GET https://api.perplexity.ai/v1/models, `pricing` in usd_per_1m_tokens).
-        Perplexity's model id already starts with `perplexity/`, so the cost-map key
-        doubles the prefix. Regression: glm-5.2 shipped glm-5.3's 0.26 cache-read rate,
-        copied from the neighbouring catalog row, an 86% overcharge on cached input.
-        """
-        info = get_model_info(
-            model=f"perplexity/{model_id}", custom_llm_provider="perplexity"
-        )
-
-        assert info["key"] == f"perplexity/perplexity/{model_id}"
-        assert info["litellm_provider"] == "perplexity"
-        assert info["mode"] == "responses"
-        assert math.isclose(info["input_cost_per_token"], usd_per_1m_input / 1e6, rel_tol=1e-9)
-        assert math.isclose(info["output_cost_per_token"], usd_per_1m_output / 1e6, rel_tol=1e-9)
-        assert math.isclose(
-            info["cache_read_input_token_cost"], usd_per_1m_cache_read / 1e6, rel_tol=1e-9
-        )
 
     def test_agent_api_fallback_rates_price_a_response_without_metered_cost(self):
         """Perplexity meters cost on the response, but when `usage.cost` is absent the
diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py
index 6ba8706b0d8..3c9112efb87 100644
--- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py
@@ -136,18 +136,6 @@ class TestVertexAIVideoConfig:
         # Should NOT include endpoint
         assert not url.endswith(":predictLongRunning")
 
-    def test_veo_31_lite_model_cost_entries_match_pricing(self):
-        for path in (ROOT_MODEL_COST_PATH, BACKUP_MODEL_COST_PATH):
-            model_cost = _load_model_cost_map(path)
-            info = model_cost.get(VEO_31_LITE_VERTEX_MODEL)
-
-            assert info is not None, f"{VEO_31_LITE_VERTEX_MODEL} missing from {path}"
-            assert info["litellm_provider"] == "vertex_ai-video-models"
-            assert info["mode"] == "video_generation"
-            assert info["max_input_tokens"] == 1024
-            assert info["output_cost_per_second"] == 0.05
-            assert info["output_cost_per_second_1080p"] == 0.08
-            assert info["supported_modalities"] == ["text", "image"]
 
     def test_veo_31_lite_provider_routing_from_local_model_map(
         self, monkeypatch: pytest.MonkeyPatch
diff --git a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
index 83c3bf1ecef..1e410e41c33 100644
--- a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
+++ b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
@@ -63,11 +63,6 @@ TIER_COST_FIELDS = (
     "output_cost_per_token_above_200k_tokens",
     "cache_read_input_token_cost_above_200k_tokens",
 )
-STALE_TIER_FIELDS = (
-    "input_cost_per_token_above_128k_tokens",
-    "output_cost_per_token_above_128k_tokens",
-    "cache_read_input_token_cost_above_128k_tokens",
-)
 
 
 def expected_retirement_date(slug: str) -> str:
@@ -102,11 +97,6 @@ def test_redirected_slug_keeps_its_retirement_date(cost_map: dict, slug: str):
     assert cost_map[slug]["deprecation_date"] == expected_retirement_date(slug)
 
 
-@pytest.mark.parametrize("slug", REDIRECTED_SLUGS)
-def test_no_slug_keeps_the_superseded_128k_tier(cost_map: dict, slug: str):
-    """The 128k tier belonged to the retired model; grok-4.3 tiers at 200k."""
-    for field in STALE_TIER_FIELDS:
-        assert field not in cost_map[slug], field
 
 
 @pytest.mark.parametrize("slug", REDIRECTED_SLUGS)
@@ -118,10 +108,6 @@ def test_redirected_slug_carries_the_target_tier_rates(cost_map: dict, slug: str
         assert entry[field] == target[field], field
 
 
-def test_a_live_xai_model_is_untouched(cost_map: dict):
-    """Guard against the repricing leaking onto models xAI still serves directly."""
-    assert cost_map["xai/grok-4.6"]["input_cost_per_token"] != cost_map[REDIRECT_TARGET]["input_cost_per_token"]
-    assert "deprecation_date" not in cost_map["xai/grok-4.6"]
 
 
 def test_both_cost_maps_agree_on_the_redirected_slugs():
diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py
index 38ddac8d510..8d3744a00e0 100644
--- a/tests/test_litellm/llms/zai/test_zai_provider.py
+++ b/tests/test_litellm/llms/zai/test_zai_provider.py
@@ -2,11 +2,9 @@
 Tests for Z.AI (Zhipu AI) provider - GLM models
 """
 
-import json
 import math
 
 import pytest
-import respx
 
 import litellm
 from litellm import completion
@@ -57,31 +55,11 @@ def test_zai_in_provider_lists():
     assert "zai" in litellm.provider_list
 
 
-def test_zai_models_in_model_cost(local_model_cost_map):
-    """Test that ZAI models are in the model cost map"""
-
-    zai_models = [
-        "zai/glm-4.7",
-        "zai/glm-4.6",
-        "zai/glm-4.5",
-        "zai/glm-4.5v",
-        "zai/glm-4.5-x",
-        "zai/glm-4.5-air",
-        "zai/glm-4.5-airx",
-        "zai/glm-4-32b-0414-128k",
-        "zai/glm-4.5-flash",
-    ]
-
-    for model in zai_models:
-        assert model in litellm.model_cost, f"Model {model} not found in model_cost"
-        assert litellm.model_cost[model]["litellm_provider"] == "zai"
 
 
 def test_zai_glm46_cost_calculation(local_model_cost_map):
     """Test the cost calculation for glm-4.6"""
 
-    key = "zai/glm-4.6"
-    info = litellm.model_cost[key]
 
     prompt_cost, completion_cost = cost_per_token(
         model="zai/glm-4.6",
@@ -94,24 +72,8 @@ def test_zai_glm46_cost_calculation(local_model_cost_map):
     assert math.isclose(completion_cost, 2.2, rel_tol=1e-6)
 
 
-def test_zai_flash_model_is_free(local_model_cost_map):
-    """Test that glm-4.5-flash has zero cost"""
-
-    key = "zai/glm-4.5-flash"
-    info = litellm.model_cost[key]
-
-    assert info["input_cost_per_token"] == 0
-    assert info["output_cost_per_token"] == 0
 
 
-def test_glm47_supports_reasoning(local_model_cost_map):
-    """Test that GLM-4.7 supports reasoning"""
-
-    key = "zai/glm-4.7"
-    assert key in litellm.model_cost, f"Model {key} not found in model_cost"
-
-    info = litellm.model_cost[key]
-    assert info["supports_reasoning"] is True
 
 
 def test_glm47_cost_calculation(local_model_cost_map):
diff --git a/tests/test_litellm/test_bedrock_extended_beta_models.py b/tests/test_litellm/test_bedrock_extended_beta_models.py
index ebbbd6cab5c..d55aac762fa 100644
--- a/tests/test_litellm/test_bedrock_extended_beta_models.py
+++ b/tests/test_litellm/test_bedrock_extended_beta_models.py
@@ -91,20 +91,6 @@ MODEL_CONFIGS = [
 class TestBedrockNewModels:
     """Unified test suite for all new Bedrock models"""
 
-    @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS)
-    def test_model_info_primary_region(
-        self, model_name, regions, max_input, max_output
-    ):
-        """Test model configuration in primary region (us-east-1)"""
-        model = f"bedrock/us-east-1/{model_name}"
-        model_info = get_model_info(model)
-
-        assert model_info is not None, f"Model {model_name} not found"
-        assert model_info["max_input_tokens"] == max_input
-        assert model_info["max_output_tokens"] == max_output
-        assert model_info["litellm_provider"] == "bedrock"
-        assert model_info["mode"] == "chat"
-        assert model_info["supports_function_calling"] is True
 
     @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS)
     def test_pricing_configured(self, model_name, regions, max_input, max_output):
@@ -128,43 +114,3 @@ class TestBedrockNewModels:
             assert model_info is not None, f"Model {model_name} not found in {region}"
             assert model_info["max_input_tokens"] == max_input
             assert model_info["max_output_tokens"] == max_output
-
-    @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS)
-    def test_sample_regional_variants(self, model_name, regions, max_input, max_output):
-        """Test sample regional variants (us-east-1, eu-west-1, ap-northeast-1)"""
-        for region in ["us-east-1", "ap-northeast-1"]:
-            if region in regions:
-                model = f"bedrock/{region}/{model_name}"
-                model_info = get_model_info(model)
-                assert (
-                    model_info is not None
-                ), f"Model {model_name} not found in {region}"
-                assert model_info["max_input_tokens"] == max_input
-                assert model_info["litellm_provider"] == "bedrock"
-
-
-class TestModelSpecificFeatures:
-    """Model-specific capability tests"""
-
-    def test_deepseek_v3_2_context_window(self):
-        """DeepSeek V3.2 has 163K context window"""
-        model_info = get_model_info("bedrock/us-east-1/deepseek.v3.2")
-        assert model_info["max_input_tokens"] == 163840
-
-    def test_minimax_m2_1_context_window(self):
-        """Minimax M2.1 has 196K input, 8K output"""
-        model_info = get_model_info("bedrock/us-east-1/minimax.minimax-m2.1")
-        assert model_info["max_input_tokens"] == 196000
-        assert model_info["max_output_tokens"] == 8192
-
-    def test_moonshotai_kimi_k2_5_context_window(self):
-        """Moonshot AI Kimi K2.5 has 256K context window"""
-        model_info = get_model_info("bedrock/us-east-1/moonshotai.kimi-k2.5")
-        assert model_info["max_input_tokens"] == 262144
-        assert model_info["max_output_tokens"] == 262144
-
-    def test_qwen3_coder_next_context_window(self):
-        """Qwen3 Coder Next has 256K input, 8K output"""
-        model_info = get_model_info("bedrock/us-east-1/qwen.qwen3-coder-next")
-        assert model_info["max_input_tokens"] == 262144
-        assert model_info["max_output_tokens"] == 8192
diff --git a/tests/test_litellm/test_bedrock_nemotron_super.py b/tests/test_litellm/test_bedrock_nemotron_super.py
deleted file mode 100644
index 969db890e84..00000000000
--- a/tests/test_litellm/test_bedrock_nemotron_super.py
+++ /dev/null
@@ -1,51 +0,0 @@
-"""
-Test suite for NVIDIA Nemotron Super 3 120B on AWS Bedrock
-Verifies model configuration, pricing, and regional availability.
-"""
-
-import os
-
-os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "true"
-
-import pytest
-
-from litellm import get_model_info
-
-
-MODEL_NAME = "nvidia.nemotron-super-3-120b"
-
-
-class TestNemotronSuper3120B:
-    """Test model definition for nvidia.nemotron-super-3-120b"""
-
-    def test_model_info_primary_region(self):
-        """Test model resolves in us-east-1"""
-        model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}")
-
-        assert model_info is not None, f"Model {MODEL_NAME} not found"
-        assert model_info["max_input_tokens"] == 256000
-        assert model_info["max_output_tokens"] == 32768
-        assert model_info["litellm_provider"] == "bedrock_converse"
-        assert model_info["mode"] == "chat"
-        assert model_info["supports_function_calling"] is True
-
-    def test_pricing_configured(self):
-        """Verify pricing matches AWS Bedrock rates"""
-        model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}")
-
-        assert model_info["input_cost_per_token"] == 1.5e-07
-        assert model_info["output_cost_per_token"] == 6.5e-07
-
-    def test_context_window(self):
-        """Nemotron Super 3 120B has 256K input, 32K output on Bedrock"""
-        model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}")
-
-        assert model_info["max_input_tokens"] == 256000
-        assert model_info["max_output_tokens"] == 32768
-
-    def test_resolves_without_region(self):
-        """Test model resolves with just bedrock/ prefix"""
-        model_info = get_model_info(f"bedrock/{MODEL_NAME}")
-
-        assert model_info is not None, f"Model {MODEL_NAME} not found without region"
-        assert model_info["max_input_tokens"] == 256000
diff --git a/tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py b/tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py
deleted file mode 100644
index 1312aa110d3..00000000000
--- a/tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py
+++ /dev/null
@@ -1,47 +0,0 @@
-"""
-Validate that AWS GovCloud (Bedrock us-gov-*) Haiku 4.5 entries carry
-the 1-hour cache write tier.
-
-AWS Bedrock GovCloud pricing applies a +20% premium over global
-Anthropic rates. Global Haiku 4.5 1h cache write is $2.00/MTok; us-gov
-is therefore $2.40/MTok — exactly 1.6x the 5-minute rate of $1.50/MTok.
-
-Source: https://aws.amazon.com/bedrock/pricing/
-"""
-
-import json
-import os
-
-import pytest
-
-
-@pytest.fixture(scope="module")
-def model_data():
-    json_path = os.path.join(
-        os.path.dirname(__file__), "../../model_prices_and_context_window.json"
-    )
-    with open(json_path) as f:
-        return json.load(f)
-
-
-HAIKU_USGOV_KEYS = [
-    "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0",
-    "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0",
-]
-
-
-@pytest.mark.parametrize("model_key", HAIKU_USGOV_KEYS)
-def test_usgov_haiku_4_5_1hr_cache_write(model_data, model_key):
-    assert model_key in model_data, f"Missing model entry: {model_key}"
-    info = model_data[model_key]
-    assert (
-        info["cache_creation_input_token_cost"] == 1.5e-06
-    ), f"{model_key}: 5m cache write should be $1.50/MTok"
-    assert (
-        info["cache_creation_input_token_cost_above_1hr"] == 2.4e-06
-    ), f"{model_key}: 1h cache write should be $2.40/MTok"
-    ratio = (
-        info["cache_creation_input_token_cost_above_1hr"]
-        / info["cache_creation_input_token_cost"]
-    )
-    assert abs(ratio - 1.6) < 1e-9, f"{model_key}: 1h/5m ratio is {ratio}, expected 1.6"
diff --git a/tests/test_litellm/test_bedrock_usgov_pricing.py b/tests/test_litellm/test_bedrock_usgov_pricing.py
index 3576834dd27..1469ec6a1bb 100644
--- a/tests/test_litellm/test_bedrock_usgov_pricing.py
+++ b/tests/test_litellm/test_bedrock_usgov_pricing.py
@@ -31,32 +31,8 @@ def model_data():
         return json.load(f)
 
 
-SONNET_4_5_USGOV_KEYS = [
-    "bedrock/us-gov-east-1/anthropic.claude-sonnet-4-5-20250929-v1:0",
-    "bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0",
-    "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0",
-    "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0",
-    "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0",
-]
 
 
-@pytest.mark.parametrize("model_key", SONNET_4_5_USGOV_KEYS)
-def test_usgov_sonnet_4_5_pricing(model_data, model_key):
-    """Each us-gov sonnet-4-5 entry must carry the +20%-over-global rates
-    that AWS publishes on the GovCloud pricing page.
-    """
-    assert model_key in model_data, f"Missing model entry: {model_key}"
-    info = model_data[model_key]
-
-    assert info["input_cost_per_token"] == 3.6e-06, (
-        f"{model_key}: input_cost_per_token should be $3.60/MTok (got {info['input_cost_per_token']})"
-    )
-    assert info["output_cost_per_token"] == 1.8e-05, f"{model_key}: output_cost_per_token should be $18.00/MTok"
-    assert info["cache_creation_input_token_cost"] == 4.5e-06, f"{model_key}: 5m cache write should be $4.50/MTok"
-    assert info["cache_creation_input_token_cost_above_1hr"] == 7.2e-06, (
-        f"{model_key}: 1h cache write should be $7.20/MTok"
-    )
-    assert info["cache_read_input_token_cost"] == 3.6e-07, f"{model_key}: cache read should be $0.36/MTok"
 
 
 def test_usgov_carries_20_percent_premium_over_global(model_data):
@@ -117,165 +93,24 @@ def test_usgov_cross_region_above_200k_ratio_to_global(model_data):
         assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2"
 
 
-CLAUDE_GOV_EXPECTED = {
-    "anthropic.claude-sonnet-5": {
-        "input_cost_per_token": 2.4e-06,
-        "output_cost_per_token": 1.2e-05,
-        "cache_creation_input_token_cost": 3e-06,
-        "cache_creation_input_token_cost_above_1hr": 4.8e-06,
-        "cache_read_input_token_cost": 2.4e-07,
-    },
-    "anthropic.claude-opus-4-8": {
-        "input_cost_per_token": 6e-06,
-        "output_cost_per_token": 3e-05,
-        "cache_creation_input_token_cost": 7.5e-06,
-        "cache_creation_input_token_cost_above_1hr": 1.2e-05,
-        "cache_read_input_token_cost": 6e-07,
-    },
-    "anthropic.claude-opus-5": {
-        "input_cost_per_token": 6e-06,
-        "output_cost_per_token": 3e-05,
-        "cache_creation_input_token_cost": 7.5e-06,
-        "cache_creation_input_token_cost_above_1hr": 1.2e-05,
-        "cache_read_input_token_cost": 6e-07,
-    },
-    "anthropic.claude-fable-5-1": {
-        "input_cost_per_token": 1.2e-05,
-        "output_cost_per_token": 6e-05,
-        "cache_creation_input_token_cost": 1.5e-05,
-        "cache_creation_input_token_cost_above_1hr": 2.4e-05,
-        "cache_read_input_token_cost": 3e-07,
-    },
-}
 
 
-USGOV_CLAUDE_KEY_TEMPLATES = {
-    "bedrock/us-gov-east-1/{base_key}": "bedrock",
-    "bedrock/us-gov-west-1/{base_key}": "bedrock",
-    "us-gov.{base_key}": "bedrock_converse",
-}
 
 
-@pytest.mark.parametrize("base_key", CLAUDE_GOV_EXPECTED)
-@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items())
-def test_usgov_claude_pricing(model_data, key_template, expected_provider, base_key):
-    """Sonnet 5, Opus 4.8, Opus 5, and Fable 5.1 gov entries, both in-region keys
-    and the us-gov. geo inference profile the model cards list for GovCloud, must
-    carry the 1.2x GovCloud premium over the global anthropic.* rates. No public
-    AWS source (offer files, pricing page) lists Claude GovCloud rows; the premium
-    is the one AWS quotes for Opus 4.8 in GovCloud ($6/$30 per million).
-    """
-    gov_key = key_template.format(base_key=base_key)
-    assert gov_key in model_data, f"Missing model entry: {gov_key}"
-    info = model_data[gov_key]
-    assert info["litellm_provider"] == expected_provider
-    assert "search_context_cost_per_query" not in info
-    for field, expected in CLAUDE_GOV_EXPECTED[base_key].items():
-        assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})"
-        ratio = info[field] / model_data[base_key][field]
-        assert abs(ratio - 1.2) < 1e-9, f"{gov_key}: {field} gov/global ratio is {ratio}, expected 1.2"
 
 
-CONVERSE_GOV_EXPECTED = {
-    "nvidia.nemotron-nano-3-30b": (7.2e-08, 2.88e-07),
-    "nvidia.nemotron-nano-9b-v2": (7.2e-08, 2.76e-07),
-    "nvidia.nemotron-nano-12b-v2": (2.4e-07, 7.2e-07),
-    "nvidia.nemotron-super-3-120b": (1.8e-07, 7.8e-07),
-    "openai.gpt-oss-20b-1:0": (8.4e-08, 3.6e-07),
-    "openai.gpt-oss-120b-1:0": (1.8e-07, 7.2e-07),
-}
 
 
-@pytest.mark.parametrize("base_key", CONVERSE_GOV_EXPECTED)
-@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items())
-def test_usgov_converse_model_pricing(model_data, key_template, expected_provider, base_key):
-    """Nemotron and gpt-oss gov entries, in-region and the us-gov. geo inference
-    profile both GovCloud regions list as ACTIVE, must match the AWS Bedrock
-    offer file, which prices both regions identically at 1.2x commercial.
-    """
-    gov_key = key_template.format(base_key=base_key)
-    assert gov_key in model_data, f"Missing model entry: {gov_key}"
-    info = model_data[gov_key]
-    expected_input, expected_output = CONVERSE_GOV_EXPECTED[base_key]
-    assert info["input_cost_per_token"] == expected_input
-    assert info["output_cost_per_token"] == expected_output
-    assert info["litellm_provider"] == expected_provider
-    base = model_data[base_key]
-    assert abs(info["input_cost_per_token"] / base["input_cost_per_token"] - 1.2) < 1e-9
-    assert abs(info["output_cost_per_token"] / base["output_cost_per_token"] - 1.2) < 1e-9
 
 
-def test_usgov_west_llama3_8b_output_price_fixed(model_data):
-    """The us-gov-west-1 llama3-8b entry carried the 70B output rate ($2.65/MTok);
-    the AWS Bedrock offer file prices output at $0.60/MTok. AWS lists the model
-    in us-gov-west-1 only, so there is no east entry to check.
-    """
-    info = model_data["bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0"]
-    assert info["input_cost_per_token"] == 3e-07
-    assert info["output_cost_per_token"] == 6e-07
 
 
-MANTLE_GOV_TIERED_EXPECTED = {
-    "openai.gpt-5.6-luna": {
-        "input_cost_per_token": 2.64e-07,
-        "input_cost_per_token_above_272k_tokens": 5.28e-07,
-        "cache_creation_input_token_cost": 3.3e-07,
-        "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07,
-        "cache_read_input_token_cost": 2.64e-08,
-        "cache_read_input_token_cost_above_272k_tokens": 5.28e-08,
-        "output_cost_per_token": 1.584e-06,
-        "output_cost_per_token_above_272k_tokens": 2.376e-06,
-    },
-    "openai.gpt-5.6-terra": {
-        "input_cost_per_token": 2.64e-06,
-        "input_cost_per_token_above_272k_tokens": 5.28e-06,
-        "cache_creation_input_token_cost": 3.3e-06,
-        "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06,
-        "cache_read_input_token_cost": 2.64e-07,
-        "cache_read_input_token_cost_above_272k_tokens": 5.28e-07,
-        "output_cost_per_token": 1.584e-05,
-        "output_cost_per_token_above_272k_tokens": 2.376e-05,
-    },
-}
 
 
-@pytest.mark.parametrize("model", MANTLE_GOV_TIERED_EXPECTED)
-def test_usgov_west_mantle_terra_luna_pricing(model_data, model):
-    """Terra and Luna carry 1.2x commercial across every tier in the
-    us-gov-west-1 offer file; the us-gov-east-1 offer file has no SKUs for them.
-    """
-    gov_key = f"bedrock_mantle/us-gov-west-1/{model}"
-    assert gov_key in model_data, f"Missing model entry: {gov_key}"
-    info = model_data[gov_key]
-    for field, expected in MANTLE_GOV_TIERED_EXPECTED[model].items():
-        assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})"
-    assert info["litellm_provider"] == "bedrock_mantle"
-    assert f"bedrock_mantle/us-gov-east-1/{model}" not in model_data
 
 
-@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"])
-def test_usgov_mantle_gpt_5_4_pricing_has_no_long_context_tier(model_data, region):
-    """gpt-5.4 gov rates come from the offer file, which publishes only the
-    standard tier in GovCloud: no long-context SKUs exist there, unlike commercial.
-    """
-    gov_key = f"bedrock_mantle/{region}/openai.gpt-5.4"
-    assert gov_key in model_data, f"Missing model entry: {gov_key}"
-    info = model_data[gov_key]
-    assert info["input_cost_per_token"] == 3.3e-06
-    assert info["cache_read_input_token_cost"] == 3.3e-07
-    assert info["output_cost_per_token"] == 1.98e-05
-    assert not any(field.endswith("_above_272k_tokens") for field in info)
 
 
-def test_usgov_mantle_grok_4_3_west_only(model_data):
-    """grok-4.3 is priced in the us-gov-west-1 offer file only; the east offer
-    file carries grok-4.6 instead.
-    """
-    info = model_data["bedrock_mantle/us-gov-west-1/xai.grok-4.3"]
-    assert info["input_cost_per_token"] == 1.5e-06
-    assert info["output_cost_per_token"] == 3e-06
-    assert info["cache_read_input_token_cost"] == 2.4e-07
-    assert "bedrock_mantle/us-gov-east-1/xai.grok-4.3" not in model_data
 
 
 def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data):
@@ -290,94 +125,18 @@ def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data):
     }
 
 
-GROK_4_6_GOV_KEYS = {
-    "us-gov.xai.grok-4.6": ("us.xai.grok-4.6", "bedrock_converse"),
-    "bedrock_mantle/us-gov-west-1/xai.grok-4.6": ("bedrock_mantle/xai.grok-4.6", "bedrock_mantle"),
-    "bedrock_mantle/us-gov-east-1/xai.grok-4.6": ("bedrock_mantle/xai.grok-4.6", "bedrock_mantle"),
-}
 
 
-@pytest.mark.parametrize("gov_key", GROK_4_6_GOV_KEYS)
-def test_usgov_grok_4_6_pricing(model_data, gov_key):
-    """Both GovCloud regions serve grok-4.6 through the us-gov. profile only, and
-    both offer files price its standard SKU at 1.2x the commercial US rate.
-    """
-    base_key, expected_provider = GROK_4_6_GOV_KEYS[gov_key]
-    assert gov_key in model_data, f"Missing model entry: {gov_key}"
-    info = model_data[gov_key]
-    assert info["litellm_provider"] == expected_provider
-    assert info["input_cost_per_token"] == 2.64e-06
-    assert info["output_cost_per_token"] == 7.92e-06
-    assert info["cache_read_input_token_cost"] == 6.6e-07
-    for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"):
-        assert abs(info[field] / model_data[base_key][field] - 1.2) < 1e-9
 
 
-NOVA_GOV_WEST_EXPECTED = {
-    "amazon.nova-lite-v1:0": (7.2e-08, 2.88e-07),
-    "amazon.nova-micro-v1:0": (4.2e-08, 1.68e-07),
-}
 
 
-@pytest.mark.parametrize("base_key", NOVA_GOV_WEST_EXPECTED)
-def test_usgov_west_nova_lite_micro_pricing(model_data, base_key):
-    """Nova Lite and Micro are on-demand in us-gov-west-1 only; the offer file
-    prices them at 1.2x commercial, like the Nova Pro row that was already there.
-    """
-    gov_key = f"bedrock/us-gov-west-1/{base_key}"
-    assert gov_key in model_data, f"Missing model entry: {gov_key}"
-    info = model_data[gov_key]
-    expected_input, expected_output = NOVA_GOV_WEST_EXPECTED[base_key]
-    assert info["litellm_provider"] == "bedrock"
-    assert info["input_cost_per_token"] == expected_input
-    assert info["output_cost_per_token"] == expected_output
-    assert abs(info["input_cost_per_token"] / model_data[base_key]["input_cost_per_token"] - 1.2) < 1e-9
-    assert abs(info["output_cost_per_token"] / model_data[base_key]["output_cost_per_token"] - 1.2) < 1e-9
-    assert f"bedrock/us-gov-east-1/{base_key}" not in model_data
 
 
-def test_usgov_west_nova_2_multimodal_embeddings_pricing(model_data):
-    """Every meter of the multimodal embedding model (tokens, images, audio and
-    video seconds) carries the 1.2x uplift the us-gov-west-1 offer file lists.
-    """
-    gov_key = "bedrock/us-gov-west-1/amazon.nova-2-multimodal-embeddings-v1:0"
-    assert gov_key in model_data, f"Missing model entry: {gov_key}"
-    info = model_data[gov_key]
-    assert info["litellm_provider"] == "bedrock"
-    assert info["mode"] == "embedding"
-    assert info["input_cost_per_token"] == 1.62e-07
-    assert info["input_cost_per_image"] == 7.2e-05
-    assert info["input_cost_per_audio_per_second"] == 0.000168
-    assert info["input_cost_per_video_per_second"] == 0.00084
-    assert "bedrock/us-gov-east-1/amazon.nova-2-multimodal-embeddings-v1:0" not in model_data
 
 
-MANTLE_GOV_FLAT_EXPECTED = {
-    "google.gemma-4-e2b": (4.8e-08, 9.6e-08, ("us-gov-west-1",)),
-    "google.gemma-4-26b-a4b": (1.56e-07, 4.8e-07, ("us-gov-west-1",)),
-    "google.gemma-4-31b": (1.68e-07, 4.8e-07, ("us-gov-west-1",)),
-    "openai.gpt-oss-20b": (8.4e-08, 3.6e-07, ("us-gov-west-1", "us-gov-east-1")),
-    "openai.gpt-oss-120b": (1.8e-07, 7.2e-07, ("us-gov-west-1", "us-gov-east-1")),
-}
 
 
-@pytest.mark.parametrize("model", MANTLE_GOV_FLAT_EXPECTED)
-def test_usgov_mantle_gemma_and_gpt_oss_pricing(model_data, model):
-    """Gemma 4 is priced in the us-gov-west-1 offer file only and gpt-oss in both;
-    each Mantle gov row carries the offer file's standard SKU, and no row exists
-    for a region whose offer file has no SKU.
-    """
-    expected_input, expected_output, regions = MANTLE_GOV_FLAT_EXPECTED[model]
-    for region in ("us-gov-west-1", "us-gov-east-1"):
-        gov_key = f"bedrock_mantle/{region}/{model}"
-        if region not in regions:
-            assert gov_key not in model_data
-            continue
-        assert gov_key in model_data, f"Missing model entry: {gov_key}"
-        info = model_data[gov_key]
-        assert info["litellm_provider"] == "bedrock_mantle"
-        assert info["input_cost_per_token"] == expected_input
-        assert info["output_cost_per_token"] == expected_output
 
 
 GOV_ROW_SOURCES = {
@@ -417,33 +176,3 @@ def test_usgov_rows_keep_commercial_limits_and_capabilities(model_data, gov_key)
     assert _non_pricing_fields(gov) == _non_pricing_fields(model_data[GOV_ROW_SOURCES[gov_key]])
     assert "search_context_cost_per_query" not in gov
     assert "source" not in gov
-
-
-AZURE_GOV_EXPECTED = {
-    "azure/us-gov/gpt-5.1": {
-        "input_cost_per_token": 1.71875e-06,
-        "cache_read_input_token_cost": 1.71875e-07,
-        "output_cost_per_token": 1.375e-05,
-    },
-    "azure/us-gov/o3-mini": {
-        "input_cost_per_token": 1.513e-06,
-        "cache_read_input_token_cost": 7.57e-07,
-        "output_cost_per_token": 6.05e-06,
-    },
-    "azure/us-gov/text-embedding-3-large": {"input_cost_per_token": 1.63e-07},
-    "azure/us-gov/text-embedding-3-small": {"input_cost_per_token": 2.5e-08},
-}
-
-
-@pytest.mark.parametrize("gov_key", AZURE_GOV_EXPECTED)
-def test_azure_usgov_pricing(model_data, gov_key):
-    """Azure Government meters from the Azure retail prices API
-    (usgovvirginia/usgovarizona, serviceName 'Foundry Models'). No Government
-    retirement schedule is published, so these entries carry no deprecation_date.
-    """
-    assert gov_key in model_data, f"Missing model entry: {gov_key}"
-    info = model_data[gov_key]
-    for field, expected in AZURE_GOV_EXPECTED[gov_key].items():
-        assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})"
-    assert info["litellm_provider"] == "azure"
-    assert "deprecation_date" not in info
diff --git a/tests/test_litellm/test_claude_fable_5_config.py b/tests/test_litellm/test_claude_fable_5_config.py
index 3ecf94602d9..52d3dccddc8 100644
--- a/tests/test_litellm/test_claude_fable_5_config.py
+++ b/tests/test_litellm/test_claude_fable_5_config.py
@@ -14,7 +14,6 @@ import os
 
 import pytest
 
-import litellm
 from litellm.constants import BEDROCK_CONVERSE_MODELS
 from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
 
@@ -28,86 +27,8 @@ def _load_root_cost_map() -> dict:
 
 
 
-def test_fable_5_model_pricing_and_capabilities():
-    model_data = _load_root_cost_map()
-
-    expected_models = [
-        ("claude-fable-5", "anthropic"),
-        ("anthropic.claude-fable-5", "bedrock_converse"),
-        ("vertex_ai/claude-fable-5", "vertex_ai-anthropic_models"),
-        # Unlike Opus 4.8 (200k on Foundry), Fable 5 has the full 1M context
-        # window on Microsoft Foundry.
-        ("azure_ai/claude-fable-5", "azure_ai"),
-    ]
-
-    for model_name, provider in expected_models:
-        assert model_name in model_data, f"Missing model entry: {model_name}"
-        info = model_data[model_name]
-
-        assert info["litellm_provider"] == provider
-        assert info["mode"] == "chat"
-        assert info["max_input_tokens"] == 1000000
-        assert info["max_output_tokens"] == 128000
-        assert info["max_tokens"] == 128000
-
-        # $10 / $50 per MTok (2x Opus 4.8), with the standard 1.25x 5m
-        # cache-write, 2x 1h cache-write, and 0.1x cache-read multipliers.
-        assert info["input_cost_per_token"] == 1e-05
-        assert info["output_cost_per_token"] == 5e-05
-        assert info["cache_creation_input_token_cost"] == 1.25e-05
-        assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05
-        assert info["cache_read_input_token_cost"] == 1e-06
-
-        # Flat-rate across the full 1M context window.
-        assert "input_cost_per_token_above_200k_tokens" not in info
-        assert "output_cost_per_token_above_200k_tokens" not in info
-
-        assert info["supports_assistant_prefill"] is False
-        assert info["supports_function_calling"] is True
-        assert info["supports_prompt_caching"] is True
-        assert info["supports_reasoning"] is True
-        assert info["supports_tool_choice"] is True
-        assert info["supports_vision"] is True
-        assert info["supports_xhigh_reasoning_effort"] is True
-        assert info["supports_max_reasoning_effort"] is True
 
 
-def test_fable_5_bedrock_regional_model_pricing():
-    model_data = _load_root_cost_map()
-
-    # Fable 5 launched with us/eu geo inference profiles plus a global profile
-    # (no au/apac/jp). Global uses base pricing; geo profiles carry the
-    # standard 10% regional premium.
-    expected_models = {
-        "global.anthropic.claude-fable-5": {
-            "input_cost_per_token": 1e-05,
-            "output_cost_per_token": 5e-05,
-            "cache_creation_input_token_cost": 1.25e-05,
-            "cache_read_input_token_cost": 1e-06,
-        },
-        "us.anthropic.claude-fable-5": {
-            "input_cost_per_token": 1.1e-05,
-            "output_cost_per_token": 5.5e-05,
-            "cache_creation_input_token_cost": 1.375e-05,
-            "cache_read_input_token_cost": 1.1e-06,
-        },
-        "eu.anthropic.claude-fable-5": {
-            "input_cost_per_token": 1.1e-05,
-            "output_cost_per_token": 5.5e-05,
-            "cache_creation_input_token_cost": 1.375e-05,
-            "cache_read_input_token_cost": 1.1e-06,
-        },
-    }
-
-    for model_name, expected in expected_models.items():
-        assert model_name in model_data, f"Missing model entry: {model_name}"
-        info = model_data[model_name]
-        assert info["litellm_provider"] == "bedrock_converse"
-        assert info["max_input_tokens"] == 1000000
-        assert info["max_output_tokens"] == 128000
-        assert info["bedrock_output_config_effort_ceiling"] == "xhigh"
-        for key, value in expected.items():
-            assert info[key] == value
 
 
 def test_fable_5_geo_multiplier_without_fast_mode():
@@ -144,11 +65,6 @@ def test_fable_5_registered_for_bedrock_converse():
     assert "anthropic.claude-fable-5" in BEDROCK_CONVERSE_MODELS
 
 
-def test_fable_5_provider_resolves_via_model_info(local_model_cost_map):
-    info = litellm.get_model_info(model="claude-fable-5")
-    assert info["litellm_provider"] == "anthropic"
-    assert info["max_input_tokens"] == 1000000
-    assert info["max_output_tokens"] == 128000
 
 
 @pytest.mark.parametrize(
@@ -222,44 +138,6 @@ FABLE_5_1_VARIANTS = (
 )
 
 
-def test_fable_5_1_model_pricing_and_capabilities():
-    model_data = _load_root_cost_map()
-
-    expected_models = [
-        ("claude-fable-5-1", "anthropic"),
-        ("anthropic.claude-fable-5-1", "bedrock_converse"),
-        ("vertex_ai/claude-fable-5-1", "vertex_ai-anthropic_models"),
-        ("azure_ai/claude-fable-5-1", "azure_ai"),
-    ]
-
-    for model_name, provider in expected_models:
-        assert model_name in model_data, f"Missing model entry: {model_name}"
-        info = model_data[model_name]
-
-        assert info["litellm_provider"] == provider
-        assert info["mode"] == "chat"
-        assert info["max_input_tokens"] == 1000000
-        assert info["max_output_tokens"] == 128000
-        assert info["max_tokens"] == 128000
-
-        assert info["input_cost_per_token"] == 1e-05
-        assert info["output_cost_per_token"] == 5e-05
-        assert info["cache_creation_input_token_cost"] == 1.25e-05
-        assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05
-
-        assert "input_cost_per_token_above_200k_tokens" not in info
-        assert "output_cost_per_token_above_200k_tokens" not in info
-
-        assert info["supports_assistant_prefill"] is False
-        assert info["supports_forced_tool_use"] is False
-        assert info["supports_function_calling"] is True
-        assert info["supports_prompt_caching"] is True
-        assert info["supports_reasoning"] is True
-        assert info["supports_tool_choice"] is True
-        assert info["supports_vision"] is True
-        assert info["supports_xhigh_reasoning_effort"] is True
-        assert info["supports_max_reasoning_effort"] is True
-        assert info["prompt_cache_min_tokens"] == 512
 
 
 @pytest.mark.parametrize(
@@ -280,46 +158,8 @@ def test_fable_5_1_cache_reads_cost_a_quarter_of_fable_5(cost_map):
         ), model_name
 
 
-def test_fable_5_1_bedrock_regional_model_pricing():
-    model_data = _load_root_cost_map()
-
-    expected_models = {
-        "global.anthropic.claude-fable-5-1": {
-            "input_cost_per_token": 1e-05,
-            "output_cost_per_token": 5e-05,
-            "cache_creation_input_token_cost": 1.25e-05,
-            "cache_read_input_token_cost": 2.5e-07,
-        },
-        "us.anthropic.claude-fable-5-1": {
-            "input_cost_per_token": 1.1e-05,
-            "output_cost_per_token": 5.5e-05,
-            "cache_creation_input_token_cost": 1.375e-05,
-            "cache_read_input_token_cost": 2.75e-07,
-        },
-        "eu.anthropic.claude-fable-5-1": {
-            "input_cost_per_token": 1.1e-05,
-            "output_cost_per_token": 5.5e-05,
-            "cache_creation_input_token_cost": 1.375e-05,
-            "cache_read_input_token_cost": 2.75e-07,
-        },
-    }
-
-    for model_name, expected in expected_models.items():
-        assert model_name in model_data, f"Missing model entry: {model_name}"
-        info = model_data[model_name]
-        assert info["litellm_provider"] == "bedrock_converse"
-        assert info["max_input_tokens"] == 1000000
-        assert info["max_output_tokens"] == 128000
-        assert info["bedrock_output_config_effort_ceiling"] == "xhigh"
-        for key, value in expected.items():
-            assert info[key] == value
 
 
-def test_fable_5_1_geo_multiplier_without_fast_mode():
-    """Fable 5.1 has no fast mode, so a ``fast`` key here would misprice
-    ``speed='fast'`` requests."""
-    model_data = _load_root_cost_map()
-    assert model_data["claude-fable-5-1"]["provider_specific_entry"] == {"us": 1.1}
 
 
 def test_fable_5_1_present_in_bundled_backup():
@@ -334,11 +174,6 @@ def test_fable_5_1_registered_for_bedrock_converse():
     assert "anthropic.claude-fable-5-1" in BEDROCK_CONVERSE_MODELS
 
 
-def test_fable_5_1_provider_resolves_via_model_info(local_model_cost_map):
-    info = litellm.get_model_info(model="claude-fable-5-1")
-    assert info["litellm_provider"] == "anthropic"
-    assert info["max_input_tokens"] == 1000000
-    assert info["max_output_tokens"] == 128000
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_litellm/test_claude_haiku_4_5_config.py b/tests/test_litellm/test_claude_haiku_4_5_config.py
index 8755e5d156f..ab99a34d378 100644
--- a/tests/test_litellm/test_claude_haiku_4_5_config.py
+++ b/tests/test_litellm/test_claude_haiku_4_5_config.py
@@ -7,55 +7,6 @@ import json
 import os
 
 
-def test_bedrock_haiku_4_5_configuration():
-    """Test that all Bedrock Claude Haiku 4.5 models use bedrock_converse provider"""
-    # Load model configuration
-    json_path = os.path.join(
-        os.path.dirname(__file__), "../../model_prices_and_context_window.json"
-    )
-    with open(json_path) as f:
-        model_data = json.load(f)
-
-    # All Bedrock Haiku 4.5 variants that should use bedrock_converse
-    bedrock_haiku_models = [
-        "anthropic.claude-haiku-4-5-20251001-v1:0",
-        "anthropic.claude-haiku-4-5@20251001",
-        "us.anthropic.claude-haiku-4-5-20251001-v1:0",
-        "eu.anthropic.claude-haiku-4-5-20251001-v1:0",
-        "apac.anthropic.claude-haiku-4-5-20251001-v1:0",
-        "jp.anthropic.claude-haiku-4-5-20251001-v1:0",
-        "global.anthropic.claude-haiku-4-5-20251001-v1:0",
-        "au.anthropic.claude-haiku-4-5-20251001-v1:0",
-    ]
-
-    for model in bedrock_haiku_models:
-        assert model in model_data, f"Model {model} not found in config"
-        model_info = model_data[model]
-
-        # Verify uses bedrock_converse (not legacy bedrock provider)
-        assert (
-            model_info["litellm_provider"] == "bedrock_converse"
-        ), f"{model} should use bedrock_converse provider, got {model_info['litellm_provider']}"
-
-        # Verify supports vision (key missing capability)
-        assert (
-            model_info.get("supports_vision") is True
-        ), f"{model} should support vision"
-
-        # Verify core capabilities
-        assert model_info.get("supports_computer_use") is True
-        assert model_info.get("supports_function_calling") is True
-        assert model_info.get("supports_tool_choice") is True
-        assert model_info.get("supports_prompt_caching") is True
-        assert model_info.get("supports_response_schema") is True
-        assert model_info.get("supports_pdf_input") is True
-        assert model_info.get("supports_assistant_prefill") is True
-        assert model_info.get("supports_reasoning") is True
-
-        # Verify token limits
-        assert model_info["max_input_tokens"] == 200000
-        assert model_info["max_output_tokens"] == 64000
-        assert model_info["mode"] == "chat"
 
 
 def test_bedrock_haiku_4_5_matches_sonnet_capabilities():
@@ -97,36 +48,3 @@ def test_bedrock_haiku_4_5_matches_sonnet_capabilities():
         assert haiku_info.get(capability) == sonnet_info.get(
             capability
         ), f"Capability {capability} mismatch: Haiku={haiku_info.get(capability)}, Sonnet={sonnet_info.get(capability)}"
-
-
-def test_anthropic_api_haiku_4_5_configuration():
-    """Test that Anthropic API Claude Haiku 4.5 has correct configuration"""
-    # Load model configuration
-    json_path = os.path.join(
-        os.path.dirname(__file__), "../../model_prices_and_context_window.json"
-    )
-    with open(json_path) as f:
-        model_data = json.load(f)
-
-    # Anthropic API models (not Bedrock)
-    anthropic_models = [
-        "claude-haiku-4-5-20251001",
-        "claude-haiku-4-5",
-    ]
-
-    for model in anthropic_models:
-        assert model in model_data, f"Model {model} not found in config"
-        model_info = model_data[model]
-
-        # Should use anthropic provider (not bedrock)
-        assert (
-            model_info["litellm_provider"] == "anthropic"
-        ), f"{model} should use anthropic provider"
-
-        # Should support vision
-        assert (
-            model_info.get("supports_vision") is True
-        ), f"{model} should support vision"
-
-        # Should have larger output token limit (64K for Anthropic API)
-        assert model_info["max_output_tokens"] == 64000
diff --git a/tests/test_litellm/test_claude_opus_4_6_config.py b/tests/test_litellm/test_claude_opus_4_6_config.py
index 89d2cd916e0..a29901adfc3 100644
--- a/tests/test_litellm/test_claude_opus_4_6_config.py
+++ b/tests/test_litellm/test_claude_opus_4_6_config.py
@@ -71,123 +71,8 @@ def test_claude_4_6_australia_region_uses_au_prefix_not_apac():
     ), "apac.anthropic.claude-sonnet-4-6 should not be in bedrock_converse_models"
 
 
-def test_opus_4_6_model_pricing_and_capabilities():
-    json_path = os.path.join(
-        os.path.dirname(__file__), "../../model_prices_and_context_window.json"
-    )
-    with open(json_path) as f:
-        model_data = json.load(f)
-
-    expected_models = {
-        "claude-opus-4-6": {
-            "provider": "anthropic",
-            "has_long_context_pricing": False,
-            "max_input_tokens": 1000000,
-        },
-        "claude-opus-4-6-20260205": {
-            "provider": "anthropic",
-            "has_long_context_pricing": False,
-            "max_input_tokens": 1000000,
-        },
-        "anthropic.claude-opus-4-6-v1": {
-            "provider": "bedrock_converse",
-            "has_long_context_pricing": False,
-            "max_input_tokens": 1000000,
-        },
-        "vertex_ai/claude-opus-4-6": {
-            "provider": "vertex_ai-anthropic_models",
-            "has_long_context_pricing": False,
-            "max_input_tokens": 1000000,
-        },
-        "azure_ai/claude-opus-4-6": {
-            "provider": "azure_ai",
-            "has_long_context_pricing": False,
-            "max_input_tokens": 1000000,
-        },
-    }
-
-    for model_name, config in expected_models.items():
-        assert model_name in model_data, f"Missing model entry: {model_name}"
-        info = model_data[model_name]
-
-        assert info["litellm_provider"] == config["provider"]
-        assert info["mode"] == "chat"
-        assert info["max_input_tokens"] == config["max_input_tokens"]
-        assert info["max_output_tokens"] == 128000
-        assert info["max_tokens"] == 128000
-
-        assert info["input_cost_per_token"] == 5e-06
-        assert info["output_cost_per_token"] == 2.5e-05
-        assert info["cache_creation_input_token_cost"] == 6.25e-06
-        assert info["cache_read_input_token_cost"] == 5e-07
-
-        if config["has_long_context_pricing"]:
-            assert info["input_cost_per_token_above_200k_tokens"] == 1e-05
-            assert info["output_cost_per_token_above_200k_tokens"] == 3.75e-05
-            assert info["cache_creation_input_token_cost_above_200k_tokens"] == 1.25e-05
-            assert info["cache_read_input_token_cost_above_200k_tokens"] == 1e-06
-        else:
-            assert "input_cost_per_token_above_200k_tokens" not in info
-            assert "output_cost_per_token_above_200k_tokens" not in info
-            assert "cache_creation_input_token_cost_above_200k_tokens" not in info
-            assert "cache_read_input_token_cost_above_200k_tokens" not in info
-
-        assert info["supports_assistant_prefill"] is False
-        assert info["supports_function_calling"] is True
-        assert info["supports_prompt_caching"] is True
-        assert info["supports_reasoning"] is True
-        assert info["supports_tool_choice"] is True
-        assert info["supports_vision"] is True
 
 
-def test_opus_4_6_bedrock_regional_model_pricing():
-    json_path = os.path.join(
-        os.path.dirname(__file__), "../../model_prices_and_context_window.json"
-    )
-    with open(json_path) as f:
-        model_data = json.load(f)
-
-    expected_models = {
-        "global.anthropic.claude-opus-4-6-v1": {
-            "input_cost_per_token": 5e-06,
-            "output_cost_per_token": 2.5e-05,
-            "cache_creation_input_token_cost": 6.25e-06,
-            "cache_read_input_token_cost": 5e-07,
-        },
-        "us.anthropic.claude-opus-4-6-v1": {
-            "input_cost_per_token": 5.5e-06,
-            "output_cost_per_token": 2.75e-05,
-            "cache_creation_input_token_cost": 6.875e-06,
-            "cache_read_input_token_cost": 5.5e-07,
-        },
-        "eu.anthropic.claude-opus-4-6-v1": {
-            "input_cost_per_token": 5.5e-06,
-            "output_cost_per_token": 2.75e-05,
-            "cache_creation_input_token_cost": 6.875e-06,
-            "cache_read_input_token_cost": 5.5e-07,
-        },
-        "au.anthropic.claude-opus-4-6-v1": {
-            "input_cost_per_token": 5.5e-06,
-            "output_cost_per_token": 2.75e-05,
-            "cache_creation_input_token_cost": 6.875e-06,
-            "cache_read_input_token_cost": 5.5e-07,
-        },
-    }
-
-    for model_name, expected in expected_models.items():
-        assert model_name in model_data, f"Missing model entry: {model_name}"
-        info = model_data[model_name]
-        assert info["litellm_provider"] == "bedrock_converse"
-        assert info["max_input_tokens"] == 1000000
-        assert info["max_output_tokens"] == 128000
-        assert info["max_tokens"] == 128000
-        assert info["supports_assistant_prefill"] is False
-        assert "input_cost_per_token_above_200k_tokens" not in info
-        assert "output_cost_per_token_above_200k_tokens" not in info
-        assert "cache_creation_input_token_cost_above_200k_tokens" not in info
-        assert "cache_read_input_token_cost_above_200k_tokens" not in info
-        for key, value in expected.items():
-            assert info[key] == value
 
 
 def test_opus_4_6_alias_and_dated_metadata_match():
diff --git a/tests/test_litellm/test_claude_opus_4_8_config.py b/tests/test_litellm/test_claude_opus_4_8_config.py
index 760512ad31b..7173b4a0e5b 100644
--- a/tests/test_litellm/test_claude_opus_4_8_config.py
+++ b/tests/test_litellm/test_claude_opus_4_8_config.py
@@ -16,7 +16,6 @@ import os
 
 import pytest
 
-import litellm
 from litellm.constants import BEDROCK_CONVERSE_MODELS
 from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
 
@@ -30,99 +29,8 @@ def _load_root_cost_map() -> dict:
 
 
 
-def test_opus_4_8_model_pricing_and_capabilities():
-    model_data = _load_root_cost_map()
-
-    expected_models = {
-        "claude-opus-4-8": {
-            "provider": "anthropic",
-            "max_input_tokens": 1000000,
-        },
-        "anthropic.claude-opus-4-8": {
-            "provider": "bedrock_converse",
-            "max_input_tokens": 1000000,
-        },
-        "vertex_ai/claude-opus-4-8": {
-            "provider": "vertex_ai-anthropic_models",
-            "max_input_tokens": 1000000,
-        },
-        "azure_ai/claude-opus-4-8": {
-            "provider": "azure_ai",
-            "max_input_tokens": 1000000,
-        },
-    }
-
-    for model_name, config in expected_models.items():
-        assert model_name in model_data, f"Missing model entry: {model_name}"
-        info = model_data[model_name]
-
-        assert info["litellm_provider"] == config["provider"]
-        assert info["mode"] == "chat"
-        assert info["max_input_tokens"] == config["max_input_tokens"]
-        assert info["max_output_tokens"] == 128000
-        assert info["max_tokens"] == 128000
-
-        # Base pricing matches Opus 4.7: $5 / $25 per MTok, with the standard
-        # 1.25x cache-write and 0.1x cache-read multipliers.
-        assert info["input_cost_per_token"] == 5e-06
-        assert info["output_cost_per_token"] == 2.5e-05
-        assert info["cache_creation_input_token_cost"] == 6.25e-06
-        assert info["cache_read_input_token_cost"] == 5e-07
-
-        # Opus 4.x flagships are flat-rate across the full context window.
-        assert "input_cost_per_token_above_200k_tokens" not in info
-        assert "output_cost_per_token_above_200k_tokens" not in info
-
-        assert info["supports_assistant_prefill"] is False
-        assert info["supports_function_calling"] is True
-        assert info["supports_prompt_caching"] is True
-        assert info["supports_reasoning"] is True
-        assert info["supports_tool_choice"] is True
-        assert info["supports_vision"] is True
-
-    assert model_data["claude-opus-4-8"]["supports_native_structured_output"] is True
 
 
-def test_opus_4_8_bedrock_regional_model_pricing():
-    model_data = _load_root_cost_map()
-
-    # Global endpoints use base pricing; regional endpoints carry a 10% premium.
-    expected_models = {
-        "global.anthropic.claude-opus-4-8": {
-            "input_cost_per_token": 5e-06,
-            "output_cost_per_token": 2.5e-05,
-            "cache_creation_input_token_cost": 6.25e-06,
-            "cache_read_input_token_cost": 5e-07,
-        },
-        "us.anthropic.claude-opus-4-8": {
-            "input_cost_per_token": 5.5e-06,
-            "output_cost_per_token": 2.75e-05,
-            "cache_creation_input_token_cost": 6.875e-06,
-            "cache_read_input_token_cost": 5.5e-07,
-        },
-        "eu.anthropic.claude-opus-4-8": {
-            "input_cost_per_token": 5.5e-06,
-            "output_cost_per_token": 2.75e-05,
-            "cache_creation_input_token_cost": 6.875e-06,
-            "cache_read_input_token_cost": 5.5e-07,
-        },
-        "au.anthropic.claude-opus-4-8": {
-            "input_cost_per_token": 5.5e-06,
-            "output_cost_per_token": 2.75e-05,
-            "cache_creation_input_token_cost": 6.875e-06,
-            "cache_read_input_token_cost": 5.5e-07,
-        },
-    }
-
-    for model_name, expected in expected_models.items():
-        assert model_name in model_data, f"Missing model entry: {model_name}"
-        info = model_data[model_name]
-        assert info["litellm_provider"] == "bedrock_converse"
-        assert info["max_input_tokens"] == 1000000
-        assert info["max_output_tokens"] == 128000
-        assert info["bedrock_output_config_effort_ceiling"] == "xhigh"
-        for key, value in expected.items():
-            assert info[key] == value
 
 
 def test_opus_4_8_fast_mode_multiplier():
@@ -134,42 +42,12 @@ def test_opus_4_8_fast_mode_multiplier():
     assert entry["fast"] == 2.0
 
 
-def test_opus_4_8_present_in_bundled_backup():
-    """The bundled backup is the runtime fallback (and what tests load with
-    ``LITELLM_LOCAL_MODEL_COST_MAP=True``) — it must carry the same entries as
-    the root cost map, otherwise the model resolves on one path but not the
-    other."""
-    backup = GetModelCostMap.load_local_model_cost_map()
-    for model_name in (
-        "claude-opus-4-8",
-        "anthropic.claude-opus-4-8",
-        "global.anthropic.claude-opus-4-8",
-        "us.anthropic.claude-opus-4-8",
-        "eu.anthropic.claude-opus-4-8",
-        "au.anthropic.claude-opus-4-8",
-        "vertex_ai/claude-opus-4-8",
-        "vertex_ai/claude-opus-4-8@default",
-        "azure_ai/claude-opus-4-8",
-    ):
-        assert model_name in backup, f"Missing from backup cost map: {model_name}"
-    assert backup["claude-opus-4-8"]["supports_native_structured_output"] is True
 
 
 def test_opus_4_8_registered_for_bedrock_converse():
     assert "anthropic.claude-opus-4-8" in BEDROCK_CONVERSE_MODELS
 
 
-def test_opus_4_8_provider_resolves_via_model_info(local_model_cost_map):
-    """Regression: ``claude-opus-4-8`` must resolve to provider ``anthropic``.
-
-    Before the cost-map entry existed, the model was unknown to LiteLLM, so it
-    could not be tied to the ``anthropic`` provider and an ``anthropic/*``
-    wildcard deployment would not match it.
-    """
-    info = litellm.get_model_info(model="claude-opus-4-8")
-    assert info["litellm_provider"] == "anthropic"
-    assert info["max_input_tokens"] == 1000000
-    assert info["max_output_tokens"] == 128000
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_litellm/test_claude_opus_5_config.py b/tests/test_litellm/test_claude_opus_5_config.py
index 34744aad17b..beb148f5e1b 100644
--- a/tests/test_litellm/test_claude_opus_5_config.py
+++ b/tests/test_litellm/test_claude_opus_5_config.py
@@ -17,7 +17,6 @@ import os
 
 import pytest
 
-import litellm
 from litellm.constants import BEDROCK_CONVERSE_MODELS
 from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
 
@@ -53,88 +52,8 @@ def _load_root_cost_map() -> dict:
 
 
 
-def test_opus_5_pricing_and_capabilities():
-    model_data = _load_root_cost_map()
-
-    expected_providers = {
-        "claude-opus-5": "anthropic",
-        "anthropic.claude-opus-5": "bedrock_converse",
-        "vertex_ai/claude-opus-5": "vertex_ai-anthropic_models",
-        "azure_ai/claude-opus-5": "azure_ai",
-    }
-
-    for model_name, provider in expected_providers.items():
-        assert model_name in model_data, f"Missing model entry: {model_name}"
-        info = model_data[model_name]
-
-        assert info["litellm_provider"] == provider
-        assert info["mode"] == "chat"
-        assert info["max_input_tokens"] == 1000000
-        assert info["max_output_tokens"] == 128000
-        assert info["max_tokens"] == 128000
-
-        # Opus 5 ships at Opus 4.8's rates: $5 / $25 per MTok, with the standard
-        # 1.25x cache-write, 2x 1-hour cache-write, and 0.1x cache-read multipliers.
-        assert info["input_cost_per_token"] == 5e-06
-        assert info["output_cost_per_token"] == 2.5e-05
-        assert info["cache_creation_input_token_cost"] == 6.25e-06
-        assert info["cache_creation_input_token_cost_above_1hr"] == 1e-05
-        assert info["cache_read_input_token_cost"] == 5e-07
-
-        # Flat rate across the full 1M window, no long-context premium.
-        assert "input_cost_per_token_above_200k_tokens" not in info
-        assert "output_cost_per_token_above_200k_tokens" not in info
-
-        # gen-5 adaptive-thinking profile: effort-driven, no sampling params, no
-        # assistant prefill.
-        assert info["supports_adaptive_thinking"] is True
-        assert info["supports_reasoning"] is True
-        assert info["supports_sampling_params"] is False
-        assert info["supports_assistant_prefill"] is False
-        assert info["supports_xhigh_reasoning_effort"] is True
-        assert info["supports_max_reasoning_effort"] is True
-
-        assert info["supports_function_calling"] is True
-        assert info["supports_prompt_caching"] is True
-        assert info["supports_tool_choice"] is True
-        assert info["supports_vision"] is True
 
 
-def test_opus_5_bedrock_regional_pricing():
-    """Global/base endpoints use base pricing; the us./eu./au./jp. regional
-    cross-region inference profiles carry a 10% premium."""
-    model_data = _load_root_cost_map()
-
-    base_pricing = {
-        "input_cost_per_token": 5e-06,
-        "output_cost_per_token": 2.5e-05,
-        "cache_creation_input_token_cost": 6.25e-06,
-        "cache_creation_input_token_cost_above_1hr": 1e-05,
-        "cache_read_input_token_cost": 5e-07,
-    }
-    regional_pricing = {
-        "input_cost_per_token": 5.5e-06,
-        "output_cost_per_token": 2.75e-05,
-        "cache_creation_input_token_cost": 6.875e-06,
-        "cache_creation_input_token_cost_above_1hr": 1.1e-05,
-        "cache_read_input_token_cost": 5.5e-07,
-    }
-
-    expected = {
-        "anthropic.claude-opus-5": base_pricing,
-        "global.anthropic.claude-opus-5": base_pricing,
-        "us.anthropic.claude-opus-5": regional_pricing,
-        "eu.anthropic.claude-opus-5": regional_pricing,
-        "au.anthropic.claude-opus-5": regional_pricing,
-        "jp.anthropic.claude-opus-5": regional_pricing,
-    }
-
-    for model_name, pricing in expected.items():
-        assert model_name in model_data, f"Missing model entry: {model_name}"
-        info = model_data[model_name]
-        assert info["litellm_provider"] == "bedrock_converse"
-        for key, value in pricing.items():
-            assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}"
 
 
 @pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS)
@@ -216,16 +135,6 @@ def test_opus_5_registered_for_bedrock_converse():
     assert "anthropic.claude-opus-5" in BEDROCK_CONVERSE_MODELS
 
 
-def test_opus_5_provider_resolves_via_model_info(local_model_cost_map):
-    """Regression: ``claude-opus-5`` must resolve to provider ``anthropic``.
-
-    Without the cost-map entry the model is unknown to LiteLLM, so it cannot be
-    tied to the ``anthropic`` provider and an ``anthropic/*`` wildcard deployment
-    would not match it."""
-    info = litellm.get_model_info(model="claude-opus-5")
-    assert info["litellm_provider"] == "anthropic"
-    assert info["max_input_tokens"] == 1000000
-    assert info["max_output_tokens"] == 128000
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_litellm/test_claude_sonnet_5_config.py b/tests/test_litellm/test_claude_sonnet_5_config.py
index 8504326cd21..bdc3bb64706 100644
--- a/tests/test_litellm/test_claude_sonnet_5_config.py
+++ b/tests/test_litellm/test_claude_sonnet_5_config.py
@@ -15,7 +15,6 @@ import os
 
 import pytest
 
-import litellm
 from litellm.constants import BEDROCK_CONVERSE_MODELS
 from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
 
@@ -42,93 +41,8 @@ def _load_root_cost_map() -> dict:
 
 
 
-def test_sonnet_5_pricing_and_capabilities():
-    model_data = _load_root_cost_map()
-
-    expected_providers = {
-        "claude-sonnet-5": "anthropic",
-        "anthropic.claude-sonnet-5": "bedrock_converse",
-        "vertex_ai/claude-sonnet-5": "vertex_ai-anthropic_models",
-        "azure_ai/claude-sonnet-5": "azure_ai",
-    }
-
-    for model_name, provider in expected_providers.items():
-        assert model_name in model_data, f"Missing model entry: {model_name}"
-        info = model_data[model_name]
-
-        assert info["litellm_provider"] == provider
-        assert info["mode"] == "chat"
-        assert info["max_input_tokens"] == 1000000
-        assert info["max_output_tokens"] == 128000
-        assert info["max_tokens"] == 128000
-
-        # Introductory Sonnet 5 pricing through 2026-08-31: $2 / $10 per MTok,
-        # with the 1.25x cache-write and 0.1x cache-read multipliers. On
-        # 2026-09-01 flip these five fields back to the sticker rate, here and
-        # in both cost-map JSON files (all ten claude-sonnet-5 entries):
-        #   input_cost_per_token: 3e-06
-        #   output_cost_per_token: 1.5e-05
-        #   cache_creation_input_token_cost: 3.75e-06
-        #   cache_creation_input_token_cost_above_1hr: 6e-06
-        #   cache_read_input_token_cost: 3e-07
-        # Regional Bedrock profiles (us./eu./au./jp.) stay at 1.1x those values:
-        # 3.3e-06 / 1.65e-05 / 4.125e-06 / 6.6e-06 / 3.3e-07 (see
-        # test_sonnet_5_bedrock_regional_pricing below).
-        assert info["input_cost_per_token"] == 2e-06
-        assert info["output_cost_per_token"] == 1e-05
-        assert info["cache_creation_input_token_cost"] == 2.5e-06
-        assert info["cache_creation_input_token_cost_above_1hr"] == 4e-06
-        assert info["cache_read_input_token_cost"] == 2e-07
-
-        # gen-5 adaptive-thinking profile: effort-driven, no sampling params, no
-        # assistant prefill.
-        assert info["supports_adaptive_thinking"] is True
-        assert info["supports_reasoning"] is True
-        assert info["supports_sampling_params"] is False
-        assert info["supports_assistant_prefill"] is False
-
-        assert info["supports_function_calling"] is True
-        assert info["supports_prompt_caching"] is True
-        assert info["supports_tool_choice"] is True
-        assert info["supports_vision"] is True
 
 
-def test_sonnet_5_bedrock_regional_pricing():
-    """Global/base endpoints use base pricing; the us./eu./au./jp. regional
-    cross-region inference profiles carry a 10% premium."""
-    model_data = _load_root_cost_map()
-
-    base_pricing = {
-        "input_cost_per_token": 2e-06,
-        "output_cost_per_token": 1e-05,
-        "cache_creation_input_token_cost": 2.5e-06,
-        "cache_creation_input_token_cost_above_1hr": 4e-06,
-        "cache_read_input_token_cost": 2e-07,
-    }
-    regional_pricing = {
-        "input_cost_per_token": 2.2e-06,
-        "output_cost_per_token": 1.1e-05,
-        "cache_creation_input_token_cost": 2.75e-06,
-        "cache_creation_input_token_cost_above_1hr": 4.4e-06,
-        "cache_read_input_token_cost": 2.2e-07,
-    }
-
-    expected = {
-        "anthropic.claude-sonnet-5": base_pricing,
-        "global.anthropic.claude-sonnet-5": base_pricing,
-        "us.anthropic.claude-sonnet-5": regional_pricing,
-        "eu.anthropic.claude-sonnet-5": regional_pricing,
-        "au.anthropic.claude-sonnet-5": regional_pricing,
-        "jp.anthropic.claude-sonnet-5": regional_pricing,
-    }
-
-    for model_name, pricing in expected.items():
-        assert model_name in model_data, f"Missing model entry: {model_name}"
-        info = model_data[model_name]
-        assert info["litellm_provider"] == "bedrock_converse"
-        assert info["bedrock_output_config_effort_ceiling"] == "xhigh"
-        for key, value in pricing.items():
-            assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}"
 
 
 def test_sonnet_5_present_in_bundled_backup():
@@ -144,16 +58,6 @@ def test_sonnet_5_registered_for_bedrock_converse():
     assert "anthropic.claude-sonnet-5" in BEDROCK_CONVERSE_MODELS
 
 
-def test_sonnet_5_provider_resolves_via_model_info(local_model_cost_map):
-    """Regression: ``claude-sonnet-5`` must resolve to provider ``anthropic``.
-
-    Before the cost-map entry existed, the model was unknown to LiteLLM, so it
-    could not be tied to the ``anthropic`` provider and an ``anthropic/*``
-    wildcard deployment would not match it."""
-    info = litellm.get_model_info(model="claude-sonnet-5")
-    assert info["litellm_provider"] == "anthropic"
-    assert info["max_input_tokens"] == 1000000
-    assert info["max_output_tokens"] == 128000
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py
index e33bcfb8378..be6be865665 100644
--- a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py
+++ b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py
@@ -27,15 +27,6 @@ BACKUP_MAP = os.path.join(
 )
 
 
-@pytest.fixture(autouse=True)
-def _use_local_model_cost_map(monkeypatch):
-    original_model_cost = litellm.model_cost
-    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-    try:
-        yield
-    finally:
-        litellm.model_cost = original_model_cost
 
 
 def _load(path: str) -> dict:
@@ -47,48 +38,12 @@ def _cloudflare_keys(data: dict) -> set:
     return {k for k in data if k.startswith("cloudflare/")}
 
 
-def test_glm_5_2_entry_is_present_and_well_formed():
-    entry = litellm.model_cost["cloudflare/@cf/zai-org/glm-5.2"]
-    assert entry["litellm_provider"] == "cloudflare"
-    assert entry["mode"] == "chat"
-    assert entry["supports_function_calling"] is True
-    assert entry["input_cost_per_token"] > 0
-    assert entry["output_cost_per_token"] > 0
 
 
-def test_vision_model_is_flagged_supports_vision():
-    entry = litellm.model_cost["cloudflare/@cf/meta/llama-3.2-11b-vision-instruct"]
-    assert entry["litellm_provider"] == "cloudflare"
-    assert entry.get("supports_vision") is True
 
 
-def test_additional_current_models_are_present():
-    for key in (
-        "cloudflare/@cf/openai/gpt-oss-120b",
-        "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast",
-    ):
-        entry = litellm.model_cost[key]
-        assert entry["litellm_provider"] == "cloudflare"
-        assert entry["mode"] == "chat"
-        assert entry["supports_function_calling"] is True
-        assert entry["input_cost_per_token"] > 0
-        assert entry["output_cost_per_token"] > 0
 
 
-@pytest.mark.parametrize(
-    "key, published_price_per_audio_minute",
-    [
-        ("cloudflare/@cf/openai/whisper", 0.00045),
-        ("cloudflare/@cf/openai/whisper-large-v3-turbo", 0.00051),
-    ],
-)
-def test_whisper_transcription_pricing_is_stored_per_second(key, published_price_per_audio_minute):
-    entry = litellm.model_cost[key]
-    assert entry["litellm_provider"] == "cloudflare"
-    assert entry["mode"] == "audio_transcription"
-    assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"]
-    assert entry["output_cost_per_second"] == 0.0
-    assert entry["input_cost_per_second"] == pytest.approx(published_price_per_audio_minute / 60)
 
 
 def test_root_and_backup_have_identical_cloudflare_keys():
diff --git a/tests/test_litellm/test_daybreak_model_metadata.py b/tests/test_litellm/test_daybreak_model_metadata.py
index dbb7ecdffac..391391444a1 100644
--- a/tests/test_litellm/test_daybreak_model_metadata.py
+++ b/tests/test_litellm/test_daybreak_model_metadata.py
@@ -32,20 +32,6 @@ def _load(path):
         return json.load(f)
 
 
-@pytest.mark.parametrize("model", DAYBREAK_MODELS)
-def test_daybreak_capability_contract(model):
-    info = _load(MAIN_PATH).get(model)
-    assert info is not None, f"{model} missing from model_prices_and_context_window.json"
-
-    assert info["litellm_provider"] == "openai"
-    assert info["mode"] == "chat"
-    assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"]
-
-    assert info["supports_computer_use"] is True
-    assert info["supports_parallel_function_calling"] is True
-    assert info["supports_function_calling"] is True
-    assert info["supports_reasoning"] is True
-    assert info["supports_vision"] is True
 
 
 def test_blue_alias_matches_its_snapshot_computer_use():
diff --git a/tests/test_litellm/test_deepseek_model_metadata.py b/tests/test_litellm/test_deepseek_model_metadata.py
index b9eb33f0972..a90ecd0ca59 100644
--- a/tests/test_litellm/test_deepseek_model_metadata.py
+++ b/tests/test_litellm/test_deepseek_model_metadata.py
@@ -39,25 +39,9 @@ class TestDeepSeekModelCostEntries:
     """Verify that provider-prefixed DeepSeek entries contain the same
     capability flags as their bare-name counterparts in the JSON files."""
 
-    def test_deepseek_chat_supports_response_schema_in_backup(self):
-        data = _load_backup_json()
-        entry = data.get("deepseek/deepseek-chat", {})
-        assert entry.get("supports_response_schema") is True
 
-    def test_deepseek_reasoner_supports_response_schema_in_backup(self):
-        data = _load_backup_json()
-        entry = data.get("deepseek/deepseek-reasoner", {})
-        assert entry.get("supports_response_schema") is True
 
-    def test_deepseek_chat_supports_system_messages_in_backup(self):
-        data = _load_backup_json()
-        entry = data.get("deepseek/deepseek-chat", {})
-        assert entry.get("supports_system_messages") is True
 
-    def test_deepseek_reasoner_supports_system_messages_in_backup(self):
-        data = _load_backup_json()
-        entry = data.get("deepseek/deepseek-reasoner", {})
-        assert entry.get("supports_system_messages") is True
 
     def test_deepseek_chat_max_input_tokens_matches_bare_in_backup(self):
         data = _load_backup_json()
@@ -71,25 +55,7 @@ class TestDeepSeekModelCostEntries:
         prefixed = data.get("deepseek/deepseek-reasoner", {})
         assert prefixed.get("max_output_tokens") == bare.get("max_output_tokens")
 
-    def test_main_json_deepseek_chat_supports_response_schema(self):
-        main_path = os.path.join(
-            os.path.dirname(os.path.dirname(litellm.__file__)),
-            "model_prices_and_context_window.json",
-        )
-        with open(main_path, encoding="utf-8") as f:
-            data = json.load(f)
-        entry = data.get("deepseek/deepseek-chat", {})
-        assert entry.get("supports_response_schema") is True
 
-    def test_main_json_deepseek_reasoner_supports_response_schema(self):
-        main_path = os.path.join(
-            os.path.dirname(os.path.dirname(litellm.__file__)),
-            "model_prices_and_context_window.json",
-        )
-        with open(main_path, encoding="utf-8") as f:
-            data = json.load(f)
-        entry = data.get("deepseek/deepseek-reasoner", {})
-        assert entry.get("supports_response_schema") is True
 
 
 # ---------------------------------------------------------------------------
diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py
index a7a9e0fc37d..858c9983221 100644
--- a/tests/test_litellm/test_fireworks_serverless_model_costs.py
+++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py
@@ -14,24 +14,9 @@ import os
 
 import pytest
 
-import litellm
 from litellm.utils import get_model_info
 
 
-@pytest.fixture(scope="module", autouse=True)
-def _local_model_cost_map():
-    """
-    Point litellm at the bundled cost map for the duration of this module
-    only. ``mp.undo()`` restores both the environment variable and
-    ``litellm.model_cost`` so nothing leaks into later tests.
-    """
-    mp = pytest.MonkeyPatch()
-    mp.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-    mp.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
-    get_model_info.cache_clear()
-    yield
-    mp.undo()
-    get_model_info.cache_clear()
 
 
 NEW_ENTRIES = {
@@ -54,17 +39,6 @@ def model_data():
         return json.load(f)
 
 
-def test_fireworks_serverless_entries_exist(model_data):
-    """The new prefixed entry carries the pricing and metadata from #37274."""
-    for key, expected in NEW_ENTRIES.items():
-        assert key in model_data, f"{key} is missing from model_prices_and_context_window.json"
-        entry = model_data[key]
-        for field, value in expected.items():
-            assert entry[field] == pytest.approx(value), f"{key}.{field}"
-        assert entry["litellm_provider"] == "fireworks_ai"
-        assert entry["mode"] == "chat"
-        assert entry["supports_function_calling"] is True
-        assert entry["supports_vision"] is False
 
 
 def test_bare_fireworks_ids_resolve_through_prefixed_entries():
diff --git a/tests/test_litellm/test_gpt_5_5_model_metadata.py b/tests/test_litellm/test_gpt_5_5_model_metadata.py
index a60fa9466e6..32fbfc533b8 100644
--- a/tests/test_litellm/test_gpt_5_5_model_metadata.py
+++ b/tests/test_litellm/test_gpt_5_5_model_metadata.py
@@ -1,53 +1,9 @@
 import json
 from pathlib import Path
 
-import pytest
-
-from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
 
 
-@pytest.mark.parametrize("model", ["azure_ai/gpt-5.5", "azure_ai/gpt-5.5-2026-04-23"])
-def test_azure_ai_gpt_5_5_model_info(model):
-    json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json"
-    with open(json_path) as f:
-        model_cost = json.load(f)
 
-    info = model_cost.get(model)
-    assert (
-        info is not None
-    ), f"{model} not found in model_prices_and_context_window.json"
-
-    assert info["litellm_provider"] == "azure_ai"
-    assert info["mode"] == "chat"
-
-    assert info["input_cost_per_token"] == 5e-06
-    assert info["output_cost_per_token"] == 3e-05
-    assert info["cache_read_input_token_cost"] == 5e-07
-
-    assert info["input_cost_per_token_above_272k_tokens"] == 1e-05
-    assert info["output_cost_per_token_above_272k_tokens"] == 4.5e-05
-    assert info["cache_read_input_token_cost_above_272k_tokens"] == 1e-06
-
-    assert info["input_cost_per_token_priority"] == 1e-05
-    assert info["output_cost_per_token_priority"] == 6e-05
-
-    assert info["max_input_tokens"] == 1050000
-    assert info["max_output_tokens"] == 128000
-    assert info["max_tokens"] == 128000
-
-    assert info["supports_function_calling"] is True
-    assert info["supports_prompt_caching"] is True
-    assert info["supports_reasoning"] is True
-    assert info["supports_response_schema"] is True
-    assert info["supports_tool_choice"] is True
-    assert info["supports_vision"] is True
-    assert info["supports_web_search"] is True
-    # gpt-5.5 dropped minimal reasoning effort support (true on gpt-5.4)
-    assert info["supports_minimal_reasoning_effort"] is False
-
-    routed_model, provider, _, _ = get_llm_provider(model=model)
-    assert routed_model == model.split("/", 1)[1]
-    assert provider == "azure_ai"
 
 
 def test_azure_ai_gpt_5_5_backup_matches_main():
diff --git a/tests/test_litellm/test_gpt_realtime_mode.py b/tests/test_litellm/test_gpt_realtime_mode.py
index 314fd63c4cc..8b730d737b3 100644
--- a/tests/test_litellm/test_gpt_realtime_mode.py
+++ b/tests/test_litellm/test_gpt_realtime_mode.py
@@ -1,10 +1,8 @@
 import json
 from pathlib import Path
 
-import pytest
 from typing_extensions import get_args, get_type_hints
 
-import litellm
 from litellm.types.utils import ModelInfoBase
 
 REALTIME_ONLY_GPT_MODELS = (
@@ -43,10 +41,6 @@ REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS = (
 ALL_REALTIME_ONLY_GPT_MODELS = REALTIME_ONLY_GPT_MODELS + REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS
 
 
-def _load_cost_map() -> dict:
-    json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json"
-    with open(json_path) as f:
-        return json.load(f)
 
 
 def test_realtime_is_a_valid_mode_literal():
@@ -54,31 +48,10 @@ def test_realtime_is_a_valid_mode_literal():
     assert "realtime" in get_args(hints["mode"])
 
 
-@pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS)
-def test_realtime_only_gpt_models_are_mode_realtime(model):
-    """These models only serve /v1/realtime and are rejected by /v1/chat/completions
-    ("This is not a chat model ..."), so they must not be tagged mode=chat."""
-    info = _load_cost_map()[model]
-    assert info["supported_endpoints"] == ["/v1/realtime"]
-    assert info["mode"] == "realtime"
 
 
-@pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS)
-def test_realtime_only_gpt_4o_models_are_mode_realtime(model):
-    """gpt-4o(-mini)-realtime-preview are realtime-only and must not be mode=chat."""
-    assert _load_cost_map()[model]["mode"] == "realtime"
 
 
-def test_get_model_info_reports_realtime_mode(monkeypatch):
-    """get_model_info must resolve the retag against the bundled cost map, not the
-    hosted map fetched from main, which lags this repo until the next promotion."""
-    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-    monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
-    litellm.get_model_info.cache_clear()
-    try:
-        assert litellm.get_model_info("gpt-realtime-mini")["mode"] == "realtime"
-    finally:
-        litellm.get_model_info.cache_clear()
 
 
 def test_backup_matches_main_for_realtime_models():
diff --git a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py
index 6f1ba702d8d..945b6e19897 100644
--- a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py
+++ b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py
@@ -3,8 +3,6 @@ from pathlib import Path
 
 import pytest
 
-import litellm
-from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
 
 REPO_ROOT = Path(__file__).parents[2]
 MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json"
@@ -28,53 +26,10 @@ def _load(path):
 
 
 
-@pytest.mark.parametrize("model", MEDIUM_3_5_MODELS)
-def test_medium_3_5_specs(model):
-    info = _load(MAIN_PATH).get(model)
-    assert info is not None, f"{model} missing from model_prices_and_context_window.json"
-
-    assert info["litellm_provider"] == "mistral"
-    assert info["mode"] == "chat"
-
-    assert info["input_cost_per_token"] == 1.5e-06
-    assert info["output_cost_per_token"] == 7.5e-06
-
-    assert info["max_input_tokens"] == 262144
-    assert info["max_output_tokens"] == 262144
-    assert info["max_tokens"] == 262144
-
-    assert info["supports_reasoning"] is True
-    assert info["supports_vision"] is True
-    assert info["supports_function_calling"] is True
-    assert info["supports_response_schema"] is True
-    assert info["supports_tool_choice"] is True
-    assert info["supports_assistant_prefill"] is True
-
-    routed_model, provider, _, _ = get_llm_provider(model=model)
-    assert routed_model == model.split("/", 1)[1]
-    assert provider == "mistral"
 
 
-def test_mistral_medium_latest_resolves_to_medium_3_5(local_model_cost_map):
-    """LIT-3883: the -latest alias was retargeted to Medium 3.5; get_model_info must
-    return the 3.5 pricing/context/reasoning, not the stale Medium 3.1 values."""
-    info = litellm.get_model_info(model="mistral/mistral-medium-latest")
-
-    assert info["input_cost_per_token"] == 1.5e-06
-    assert info["output_cost_per_token"] == 7.5e-06
-    assert info["max_input_tokens"] == 262144
-    assert info["supports_reasoning"] is True
 
 
-def test_mistral_medium_2508_keeps_medium_3_1_specs():
-    """The date-pinned 2508 alias is Medium 3.1 and must not inherit 3.5 pricing."""
-    info = _load(MAIN_PATH).get("mistral/mistral-medium-2508")
-    assert info is not None, "mistral/mistral-medium-2508 missing from cost map"
-
-    assert info["input_cost_per_token"] == 4e-07
-    assert info["output_cost_per_token"] == 2e-06
-    assert info["max_input_tokens"] == 131072
-    assert info.get("supports_reasoning") is not True
 
 
 @pytest.mark.parametrize("model", SYNCED_MODELS)
diff --git a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py b/tests/test_litellm/test_mistral_small_4_0_model_metadata.py
index 0442321ba0b..16b126a017c 100644
--- a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py
+++ b/tests/test_litellm/test_mistral_small_4_0_model_metadata.py
@@ -18,27 +18,6 @@ def _load(path):
         return json.load(f)
 
 
-@pytest.mark.parametrize("model", SMALL_4_0_MODELS)
-def test_small_4_0_specs(model):
-    info = _load(MAIN_PATH).get(model)
-    assert info is not None, f"{model} missing from model_prices_and_context_window.json"
-
-    assert info["litellm_provider"] == "mistral"
-    assert info["mode"] == "chat"
-
-    assert info["input_cost_per_token"] == 1.5e-07
-    assert info["output_cost_per_token"] == 6e-07
-
-    assert info["max_input_tokens"] == 262144
-    assert info["max_output_tokens"] == 262144
-    assert info["max_tokens"] == 262144
-
-    assert info["supports_reasoning"] is True
-    assert info["supports_vision"] is True
-    assert info["supports_function_calling"] is True
-    assert info["supports_response_schema"] is True
-    assert info["supports_tool_choice"] is True
-    assert info["supports_assistant_prefill"] is True
 
 
 @pytest.mark.parametrize("model", SMALL_4_0_MODELS)
diff --git a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py
index 0587883aa44..fd2224d7720 100644
--- a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py
+++ b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py
@@ -24,43 +24,6 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di
 
 
 
-@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING)
-def test_muse_spark_1_2_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float):
-    info = _load_cost_map().get(model)
-    assert info is not None, f"{model} not found in model_prices_and_context_window.json"
-
-    assert info["litellm_provider"] == "meta"
-    assert info["mode"] == "chat"
-
-    assert info["input_cost_per_token"] == input_cost
-    assert info["output_cost_per_token"] == output_cost
-    assert info["cache_read_input_token_cost"] == cached_cost
-
-    assert info["max_input_tokens"] == 1048576
-    assert info["max_output_tokens"] == 131072
-    assert info["max_tokens"] == 131072
-
-    assert info["supports_function_calling"] is True
-    assert info["supports_parallel_function_calling"] is True
-    assert info["supports_prompt_caching"] is True
-    assert info["supports_reasoning"] is True
-    assert info["supports_response_schema"] is True
-    assert info["supports_tool_choice"] is True
-    assert info["supports_vision"] is True
-    assert info["supports_pdf_input"] is True
-    assert info["supports_web_search"] is True
-    assert info["supports_minimal_reasoning_effort"] is True
-    assert info["supports_xhigh_reasoning_effort"] is True
-
-    assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"]
-    assert info["supported_modalities"] == ["text", "image", "video"]
-    assert info["supported_output_modalities"] == ["text"]
-
-    assert info["search_context_cost_per_query"] == {
-        "search_context_size_high": WEB_SEARCH_COST_PER_QUERY,
-        "search_context_size_low": WEB_SEARCH_COST_PER_QUERY,
-        "search_context_size_medium": WEB_SEARCH_COST_PER_QUERY,
-    }
 
 
 @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING)
diff --git a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py b/tests/test_litellm/test_muse_spark_1_3_model_metadata.py
index 1ecd9490f78..3328e916af5 100644
--- a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py
+++ b/tests/test_litellm/test_muse_spark_1_3_model_metadata.py
@@ -24,43 +24,6 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di
 
 
 
-@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING)
-def test_muse_spark_1_3_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float):
-    info = _load_cost_map().get(model)
-    assert info is not None, f"{model} not found in model_prices_and_context_window.json"
-
-    assert info["litellm_provider"] == "meta"
-    assert info["mode"] == "chat"
-
-    assert info["input_cost_per_token"] == input_cost
-    assert info["output_cost_per_token"] == output_cost
-    assert info["cache_read_input_token_cost"] == cached_cost
-
-    assert info["max_input_tokens"] == 1048576
-    assert info["max_output_tokens"] == 131072
-    assert info["max_tokens"] == 131072
-
-    assert info["supports_function_calling"] is True
-    assert info["supports_parallel_function_calling"] is True
-    assert info["supports_prompt_caching"] is True
-    assert info["supports_reasoning"] is True
-    assert info["supports_response_schema"] is True
-    assert info["supports_tool_choice"] is True
-    assert info["supports_vision"] is True
-    assert info["supports_pdf_input"] is True
-    assert info["supports_web_search"] is True
-    assert info["supports_minimal_reasoning_effort"] is True
-    assert info["supports_xhigh_reasoning_effort"] is True
-
-    assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"]
-    assert info["supported_modalities"] == ["text", "image", "video"]
-    assert info["supported_output_modalities"] == ["text"]
-
-    assert info["search_context_cost_per_query"] == {
-        "search_context_size_high": WEB_SEARCH_COST_PER_QUERY,
-        "search_context_size_low": WEB_SEARCH_COST_PER_QUERY,
-        "search_context_size_medium": WEB_SEARCH_COST_PER_QUERY,
-    }
 
 
 @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING)
diff --git a/tests/test_litellm/test_replicate_model_key_format.py b/tests/test_litellm/test_replicate_model_key_format.py
index 8c2b72f8ed2..8c52ae3603e 100644
--- a/tests/test_litellm/test_replicate_model_key_format.py
+++ b/tests/test_litellm/test_replicate_model_key_format.py
@@ -20,12 +20,6 @@ def test_replicate_models_have_valid_key_prefix(model_cost: dict[str, Any]) -> N
     )
 
 
-def test_replicate_openai_gpt_oss_20b_key_exists(model_cost: dict[str, Any]) -> None:
-    assert "replicate/openai/gpt-oss-20b" in model_cost
-    info = model_cost["replicate/openai/gpt-oss-20b"]
-    assert info["litellm_provider"] == "replicate"
-    assert info["mode"] == "chat"
-    assert info["supports_function_calling"] is True
 
 
 def test_replicate_backup_matches_main() -> None:
diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/test_litellm/test_together_ai_model_metadata.py
index c9e2863d240..fd572733466 100644
--- a/tests/test_litellm/test_together_ai_model_metadata.py
+++ b/tests/test_litellm/test_together_ai_model_metadata.py
@@ -5,7 +5,6 @@ from typing import Final
 import pytest
 from pydantic import TypeAdapter
 
-from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
 
 REPO_ROOT: Final = Path(__file__).parents[2]
 
@@ -77,57 +76,12 @@ def cost_map() -> CostMap:
         return COST_MAP_ADAPTER.validate_python(json.load(f))
 
 
-@pytest.mark.parametrize("model", SERVERLESS_CHAT_MODELS)
-def test_together_serverless_chat_model_is_mapped(cost_map: CostMap, model: str):
-    info = cost_map.get(model)
-    assert info is not None, f"{model} missing from model_prices_and_context_window.json"
-    assert info["litellm_provider"] == "together_ai"
-    assert info["mode"] == "chat"
-    assert info["input_cost_per_token"] >= 0
-    assert info["output_cost_per_token"] >= info["input_cost_per_token"]
-    assert "deprecation_date" not in info
-
-    routed_model, provider, _, _ = get_llm_provider(model=model)
-    assert routed_model == model.removeprefix("together_ai/")
-    assert provider == "together_ai"
 
 
-def test_together_kimi_k3_pricing_and_capabilities(cost_map: CostMap):
-    info = cost_map["together_ai/moonshotai/Kimi-K3"]
-    assert info["input_cost_per_token"] == 3e-06
-    assert info["output_cost_per_token"] == 1.5e-05
-    assert info["max_input_tokens"] == 1048576
-    assert info["supports_function_calling"] is True
-    assert info["supports_tool_choice"] is True
-    assert info["supports_response_schema"] is True
-    assert info["supports_vision"] is True
-    assert info["supports_reasoning"] is True
 
 
-def test_together_glm_52_pricing(cost_map: CostMap):
-    info = cost_map["together_ai/zai-org/GLM-5.2"]
-    assert info["input_cost_per_token"] == 1.4e-06
-    assert info["output_cost_per_token"] == 4.4e-06
-    assert info["max_input_tokens"] == 1048575
-    assert info["max_output_tokens"] == 128000
-    assert info["supports_function_calling"] is True
-    assert info["supports_reasoning"] is True
 
 
-def test_together_glm_53_flash_pricing_and_capabilities(cost_map: CostMap):
-    info = cost_map["together_ai/zai-org/GLM-5.3-Flash"]
-    assert info["input_cost_per_token"] == 1.5e-07
-    assert info["output_cost_per_token"] == 5e-07
-    assert info["cache_read_input_token_cost"] == 3e-08
-    assert info["max_input_tokens"] == 1048575
-    assert info["max_output_tokens"] == 128000
-    assert info["supports_function_calling"] is True
-    assert info["supports_parallel_function_calling"] is True
-    assert info["supports_prompt_caching"] is True
-    assert info["supports_tool_choice"] is True
-    assert info["supports_response_schema"] is True
-    assert info["supports_vision"] is True
-    assert info["supports_reasoning"] is True
 
 
 def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost_map: CostMap):
@@ -142,19 +96,8 @@ def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost
     assert inflated == []
 
 
-def test_together_multilingual_e5_embedding_entry(cost_map: CostMap):
-    info = cost_map["together_ai/intfloat/multilingual-e5-large-instruct"]
-    assert info["mode"] == "embedding"
-    assert info["input_cost_per_token"] == 2e-08
-    assert info["max_input_tokens"] == 514
-    assert info["output_vector_size"] == 1024
 
 
-def test_together_llama_33_70b_repriced_to_current_together_rate(cost_map: CostMap):
-    info = cost_map["together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo"]
-    assert info["input_cost_per_token"] == 1.04e-06
-    assert info["output_cost_per_token"] == 1.04e-06
-    assert info["max_input_tokens"] == 131072
 
 
 @pytest.mark.parametrize("model", sorted(DEPRECATED_MODELS))
@@ -210,32 +153,9 @@ CACHED_INPUT_MODELS: Final = (
 )
 
 
-@pytest.mark.parametrize("model", CACHED_INPUT_MODELS)
-def test_together_cached_input_model_carries_cache_read_pricing(cost_map: CostMap, model: str):
-    info = cost_map.get(model)
-    assert info is not None, f"{model} missing from model_prices_and_context_window.json"
-    assert info.get("supports_prompt_caching") is True
-    cache_read = info.get("cache_read_input_token_cost")
-    assert isinstance(cache_read, float)
-    assert 0 < cache_read < info["input_cost_per_token"]
-    assert "cache_creation_input_token_cost" not in info
 
 
 def test_together_prompt_caching_flag_implies_cache_read_rate(cost_map: CostMap):
     for model, info in cost_map.items():
         if model.startswith("together_ai/") and info.get("supports_prompt_caching"):
             assert "cache_read_input_token_cost" in info, f"{model} flags caching without a cache read rate"
-
-
-def test_together_deepseek_v4_flash_cache_read_rate(cost_map: CostMap):
-    info = cost_map["together_ai/deepseek-ai/DeepSeek-V4-Flash-0731"]
-    assert info["input_cost_per_token"] == 1.4e-07
-    assert info["cache_read_input_token_cost"] == 3e-08
-    assert info["output_cost_per_token"] == 2.8e-07
-
-
-def test_together_qwen_37_max_repriced_to_current_together_rate(cost_map: CostMap):
-    info = cost_map["together_ai/Qwen/Qwen3.7-Max"]
-    assert info["input_cost_per_token"] == 2.5e-06
-    assert info["output_cost_per_token"] == 7.5e-06
-    assert info["cache_read_input_token_cost"] == 5e-07

From ac573fd66e855d834606316a26dca61d37305f2f Mon Sep 17 00:00:00 2001
From: mateo <mateo@berri.ai>
Date: Tue, 8 Sep 2026 02:20:39 +0000
Subject: [PATCH 154/164] test: remove remaining static cost assertions

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../test_bedrock_extended_beta_models.py      | 116 ------------------
 .../test_bedrock_usgov_pricing.py             |  12 --
 2 files changed, 128 deletions(-)
 delete mode 100644 tests/test_litellm/test_bedrock_extended_beta_models.py

diff --git a/tests/test_litellm/test_bedrock_extended_beta_models.py b/tests/test_litellm/test_bedrock_extended_beta_models.py
deleted file mode 100644
index d55aac762fa..00000000000
--- a/tests/test_litellm/test_bedrock_extended_beta_models.py
+++ /dev/null
@@ -1,116 +0,0 @@
-"""
-Test suite for AWS Bedrock extended beta model support
-Tests model configuration, pricing, and regional availability for:
-- DeepSeek V3.2
-- Minimax M2.1
-- Moonshot AI Kimi K2.5
-- Qwen3 Coder Next
-"""
-
-import os
-
-# Set env var to use local model cost map instead of fetching from remote
-os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "true"
-
-import pytest
-
-from litellm import get_model_info
-
-# Model configurations: (model_name, regions, max_input, max_output)
-MODEL_CONFIGS = [
-    (
-        "deepseek.v3.2",
-        [
-            "ap-northeast-1",
-            "ap-south-1",
-            "ap-southeast-3",
-            "eu-north-1",
-            "sa-east-1",
-            "us-east-1",
-            "us-east-2",
-            "us-west-2",
-        ],
-        163840,
-        163840,
-    ),
-    (
-        "minimax.minimax-m2.1",
-        [
-            "ap-northeast-1",
-            "ap-south-1",
-            "ap-southeast-3",
-            "eu-central-1",
-            "eu-north-1",
-            "eu-south-1",
-            "eu-west-1",
-            "eu-west-2",
-            "sa-east-1",
-            "us-east-1",
-            "us-east-2",
-            "us-west-2",
-        ],
-        196000,
-        8192,
-    ),
-    (
-        "moonshotai.kimi-k2.5",
-        [
-            "ap-northeast-1",
-            "ap-south-1",
-            "ap-southeast-3",
-            "eu-north-1",
-            "sa-east-1",
-            "us-east-1",
-            "us-east-2",
-            "us-west-2",
-        ],
-        262144,
-        262144,
-    ),
-    (
-        "qwen.qwen3-coder-next",
-        [
-            "ap-northeast-1",
-            "ap-south-1",
-            "ap-southeast-3",
-            "eu-central-1",
-            "eu-south-1",
-            "eu-west-1",
-            "eu-west-2",
-            "sa-east-1",
-            "us-east-1",
-            "us-east-2",
-            "us-west-2",
-        ],
-        262144,
-        8192,
-    ),
-]
-
-
-class TestBedrockNewModels:
-    """Unified test suite for all new Bedrock models"""
-
-
-    @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS)
-    def test_pricing_configured(self, model_name, regions, max_input, max_output):
-        """Verify pricing is set for all models"""
-        model = f"bedrock/us-east-1/{model_name}"
-        model_info = get_model_info(model)
-
-        assert (
-            model_info["input_cost_per_token"] > 0
-        ), f"Missing input cost for {model_name}"
-        assert (
-            model_info["output_cost_per_token"] > 0
-        ), f"Missing output cost for {model_name}"
-
-    @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS)
-    def test_region_count(self, model_name, regions, max_input, max_output):
-        """Verify each bedrock/{region}/{model_name} resolves via get_model_info"""
-        for region in regions:
-            model = f"bedrock/{region}/{model_name}"
-            model_info = get_model_info(model)
-            assert model_info is not None, f"Model {model_name} not found in {region}"
-            assert model_info["max_input_tokens"] == max_input
-            assert model_info["max_output_tokens"] == max_output
diff --git a/tests/test_litellm/test_bedrock_usgov_pricing.py b/tests/test_litellm/test_bedrock_usgov_pricing.py
index 1469ec6a1bb..9e7e1e5c9f6 100644
--- a/tests/test_litellm/test_bedrock_usgov_pricing.py
+++ b/tests/test_litellm/test_bedrock_usgov_pricing.py
@@ -68,18 +68,6 @@ EXPECTED_USGOV_ABOVE_200K = {
 }
 
 
-@pytest.mark.parametrize("field,expected", EXPECTED_USGOV_ABOVE_200K.items())
-def test_usgov_cross_region_above_200k_carries_gov_premium(model_data, field, expected):
-    """The `_above_200k_tokens` tier on the us-gov cross-region inference
-    profile must also carry the +20% GovCloud uplift. The original PR
-    corrected the base rates but left the 200k-tier fields at the +10%
-    commercial-US rates, undercharging long-context requests.
-    """
-    info = model_data[USGOV_CROSS_REGION_KEY]
-    assert field in info, f"{USGOV_CROSS_REGION_KEY}: missing field {field}"
-    assert info[field] == expected, f"{USGOV_CROSS_REGION_KEY}: {field} should be {expected} (got {info[field]})"
-
-
 def test_usgov_cross_region_above_200k_ratio_to_global(model_data):
     """Cross-check via the property-based invariant: every `_above_200k_tokens`
     field on the us-gov cross-region profile must equal 1.2x the global

From 5cfe20a68d541580fd1d7198a136ec56e88c2255 Mon Sep 17 00:00:00 2001
From: mateo <mateo@berri.ai>
Date: Tue, 8 Sep 2026 02:21:55 +0000
Subject: [PATCH 155/164] test: collapse blank lines left by removed tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../llm_translation/test_bedrock_govcloud.py  |  3 --
 tests/llm_translation/test_hyperbolic.py      |  3 --
 tests/llm_translation/test_lambda_ai.py       |  2 -
 tests/llm_translation/test_morph.py           |  2 -
 tests/llm_translation/test_openai_o1.py       |  3 --
 tests/local_testing/test_get_model_info.py    |  8 ----
 .../test_xai_oauth_routing.py                 |  2 -
 .../test_mai_image_generation.py              |  1 -
 .../test_azure_ai_fw_models_metadata.py       |  5 ---
 .../test_azure_ai_kimi_k26_metadata.py        |  4 --
 ..._cross_region_inference_profile_mapping.py |  8 ----
 ...bedrock_mantle_responses_transformation.py |  4 --
 .../test_bedrock_mantle_transformation.py     |  7 ----
 .../test_fireworks_ai_chat_transformation.py  |  4 --
 .../test_fireworks_ai_kimi_model_metadata.py  |  2 -
 .../test_gemini_realtime_transformation.py    |  4 --
 .../test_inception_chat_transformation.py     |  2 -
 ...est_inception_completion_transformation.py |  2 -
 .../test_moonshot_chat_transformation.py      |  6 ---
 .../llms/openai_like/test_json_providers.py   |  1 -
 .../openai_like/test_libertai_provider.py     |  2 -
 .../test_perplexity_cost_calculator.py        |  2 -
 .../test_vertex_video_transformation.py       |  1 -
 .../xai/test_xai_redirected_slug_pricing.py   |  4 --
 .../llms/zai/test_zai_provider.py             |  7 ----
 .../test_bedrock_usgov_pricing.py             | 38 -------------------
 .../test_claude_fable_5_config.py             | 15 --------
 .../test_claude_haiku_4_5_config.py           |  2 -
 .../test_claude_opus_4_6_config.py            |  4 --
 .../test_claude_opus_4_8_config.py            |  9 -----
 .../test_litellm/test_claude_opus_5_config.py |  7 ----
 .../test_claude_sonnet_5_config.py            |  7 ----
 ...st_cloudflare_workers_ai_model_metadata.py | 10 -----
 .../test_daybreak_model_metadata.py           |  2 -
 .../test_deepseek_model_metadata.py           |  6 ---
 .../test_fireworks_serverless_model_costs.py  |  4 --
 .../test_gpt_5_5_model_metadata.py            |  4 --
 tests/test_litellm/test_gpt_realtime_mode.py  |  8 ----
 .../test_mistral_medium_3_5_model_metadata.py |  7 ----
 .../test_mistral_small_4_0_model_metadata.py  |  2 -
 .../test_muse_spark_1_2_model_metadata.py     |  3 --
 .../test_muse_spark_1_3_model_metadata.py     |  3 --
 .../test_replicate_model_key_format.py        |  2 -
 .../test_together_ai_model_metadata.py        | 14 -------
 44 files changed, 236 deletions(-)

diff --git a/tests/llm_translation/test_bedrock_govcloud.py b/tests/llm_translation/test_bedrock_govcloud.py
index a69b786fd45..3ac1fa7cf2e 100644
--- a/tests/llm_translation/test_bedrock_govcloud.py
+++ b/tests/llm_translation/test_bedrock_govcloud.py
@@ -40,7 +40,6 @@ class TestBedrockGovCloudSupport:
         assert "us-gov-east-1" in all_regions
         assert "us-gov-west-1" in all_regions
 
-
     def test_govcloud_model_routing(self):
         """Test that GovCloud models are routed correctly"""
         # Test Claude model routing
@@ -117,8 +116,6 @@ class TestBedrockGovCloudSupport:
         assert not any("us-gov-east-1" in model for model in litellm.bedrock_models)
         assert not any("us-gov-west-1" in model for model in litellm.bedrock_models)
 
-
-
     @patch("litellm.completion")
     def test_govcloud_completion_cost_calculation(self, mock_completion):
         """Test that completion requests use correct pricing for GovCloud models"""
diff --git a/tests/llm_translation/test_hyperbolic.py b/tests/llm_translation/test_hyperbolic.py
index 0dd1c4924c0..b7206e40a4e 100644
--- a/tests/llm_translation/test_hyperbolic.py
+++ b/tests/llm_translation/test_hyperbolic.py
@@ -1,6 +1,5 @@
 
 
-
 import litellm
 from litellm import get_llm_provider
 
@@ -65,8 +64,6 @@ def test_hyperbolic_in_provider_lists():
     assert "https://api.hyperbolic.xyz/v1" in openai_compatible_endpoints
 
 
-
-
 def test_hyperbolic_supported_params():
     """Test that supported OpenAI parameters are correctly configured"""
     from litellm.llms.hyperbolic.chat.transformation import HyperbolicChatConfig
diff --git a/tests/llm_translation/test_lambda_ai.py b/tests/llm_translation/test_lambda_ai.py
index b2fb72f8412..edba459b352 100644
--- a/tests/llm_translation/test_lambda_ai.py
+++ b/tests/llm_translation/test_lambda_ai.py
@@ -102,8 +102,6 @@ async def test_lambda_ai_completion_call():
             raise
 
 
-
-
 def test_lambda_ai_model_list_populated():
     """Test that lambda_ai_models list is populated correctly"""
     # Ensure we're using local model cost map and repopulate models
diff --git a/tests/llm_translation/test_morph.py b/tests/llm_translation/test_morph.py
index 47ad3a1749b..752fb3b9083 100644
--- a/tests/llm_translation/test_morph.py
+++ b/tests/llm_translation/test_morph.py
@@ -68,8 +68,6 @@ def test_morph_in_provider_lists():
     )
 
 
-
-
 def test_morph_supported_params():
     """Test that MorphChatConfig returns correct supported parameters."""
     config = MorphChatConfig()
diff --git a/tests/llm_translation/test_openai_o1.py b/tests/llm_translation/test_openai_o1.py
index 9de5d5d9431..fd25e04d67d 100644
--- a/tests/llm_translation/test_openai_o1.py
+++ b/tests/llm_translation/test_openai_o1.py
@@ -2,7 +2,6 @@ import os
 from unittest.mock import patch
 
 
-
 import pytest
 
 import litellm
@@ -182,8 +181,6 @@ class TestOpenAIO3(BaseOSeriesModelsTest, BaseLLMChatTest):
         pass
 
 
-
-
 def test_o3_reasoning_effort():
     resp = litellm.completion(
         model="o3-mini",
diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py
index 562ed240b9c..38ccfd91f95 100644
--- a/tests/local_testing/test_get_model_info.py
+++ b/tests/local_testing/test_get_model_info.py
@@ -47,14 +47,6 @@ def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch):
     assert model_info["input_cost_per_token"] == 0.0
 
 
-
-
-
-
-
-
-
-
 def test_get_model_info_gemini_pro():
     info = litellm.get_model_info("gemini-2.0-flash")
     print("info", info)
diff --git a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py
index 83ede898e49..d24e2b58db8 100644
--- a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py
+++ b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py
@@ -45,8 +45,6 @@ def test_xai_openai_compatible_provider_info():
     assert dynamic_api_key == "api-key"
 
 
-
-
 def test_xai_validate_environment_reads_api_key(monkeypatch):
     monkeypatch.setenv("XAI_API_KEY", "api-key")
 
diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
index 669c566f96b..9bdc79919d2 100644
--- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
+++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
@@ -37,7 +37,6 @@ class TestAzureMAIImageGeneration:
         assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro")
         assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-DS-R1")
 
-
     def test_get_mai_image_generation_url(self):
         url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url(
             api_base="https://my-resource.services.ai.azure.com",
diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py
index d9b948e212a..1b2ca298694 100644
--- a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py
+++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py
@@ -13,7 +13,6 @@ from importlib.resources import files
 import pytest
 
 
-
 @pytest.fixture(scope="module")
 def use_local_model_cost_map():
     monkeypatch = pytest.MonkeyPatch()
@@ -39,8 +38,6 @@ def use_local_model_cost_map():
         monkeypatch.undo()
 
 
-
-
 @pytest.mark.parametrize(
     "model_name,expected_prompt,expected_completion",
     [
@@ -72,8 +69,6 @@ def test_azure_ai_fw_cost_per_token(
     assert completion_cost == pytest.approx(expected_completion)
 
 
-
-
 def test_azure_ai_fw_nemotron_lightning_supports_tool_choice(use_local_model_cost_map):
     from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig
 
diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py
index 18bdf60e9a0..cbcc2a94043 100644
--- a/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py
+++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py
@@ -33,10 +33,6 @@ def use_local_model_cost_map():
         monkeypatch.undo()
 
 
-
-
-
-
 def test_azure_ai_kimi_k26_cost_per_token(use_local_model_cost_map):
     from litellm.llms.azure_ai.cost_calculator import cost_per_token
     from litellm.types.utils import Usage
diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py
index c697bcb24b0..fda3c8ceb8f 100644
--- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py
+++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py
@@ -99,8 +99,6 @@ GPT_5_6_PROFILES = [
 ]
 
 
-
-
 def _bedrock_response(model, usage):
     return ModelResponse(
         id="test",
@@ -118,8 +116,6 @@ def _bedrock_response(model, usage):
     )
 
 
-
-
 def test_proxy_cost_calculation_scenario():
     """Test exact GitHub issue scenario: proxy cost calculation"""
     model = "litellm_proxy/bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0"
@@ -159,8 +155,6 @@ def test_bedrock_gpt_5_6_profiles_route_to_converse(profile, local_model_cost_ma
     assert BedrockModelInfo.get_bedrock_route(f"bedrock/{profile.model_id}") == "converse"
 
 
-
-
 def test_bedrock_gpt_5_6_above_272k_tier_applies_to_cost(local_model_cost_map):
     """A prompt over 272K tokens is billed at the long-context rate, not the base rate."""
     response = _bedrock_response(
@@ -220,8 +214,6 @@ def test_bedrock_gpt_5_6_bills_cache_write_tokens(local_model_cost_map):
     assert cost == pytest.approx(expected, rel=1e-9)
 
 
-
-
 @pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id)
 def test_bedrock_gpt_5_6_offers_tools_and_reasoning_effort_but_not_thinking(profile, local_model_cost_map):
     """GPT-5.x on Converse maps reasoning_effort to reasoning.effort, so reasoning_effort
diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
index 5994de28ba8..9457d5faaff 100644
--- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
+++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
@@ -157,7 +157,6 @@ class TestBedrockMantleResponsesURL:
         assert url == "https://bedrock-mantle.us-east-2.api.aws/v1/responses"
         assert url.count("/responses") == 1
 
-
     def test_url_aws_region_name_overrides_stale_api_base(self, monkeypatch):
         monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
         monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
@@ -1776,9 +1775,6 @@ class TestBedrockMantleResponsesSigV4:
 
 class TestBedrockMantleResponsesPricing:
 
-
-
-
     @pytest.mark.parametrize(
         "model, input_cost, output_cost",
         [
diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py
index e370cb22ce7..1be94d4daa2 100644
--- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py
+++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py
@@ -684,9 +684,6 @@ class TestBedrockMantleProviderResolution:
 class TestBedrockMantlePricing:
     """Tests that verify Bedrock Mantle uses correct AWS Bedrock pricing, not OpenAI pricing."""
 
-
-
-
     def test_safeguard_models_have_larger_output_tokens(self, monkeypatch):
         monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
         litellm.add_known_models()
@@ -697,10 +694,6 @@ class TestBedrockMantlePricing:
         assert info_safeguard["max_output_tokens"] > info_120b["max_output_tokens"]
 
 
-
-
-
-
 @pytest.mark.parametrize(
     "model_id",
     [
diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
index a79baef5ee5..d4ef4282b27 100644
--- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
+++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
@@ -16,8 +16,6 @@ from litellm.types.utils import (
 )
 
 
-
-
 def test_validate_environment_sets_session_affinity_from_litellm_session_id():
     config = FireworksAIConfig()
 
@@ -395,8 +393,6 @@ def test_get_supported_openai_params_parallel_tool_calls_without_tool_choice(
     assert "tool_choice" not in supported_params
 
 
-
-
 def test_get_provider_info_omits_false_supports_reasoning(monkeypatch):
     """Test that Fireworks only overrides supports_reasoning for supported models."""
     config = FireworksAIConfig()
diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py
index ba40f02ddc1..41f6ad9d99d 100644
--- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py
+++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py
@@ -56,8 +56,6 @@ def use_local_model_cost_map():
         monkeypatch.undo()
 
 
-
-
 @pytest.mark.parametrize("alias", KIMI_ALIASES)
 def test_fireworks_kimi_get_model_info_limits(use_local_model_cost_map, alias):
     model_info = use_local_model_cost_map.get_model_info(model=alias)
diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py
index 8295cf72524..2b3b6343fad 100644
--- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py
+++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py
@@ -306,8 +306,6 @@ def test_gemini_realtime_transformation_generation_complete():
     assert contains_audio_done_event, "Expected audio done event"
 
 
-
-
 def test_gemini_realtime_tool_call_transformation():
     """Test transformation of Gemini toolCall to OpenAI function_call_arguments.done format."""
     config = GeminiRealtimeConfig()
@@ -1831,8 +1829,6 @@ def test_is_audio_only_live_model_uses_cost_map(model, expected, patch_gemini_au
     assert GeminiRealtimeConfig._is_audio_only_live_model(model) == expected
 
 
-
-
 def test_is_setup_message_and_is_content_message():
     config = GeminiRealtimeConfig()
     assert config.is_setup_message({"setup": {}}) is True
diff --git a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py
index fff352a2f6c..4c0f5969249 100644
--- a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py
+++ b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py
@@ -231,8 +231,6 @@ def test_inception_in_provider_lists():
     assert "https://api.inceptionlabs.ai/v1" in litellm.openai_compatible_endpoints
 
 
-
-
 def test_inception_model_list_populated(monkeypatch):
     monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
     litellm.model_cost = litellm.get_model_cost_map(url="")
diff --git a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py
index 347cfe4cfc5..ed3f34fc744 100644
--- a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py
+++ b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py
@@ -143,8 +143,6 @@ async def test_inception_fim_async():
     assert r.choices[0].text == "a + b"
 
 
-
-
 def test_inception_fim_targets_fim_endpoint():
     """
     End-to-end: a FIM request must hit `/v1/fim/completions` (NOT
diff --git a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py
index 2d6751fca63..d484fa437ae 100644
--- a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py
+++ b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py
@@ -709,11 +709,6 @@ class TestKimiK26ModelRegistry:
         return GetModelCostMap.load_local_model_cost_map()
 
 
-
-
-
-
-
 class TestMoonshotResponseSchemaSupport:
     """Every model currently live on api.moonshot.ai supports json_schema
     response_format, which gates discovery via litellm.responses(). The flag
@@ -735,7 +730,6 @@ class TestMoonshotResponseSchemaSupport:
     def model_cost_map(self):
         return GetModelCostMap.load_local_model_cost_map()
 
-
     def test_supports_response_schema_utility_reports_true(self, model_cost_map, monkeypatch):
         monkeypatch.setattr(litellm, "model_cost", model_cost_map)
         assert litellm.utils.supports_response_schema(model="moonshot/kimi-k2.5") is True
diff --git a/tests/test_litellm/llms/openai_like/test_json_providers.py b/tests/test_litellm/llms/openai_like/test_json_providers.py
index fb5d28b8d3b..d84cc8d3237 100644
--- a/tests/test_litellm/llms/openai_like/test_json_providers.py
+++ b/tests/test_litellm/llms/openai_like/test_json_providers.py
@@ -318,7 +318,6 @@ class TestDarkbloom:
         assert config.custom_llm_provider == "darkbloom"
 
 
-
 class TestPublicAIIntegration:
     """Integration tests for PublicAI provider"""
 
diff --git a/tests/test_litellm/llms/openai_like/test_libertai_provider.py b/tests/test_litellm/llms/openai_like/test_libertai_provider.py
index dc7d5d18f36..c17eaf7c87f 100644
--- a/tests/test_litellm/llms/openai_like/test_libertai_provider.py
+++ b/tests/test_litellm/llms/openai_like/test_libertai_provider.py
@@ -59,7 +59,6 @@ class TestLibertAIProviderConfig:
         assert api_base == "https://custom.example.com/v1"
         assert api_key == "sk-test"
 
-
     def test_libertai_router_config(self):
         """Test that libertai can be used in Router configuration"""
         from litellm import Router
@@ -79,7 +78,6 @@ class TestLibertAIProviderConfig:
         assert len(router.model_list) == 1
         assert router.model_list[0]["model_name"] == "libertai-chat"
 
-
     def test_libertai_supported_endpoints_matrix(self):
         """The runtime-served backup matrix (GET /public/supported_endpoints) lists libertai."""
         import json
diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py
index 921022ce562..7556b215e66 100644
--- a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py
+++ b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py
@@ -316,7 +316,6 @@ class TestPerplexityCostCalculator:
 
         assert math.isclose(total_cost, expected_total, rel_tol=1e-6)
 
-
     @pytest.mark.parametrize("citation_tokens", [0, 10, 25, 100])
     @pytest.mark.parametrize("search_queries", [0, 1, 5, 10])
     @pytest.mark.parametrize("reasoning_tokens", [0, 15, 30])
@@ -462,7 +461,6 @@ class TestPerplexityCostCalculator:
         assert math.isclose(prompt_cost, expected_prompt, rel_tol=1e-9)
         assert math.isclose(completion_cost, expected_completion, rel_tol=1e-9)
 
-
     def test_agent_api_fallback_rates_price_a_response_without_metered_cost(self):
         """Perplexity meters cost on the response, but when `usage.cost` is absent the
         calculator falls back to the mapped per-token rates. Regression: that fallback
diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py
index 3c9112efb87..04e46eab1b7 100644
--- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py
@@ -136,7 +136,6 @@ class TestVertexAIVideoConfig:
         # Should NOT include endpoint
         assert not url.endswith(":predictLongRunning")
 
-
     def test_veo_31_lite_provider_routing_from_local_model_map(
         self, monkeypatch: pytest.MonkeyPatch
     ):
diff --git a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
index 1e410e41c33..a6b1c9a92ce 100644
--- a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
+++ b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
@@ -97,8 +97,6 @@ def test_redirected_slug_keeps_its_retirement_date(cost_map: dict, slug: str):
     assert cost_map[slug]["deprecation_date"] == expected_retirement_date(slug)
 
 
-
-
 @pytest.mark.parametrize("slug", REDIRECTED_SLUGS)
 def test_redirected_slug_carries_the_target_tier_rates(cost_map: dict, slug: str):
     """The request executes as grok-4.3, so it is tiered at grok-4.3's 200k boundary."""
@@ -108,8 +106,6 @@ def test_redirected_slug_carries_the_target_tier_rates(cost_map: dict, slug: str
         assert entry[field] == target[field], field
 
 
-
-
 def test_both_cost_maps_agree_on_the_redirected_slugs():
     prices = json.loads(PRICES_PATH.read_text(encoding="utf-8"))
     backup = json.loads(BACKUP_PRICES_PATH.read_text(encoding="utf-8"))
diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py
index 8d3744a00e0..069ac5727f6 100644
--- a/tests/test_litellm/llms/zai/test_zai_provider.py
+++ b/tests/test_litellm/llms/zai/test_zai_provider.py
@@ -55,12 +55,9 @@ def test_zai_in_provider_lists():
     assert "zai" in litellm.provider_list
 
 
-
-
 def test_zai_glm46_cost_calculation(local_model_cost_map):
     """Test the cost calculation for glm-4.6"""
 
-
     prompt_cost, completion_cost = cost_per_token(
         model="zai/glm-4.6",
         prompt_tokens=1000000,  # 1M tokens
@@ -72,10 +69,6 @@ def test_zai_glm46_cost_calculation(local_model_cost_map):
     assert math.isclose(completion_cost, 2.2, rel_tol=1e-6)
 
 
-
-
-
-
 def test_glm47_cost_calculation(local_model_cost_map):
     """Test cost calculation for GLM-4.7"""
 
diff --git a/tests/test_litellm/test_bedrock_usgov_pricing.py b/tests/test_litellm/test_bedrock_usgov_pricing.py
index 9e7e1e5c9f6..3dfd7350a06 100644
--- a/tests/test_litellm/test_bedrock_usgov_pricing.py
+++ b/tests/test_litellm/test_bedrock_usgov_pricing.py
@@ -31,10 +31,6 @@ def model_data():
         return json.load(f)
 
 
-
-
-
-
 def test_usgov_carries_20_percent_premium_over_global(model_data):
     """The us-gov rates must equal 1.2x the global anthropic.* rates,
     matching AWS's documented GovCloud uplift.
@@ -80,26 +76,6 @@ def test_usgov_cross_region_above_200k_ratio_to_global(model_data):
         ratio = usgov_info[field] / global_info[field]
         assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2"
 
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
 
 def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data):
     """us-gov-east-1 serves claude-3-haiku through the us-gov. inference profile
@@ -113,20 +89,6 @@ def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data):
     }
 
 
-
-
-
-
-
-
-
-
-
-
-
-
-
-
 GOV_ROW_SOURCES = {
     "us-gov.anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1",
     "bedrock/us-gov-west-1/anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1",
diff --git a/tests/test_litellm/test_claude_fable_5_config.py b/tests/test_litellm/test_claude_fable_5_config.py
index 52d3dccddc8..0473161faac 100644
--- a/tests/test_litellm/test_claude_fable_5_config.py
+++ b/tests/test_litellm/test_claude_fable_5_config.py
@@ -26,11 +26,6 @@ def _load_root_cost_map() -> dict:
         return json.load(f)
 
 
-
-
-
-
-
 def test_fable_5_geo_multiplier_without_fast_mode():
     """First-party ``inference_geo='us'`` carries the 1.1x premium, but unlike
     the Opus line there is no fast-mode variant for Fable 5; a ``fast`` key
@@ -65,8 +60,6 @@ def test_fable_5_registered_for_bedrock_converse():
     assert "anthropic.claude-fable-5" in BEDROCK_CONVERSE_MODELS
 
 
-
-
 @pytest.mark.parametrize(
     "cost_map",
     [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
@@ -138,8 +131,6 @@ FABLE_5_1_VARIANTS = (
 )
 
 
-
-
 @pytest.mark.parametrize(
     "cost_map",
     [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
@@ -158,10 +149,6 @@ def test_fable_5_1_cache_reads_cost_a_quarter_of_fable_5(cost_map):
         ), model_name
 
 
-
-
-
-
 def test_fable_5_1_present_in_bundled_backup():
     backup = GetModelCostMap.load_local_model_cost_map()
     root = _load_root_cost_map()
@@ -174,8 +161,6 @@ def test_fable_5_1_registered_for_bedrock_converse():
     assert "anthropic.claude-fable-5-1" in BEDROCK_CONVERSE_MODELS
 
 
-
-
 @pytest.mark.parametrize(
     "model",
     [
diff --git a/tests/test_litellm/test_claude_haiku_4_5_config.py b/tests/test_litellm/test_claude_haiku_4_5_config.py
index ab99a34d378..9172b6479a5 100644
--- a/tests/test_litellm/test_claude_haiku_4_5_config.py
+++ b/tests/test_litellm/test_claude_haiku_4_5_config.py
@@ -7,8 +7,6 @@ import json
 import os
 
 
-
-
 def test_bedrock_haiku_4_5_matches_sonnet_capabilities():
     """
     Test that Haiku 4.5 has same capabilities as Sonnet 4.5
diff --git a/tests/test_litellm/test_claude_opus_4_6_config.py b/tests/test_litellm/test_claude_opus_4_6_config.py
index a29901adfc3..9a8632924f2 100644
--- a/tests/test_litellm/test_claude_opus_4_6_config.py
+++ b/tests/test_litellm/test_claude_opus_4_6_config.py
@@ -71,10 +71,6 @@ def test_claude_4_6_australia_region_uses_au_prefix_not_apac():
     ), "apac.anthropic.claude-sonnet-4-6 should not be in bedrock_converse_models"
 
 
-
-
-
-
 def test_opus_4_6_alias_and_dated_metadata_match():
     json_path = os.path.join(
         os.path.dirname(__file__), "../../model_prices_and_context_window.json"
diff --git a/tests/test_litellm/test_claude_opus_4_8_config.py b/tests/test_litellm/test_claude_opus_4_8_config.py
index 7173b4a0e5b..e75fdba54ed 100644
--- a/tests/test_litellm/test_claude_opus_4_8_config.py
+++ b/tests/test_litellm/test_claude_opus_4_8_config.py
@@ -28,11 +28,6 @@ def _load_root_cost_map() -> dict:
         return json.load(f)
 
 
-
-
-
-
-
 def test_opus_4_8_fast_mode_multiplier():
     """Opus 4.8 dropped fast-mode pricing to 2x base ($10/$50 per MTok);
     Opus 4.7 was 6x ($30/$150)."""
@@ -42,14 +37,10 @@ def test_opus_4_8_fast_mode_multiplier():
     assert entry["fast"] == 2.0
 
 
-
-
 def test_opus_4_8_registered_for_bedrock_converse():
     assert "anthropic.claude-opus-4-8" in BEDROCK_CONVERSE_MODELS
 
 
-
-
 @pytest.mark.parametrize(
     "cost_map",
     [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
diff --git a/tests/test_litellm/test_claude_opus_5_config.py b/tests/test_litellm/test_claude_opus_5_config.py
index beb148f5e1b..285d556ef2b 100644
--- a/tests/test_litellm/test_claude_opus_5_config.py
+++ b/tests/test_litellm/test_claude_opus_5_config.py
@@ -51,11 +51,6 @@ def _load_root_cost_map() -> dict:
         return json.load(f)
 
 
-
-
-
-
-
 @pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS)
 def test_opus_5_bedrock_entries_declare_no_effort_ceiling(model_name):
     """Bedrock accepts every effort level for Opus 5, so no clamp belongs here.
@@ -135,8 +130,6 @@ def test_opus_5_registered_for_bedrock_converse():
     assert "anthropic.claude-opus-5" in BEDROCK_CONVERSE_MODELS
 
 
-
-
 @pytest.mark.parametrize(
     "cost_map",
     [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
diff --git a/tests/test_litellm/test_claude_sonnet_5_config.py b/tests/test_litellm/test_claude_sonnet_5_config.py
index bdc3bb64706..8c6d2cd1851 100644
--- a/tests/test_litellm/test_claude_sonnet_5_config.py
+++ b/tests/test_litellm/test_claude_sonnet_5_config.py
@@ -40,11 +40,6 @@ def _load_root_cost_map() -> dict:
         return json.load(f)
 
 
-
-
-
-
-
 def test_sonnet_5_present_in_bundled_backup():
     """The bundled backup is the runtime fallback (and what tests load with
     ``LITELLM_LOCAL_MODEL_COST_MAP=True``); it must carry the same entries as the
@@ -58,8 +53,6 @@ def test_sonnet_5_registered_for_bedrock_converse():
     assert "anthropic.claude-sonnet-5" in BEDROCK_CONVERSE_MODELS
 
 
-
-
 @pytest.mark.parametrize(
     "cost_map",
     [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
diff --git a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py
index be6be865665..4e770be7c3e 100644
--- a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py
+++ b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py
@@ -27,8 +27,6 @@ BACKUP_MAP = os.path.join(
 )
 
 
-
-
 def _load(path: str) -> dict:
     with open(path, encoding="utf-8") as f:
         return json.load(f)
@@ -38,14 +36,6 @@ def _cloudflare_keys(data: dict) -> set:
     return {k for k in data if k.startswith("cloudflare/")}
 
 
-
-
-
-
-
-
-
-
 def test_root_and_backup_have_identical_cloudflare_keys():
     if not os.path.exists(ROOT_MAP):
         pytest.skip("root cost map only ships in source checkouts")
diff --git a/tests/test_litellm/test_daybreak_model_metadata.py b/tests/test_litellm/test_daybreak_model_metadata.py
index 391391444a1..c3bac14dbbd 100644
--- a/tests/test_litellm/test_daybreak_model_metadata.py
+++ b/tests/test_litellm/test_daybreak_model_metadata.py
@@ -32,8 +32,6 @@ def _load(path):
         return json.load(f)
 
 
-
-
 def test_blue_alias_matches_its_snapshot_computer_use():
     cost_map = _load(MAIN_PATH)
 
diff --git a/tests/test_litellm/test_deepseek_model_metadata.py b/tests/test_litellm/test_deepseek_model_metadata.py
index a90ecd0ca59..9cbd14ebd1e 100644
--- a/tests/test_litellm/test_deepseek_model_metadata.py
+++ b/tests/test_litellm/test_deepseek_model_metadata.py
@@ -39,10 +39,6 @@ class TestDeepSeekModelCostEntries:
     """Verify that provider-prefixed DeepSeek entries contain the same
     capability flags as their bare-name counterparts in the JSON files."""
 
-
-
-
-
     def test_deepseek_chat_max_input_tokens_matches_bare_in_backup(self):
         data = _load_backup_json()
         bare = data.get("deepseek-chat", {})
@@ -56,8 +52,6 @@ class TestDeepSeekModelCostEntries:
         assert prefixed.get("max_output_tokens") == bare.get("max_output_tokens")
 
 
-
-
 # ---------------------------------------------------------------------------
 # API-level tests – verify supports_response_schema returns True
 # ---------------------------------------------------------------------------
diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py
index 858c9983221..701938f5677 100644
--- a/tests/test_litellm/test_fireworks_serverless_model_costs.py
+++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py
@@ -17,8 +17,6 @@ import pytest
 from litellm.utils import get_model_info
 
 
-
-
 NEW_ENTRIES = {
     "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": {
         "input_cost_per_token": 1.32e-06,
@@ -39,8 +37,6 @@ def model_data():
         return json.load(f)
 
 
-
-
 def test_bare_fireworks_ids_resolve_through_prefixed_entries():
     """Bare IDs from #37274 resolve via the provider-prefix lookup path."""
     for bare_id, prefixed_key in [
diff --git a/tests/test_litellm/test_gpt_5_5_model_metadata.py b/tests/test_litellm/test_gpt_5_5_model_metadata.py
index 32fbfc533b8..e07efbcc913 100644
--- a/tests/test_litellm/test_gpt_5_5_model_metadata.py
+++ b/tests/test_litellm/test_gpt_5_5_model_metadata.py
@@ -2,10 +2,6 @@ import json
 from pathlib import Path
 
 
-
-
-
-
 def test_azure_ai_gpt_5_5_backup_matches_main():
     """Ensure the bundled model cost map stays in sync with the canonical file."""
     repo_root = Path(__file__).parents[2]
diff --git a/tests/test_litellm/test_gpt_realtime_mode.py b/tests/test_litellm/test_gpt_realtime_mode.py
index 8b730d737b3..8c41e474486 100644
--- a/tests/test_litellm/test_gpt_realtime_mode.py
+++ b/tests/test_litellm/test_gpt_realtime_mode.py
@@ -41,19 +41,11 @@ REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS = (
 ALL_REALTIME_ONLY_GPT_MODELS = REALTIME_ONLY_GPT_MODELS + REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS
 
 
-
-
 def test_realtime_is_a_valid_mode_literal():
     hints = get_type_hints(ModelInfoBase, include_extras=False)
     assert "realtime" in get_args(hints["mode"])
 
 
-
-
-
-
-
-
 def test_backup_matches_main_for_realtime_models():
     repo_root = Path(__file__).parents[2]
     with open(repo_root / "model_prices_and_context_window.json") as f:
diff --git a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py
index 945b6e19897..d73311baae9 100644
--- a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py
+++ b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py
@@ -25,13 +25,6 @@ def _load(path):
         return json.load(f)
 
 
-
-
-
-
-
-
-
 @pytest.mark.parametrize("model", SYNCED_MODELS)
 def test_backup_matches_main(model):
     """Ensure the bundled (backup) cost map stays in sync with the canonical file."""
diff --git a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py b/tests/test_litellm/test_mistral_small_4_0_model_metadata.py
index 16b126a017c..182c444bac9 100644
--- a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py
+++ b/tests/test_litellm/test_mistral_small_4_0_model_metadata.py
@@ -18,8 +18,6 @@ def _load(path):
         return json.load(f)
 
 
-
-
 @pytest.mark.parametrize("model", SMALL_4_0_MODELS)
 def test_backup_matches_main(model):
     main_cost = _load(MAIN_PATH)
diff --git a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py
index fd2224d7720..02527a98711 100644
--- a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py
+++ b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py
@@ -23,9 +23,6 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di
         return json.load(f)
 
 
-
-
-
 @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING)
 def test_muse_spark_1_2_cost_per_token(
     local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float
diff --git a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py b/tests/test_litellm/test_muse_spark_1_3_model_metadata.py
index 3328e916af5..92b099fc780 100644
--- a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py
+++ b/tests/test_litellm/test_muse_spark_1_3_model_metadata.py
@@ -23,9 +23,6 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di
         return json.load(f)
 
 
-
-
-
 @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING)
 def test_muse_spark_1_3_cost_per_token(
     local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float
diff --git a/tests/test_litellm/test_replicate_model_key_format.py b/tests/test_litellm/test_replicate_model_key_format.py
index 8c52ae3603e..77ae5e1b069 100644
--- a/tests/test_litellm/test_replicate_model_key_format.py
+++ b/tests/test_litellm/test_replicate_model_key_format.py
@@ -20,8 +20,6 @@ def test_replicate_models_have_valid_key_prefix(model_cost: dict[str, Any]) -> N
     )
 
 
-
-
 def test_replicate_backup_matches_main() -> None:
     repo_root = Path(__file__).parents[2]
     main_path = repo_root / "model_prices_and_context_window.json"
diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/test_litellm/test_together_ai_model_metadata.py
index fd572733466..b9764eca2f8 100644
--- a/tests/test_litellm/test_together_ai_model_metadata.py
+++ b/tests/test_litellm/test_together_ai_model_metadata.py
@@ -76,14 +76,6 @@ def cost_map() -> CostMap:
         return COST_MAP_ADAPTER.validate_python(json.load(f))
 
 
-
-
-
-
-
-
-
-
 def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost_map: CostMap):
     inflated = sorted(
         model
@@ -96,10 +88,6 @@ def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost
     assert inflated == []
 
 
-
-
-
-
 @pytest.mark.parametrize("model", sorted(DEPRECATED_MODELS))
 def test_together_deprecated_model_carries_deprecation_date(cost_map: CostMap, model: str):
     info = cost_map.get(model)
@@ -153,8 +141,6 @@ CACHED_INPUT_MODELS: Final = (
 )
 
 
-
-
 def test_together_prompt_caching_flag_implies_cache_read_rate(cost_map: CostMap):
     for model, info in cost_map.items():
         if model.startswith("together_ai/") and info.get("supports_prompt_caching"):

From dc035cba624d32baa77b3c3e77cdd4fbf15d03ea Mon Sep 17 00:00:00 2001
From: mateo <mateo@berri.ai>
Date: Tue, 8 Sep 2026 02:33:04 +0000
Subject: [PATCH 156/164] test: preserve live xai pricing invariant

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../llms/xai/test_xai_redirected_slug_pricing.py            | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
index a6b1c9a92ce..4b7ba3f75f3 100644
--- a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
+++ b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
@@ -97,6 +97,12 @@ def test_redirected_slug_keeps_its_retirement_date(cost_map: dict, slug: str):
     assert cost_map[slug]["deprecation_date"] == expected_retirement_date(slug)
 
 
+def test_a_live_xai_model_is_untouched(cost_map: dict):
+    """Guard against the repricing leaking onto models xAI still serves directly."""
+    assert cost_map["xai/grok-4.6"]["input_cost_per_token"] != cost_map[REDIRECT_TARGET]["input_cost_per_token"]
+    assert "deprecation_date" not in cost_map["xai/grok-4.6"]
+
+
 @pytest.mark.parametrize("slug", REDIRECTED_SLUGS)
 def test_redirected_slug_carries_the_target_tier_rates(cost_map: dict, slug: str):
     """The request executes as grok-4.3, so it is tiered at grok-4.3's 200k boundary."""

From adcfe8cb7f2eec44d79371a624c3435245950970 Mon Sep 17 00:00:00 2001
From: mateo <mateo@berri.ai>
Date: Tue, 8 Sep 2026 02:44:52 +0000
Subject: [PATCH 157/164] test: pin redirected xai slugs to the target's tier
 field set

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
index 4b7ba3f75f3..e591c1ae682 100644
--- a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
+++ b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
@@ -110,6 +110,7 @@ def test_redirected_slug_carries_the_target_tier_rates(cost_map: dict, slug: str
     entry = cost_map[slug]
     for field in TIER_COST_FIELDS:
         assert entry[field] == target[field], field
+    assert {k for k in entry if "_above_" in k} == {k for k in target if "_above_" in k}
 
 
 def test_both_cost_maps_agree_on_the_redirected_slugs():

From 1a6aa98230571db22ccd37e5db2c6011a4f0c4c4 Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Mon, 7 Sep 2026 23:29:42 -0700
Subject: [PATCH 158/164] fix(spend): compare auto-router targets by deployment
 identity (#40206)

Preserve deployment identity through savings calculation, with canonical model fallback only when either ID is absent. Cover negotiated rates, unchanged deployments, alias/base-model cache accounting and missing IDs.

Fixes #38811. Based on the deployment-identity approach proposed by @QuantumBreakz in #38834.

Co-authored-by: Claude Code <noreply@anthropic.com>
---
 litellm/proxy/spend_tracking/savings.py       | 15 ++-
 .../proxy/spend_tracking/test_savings.py      | 98 +++++++++++++++++++
 2 files changed, 108 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py
index c541b9b40e5..950fcca2039 100644
--- a/litellm/proxy/spend_tracking/savings.py
+++ b/litellm/proxy/spend_tracking/savings.py
@@ -299,6 +299,8 @@ def compute_autorouter_savings(
     selected_info: ModelInfo | None = None,
     baseline_info: ModelInfo | None = None,
     cost_breakdown: Mapping[str, object] | None = None,
+    baseline_deployment_id: str | None = None,
+    selected_deployment_id: str | None = None,
 ) -> float:
     """Net dollars the router saved, or cost, by serving this request on ``selected_model``.
 
@@ -334,11 +336,12 @@ def compute_autorouter_savings(
     selected: Final = _resolve_model(selected_model, selected_provider)
     if baseline is None or selected is None:
         return 0.0
-    # Same model is only the same cost when it is also the same deployment. Two
-    # deployments of one model can carry different negotiated rates, and routing from
-    # the dear one to the cheap one is a real saving that short-circuiting on the model
-    # name alone reports as zero.
-    if baseline == selected:
+    same_target: Final = (
+        baseline_deployment_id == selected_deployment_id
+        if baseline_deployment_id and selected_deployment_id
+        else baseline == selected
+    )
+    if same_target:
         return 0.0
     basis: Final = _pricing_basis(cost_breakdown)
     effective_baseline_info: Final = baseline_info if baseline_info is not None else _model_info(baseline)
@@ -517,6 +520,8 @@ def autorouter_savings_for_request(
         selected_info=_effective_model_info(router_instance, model_id, model or ""),
         baseline_info=_effective_model_info(router_instance, baseline_id, baseline_model or ""),
         cost_breakdown=cost_breakdown,
+        baseline_deployment_id=baseline_id,
+        selected_deployment_id=model_id,
     )
     classifier_cost: Final = classifier_cost_from_decision(decision)
     return gross if classifier_cost is None else gross - classifier_cost
diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py
index cc8fdeb0160..e466edab131 100644
--- a/tests/test_litellm/proxy/spend_tracking/test_savings.py
+++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py
@@ -1162,6 +1162,104 @@ def test_prompt_caching_prices_at_the_deployment_rate_not_the_public_one():
     assert result.prompt_caching > at_public_rates.prompt_caching
 
 
+@pytest.mark.parametrize(
+    "baseline_id, selected_id, selected_multiplier, billed_input, classifier_cost, expected",
+    [
+        ("baseline", "selected", 0.1, None, 0.0, 0.0135),
+        ("baseline", "selected", 2.0, None, 0.0, -0.015),
+        ("baseline", "selected", 1.0, None, 0.0, 0.0),
+        ("baseline", "selected", 0.1, 0.004, 0.001, 0.01),
+        ("baseline", "baseline", 0.1, 0.004, 0.001, -0.001),
+        (None, "selected", 0.1, None, 0.0, 0.0),
+        ("baseline", None, 0.1, None, 0.0, 0.0),
+        (None, None, 0.1, None, 0.0, 0.0),
+        ("", "selected", 0.1, None, 0.0, 0.0),
+        ("baseline", "", 0.1, None, 0.0, 0.0),
+    ],
+)
+def test_autorouter_savings_distinguishes_priced_deployments(
+    baseline_id: str | None,
+    selected_id: str | None,
+    selected_multiplier: float,
+    billed_input: float | None,
+    classifier_cost: float,
+    expected: float,
+) -> None:
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": name,
+                "litellm_params": {
+                    "model": "anthropic/claude-opus-5",
+                    "api_key": "test-key",
+                    "input_cost_per_token": 1e-5 * multiplier,
+                    "output_cost_per_token": 5e-5 * multiplier,
+                },
+                "model_info": {"id": name},
+            }
+            for name, multiplier in (("baseline", 1.0), ("selected", selected_multiplier))
+        ]
+    )
+    result: Final = compute_savings_spend(
+        model="claude-opus-5",
+        custom_llm_provider="anthropic",
+        compression_saved_tokens=0,
+        gateway_injected_cache=False,
+        model_id=selected_id,
+        llm_router=lambda: router,
+        routing_decision={
+            "savings_baseline_model": "anthropic/claude-opus-5",
+            "savings_baseline_deployment_id": baseline_id,
+            "conversation_continuing": False,
+            "classifier_cost": classifier_cost,
+        },
+        usage_object={"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100},
+        cost_breakdown=None if billed_input is None else {"input_cost": billed_input, "output_cost": 0.0},
+    )
+    assert result.autorouter == pytest.approx(expected)
+
+
+@pytest.mark.parametrize("selected_model", ["azure/contract-deployment", "contract-deployment"])
+def test_autorouter_savings_recognizes_one_deployment_under_its_base_model(selected_model: str) -> None:
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "contract",
+                "litellm_params": {
+                    "model": "azure/contract-deployment",
+                    "api_key": "test-key",
+                    "api_base": "https://example.openai.azure.com",
+                    "input_cost_per_token": 0.0001,
+                    "output_cost_per_token": 0.0002,
+                    "cache_read_input_token_cost": 0.00001,
+                },
+                "model_info": {"id": "contract", "base_model": "azure/gpt-5.5"},
+            }
+        ]
+    )
+    result: Final = compute_savings_spend(
+        model=selected_model,
+        custom_llm_provider="azure",
+        compression_saved_tokens=0,
+        gateway_injected_cache=False,
+        model_id="contract",
+        llm_router=lambda: router,
+        routing_decision={
+            "savings_baseline_model": "azure/gpt-5.5",
+            "savings_baseline_deployment_id": "contract",
+            "conversation_continuing": True,
+        },
+        usage_object={
+            "prompt_tokens": 21000,
+            "completion_tokens": 100,
+            "total_tokens": 21100,
+            "prompt_tokens_details": {"text_tokens": 1000, "cached_tokens": 0, "cache_creation_tokens": 20000},
+        },
+        cost_breakdown={"input_cost": 2.1, "output_cost": 0.02},
+    )
+    assert result.autorouter == 0.0
+
+
 def test_a_recorded_baseline_deployment_prices_at_its_configured_rate():
     """A hardest-tier deployment with a negotiated rate is what the traffic would
     really have cost; pricing its model publicly misstates the saving."""

From 9a9b4c4c2538bf0df6d68eaabe48cefbb4824e7d Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Mon, 7 Sep 2026 23:37:42 -0700
Subject: [PATCH 159/164] feat(ui): show auto-router classification rate
 (#40192)

---
 .../AutoRouterBenchmarksTab.test.tsx          | 12 ++++-----
 .../_components/AutoRouterBenchmarksTab.tsx   | 26 ++++++++++++++-----
 .../_components/costOptimizationUtils.test.ts | 18 +++++++++++++
 .../_components/costOptimizationUtils.ts      |  7 +++++
 ...KeyAutoRouterUsageTab.integration.test.tsx |  2 +-
 5 files changed, 52 insertions(+), 13 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx
index 006da4f2725..2820a9dce83 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx
@@ -187,10 +187,10 @@ describe("AutoRouterBenchmarksTab", () => {
   });
 
   it.each([
-    { spend: 20665.28, classifier_cost: 342.18, turns: 140815, llm: "$20,323.10", cost: "$342.18" },
-    { spend: 0, classifier_cost: 0, turns: 0, llm: "$0.00", cost: "$0.00" },
-    { spend: 0.002, classifier_cost: 0.0004, turns: 100, llm: "$0.0016", cost: "$0.0004" },
-  ])("shows total classification cost across $turns turns without a per-turn rate", ({ llm, cost, ...values }) => {
+    { spend: 20665.28, classifier_cost: 342.18, turns: 140815, llm: "$20,323.10", cost: "$342.18", rate: "$2.43" },
+    { spend: 0, classifier_cost: 0, turns: 0, llm: "$0.00", cost: "$0.00", rate: "$0.00" },
+    { spend: 0.002, classifier_cost: 0.0004, turns: 100, llm: "$0.0016", cost: "$0.0004", rate: "$0.0040" },
+  ])("shows total classification cost and its rate across $turns turns", ({ llm, cost, rate, ...values }) => {
     const stats = totals({ ...values, saved_spend: 10126.28, baseline_spend: values.spend + 10126.28 });
     mockHook({ data: response([group(stats)], stats) });
     renderTab();
@@ -201,7 +201,7 @@ describe("AutoRouterBenchmarksTab", () => {
         .map((node) => node.textContent)
         .slice(1, 3),
     ).toEqual([llm, cost]);
-    expect(screen.queryByText(/1K turns/)).not.toBeInTheDocument();
+    expect(screen.getByText(`(${rate} / 1K turns)`)).toBeInTheDocument();
     expect(screen.getAllByText("$10,126.28").length).toBeGreaterThan(0);
   });
 
@@ -237,7 +237,7 @@ describe("AutoRouterBenchmarksTab", () => {
     expect(terms).toEqual([
       "Actual auto-router spend",
       "LLM spend",
-      "Classification cost",
+      "Classification cost($2.00 / 1K turns)",
       "Estimated spend at highest-tier model",
     ]);
     expect(values).toEqual(["$359.86", "$353.71", "$6.15", "$2,534.45"]);
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx
index 33ad1bfe555..ce5ab1c6776 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx
@@ -30,7 +30,7 @@ import {
   type BenchmarkView,
   type BucketRow,
 } from "./autoRouterBenchmarks";
-import { formatRangeLabel, usd } from "./costOptimizationUtils";
+import { classificationRatePer1kTurns, formatRangeLabel, usd } from "./costOptimizationUtils";
 import ShadowEvalSection from "./ShadowEvalSection";
 import TierTurnsChart from "./TierTurnsChart";
 import { useAutoRouterBenchmarks } from "./useAutoRouterBenchmarks";
@@ -52,9 +52,17 @@ const Metric: React.FC<{ label: string; value: string; hint?: string }> = ({ lab
   </Card>
 );
 
-const SpendRow: React.FC<{ label: string; value: string; subdued?: boolean }> = ({ label, value, subdued }) => (
+const SpendRow: React.FC<{ label: string; value: string; hint?: string; subdued?: boolean }> = ({
+  label,
+  value,
+  hint,
+  subdued,
+}) => (
   <dl className="flex flex-wrap items-baseline justify-between gap-x-6 gap-y-1 py-2">
-    <dt className="min-w-0 text-sm text-muted-foreground">{label}</dt>
+    <dt className="flex min-w-0 flex-wrap items-baseline gap-x-2 text-sm text-muted-foreground">
+      {label}
+      {hint && <span className="text-xs">{hint}</span>}
+    </dt>
     <dd
       className={`min-w-0 break-all tabular-nums ${subdued ? "text-sm font-normal text-muted-foreground" : "text-base font-semibold text-foreground"}`}
     >
@@ -99,6 +107,11 @@ const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => {
               subdued
               label="Classification cost"
               value={stats.classifier_cost == null ? "Unavailable" : usd(stats.classifier_cost)}
+              hint={
+                stats.classifier_cost == null
+                  ? undefined
+                  : classificationRatePer1kTurns(stats.classifier_cost, stats.turns)
+              }
             />
           </div>
           {stats.classifier_cost == null && (
@@ -277,9 +290,10 @@ const BenchmarksBody: React.FC<BenchmarksBodyProps> = ({ isPending, error, data,
       <p className="text-xs text-muted-foreground">
         Compares your actual routed spend with the estimated cost of using only the most expensive model configured in
         the auto-router. It accounts for both the cache savings from staying on one model and the added cache costs from
-        switching models. Savings are net of recorded LLM classification cost, which is included in actual spend. The
-        range counts whole sessions that overlap it, so totals can differ slightly from the Overall tab, which buckets
-        savings by UTC day.
+        switching models. Savings are net of recorded LLM classification cost, which is included in actual spend.
+        Classification cost per 1K turns is averaged over all auto-router turns, including those that skip
+        classification. The range counts whole sessions that overlap it, so totals can differ slightly from the Overall
+        tab, which buckets savings by UTC day.
       </p>
 
       <div className="space-y-4">
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts
index 9a2d7a0b0ec..5d2c48e6440 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts
@@ -7,6 +7,7 @@ import {
   SAVINGS_DRIVERS,
   SAVINGS_SERIES,
   buildDailyToolSeries,
+  classificationRatePer1kTurns,
   computeCacheLeakage,
   formatRangeLabel,
   isAnthropicModel,
@@ -401,6 +402,23 @@ describe("usd", () => {
   });
 });
 
+describe("classificationRatePer1kTurns", () => {
+  it("normalizes total classification cost to one thousand turns", () => {
+    expect(classificationRatePer1kTurns(342.18, 140815)).toBe("($2.43 / 1K turns)");
+    expect(classificationRatePer1kTurns(0.0004, 100)).toBe("($0.0040 / 1K turns)");
+  });
+
+  it("shows a floor instead of rounding a real cost down to zero", () => {
+    expect(classificationRatePer1kTurns(0.00001, 1000)).toBe("(<$0.0001 / 1K turns)");
+    expect(classificationRatePer1kTurns(0.0001, 1000)).toBe("($0.0001 / 1K turns)");
+  });
+
+  it("reports zero when there are no turns or no classification cost", () => {
+    expect(classificationRatePer1kTurns(0, 0)).toBe("($0.00 / 1K turns)");
+    expect(classificationRatePer1kTurns(0, 100)).toBe("($0.00 / 1K turns)");
+  });
+});
+
 describe("savings driver colours", () => {
   it("keeps a driver's colour when a driver above it is filtered out", () => {
     // Charts colour by position in the data they are given, and the donut is given
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts
index 7019b0d3301..464c779aa2b 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts
@@ -10,6 +10,13 @@ export const usd = (value: number): string => {
   return `${value < 0 ? "-" : ""}$${formatNumberWithCommas(magnitude, decimals)}`;
 };
 
+export const classificationRatePer1kTurns = (classifierCost: number, turns: number): string => {
+  if (turns <= 0) return `(${usd(0)} / 1K turns)`;
+  const rate = (classifierCost * 1000) / turns;
+  if (rate > 0 && rate < 0.0001) return "(<$0.0001 / 1K turns)";
+  return `(${usd(rate)} / 1K turns)`;
+};
+
 export const pct = (ratio: number): string => `${formatNumberWithCommas(ratio * 100, 1)}%`;
 
 export const shortDate = (iso: string): string =>
diff --git a/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx
index be95c0e600d..1f6a67b27ac 100644
--- a/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx
+++ b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx
@@ -89,7 +89,7 @@ describe("KeyAutoRouterUsageTab", () => {
     expect(screen.getByText("$1.00")).toBeInTheDocument();
     expect(screen.getByText("Classification cost")).toBeInTheDocument();
     expect(screen.getByText("$0.2500")).toBeInTheDocument();
-    expect(screen.queryByText(/1K turns/)).not.toBeInTheDocument();
+    expect(screen.getByText("($62.50 / 1K turns)")).toBeInTheDocument();
     expect(screen.getByText("Estimated spend at highest-tier model")).toBeInTheDocument();
     expect(screen.getByText("$10.00")).toBeInTheDocument();
     expect(screen.getByText("Auto-router prompt caching")).toBeInTheDocument();

From 1af7a403c66e037bec2e0ae6ea455a1c10b17b1b Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Mon, 7 Sep 2026 23:40:02 -0700
Subject: [PATCH 160/164] feat(mcp): start the named server's OAuth directly
 for a resource-scoped gateway flow (#39933)

An aggregate gateway DCR authorize whose RFC 8707 resource resolves to exactly one
gateway-managed oauth2 server sealed that server into the flow and then sent the browser to
the generic connect grid anyway, so the user had to find the server the client had already
named and click Connect.

The connect URL now carries only the flow handle. GET /authorize/flow classifies the sealed
flow as unscoped, interactive, M2M, or stale, and returns the matching state to the page.
Interactive flows require a live per-user vendor credential before minting and do not burn the
flow on an early submit. M2M flows use the gateway's configured service credential and finish
without an interactive OAuth trip. Stale flows fail closed instead of becoming unscoped.
The existing explicit Finish action and a new Cancel path preserve deliberate user intent.
---
 litellm/proxy/_experimental/mcp_server/db.py  |  21 +-
 .../mcp_server/discoverable_endpoints.py      |  69 +++--
 .../mcp_server/gateway_dcr_flow.py            | 187 +++++++++----
 litellm/proxy/_lazy_openapi_snapshot.json     |  40 +++
 .../mcp_server/test_discoverable_endpoints.py |  52 ++++
 .../mcp_server/test_gateway_dcr_flow.py       | 254 ++++++++++++++++--
 .../src/app/chat/integrations/page.tsx        |  30 +--
 .../src/app/connect/page.test.tsx             |  90 +------
 ui/litellm-dashboard/src/app/connect/page.tsx |  23 +-
 .../chat/ConnectFlowBanner.test.tsx           |  69 +++--
 .../src/components/chat/ConnectFlowBanner.tsx | 123 ++++++---
 .../chat/ConnectFlowSurface.test.tsx          | 118 ++++++++
 .../components/chat/ConnectFlowSurface.tsx    |  59 ++++
 .../src/components/chat/MCPAppsPanel.test.tsx |   9 +
 .../src/components/chat/MCPAppsPanel.tsx      |  33 ++-
 .../src/components/networking.tsx             |  18 ++
 .../src/lib/http/client.test.ts               |   9 +
 ui/litellm-dashboard/src/lib/http/client.ts   |   6 +-
 ui/litellm-dashboard/src/lib/http/schema.d.ts |  48 ++++
 19 files changed, 950 insertions(+), 308 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx
 create mode 100644 ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.tsx

diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py
index 082a90fdcfb..7379126983a 100644
--- a/litellm/proxy/_experimental/mcp_server/db.py
+++ b/litellm/proxy/_experimental/mcp_server/db.py
@@ -4,7 +4,7 @@ import hashlib
 import json
 from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast
 
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
@@ -125,6 +125,9 @@ class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False):
     server_id: str
 
 
+OAuthGrantState = Literal["valid", "refreshable", "absent"]
+
+
 class _OAuthTokenRefreshResponse(TypedDict, total=False):
     access_token: str
     refresh_token: str
@@ -1465,6 +1468,15 @@ def is_oauth_credential_expired(cred: OAuthCredentialPayload, buffer_seconds: in
         return False
 
 
+def oauth_grant_state(cred: OAuthCredentialPayload | None) -> OAuthGrantState:
+    """Classify local grant readiness without attempting a refresh or checking upstream revocation."""
+    if not cred or not cred.get("access_token"):
+        return "absent"
+    if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS):
+        return "valid"
+    return "refreshable" if cred.get("refresh_token") else "absent"
+
+
 async def get_user_oauth_credential(
     prisma_client: PrismaClient,
     user_id: str,
@@ -1727,12 +1739,11 @@ async def resolve_valid_user_oauth_token(
     dict it already holds. ``prisma_client`` is fetched lazily and only when a refresh
     actually happens, so the valid-token path never requires a DB handle.
     """
-    if not cred or not cred.get("access_token"):
+    grant: Final = oauth_grant_state(cred)
+    if cred is None or grant == "absent":
         return None
-    if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS):
+    if grant == "valid":
         return cred
-    if not cred.get("refresh_token"):
-        return None
     if prisma_client is None:
         from litellm.proxy.utils import get_prisma_client_or_throw
 
diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
index e10bfd41ed6..cab4b6c161a 100644
--- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
+++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
@@ -43,9 +43,11 @@ from litellm.proxy._experimental.mcp_server.faults import (
     render_token_fault,
 )
 from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
+    VendorCredentialState,
     aggregate_authorize,
     aggregate_token,
     complete_connect_flow,
+    describe_connect_flow,
     introspect_gateway_token,
     is_gateway_dcr_client_id,
     is_proxy_api_resource,
@@ -798,21 +800,7 @@ def _bridge_access_denied_redirect(redirect_uri: str, state: str, mcp_server: MC
     return RedirectResponse(_append_query_params(redirect_uri, params), status_code=302)
 
 
-async def _bridge_authorize_access_denial(
-    litellm_user_id: str,
-    mcp_server: MCPServer,
-    redirect_uri: str,
-    state: str,
-) -> RedirectResponse | None:
-    """The denial redirect for a signed-in user who cannot reach the target server, or None to proceed.
-
-    Admits the user exactly as MCP egress will (the same ``reload_admitted_user`` constructor and the
-    same ``get_allowed_mcp_servers`` resolver), so an envelope is minted only when the resulting
-    session can actually list and call the server's tools. Without this gate the flow completes, the
-    client shows connected, and every tool request fail-closes to an empty list with nothing telling
-    the operator why. An availability fault (5xx, e.g. a DB outage's 503) propagates; an unknown or
-    deactivated user denies like a missing grant, fail closed.
-    """
+async def _user_can_reach_mcp_server(user_id: str, server_id: str) -> bool:
     from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
         MCPRequestHandler,
     )
@@ -821,13 +809,22 @@ async def _bridge_authorize_access_denial(
     )
 
     try:
-        admitted: Final = await MCPRequestHandler.reload_admitted_user(litellm_user_id)
+        admitted: Final = await MCPRequestHandler.reload_admitted_user(user_id)
     except HTTPException as exc:
         if exc.status_code >= 500:
             raise
-        return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
-    allowed_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers(admitted)
-    if mcp_server.server_id in allowed_server_ids:
+        return False
+    return server_id in await global_mcp_server_manager.get_allowed_mcp_servers(admitted)
+
+
+async def _bridge_authorize_access_denial(
+    litellm_user_id: str,
+    mcp_server: MCPServer,
+    redirect_uri: str,
+    state: str,
+) -> RedirectResponse | None:
+    """The denial redirect for a signed-in user who cannot reach the target server, or None to proceed."""
+    if await _user_can_reach_mcp_server(litellm_user_id, mcp_server.server_id):
         return None
     return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
 
@@ -1910,6 +1907,38 @@ async def token_endpoint(
     )
 
 
+async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCredentialState:
+    """Whether the gateway itself can see a live vendor credential for this user and server.
+
+    The one reading of "authorized" the connect page displays and the finish step enforces, so
+    the button a user sees and the grant they get cannot disagree. A read fault is neither, and
+    fails the scoped grant closed."""
+    from litellm.proxy._experimental.mcp_server.db import (  # noqa: PLC0415  # circular import at module load
+        get_user_oauth_credential,
+        oauth_grant_state,
+    )
+    from litellm.proxy.proxy_server import prisma_client  # noqa: PLC0415  # circular import at module load
+
+    if prisma_client is None:
+        return "unavailable"
+    try:
+        credential: Final = await get_user_oauth_credential(prisma_client, user_id, server_id)
+    except Exception:  # noqa: BLE001  # a credential-read fault must fail the scoped grant closed
+        return "unavailable"
+    return "absent" if oauth_grant_state(credential) == "absent" else "present"
+
+
+@router.get("/authorize/flow")
+async def authorize_flow(request: Request, flow: str) -> Response:
+    return await describe_connect_flow(
+        request=request,
+        flow_handle=flow,
+        session_user_id=_session_cookie_user_id(request),
+        lookup_vendor_credential=_vendor_credential_state,
+        lookup_server_reachability=_user_can_reach_mcp_server,
+    )
+
+
 @router.post("/authorize/complete")
 async def authorize_complete(
     request: Request,
@@ -1934,6 +1963,8 @@ async def authorize_complete(
         delivery=delivery,
         team_id=team_id,
         decision=decision,
+        lookup_vendor_credential=_vendor_credential_state,
+        lookup_server_reachability=_user_can_reach_mcp_server,
     )
 
 
diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
index c7b0045dde5..3d94fa345d0 100644
--- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
+++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
@@ -152,10 +152,8 @@ _AUTH_CODE_DEBUG_KEY: Final = "gateway_authorization_code"
 
 ReloadUserFailure = Literal["unresolvable", "unavailable", "faulted", "no_active_key"]
 ReloadUser = Callable[[str], Awaitable[ReloadUserFailure | None]]
-"""Injected live-user revalidation (the token endpoint's mirror of admission):
-``None`` means the user is active; ``unavailable`` is a retryable DB outage; ``faulted`` is
-a DB fault retrying will not clear (still 503, worded so nobody just waits); anything else
-fails the grant closed."""
+VendorCredentialState = Literal["present", "absent", "unavailable"]
+"""The per-user vendor credential read has three outcomes: present, absent, or unavailable."""
 
 _DB_UNAVAILABLE_DESCRIPTION: Final = "the gateway database is unavailable; retry"
 _DB_FAULTED_DESCRIPTION: Final = (
@@ -195,6 +193,16 @@ class ConsentTeam(BaseModel):
     team_alias: str | None = None
 
 
+class LookupVendorCredential(Protocol):
+    """Injected read of a user's vendor credential for one server."""
+
+    def __call__(self, user_id: str, server_id: str, /) -> Awaitable[VendorCredentialState]: ...
+
+
+class LookupServerReachability(Protocol):
+    def __call__(self, user_id: str, server_id: str, /) -> Awaitable[bool]: ...
+
+
 class LookupConsentTeams(Protocol):
     """Injected lookup of the teams a signed-in user may bind a proxy-API credential to."""
 
@@ -205,6 +213,14 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr
     return "unresolvable"
 
 
+async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState:
+    return "unavailable"
+
+
+async def _unreachable_server(user_id: str, server_id: str) -> bool:
+    return False
+
+
 class GatewayDcrClient(BaseModel):
     """The registration record sealed into a gateway DCR ``client_id``.
 
@@ -449,7 +465,10 @@ def aggregate_authorize(
 
     A per-server RFC 8707 ``resource`` naming a gateway-managed oauth2 server scopes the
     flow to that one server: the scope is sealed into the flow, carried into the code, and
-    bound into the session token, while the connect page interlude runs exactly as before.
+    bound into the session token. The connect URL carries only the flow handle; the page
+    learns the client origin, the scoped server, and whether its vendor OAuth is done from
+    :func:`describe_connect_flow`, which reads the sealed flow, so nothing a link can carry
+    steers which server the page authorizes or names on the confirmation.
 
     Validation failures respond directly with 400 and never redirect: per RFC 6749
     section 4.1.2.1 an unvalidated redirect URI must not receive an error redirect, and
@@ -474,10 +493,7 @@ def aggregate_authorize(
         resource_server_id=scoped_server.server_id if scoped_server is not None else None,
         audience=None,
     )
-    connect_url: Final = _append_query_params(
-        f"{base_url}/ui/connect",
-        (("connect_flow", handle), ("connect_client", _origin_only(redirect_uri))),
-    )
+    connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),))
     response: Final = RedirectResponse(connect_url, status_code=303)
     _set_flow_cookie(response, request, handle, flow)
     return response
@@ -684,6 +700,99 @@ def _origin_only(url: str) -> str:
     return f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else ""
 
 
+def _open_flow_for(
+    request: Request, flow_handle: str, session_user_id: str | None, now: datetime
+) -> _ConnectFlow | Response:
+    sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle))
+    if sealed_flow is None:
+        return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
+    flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY)
+    if flow is None or now.timestamp() >= flow.exp:
+        return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
+    if session_user_id is None:
+        return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting")
+    if session_user_id != flow.user_id:
+        return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow")
+    return flow
+
+
+async def _flow_target(
+    flow: _ConnectFlow, lookup_server_reachability: LookupServerReachability
+) -> tuple[Literal["unscoped", "interactive", "m2m", "stale"], MCPServer | None]:
+    if flow.resource_server_id is None:
+        return "unscoped", None
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (  # noqa: PLC0415  # import cycle
+        MCPServerManager,
+        global_mcp_server_manager,
+    )
+
+    server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id)
+    if (
+        server is None
+        or not server.is_gateway_managed_oauth2
+        or not await lookup_server_reachability(flow.user_id, server.server_id)
+    ):
+        return "stale", None
+    state: Final = "m2m" if MCPServerManager.effective_oauth2_flow(server) == "client_credentials" else "interactive"
+    return state, server
+
+
+class ConnectFlowDescription(TypedDict):
+    """What the connect page is allowed to know about one in-flight flow."""
+
+    state: ReadOnly[Literal["unscoped", "interactive", "m2m", "stale"]]
+    client_origin: ReadOnly[str]
+    server_id: ReadOnly[str | None]
+    server_name: ReadOnly[str | None]
+    connected: ReadOnly[bool | None]
+
+
+async def _describe_opened_flow(
+    flow: _ConnectFlow,
+    lookup_vendor_credential: LookupVendorCredential,
+    lookup_server_reachability: LookupServerReachability,
+) -> ConnectFlowDescription | Response:
+    state, server = await _flow_target(flow, lookup_server_reachability)
+    if state == "interactive" and server is not None:
+        credential: Final = await lookup_vendor_credential(flow.user_id, server.server_id)
+        if credential == "unavailable":
+            return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION)
+        interactive_description: Final[ConnectFlowDescription] = {
+            "state": state,
+            "client_origin": _origin_only(flow.redirect_uri),
+            "server_id": server.server_id,
+            "server_name": server.server_name or server.alias or server.name,
+            "connected": credential == "present",
+        }
+        return interactive_description
+    described: Final[ConnectFlowDescription] = {
+        "state": state,
+        "client_origin": _origin_only(flow.redirect_uri),
+        "server_id": None if server is None else server.server_id,
+        "server_name": None if server is None else (server.server_name or server.alias or server.name),
+        "connected": state == "m2m" or None,
+    }
+    return described
+
+
+async def describe_connect_flow(
+    request: Request,
+    flow_handle: str,
+    session_user_id: str | None,
+    lookup_vendor_credential: LookupVendorCredential,
+    lookup_server_reachability: LookupServerReachability,
+) -> Response:
+    opened: Final = _open_flow_for(request, flow_handle, session_user_id, datetime.now(timezone.utc))
+    if isinstance(opened, Response):
+        return opened
+    described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability)
+    return (
+        described
+        if isinstance(described, Response)
+        else JSONResponse(content=described, headers=TOKEN_NO_CACHE_HEADERS)
+    )
+
+
 async def complete_connect_flow(
     request: Request,
     flow_handle: str,
@@ -692,56 +801,34 @@ async def complete_connect_flow(
     delivery: str | None = None,
     team_id: str | None = None,
     decision: str | None = None,
+    lookup_vendor_credential: LookupVendorCredential = _unavailable_vendor_credential,
+    lookup_server_reachability: LookupServerReachability = _unreachable_server,
 ) -> Response:
-    """The deliberate finish step of the connect flow: mint the gateway authorization
-    code and send the browser back to the client.
+    """Mint the code only after a deliberate POST by the sealed user.
 
-    Reached by POST so a cross-site GET cannot trigger it, and bound to the HttpOnly
-    per-flow cookie plus an exact match between the signed-in user and the user sealed
-    into the flow: a link crafted by another party dies here with ``access_denied``
-    instead of minting a code for the victim's identity. The flow is single-use (an atomic
-    claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in.
-
-    ``delivery`` chooses how the code reaches the client. Default (absent or
-    ``"redirect"``) is the 303 to the client's registered redirect URI. ``"manual"``
-    renders the callback URL on a page instead, for a client whose redirect URI is a
-    loopback host but which runs on a DIFFERENT machine than the browser (EC2/SSH box,
-    container): the 303 would dereference the browser machine's loopback and the code
-    would never arrive, so the user carries it over by pasting the URL into the client or
-    fetching it from the client machine's terminal. Manual delivery is honored only for
-    loopback redirect URIs; a routable redirect URI works from any browser by
-    construction, so those flows always redirect. The user who sees the page is exactly
-    the user the 303 would have carried the code to, and the same user already sees the
-    code today in the dead redirect's address bar, so the page exposes the code to no new
-    party. Unknown ``delivery`` values are rejected rather than defaulted: a client that
-    asked for manual delivery and got a dead redirect instead would silently lose its
-    code.
-
-    ``decision`` and ``team_id`` come from the native-client consent page. ``"deny"``
-    burns the flow and sends the client ``error=access_denied`` so it stops waiting;
-    ``team_id`` is sealed into the code only for proxy-API flows, where it picks which of
-    the user's teams the minted credential is attributed to.
+    A scoped flow additionally requires its sealed server to have a live vendor credential
+    before a code can be minted. The check happens before the single-use claim, so a
+    premature submit can be retried after authorization; denial deliberately bypasses it.
     """
     if delivery not in (None, "redirect", "manual"):
         return _oauth_error(400, "invalid_request", "delivery must be 'redirect' or 'manual'")
     if decision not in (None, "approve", "deny"):
         return _oauth_error(400, "invalid_request", "decision must be 'approve' or 'deny'")
-    sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle))
-    if sealed_flow is None:
-        return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
-    flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY)
-    if flow is None:
-        return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
     now: Final = datetime.now(timezone.utc)
-    if now.timestamp() >= flow.exp:
-        return _oauth_error(400, "invalid_request", "the connect flow has expired; restart the connection")
-    if session_user_id is None:
-        return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting")
-    if session_user_id != flow.user_id:
-        return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow")
+    opened: Final = _open_flow_for(request, flow_handle, session_user_id, now)
+    if isinstance(opened, Response):
+        return opened
+    if decision != "deny":
+        described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability)
+        if isinstance(described, Response):
+            return described
+        if described["state"] == "stale":
+            return _oauth_error(400, "invalid_request", "the requested MCP server is no longer available")
+        if described["connected"] is False:
+            return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing")
     flow_refusal: Final = _claim_refusal(
         await _SingleUseGuard(cache).claim(
-            f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
+            f"{_USED_FLOW_CACHE_PREFIX}{opened.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
         ),
         replayed=_oauth_error(
             400, "invalid_request", "this connect flow was already completed; restart the connection"
@@ -750,7 +837,7 @@ async def complete_connect_flow(
     if flow_refusal is not None:
         return flow_refusal
     response: Final = (
-        _denied_flow_response(flow) if decision == "deny" else _approved_flow_response(flow, delivery, team_id, now)
+        _denied_flow_response(opened) if decision == "deny" else _approved_flow_response(opened, delivery, team_id, now)
     )
     path, secure = _cookie_path_and_secure(request)
     response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax")
diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json
index dba2b2428aa..71475320c2c 100644
--- a/litellm/proxy/_lazy_openapi_snapshot.json
+++ b/litellm/proxy/_lazy_openapi_snapshot.json
@@ -19886,6 +19886,46 @@
           ]
         }
       },
+      "/authorize/flow": {
+        "get": {
+          "operationId": "authorize_flow_authorize_flow_get",
+          "parameters": [
+            {
+              "in": "query",
+              "name": "flow",
+              "required": true,
+              "schema": {
+                "title": "Flow",
+                "type": "string"
+              }
+            }
+          ],
+          "responses": {
+            "200": {
+              "content": {
+                "application/json": {
+                  "schema": {}
+                }
+              },
+              "description": "Successful Response"
+            },
+            "422": {
+              "content": {
+                "application/json": {
+                  "schema": {
+                    "$ref": "#/components/schemas/HTTPValidationError"
+                  }
+                }
+              },
+              "description": "Validation Error"
+            }
+          },
+          "summary": "Authorize Flow",
+          "tags": [
+            "mcp_discoverable"
+          ]
+        }
+      },
       "/callback": {
         "get": {
           "description": "OAuth 2.0 authorization response handler for MCP loopback clients.\n\nAccepts either:\n\n- A successful authorization response (``code`` + ``state``), which is\n  forwarded back to the validated client ``redirect_uri`` with the\n  original (un-wrapped) ``state``.\n- An error response (``error``[+``error_description``/``error_uri``]), per\n  RFC 6749 \u00a74.1.2.1. When ``state`` is present and decodes to a trusted\n  ``redirect_uri``, the error params are propagated back to the client so\n  its OAuth library can surface them. Otherwise we render an HTML error\n  page so the user is not left on an opaque 422 / blank screen.",
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
index be4206a1faf..763200c3709 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
@@ -4,6 +4,7 @@ import hashlib
 import json
 import time
 from base64 import urlsafe_b64encode
+from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -18,6 +19,57 @@ if TYPE_CHECKING:
     from litellm.types.mcp_server.mcp_server_manager import MCPServer
 
 
+def _stored_grant(access_token="access-token", refresh_token=None, expires_in_seconds=None, expires_at=None):
+    credential = {"type": "oauth2", "access_token": access_token}
+    if refresh_token is not None:
+        credential["refresh_token"] = refresh_token
+    if expires_in_seconds is not None:
+        credential["expires_at"] = (datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds)).isoformat()
+    if expires_at is not None:
+        credential["expires_at"] = expires_at
+    return credential
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("fields", "egress_has_token"),
+    [
+        (None, False),
+        ({"access_token": "", "refresh_token": "refresh-token"}, False),
+        ({}, True),
+        ({"expires_at": "never"}, True),
+        ({"expires_in_seconds": 600}, True),
+        ({"expires_in_seconds": 30}, False),
+        ({"expires_in_seconds": -300}, False),
+        ({"expires_in_seconds": -300, "refresh_token": ""}, False),
+        ({"expires_in_seconds": 30, "refresh_token": "refresh-token"}, True),
+        ({"expires_in_seconds": -300, "refresh_token": "refresh-token"}, True),
+    ],
+)
+async def test_vendor_credential_state_agrees_with_egress_token_resolution(monkeypatch, fields, egress_has_token):
+    from litellm.proxy._experimental.mcp_server import db as mcp_db
+    from litellm.proxy._experimental.mcp_server import discoverable_endpoints
+
+    monkeypatch.setattr(mcp_db, "MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", 60)
+    credential = _stored_grant(**fields) if fields is not None else None
+    read = AsyncMock(return_value=credential)
+    refresh = AsyncMock(return_value=_stored_grant(access_token="fresh-token", expires_in_seconds=3600))
+    prisma = MagicMock()
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
+    monkeypatch.setattr(mcp_db, "get_user_oauth_credential", read)
+    monkeypatch.setattr(mcp_db, "refresh_user_oauth_token", refresh)
+
+    connect = await discoverable_endpoints._vendor_credential_state("user-1", "server-1")
+    read.assert_awaited_once_with(prisma, "user-1", "server-1")
+    refresh.assert_not_awaited()
+    egress = await mcp_db.resolve_valid_user_oauth_token(
+        user_id="user-1", server=MagicMock(), cred=credential, prisma_client=prisma
+    )
+
+    assert (egress is not None) is egress_has_token
+    assert connect == ("present" if egress_has_token else "absent")
+
+
 # Fixture to mock IP address check for all MCP tests
 # This prevents tests from failing due to IP-based access control
 @pytest.fixture(autouse=True)
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
index 1670370f082..73a52a8d2e8 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
@@ -26,6 +26,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
     aggregate_authorize,
     aggregate_token,
     complete_connect_flow,
+    describe_connect_flow,
     introspect_gateway_token,
     is_gateway_dcr_client_id,
     is_proxy_api_resource,
@@ -230,7 +231,7 @@ async def test_authorize_with_session_hands_browser_to_connect_page_with_flow_co
     assert location.path == "/ui/connect"
     params = parse_qs(location.query)
     handle = params["connect_flow"][0]
-    assert params["connect_client"] == ["https://claude.ai"]
+    assert set(params) == {"connect_flow"}
     set_cookie = response.headers["set-cookie"]
     assert f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" in set_cookie
     assert "HttpOnly" in set_cookie
@@ -889,14 +890,60 @@ def _opened_principal(payload):
     return admitted.principal
 
 
-async def _finish_connect_page(response):
+class _VendorCredential:
+    def __init__(self, state="present"):
+        self.calls = []
+        self.state = state
+
+    async def __call__(self, user_id, server_id):
+        self.calls.append((user_id, server_id))
+        return self.state
+
+
+class _ServerReachability:
+    def __init__(self, reachable=True):
+        self.calls = []
+        self.reachable = reachable
+
+    async def __call__(self, user_id, server_id):
+        self.calls.append((user_id, server_id))
+        return self.reachable
+
+
+async def _complete_page(response, scoped_server=None, vendor=None, reachable=None, cache=None, **overrides):
+    from unittest.mock import patch
+
     handle, cookies = _flow_cookie_from(response)
-    completed = await complete_connect_flow(
-        request=_request("/authorize/complete", cookies=cookies, method="POST"),
-        flow_handle=handle,
-        session_user_id="u1",
-        cache=DualCache(),
-    )
+    with patch(_MANAGER_PATCH) as manager:
+        manager.get_mcp_server_by_id.return_value = scoped_server
+        return await complete_connect_flow(
+            request=_request("/authorize/complete", cookies=cookies, method="POST"),
+            flow_handle=handle,
+            session_user_id="u1",
+            cache=cache or DualCache(),
+            lookup_vendor_credential=vendor or _VendorCredential(),
+            lookup_server_reachability=reachable or _ServerReachability(),
+            **overrides,
+        )
+
+
+async def _describe_page(response, scoped_server=None, vendor=None, reachable=None, session_user_id="u1", cookies=None):
+    from unittest.mock import patch
+
+    handle, flow_cookies = _flow_cookie_from(response)
+    with patch(_MANAGER_PATCH) as manager:
+        manager.get_mcp_server_by_id.return_value = scoped_server
+        return await describe_connect_flow(
+            request=_request("/authorize/flow", cookies=flow_cookies if cookies is None else cookies),
+            flow_handle=handle,
+            session_user_id=session_user_id,
+            lookup_vendor_credential=vendor or _VendorCredential(),
+            lookup_server_reachability=reachable or _ServerReachability(),
+        )
+
+
+async def _finish_connect_page(response, scoped_server=None):
+    completed = await _complete_page(response, scoped_server=scoped_server)
     return parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
 
 
@@ -910,23 +957,43 @@ def _sealed_wire_json(sealed, prefix, debug_key):
 
 @pytest.mark.asyncio
 async def test_scoped_authorize_runs_connect_page_with_sealed_scope():
-    """LIT-4917: a per-server RFC 8707 resource naming a gateway-managed oauth2 server
-    seals that server into the flow. The connect page interlude runs exactly as before
-    (the scope restricts, it never skips consent), and the code minted at the finish step
-    and the session pair it redeems for are both scoped."""
+    """LIT-4917 plus LIT-7075: a per-server RFC 8707 resource naming a gateway-managed oauth2
+    server seals that server into the flow. The connect URL carries only the handle; the page
+    learns the scoped server and its vendor state from describe_connect_flow, and the finish
+    step refuses to mint a scoped code until that vendor credential exists, without burning
+    the flow. The code minted afterwards and the session pair it redeems for are both scoped."""
     from unittest.mock import patch
 
     client_id = (await _register([REDIRECT_URI]))["client_id"]
+    github = _scoped_mcp_server()
     with patch(_MANAGER_PATCH) as manager:
-        manager.get_mcp_server_by_name.return_value = _scoped_mcp_server()
+        manager.get_mcp_server_by_name.return_value = github
         response = _scoped_authorize(client_id, SCOPED_RESOURCE)
     assert response.status_code == 303
-    assert "/ui/connect" in response.headers["location"]
+    location = urlparse(response.headers["location"])
+    assert location.path == "/ui/connect"
+    assert set(parse_qs(location.query)) == {"connect_flow"}
     _, cookies = _flow_cookie_from(response)
     assert (
         _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")["resource_server_id"] == "github-id"
     )
-    code = await _finish_connect_page(response)
+    described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("absent"))
+    assert json.loads(described.body) == {
+        "state": "interactive",
+        "client_origin": "https://claude.ai",
+        "server_id": "github-id",
+        "server_name": "github",
+        "connected": False,
+    }
+    cache = DualCache()
+    premature = await _complete_page(response, scoped_server=github, vendor=_VendorCredential("absent"), cache=cache)
+    assert premature.status_code == 400
+    assert "authorize the requested MCP server" in json.loads(premature.body)["error_description"]
+    present = _VendorCredential("present")
+    completed = await _complete_page(response, scoped_server=github, vendor=present, cache=cache)
+    assert completed.status_code == 303
+    assert present.calls == [("u1", "github-id")]
+    code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
     assert (
         _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")["resource_server_id"]
         == "github-id"
@@ -952,9 +1019,11 @@ async def test_scoped_authorize_runs_connect_page_with_sealed_scope():
 )
 async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, resolves):
     """Every resource shape outside 'exactly one gateway-managed server' keeps today's flow:
-    connect page interlude, and NONE of the minted artifacts carry the scope key on the
-    wire, not the flow cookie, not the code, not the session JWT, so an unscoped flow
-    started on a new pod completes on a pod whose strict models predate the claim."""
+    the generic connect grid (describe names no server, the finish step never consults the
+    vendor credential), and NONE of the minted
+    artifacts carry the scope key on the wire, not the flow cookie, not the code, not the
+    session JWT, so an unscoped flow started on a new pod completes on a pod whose strict
+    models predate the claim."""
     import base64
     from unittest.mock import patch
 
@@ -963,10 +1032,18 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource,
         manager.get_mcp_server_by_name.return_value = None if resolves is None else _scoped_mcp_server()
         response = _scoped_authorize(client_id, resource)
     assert response.status_code == 303
-    assert "/ui/connect" in response.headers["location"]
+    location = urlparse(response.headers["location"])
+    assert location.path == "/ui/connect"
+    assert set(parse_qs(location.query)) == {"connect_flow"}
     _, cookies = _flow_cookie_from(response)
     assert "resource_server_id" not in _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")
-    code = await _finish_connect_page(response)
+    vendor = _VendorCredential("absent")
+    described = await _describe_page(response, vendor=vendor)
+    assert json.loads(described.body)["state"] == "unscoped"
+    assert json.loads(described.body)["server_id"] is None
+    completed = await _complete_page(response, vendor=vendor)
+    assert vendor.calls == []
+    code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
     assert "resource_server_id" not in _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")
     token_response = await _redeem(code, client_id)
     payload = json.loads(token_response.body)
@@ -979,19 +1056,150 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource,
 @pytest.mark.asyncio
 async def test_scoped_authorize_delegate_server_stays_unscoped():
     """A delegate-auth oauth2 server is outside the gateway-managed set (its keyless flow is
-    upstream PKCE via the relay), so a resource naming it never scopes the gateway flow."""
+    upstream PKCE via the relay), so a resource naming it never scopes the gateway flow and
+    never narrows the connect page to it."""
     from unittest.mock import patch
 
     client_id = (await _register([REDIRECT_URI]))["client_id"]
     with patch(_MANAGER_PATCH) as manager:
         manager.get_mcp_server_by_name.return_value = _scoped_mcp_server(delegate_auth_to_upstream=True)
         response = _scoped_authorize(client_id, SCOPED_RESOURCE)
-    assert "/ui/connect" in response.headers["location"]
+    location = urlparse(response.headers["location"])
+    assert location.path == "/ui/connect"
+    assert set(parse_qs(location.query)) == {"connect_flow"}
+    assert json.loads((await _describe_page(response)).body)["server_id"] is None
     code = await _finish_connect_page(response)
     token_response = await _redeem(code, client_id)
     assert _opened_principal(json.loads(token_response.body)).resource_server_id is None
 
 
+@pytest.mark.asyncio
+async def test_m2m_scoped_flow_mints_without_a_user_credential():
+    """A client-credentials server is already authorized by its gateway service credential, so
+    a resource-scoped flow finishes without consulting the per-user vault."""
+    from unittest.mock import patch
+
+    client_id = (await _register([REDIRECT_URI]))["client_id"]
+    m2m = _scoped_mcp_server(oauth2_flow="client_credentials")
+    with patch(_MANAGER_PATCH) as manager:
+        manager.get_mcp_server_by_name.return_value = m2m
+        response = _scoped_authorize(client_id, SCOPED_RESOURCE)
+    vendor = _VendorCredential("unavailable")
+    described = await _describe_page(response, scoped_server=m2m, vendor=vendor)
+    assert json.loads(described.body)["state"] == "m2m"
+    assert json.loads(described.body)["connected"] is True
+    assert vendor.calls == []
+    completed = await _complete_page(response, scoped_server=m2m, vendor=vendor)
+    assert completed.status_code == 303
+    assert vendor.calls == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("oauth2_flow", ["authorization_code", "client_credentials"])
+async def test_unreachable_scoped_flow_cannot_describe_or_finish(oauth2_flow):
+    from unittest.mock import patch
+
+    client_id = (await _register([REDIRECT_URI]))["client_id"]
+    server = _scoped_mcp_server(oauth2_flow=oauth2_flow)
+    with patch(_MANAGER_PATCH) as manager:
+        manager.get_mcp_server_by_name.return_value = server
+        response = _scoped_authorize(client_id, SCOPED_RESOURCE)
+    reachable = _ServerReachability(False)
+    vendor = _VendorCredential("present")
+    described = await _describe_page(response, scoped_server=server, reachable=reachable, vendor=vendor)
+    assert json.loads(described.body) == {
+        "state": "stale",
+        "client_origin": "https://claude.ai",
+        "server_id": None,
+        "server_name": None,
+        "connected": None,
+    }
+    cache = DualCache()
+    refused = await _complete_page(response, scoped_server=server, reachable=reachable, vendor=vendor, cache=cache)
+    assert refused.status_code == 400
+    assert vendor.calls == []
+    assert reachable.calls == [("u1", "github-id"), ("u1", "github-id")]
+    completed = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache)
+    assert completed.status_code == 303
+
+
+@pytest.mark.asyncio
+async def test_stale_scoped_flow_remains_distinct_from_unscoped():
+    """A server removed after authorize stays a stale scoped flow, so the page cannot offer a
+    broader unscoped grant or report a misleading Finish action."""
+    from unittest.mock import patch
+
+    client_id = (await _register([REDIRECT_URI]))["client_id"]
+    with patch(_MANAGER_PATCH) as manager:
+        manager.get_mcp_server_by_name.return_value = _scoped_mcp_server()
+        response = _scoped_authorize(client_id, SCOPED_RESOURCE)
+    described = await _describe_page(response, scoped_server=None)
+    assert json.loads(described.body)["state"] == "stale"
+    assert json.loads(described.body)["connected"] is None
+    stale = await _complete_page(response, scoped_server=None)
+    assert stale.status_code == 400
+    assert json.loads(stale.body)["error_description"] == "the requested MCP server is no longer available"
+
+
+@pytest.mark.asyncio
+async def test_scoped_flow_deny_and_stale_server_never_need_the_vendor_credential():
+    """Cancel is the escape hatch: a scoped user who cannot finish the vendor step still ends
+    the flow with access_denied and no credential lookup. A scoped server that is no longer
+    gateway-managed refuses to mint (nothing could serve that code) but also burns nothing."""
+    from unittest.mock import patch
+
+    client_id = (await _register([REDIRECT_URI]))["client_id"]
+    github = _scoped_mcp_server()
+    with patch(_MANAGER_PATCH) as manager:
+        manager.get_mcp_server_by_name.return_value = github
+        response = _scoped_authorize(client_id, SCOPED_RESOURCE)
+    cache = DualCache()
+    skipped_reachability = _ServerReachability(False)
+    stale = await _complete_page(response, scoped_server=None, reachable=skipped_reachability, cache=cache)
+    assert stale.status_code == 400
+    assert skipped_reachability.calls == []
+    assert json.loads(stale.body)["error_description"] == "the requested MCP server is no longer available"
+    described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("unavailable"))
+    assert described.status_code == 503
+    vendor = _VendorCredential("absent")
+    deny_reachability = _ServerReachability(False)
+    denied = await _complete_page(
+        response,
+        scoped_server=github,
+        vendor=vendor,
+        reachable=deny_reachability,
+        cache=cache,
+        decision="deny",
+    )
+    assert denied.status_code == 303
+    assert parse_qs(urlparse(denied.headers["location"]).query)["error"] == ["access_denied"]
+    assert vendor.calls == []
+    assert deny_reachability.calls == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "session_user_id, cookies, expected_status, expected_error",
+    [
+        ("u1", {}, 400, "invalid_request"),
+        ("u1", {"mcp_connect_flow_wrong": "garbage"}, 400, "invalid_request"),
+        (None, None, 401, "login_required"),
+        ("u2", None, 403, "access_denied"),
+    ],
+)
+async def test_describe_connect_flow_refuses_exactly_like_the_finish_step(
+    session_user_id, cookies, expected_status, expected_error
+):
+    """The page's read of the flow is gated the same way minting is: the HttpOnly cookie for
+    that handle must open and the signed-in user must be the sealed one. A lure link with a
+    made-up handle therefore learns nothing and starts nothing."""
+    client_id = (await _register([REDIRECT_URI]))["client_id"]
+    response = _authorize(client_id, session_user_id="u1")
+    described = await _describe_page(response, session_user_id=session_user_id, cookies=cookies)
+    assert described.status_code == expected_status
+    assert json.loads(described.body)["error"] == expected_error
+
+
 @pytest.mark.asyncio
 async def test_token_rejects_resource_conflicting_with_sealed_scope():
     """RFC 8707 section 2.2: redeeming a scoped code (or rotating a scoped refresh token)
@@ -1005,7 +1213,7 @@ async def test_token_rejects_resource_conflicting_with_sealed_scope():
     with patch(_MANAGER_PATCH) as manager:
         manager.get_mcp_server_by_name.return_value = github
         response = _scoped_authorize(client_id, SCOPED_RESOURCE)
-    code = await _finish_connect_page(response)
+    code = await _finish_connect_page(response, scoped_server=github)
 
     with patch(_MANAGER_PATCH) as manager:
         manager.get_mcp_server_by_name.return_value = linear
diff --git a/ui/litellm-dashboard/src/app/chat/integrations/page.tsx b/ui/litellm-dashboard/src/app/chat/integrations/page.tsx
index 30ce62d8081..a663e16c2b2 100644
--- a/ui/litellm-dashboard/src/app/chat/integrations/page.tsx
+++ b/ui/litellm-dashboard/src/app/chat/integrations/page.tsx
@@ -1,43 +1,19 @@
 "use client";
 
-import { Suspense, useEffect } from "react";
-import { useRouter, useSearchParams } from "next/navigation";
+import { Suspense } from "react";
 import { useChatShell } from "@/contexts/ChatShellContext";
-import MCPAppsPanel from "@/components/chat/MCPAppsPanel";
-import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner";
+import ConnectFlowSurface from "@/components/chat/ConnectFlowSurface";
 
 // useSearchParams() requires a Suspense boundary for static export.
 function IntegrationsPageContent() {
   const { accessToken, selectedMCPServers, setSelectedMCPServers } = useChatShell();
-  const router = useRouter();
-  const searchParams = useSearchParams();
-  const oauthReturn = searchParams.get("mcpOauthReturn");
-  // Set by the gateway DCR authorize when a DCR client sends the user here to
-  // authorize servers before finishing sign-in (see gateway_dcr_flow.py). The
-  // handle keys the sealed per-flow cookie; connect_client is the client origin
-  // for display only. connect_flow is NOT cleaned from the URL: the finish form
-  // needs it, and the sealed cookie (not the URL) is the security boundary.
-  const connectFlow = searchParams.get("connect_flow");
-  const connectClient = searchParams.get("connect_client");
-
-  // Clean up the OAuth return param after it's been consumed — real routing means
-  // we no longer need it to pick a tab, but it should not linger in the address bar.
-  useEffect(() => {
-    if (oauthReturn) {
-      const url = new URL(window.location.href);
-      url.searchParams.delete("mcpOauthReturn");
-      router.replace(url.pathname + url.search);
-    }
-  }, [oauthReturn, router]);
 
   return (
     <div className="flex-1 min-h-0 overflow-auto w-full py-8 px-8">
-      {connectFlow && <ConnectFlowBanner flowHandle={connectFlow} clientOrigin={connectClient} />}
-      <MCPAppsPanel
+      <ConnectFlowSurface
         accessToken={accessToken}
         selectedServers={selectedMCPServers}
         onChange={setSelectedMCPServers}
-        connectMode={!!connectFlow}
       />
     </div>
   );
diff --git a/ui/litellm-dashboard/src/app/connect/page.test.tsx b/ui/litellm-dashboard/src/app/connect/page.test.tsx
index 7d49a8b6a4c..6a6ba24bd87 100644
--- a/ui/litellm-dashboard/src/app/connect/page.test.tsx
+++ b/ui/litellm-dashboard/src/app/connect/page.test.tsx
@@ -2,101 +2,29 @@ import { afterEach, describe, expect, it, vi } from "vitest";
 import { render, screen } from "@testing-library/react";
 import ConnectPage from "./page";
 
-interface PanelProps {
+interface SurfaceProps {
   accessToken: string;
   selectedServers: string[];
   onChange: (servers: string[]) => void;
-  connectMode?: boolean;
 }
 
-interface BannerProps {
-  flowHandle: string;
-  clientOrigin: string | null;
-}
-
-const { mockReplace, mockPanel, mockBanner, state } = vi.hoisted(() => {
-  const state = {
-    oauthReturn: null as string | null,
-    connectFlow: null as string | null,
-    connectClient: null as string | null,
-  };
-  return {
-    state,
-    mockReplace: vi.fn(),
-    mockPanel: vi.fn((_props: PanelProps) => <div data-testid="mcp-apps-panel" />),
-    mockBanner: vi.fn((_props: BannerProps) => <div data-testid="connect-flow-banner" />),
-  };
-});
-
-vi.mock("next/navigation", () => ({
-  useRouter: () => ({ replace: mockReplace }),
-  useSearchParams: () => ({
-    get: (key: string) => {
-      if (key === "mcpOauthReturn") return state.oauthReturn;
-      if (key === "connect_flow") return state.connectFlow;
-      if (key === "connect_client") return state.connectClient;
-      return null;
-    },
-  }),
+const { mockSurface } = vi.hoisted(() => ({
+  mockSurface: vi.fn((_props: SurfaceProps) => <div data-testid="connect-flow-surface" />),
 }));
+
 vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
   default: () => ({ accessToken: "token-123" }),
 }));
-vi.mock("@/components/chat/MCPAppsPanel", () => ({ default: mockPanel }));
-vi.mock("@/components/chat/ConnectFlowBanner", () => ({ default: mockBanner }));
+vi.mock("@/components/chat/ConnectFlowSurface", () => ({ default: mockSurface }));
 
 describe("ConnectPage", () => {
   afterEach(() => {
-    state.oauthReturn = null;
-    state.connectFlow = null;
-    state.connectClient = null;
-    mockReplace.mockClear();
-    mockPanel.mockClear();
-    mockBanner.mockClear();
+    mockSurface.mockClear();
   });
 
-  it("renders the MCP connect panel with the user's access token", () => {
+  it("renders the gateway connect surface with the user's access token and an empty selection", () => {
     render(<ConnectPage />);
-    expect(screen.getByTestId("mcp-apps-panel")).toBeInTheDocument();
-    expect(mockPanel.mock.calls[0][0]).toMatchObject({ accessToken: "token-123", selectedServers: [] });
-  });
-
-  it("strips the mcpOauthReturn param from the URL after an OAuth return", () => {
-    state.oauthReturn = "apps";
-    window.history.replaceState({}, "", "/connect?mcpOauthReturn=apps");
-    render(<ConnectPage />);
-    expect(mockReplace).toHaveBeenCalledWith("/connect");
-  });
-
-  it("does not rewrite the URL when there is no OAuth return param", () => {
-    render(<ConnectPage />);
-    expect(mockReplace).not.toHaveBeenCalled();
-  });
-
-  it("mounts the gateway connect banner and puts the panel in connect mode for a DCR flow", () => {
-    state.connectFlow = "flow-handle-123";
-    state.connectClient = "https://claude.ai";
-    render(<ConnectPage />);
-    expect(screen.getByTestId("connect-flow-banner")).toBeInTheDocument();
-    expect(mockBanner.mock.calls[0][0]).toMatchObject({
-      flowHandle: "flow-handle-123",
-      clientOrigin: "https://claude.ai",
-    });
-    expect(mockPanel.mock.calls[0][0].connectMode).toBe(true);
-  });
-
-  it("shows no connect banner and leaves connect mode off for a plain visit", () => {
-    render(<ConnectPage />);
-    expect(screen.queryByTestId("connect-flow-banner")).not.toBeInTheDocument();
-    expect(mockBanner).not.toHaveBeenCalled();
-    expect(mockPanel.mock.calls[0][0].connectMode).toBe(false);
-  });
-
-  it("keeps the connect flow handle in the URL while stripping the OAuth return param", () => {
-    state.oauthReturn = "apps";
-    state.connectFlow = "flow-handle-123";
-    window.history.replaceState({}, "", "/connect?connect_flow=flow-handle-123&mcpOauthReturn=apps");
-    render(<ConnectPage />);
-    expect(mockReplace).toHaveBeenCalledWith("/connect?connect_flow=flow-handle-123");
+    expect(screen.getByTestId("connect-flow-surface")).toBeInTheDocument();
+    expect(mockSurface.mock.calls[0][0]).toMatchObject({ accessToken: "token-123", selectedServers: [] });
   });
 });
diff --git a/ui/litellm-dashboard/src/app/connect/page.tsx b/ui/litellm-dashboard/src/app/connect/page.tsx
index 3f0c269e86b..652c044197b 100644
--- a/ui/litellm-dashboard/src/app/connect/page.tsx
+++ b/ui/litellm-dashboard/src/app/connect/page.tsx
@@ -1,36 +1,19 @@
 "use client";
 
-import { Suspense, useEffect, useState } from "react";
-import { useRouter, useSearchParams } from "next/navigation";
+import { Suspense, useState } from "react";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
-import MCPAppsPanel from "@/components/chat/MCPAppsPanel";
-import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner";
+import ConnectFlowSurface from "@/components/chat/ConnectFlowSurface";
 
 function ConnectPageContent() {
   const { accessToken } = useAuthorized();
   const [selectedServers, setSelectedServers] = useState<string[]>([]);
-  const router = useRouter();
-  const searchParams = useSearchParams();
-  const oauthReturn = searchParams.get("mcpOauthReturn");
-  const connectFlow = searchParams.get("connect_flow");
-  const connectClient = searchParams.get("connect_client");
-
-  useEffect(() => {
-    if (oauthReturn) {
-      const url = new URL(window.location.href);
-      url.searchParams.delete("mcpOauthReturn");
-      router.replace(url.pathname + url.search);
-    }
-  }, [oauthReturn, router]);
 
   return (
     <div className="mx-auto w-full max-w-5xl px-8 py-8">
-      {connectFlow && <ConnectFlowBanner flowHandle={connectFlow} clientOrigin={connectClient} />}
-      <MCPAppsPanel
+      <ConnectFlowSurface
         accessToken={accessToken ?? ""}
         selectedServers={selectedServers}
         onChange={setSelectedServers}
-        connectMode={!!connectFlow}
       />
     </div>
   );
diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx
index b3cd6e229af..5caf15d1fce 100644
--- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx
+++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx
@@ -1,59 +1,61 @@
 import { afterEach, describe, expect, it, vi } from "vitest";
 import { render, screen } from "@testing-library/react";
+import type { ConnectFlowStatus } from "@/components/networking";
 import ConnectFlowBanner, { isLoopbackOrigin } from "./ConnectFlowBanner";
 
 vi.mock("@/components/networking", () => ({
   getProxyBaseUrl: () => "https://gateway.example.com",
 }));
 
+vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({
+  useUserMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle" }),
+}));
+
 afterEach(() => {
   vi.restoreAllMocks();
-  sessionStorage.clear();
 });
 
+const unscoped = (client_origin: string): ConnectFlowStatus => ({
+  state: "unscoped",
+  client_origin,
+  server_id: null,
+  server_name: null,
+  connected: null,
+});
+
+const renderBanner = (clientOrigin: string) =>
+  render(
+    <ConnectFlowBanner
+      flowHandle="flow-handle-123"
+      flow={unscoped(clientOrigin)}
+      accessToken="tok"
+      onConnected={vi.fn()}
+      failed={false}
+    />,
+  );
+
 describe("ConnectFlowBanner", () => {
-  it("posts the flow handle to the proxy /authorize/complete as a full-page form", () => {
-    const { container } = render(<ConnectFlowBanner flowHandle="flow-handle-123" clientOrigin="https://claude.ai" />);
+  it("posts only the flow handle to the proxy /authorize/complete as a full-page form", () => {
+    const { container } = renderBanner("https://claude.ai");
 
     const form = container.querySelector("form")!;
     expect(form).toHaveAttribute("method", "POST");
     expect(form).toHaveAttribute("action", "https://gateway.example.com/authorize/complete");
-
-    const hidden = form.querySelector('input[name="flow"]') as HTMLInputElement;
-    expect(hidden.value).toBe("flow-handle-123");
-    // No token, code, or secret is ever placed in the form; the sealed cookie carries them.
+    expect(screen.getByDisplayValue("flow-handle-123")).toHaveAttribute("name", "flow");
     expect(form.innerHTML).not.toContain("token");
-  });
-
-  it("shows the client origin so the user knows what they are connecting to", () => {
-    render(<ConnectFlowBanner flowHandle="h" clientOrigin="https://claude.ai" />);
-    expect(screen.getAllByText(/claude\.ai/).length).toBeGreaterThan(0);
     expect(screen.getByRole("button", { name: /finish connecting/i })).toBeInTheDocument();
+    expect(screen.queryByRole("button", { name: "Cancel" })).not.toBeInTheDocument();
   });
 
-  it("falls back to a generic label when the client origin is unknown", () => {
-    render(<ConnectFlowBanner flowHandle="h" clientOrigin={null} />);
-    expect(screen.getAllByText(/the application/).length).toBeGreaterThan(0);
-  });
-
-  it("offers manual delivery for a loopback client, posted only when checked", () => {
-    const { container } = render(
-      <ConnectFlowBanner flowHandle="flow-handle-123" clientOrigin="http://localhost:3118" />,
-    );
-
-    const checkbox = container.querySelector('input[type="checkbox"][name="delivery"]') as HTMLInputElement;
-    expect(checkbox).not.toBeNull();
+  it("offers manual delivery only for a loopback client, posted only when checked", () => {
+    const loopback = renderBanner("http://localhost:3118");
+    const checkbox = loopback.container.querySelector('input[type="checkbox"][name="delivery"]') as HTMLInputElement;
     expect(checkbox.value).toBe("manual");
     expect(checkbox.checked).toBe(false);
-    expect(screen.getByText(/remote or SSH machine/i)).toBeInTheDocument();
-  });
+    loopback.unmount();
 
-  it("does not offer manual delivery for a routable client origin or an unknown one", () => {
-    const routable = render(<ConnectFlowBanner flowHandle="h" clientOrigin="https://claude.ai" />);
+    const routable = renderBanner("https://claude.ai");
     expect(routable.container.querySelector('input[name="delivery"]')).toBeNull();
-
-    const unknown = render(<ConnectFlowBanner flowHandle="h" clientOrigin={null} />);
-    expect(unknown.container.querySelector('input[name="delivery"]')).toBeNull();
   });
 
   it("classifies loopback origins like the server does", () => {
@@ -70,12 +72,9 @@ describe("ConnectFlowBanner", () => {
   });
 
   it("does NOT complete the flow on pagehide (completion requires the explicit button)", () => {
-    // Security regression: an attacker could lure a signed-in victim to their own client's
-    // authorize URL; the victim merely closing the tab must NOT deliver a victim-bound code.
-    // Completion is a deliberate button press, never a side effect of leaving the page.
     const beaconMock = vi.fn(() => true);
     vi.stubGlobal("navigator", { ...navigator, sendBeacon: beaconMock });
-    render(<ConnectFlowBanner flowHandle="flow-xyz" clientOrigin="https://claude.ai" />);
+    renderBanner("https://claude.ai");
 
     window.dispatchEvent(new Event("pagehide"));
 
diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx
index cea42f916f8..0d6e708f734 100644
--- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx
+++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx
@@ -2,30 +2,18 @@
 
 import React from "react";
 import { CheckCircle } from "lucide-react";
-import { getProxyBaseUrl } from "@/components/networking";
+import { getProxyBaseUrl, ConnectFlowStatus } from "@/components/networking";
+import { OAuth2ConnectButton } from "@/components/chat/MCPAppsPanel";
 
 interface Props {
   flowHandle: string;
-  clientOrigin: string | null;
+  flow?: ConnectFlowStatus;
+  accessToken: string;
+  onConnected: () => void;
+  failed: boolean;
 }
 
-/**
- * The interlude shown when a DCR client (Claude Desktop, MCP Inspector) sends the user
- * through the gateway sign-in and lands them on the apps grid to authorize servers. The
- * grid below authorizes individual servers into the per-user vault; this banner is the
- * finish step that returns the user to the client.
- *
- * Finishing requires the explicit "Finish connecting" button: a native form POST to the proxy's
- * /authorize/complete, which mints the gateway authorization code and 303-redirects to the DCR
- * client's own redirect URI (the full-page navigation carries the HttpOnly per-flow cookie and
- * follows the cross-origin redirect to the client's loopback).
- *
- * The button press IS the consent gate and must not be bypassed. An earlier version auto-finished
- * on tab close via navigator.sendBeacon; that let an attacker who lured a signed-in victim to their
- * own client's authorize URL harvest a victim-bound code the moment the victim closed the tab
- * (no click). Merely visiting the authorize URL is attacker-inducible, so completion has to be a
- * deliberate user action, not a side effect of leaving the page.
- */
+/** Finish remains an explicit POST because a cross-site navigation must never mint a code. */
 export function isLoopbackOrigin(origin: string | null): boolean {
   if (!origin) return false;
   try {
@@ -36,10 +24,44 @@ export function isLoopbackOrigin(origin: string | null): boolean {
   }
 }
 
-const ConnectFlowBanner: React.FC<Props> = ({ flowHandle, clientOrigin }) => {
+const copyFor = (flow: ConnectFlowStatus | undefined, failed: boolean): readonly [string, string] => {
+  const clientLabel = flow?.client_origin ?? "the application";
+  const serverLabel = flow?.server_name ?? "the requested MCP server";
+  if (failed || flow === undefined || flow.state === "stale") {
+    return [
+      "The connection cannot continue",
+      `The gateway could not validate this connection. Cancel to return to ${clientLabel}.`,
+    ];
+  }
+  if (flow.state === "unscoped") {
+    return [
+      `Connect your MCP servers to ${clientLabel}`,
+      `Authorize the servers you want to use below, then click Finish connecting to return to ${clientLabel}.`,
+    ];
+  }
+  if (flow.state === "interactive" && !flow.connected) {
+    return [
+      `Allow ${clientLabel} to use ${serverLabel}`,
+      `Authorize ${serverLabel} below to continue, or cancel to send ${clientLabel} away.`,
+    ];
+  }
+  return [
+    `Allow ${clientLabel} to use ${serverLabel}`,
+    `Click Finish connecting to give ${clientLabel} access to ${serverLabel} as you.`,
+  ];
+};
+
+const ConnectFlowBanner: React.FC<Props> = ({ flowHandle, flow, accessToken, onConnected, failed }) => {
   const action = `${getProxyBaseUrl()}/authorize/complete`;
-  const clientLabel = clientOrigin ?? "the application";
-  const loopbackClient = isLoopbackOrigin(clientOrigin);
+  const state = failed || flow === undefined ? "stale" : flow.state;
+  const canFinish = state === "unscoped" || (state !== "stale" && flow?.connected === true);
+  const canCancel = state !== "unscoped";
+  const loopbackClient = isLoopbackOrigin(flow?.client_origin ?? null);
+  const vendorServer =
+    state === "interactive" && flow?.connected === false && flow.server_id !== null
+      ? { server_id: flow.server_id, server_name: flow.server_name }
+      : null;
+  const copy = copyFor(flow, failed);
 
   return (
     <div className="mb-6 rounded-lg border border-primary/30 bg-primary/5 px-5 py-4">
@@ -47,27 +69,48 @@ const ConnectFlowBanner: React.FC<Props> = ({ flowHandle, clientOrigin }) => {
         <div className="flex items-start gap-3 min-w-0">
           <CheckCircle className="h-5 w-5 text-primary shrink-0 mt-0.5" />
           <div className="min-w-0">
-            <p className="text-sm font-semibold text-foreground">Connect your MCP servers to {clientLabel}</p>
-            <p className="text-[13px] text-muted-foreground mt-0.5">
-              Authorize the servers you want to use below, then click Finish connecting to return to {clientLabel}.
-            </p>
+            <p className="text-sm font-semibold text-foreground">{copy[0]}</p>
+            <p className="text-[13px] text-muted-foreground mt-0.5">{copy[1]}</p>
           </div>
         </div>
-        <form method="POST" action={action} className="shrink-0">
-          <input type="hidden" name="flow" value={flowHandle} />
-          <button
-            type="submit"
-            className="h-[38px] rounded-md bg-primary px-4 text-sm font-semibold text-primary-foreground hover:bg-primary/90"
-          >
-            Finish connecting
-          </button>
-          {loopbackClient && (
-            <label className="mt-2 flex items-center gap-2 text-[13px] text-muted-foreground">
-              <input type="checkbox" name="delivery" value="manual" />
-              My client is on a remote or SSH machine
-            </label>
+        <div className="flex shrink-0 gap-2">
+          {vendorServer !== null && (
+            <OAuth2ConnectButton
+              server={vendorServer}
+              accessToken={accessToken}
+              onConnect={onConnected}
+              variant="button"
+              autoStartKey={`litellm-mcp-autostart:${flowHandle}`}
+            />
           )}
-        </form>
+          <form method="POST" action={action}>
+            <input type="hidden" name="flow" value={flowHandle} />
+            {canFinish && (
+              <button
+                type="submit"
+                className="h-[38px] rounded-md bg-primary px-4 text-sm font-semibold text-primary-foreground hover:bg-primary/90"
+              >
+                Finish connecting
+              </button>
+            )}
+            {canCancel && (
+              <button
+                type="submit"
+                name="decision"
+                value="deny"
+                className="ml-2 h-[38px] rounded-md border px-4 text-sm font-semibold text-foreground hover:bg-accent/40"
+              >
+                Cancel
+              </button>
+            )}
+            {loopbackClient && (
+              <label className="mt-2 flex items-center gap-2 text-[13px] text-muted-foreground">
+                <input type="checkbox" name="delivery" value="manual" />
+                My client is on a remote or SSH machine
+              </label>
+            )}
+          </form>
+        </div>
       </div>
     </div>
   );
diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx
new file mode 100644
index 00000000000..cf59dcfc368
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx
@@ -0,0 +1,118 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { act, render, screen, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import ConnectFlowSurface from "./ConnectFlowSurface";
+import { fetchConnectFlow } from "@/components/networking";
+
+const { startOAuthFlow, state, onSuccess } = vi.hoisted(() => ({
+  startOAuthFlow: vi.fn(),
+  onSuccess: { current: undefined as (() => void) | undefined },
+  state: { oauthReturn: null as string | null, connectFlow: null as string | null },
+}));
+
+vi.mock("next/navigation", () => ({
+  useRouter: () => ({ replace: vi.fn() }),
+  useSearchParams: () => ({
+    get: (key: string) => ({ mcpOauthReturn: state.oauthReturn, connect_flow: state.connectFlow })[key] ?? null,
+  }),
+}));
+vi.mock("@/components/networking", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@/components/networking")>()),
+  fetchConnectFlow: vi.fn(),
+  getProxyBaseUrl: () => "https://gateway.example.com",
+}));
+vi.mock("@/components/chat/MCPAppsPanel", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@/components/chat/MCPAppsPanel")>()),
+  default: () => <div data-testid="mcp-apps-panel" />,
+}));
+vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({
+  useUserMcpOAuthFlow: ({ onSuccess: success }: { onSuccess: () => void }) => {
+    onSuccess.current = success;
+    return { startOAuthFlow, status: "idle" };
+  },
+}));
+
+const flow = (state: "unscoped" | "interactive" | "m2m" | "stale", connected: boolean | null = null) => ({
+  state,
+  client_origin: "https://claude.ai",
+  server_id: state === "interactive" || state === "m2m" ? "s-design" : null,
+  server_name: state === "interactive" || state === "m2m" ? "design_tool" : null,
+  connected,
+});
+
+const renderSurface = () =>
+  render(
+    <QueryClientProvider client={new QueryClient({ defaultOptions: { queries: { retry: false } } })}>
+      <ConnectFlowSurface accessToken="token-123" selectedServers={[]} onChange={vi.fn()} />
+    </QueryClientProvider>,
+  );
+
+afterEach(() => {
+  state.oauthReturn = null;
+  state.connectFlow = null;
+  onSuccess.current = undefined;
+  sessionStorage.clear();
+  vi.clearAllMocks();
+});
+
+describe("ConnectFlowSurface", () => {
+  it.each([
+    { result: flow("unscoped"), grid: true, finish: true, cancel: false, oauthStarts: 0 },
+    { result: flow("interactive", false), grid: false, finish: false, cancel: true, oauthStarts: 1 },
+    { result: flow("interactive", true), grid: false, finish: true, cancel: true, oauthStarts: 0 },
+    { result: flow("m2m", true), grid: false, finish: true, cancel: true, oauthStarts: 0 },
+    { result: flow("stale"), grid: false, finish: false, cancel: true, oauthStarts: 0 },
+  ])(
+    "renders $result.state without widening its action surface",
+    async ({ result, grid, finish, cancel, oauthStarts }) => {
+      state.connectFlow = "flow-handle-123";
+      vi.mocked(fetchConnectFlow).mockResolvedValue(result);
+      renderSurface();
+
+      await screen.findByRole("button", { name: /finish connecting|cancel|connect/i });
+      await waitFor(() => expect(startOAuthFlow).toHaveBeenCalledTimes(oauthStarts));
+      expect(screen.queryByTestId("mcp-apps-panel") !== null).toBe(grid);
+      expect(screen.queryByRole("button", { name: /finish connecting/i }) !== null).toBe(finish);
+      expect(screen.queryByRole("button", { name: "Cancel" }) !== null).toBe(cancel);
+    },
+  );
+
+  it("keeps the grid and Finish hidden until the gateway accepts a handle", () => {
+    state.connectFlow = "flow-handle-123";
+    vi.mocked(fetchConnectFlow).mockReturnValue(new Promise(() => {}));
+    renderSurface();
+
+    expect(screen.queryByTestId("mcp-apps-panel")).not.toBeInTheDocument();
+    expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument();
+    expect(screen.getByRole("button", { name: "Cancel" })).toHaveAttribute("value", "deny");
+  });
+
+  it("keeps the grid and Finish hidden when flow validation fails", async () => {
+    state.connectFlow = "invalid-handle";
+    vi.mocked(fetchConnectFlow).mockRejectedValue(new Error("invalid flow"));
+    renderSurface();
+
+    await screen.findByRole("button", { name: "Cancel" });
+    expect(screen.queryByTestId("mcp-apps-panel")).not.toBeInTheDocument();
+    expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument();
+  });
+
+  it("refetches the sealed flow after the vendor connection completes", async () => {
+    state.connectFlow = "flow-handle-123";
+    vi.mocked(fetchConnectFlow)
+      .mockResolvedValueOnce(flow("interactive", false))
+      .mockResolvedValueOnce(flow("interactive", true));
+    renderSurface();
+
+    await waitFor(() => expect(startOAuthFlow).toHaveBeenCalledOnce());
+    await act(async () => onSuccess.current?.());
+
+    await screen.findByRole("button", { name: /finish connecting/i });
+  });
+
+  it("renders the ordinary panel without a flow handle", () => {
+    renderSurface();
+    expect(fetchConnectFlow).not.toHaveBeenCalled();
+    expect(screen.getByTestId("mcp-apps-panel")).toBeInTheDocument();
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.tsx
new file mode 100644
index 00000000000..33a61fceacf
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.tsx
@@ -0,0 +1,59 @@
+"use client";
+
+import React, { useEffect } from "react";
+import { useRouter, useSearchParams } from "next/navigation";
+import { useQuery } from "@tanstack/react-query";
+import MCPAppsPanel from "@/components/chat/MCPAppsPanel";
+import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner";
+import { fetchConnectFlow } from "@/components/networking";
+
+interface Props {
+  accessToken: string;
+  selectedServers: string[];
+  onChange: (servers: string[]) => void;
+}
+
+/** Renders the sealed gateway connect flow without trusting URL context. */
+const ConnectFlowSurface: React.FC<Props> = ({ accessToken, selectedServers, onChange }) => {
+  const router = useRouter();
+  const searchParams = useSearchParams();
+  const oauthReturn = searchParams.get("mcpOauthReturn");
+  const connectFlow = searchParams.get("connect_flow");
+
+  useEffect(() => {
+    if (oauthReturn) {
+      const url = new URL(window.location.href);
+      url.searchParams.delete("mcpOauthReturn");
+      router.replace(url.pathname + url.search);
+    }
+  }, [oauthReturn, router]);
+
+  const flowQuery = {
+    queryKey: ["gateway-connect-flow", connectFlow],
+    queryFn: () => fetchConnectFlow(connectFlow!),
+    enabled: !!connectFlow,
+    retry: false,
+  };
+  const { data: flow, isError, refetch } = useQuery(flowQuery);
+
+  if (connectFlow === null) {
+    return <MCPAppsPanel accessToken={accessToken} selectedServers={selectedServers} onChange={onChange} />;
+  }
+
+  return (
+    <>
+      <ConnectFlowBanner
+        flowHandle={connectFlow}
+        flow={flow}
+        accessToken={accessToken}
+        onConnected={refetch}
+        failed={isError}
+      />
+      {flow?.state === "unscoped" && (
+        <MCPAppsPanel accessToken={accessToken} selectedServers={selectedServers} onChange={onChange} connectMode />
+      )}
+    </>
+  );
+};
+
+export default ConnectFlowSurface;
diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx
index e8795c36bc6..c9609405676 100644
--- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx
+++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx
@@ -88,6 +88,13 @@ describe("MCPAppsPanel logos", () => {
 });
 
 const connectServers = [
+  {
+    server_id: "s-m2m",
+    server_name: "service_tool",
+    auth_type: "oauth2",
+    oauth2_flow: "client_credentials",
+    connected_app_reachable: true,
+  },
   {
     server_id: "s-reach",
     server_name: "reachable_srv",
@@ -124,6 +131,8 @@ describe("MCPAppsPanel connected-app reachability (LIT-4861)", () => {
     expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, true);
     expect(screen.queryByText("unreachable_srv")).not.toBeInTheDocument();
     expect(screen.getByText("Connected (1)")).toBeInTheDocument();
+    expect(screen.getByText("service_tool")).toBeInTheDocument();
+    expect(screen.queryByText("Connect", { exact: true })).not.toBeInTheDocument();
     const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]);
     expect(toolCountFetchedIds).toContain("s-reach");
     expect(toolCountFetchedIds).not.toContain("s-unreach");
diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx
index 38a095c067d..1fee8923e94 100644
--- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx
+++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx
@@ -13,23 +13,32 @@ import {
   getMCPOAuthUserCredentialStatus,
   listMCPTools,
 } from "../networking";
-import { AUTH_TYPE, MCPServer, MCPTool, handleTransport, isUnsupportedOnGatewayConnect } from "../mcp_tools/types";
+import {
+  getMcpOAuthMode,
+  MCPServer,
+  MCPTool,
+  handleTransport,
+  isUnsupportedOnGatewayConnect,
+} from "../mcp_tools/types";
 import { Logo } from "@/components/molecules/logo/Logo";
 import { toast } from "@/lib/toast";
 import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow";
+import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
 
 interface OAuth2ConnectButtonProps {
-  server: MCPServer;
+  server: Pick<MCPServer, "server_id" | "server_name" | "alias">;
   accessToken: string;
   onConnect: (serverId: string) => void;
   variant?: "badge" | "button";
+  autoStartKey?: string | null;
 }
 
-const OAuth2ConnectButton: React.FC<OAuth2ConnectButtonProps> = ({
+export const OAuth2ConnectButton: React.FC<OAuth2ConnectButtonProps> = ({
   server,
   accessToken,
   onConnect,
   variant = "badge",
+  autoStartKey = null,
 }) => {
   const name = server.server_name ?? server.alias ?? server.server_id;
   const { startOAuthFlow, status } = useUserMcpOAuthFlow({
@@ -39,6 +48,12 @@ const OAuth2ConnectButton: React.FC<OAuth2ConnectButtonProps> = ({
     onSuccess: useCallback(() => onConnect(server.server_id), [onConnect, server.server_id]),
   });
 
+  useEffect(() => {
+    if (autoStartKey === null || status !== "idle" || getSecureItem(autoStartKey) !== null) return;
+    setSecureItem(autoStartKey, "1");
+    startOAuthFlow();
+  }, [autoStartKey, status, startOAuthFlow]);
+
   const loading = status === "authorizing" || status === "exchanging";
 
   if (variant === "button") {
@@ -190,7 +205,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
         if (!isCurrentLoad()) return;
         const list: MCPServer[] = Array.isArray(serverData) ? serverData : serverData?.data ?? [];
         const reachable = connectMode ? list.filter((s) => s.connected_app_reachable !== false) : list;
-        const oauthServers = reachable.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2);
+        const oauthServers = reachable.filter((s) => getMcpOAuthMode(s) === "authorization_code");
         commitServers(reachable);
         setOauthChecking(new Set(oauthServers.map((s) => s.server_id)));
         setLoading(false);
@@ -274,7 +289,10 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
         <span className="text-[11px] text-muted-foreground shrink-0 whitespace-nowrap">{unavailabilityLabel}</span>
       );
     }
-    if (server.auth_type === AUTH_TYPE.OAUTH2) {
+    if (getMcpOAuthMode(server) === "m2m") {
+      return <CheckCircle className="h-3.5 w-3.5 text-success shrink-0" />;
+    }
+    if (getMcpOAuthMode(server) === "authorization_code") {
       if (oauthConnected.has(server.server_id)) {
         return <CheckCircle className="h-3.5 w-3.5 text-success shrink-0" />;
       }
@@ -339,7 +357,10 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
       if (unavailabilityLabel !== null) {
         return <span className="text-[13px] text-muted-foreground py-2.5 shrink-0">{unavailabilityLabel}</span>;
       }
-      if (detailServer.auth_type !== AUTH_TYPE.OAUTH2) {
+      if (getMcpOAuthMode(detailServer) === "m2m") {
+        return <span className="text-[13px] text-muted-foreground">Authorized</span>;
+      }
+      if (getMcpOAuthMode(detailServer) !== "authorization_code") {
         return (
           <Button
             variant={isConnected ? "outline" : "default"}
diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx
index 0e9133d6d3c..d7cf5b6110b 100644
--- a/ui/litellm-dashboard/src/components/networking.tsx
+++ b/ui/litellm-dashboard/src/components/networking.tsx
@@ -4918,6 +4918,24 @@ export const fetchDiscoverableMCPServers = async (accessToken: string) => {
   }
 };
 
+export interface ConnectFlowStatus {
+  state: "unscoped" | "interactive" | "m2m" | "stale";
+  client_origin: string;
+  server_id: string | null;
+  server_name: string | null;
+  connected: boolean | null;
+}
+
+/**
+ * What the gateway says about one in-flight connect flow, read from the sealed HttpOnly
+ * flow cookie rather than from the address bar (LIT-7075): the client asking, the server
+ * the flow is scoped to, and whether that server's vendor OAuth is already done. Sent with
+ * no access token; the flow and session cookies are the credential, exactly as they are for
+ * the finish form this page posts.
+ */
+export const fetchConnectFlow = async (flowHandle: string): Promise<ConnectFlowStatus> =>
+  apiClient.get(`/authorize/flow`, { query: { flow: flowHandle }, credentials: "include" });
+
 export const fetchMCPServers = async (accessToken: string, teamId?: string | null, connectedAppView?: boolean) => {
   try {
     return await apiClient.get(`/v1/mcp/server`, {
diff --git a/ui/litellm-dashboard/src/lib/http/client.test.ts b/ui/litellm-dashboard/src/lib/http/client.test.ts
index c3be1d48b42..0d765371e0c 100644
--- a/ui/litellm-dashboard/src/lib/http/client.test.ts
+++ b/ui/litellm-dashboard/src/lib/http/client.test.ts
@@ -41,6 +41,15 @@ describe("createApiClient", () => {
     expect(init.body).toBeUndefined();
   });
 
+  it("forwards the browser credential policy", async () => {
+    const fetchImpl = vi.fn(async () => okResponse({}));
+    const client = createApiClient({ getBaseUrl: () => "https://proxy.example", fetchImpl });
+
+    await client.get("/authorize/flow", { credentials: "include" });
+
+    expect(fetchImpl.mock.calls[0][1].credentials).toBe("include");
+  });
+
   it("JSON-serializes the body for writes", async () => {
     const fetchImpl = vi.fn(async () => okResponse({}));
     const client = createApiClient({ getBaseUrl: () => "", fetchImpl });
diff --git a/ui/litellm-dashboard/src/lib/http/client.ts b/ui/litellm-dashboard/src/lib/http/client.ts
index b8cfb81211a..c78b2b10011 100644
--- a/ui/litellm-dashboard/src/lib/http/client.ts
+++ b/ui/litellm-dashboard/src/lib/http/client.ts
@@ -25,6 +25,8 @@ export interface RequestOptions {
   query?: QueryParams;
   headers?: Record<string, string>;
   signal?: AbortSignal;
+  /** Send browser cookies with the request; needed for cookie-authenticated proxy routes. */
+  credentials?: RequestCredentials;
 }
 
 export class ApiError extends Error {
@@ -136,7 +138,7 @@ export function createApiClient(config: ApiClientConfig): ApiClient {
   const doFetch: typeof fetch = (input, init) => (fetchImpl ?? fetch)(input, init);
 
   async function request<T = any>(method: HttpMethod, path: string, options: RequestOptions = {}): Promise<T> {
-    const { accessToken, body, rawBody, query, headers: extraHeaders, signal } = options;
+    const { accessToken, body, rawBody, query, headers: extraHeaders, signal, credentials } = options;
 
     const url = appendQuery(`${getBaseUrl()}${path}`, query);
 
@@ -152,7 +154,7 @@ export function createApiClient(config: ApiClientConfig): ApiClient {
       Object.assign(headers, extraHeaders);
     }
 
-    const init: RequestInit = { method, headers, signal };
+    const init: RequestInit = { method, headers, signal, credentials };
     if (rawBody !== undefined) {
       init.body = rawBody;
     } else if (body !== undefined) {
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 6c2311a0ed9..232b366d8ce 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -1169,6 +1169,23 @@ export interface paths {
         patch?: never;
         trace?: never;
     };
+    "/authorize/flow": {
+        parameters: {
+            query?: never;
+            header?: never;
+            path?: never;
+            cookie?: never;
+        };
+        /** Authorize Flow */
+        get: operations["authorize_flow_authorize_flow_get"];
+        put?: never;
+        post?: never;
+        delete?: never;
+        options?: never;
+        head?: never;
+        patch?: never;
+        trace?: never;
+    };
     "/auto_router/benchmarks": {
         parameters: {
             query?: never;
@@ -41433,6 +41450,37 @@ export interface operations {
             };
         };
     };
+    authorize_flow_authorize_flow_get: {
+        parameters: {
+            query: {
+                flow: string;
+            };
+            header?: never;
+            path?: never;
+            cookie?: never;
+        };
+        requestBody?: never;
+        responses: {
+            /** @description Successful Response */
+            200: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": unknown;
+                };
+            };
+            /** @description Validation Error */
+            422: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": components["schemas"]["HTTPValidationError"];
+                };
+            };
+        };
+    };
     get_auto_router_benchmarks_auto_router_benchmarks_get: {
         parameters: {
             query?: {

From 95c0f9db7d92cefafa5e2d9799786b4931c9a390 Mon Sep 17 00:00:00 2001
From: Clement <clement.ho@paradigm.co>
Date: Wed, 9 Sep 2026 01:05:23 +0800
Subject: [PATCH 161/164] fix(router): rank streaming latency routing by raw
 TTFT, not TTFT per token (#40202)

* fix(router): rank streaming latency routing by raw TTFT, not TTFT per token

Latency-based routing divided time-to-first-token by completion_tokens
before storing it, so a deployment that streamed a long answer looked
faster to first token than one that answered briefly. TTFT is now stored
as plain seconds (first token time minus request start) in both the sync
and async success handlers, which is what the routing decision compares.

Non-streaming latency normalization per output token is unchanged.

Claude-Session: https://claude.ai/code/session_01Ttd5Q9ZhRPB4ch5guos3rj

* fix(router): store streaming TTFT under a seconds-only cache key

Workers on the previous release keep writing seconds-per-token samples
under "time_to_first_token" in the shared router cache during a rolling
deploy, so mixing the new raw-seconds samples into the same list averaged
incompatible units. Raw TTFT now lives under "time_to_first_token_seconds"
and routing reads only that key.

Also fix the regression test's token counts: with 50 tokens on the fast
deployment and 500 on the slow one the old per-token formula picks the
slow deployment, so the routing assertion now catches the bug.

Claude-Session: https://claude.ai/code/session_01Ttd5Q9ZhRPB4ch5guos3rj

* test(router): cover the TTFT sliding window from the unit-test shard

Move the TTFT list trimming checks from the CircleCI-only suite into the
mapped unit test file as one sync/async parametrized test, so the changed
lines in lowest_latency.py are exercised by the GitHub unit-test shard
that reports patch coverage.

Claude-Session: https://claude.ai/code/session_01Ttd5Q9ZhRPB4ch5guos3rj
---
 litellm/router_strategy/lowest_latency.py     |  54 +++----
 .../test_lowest_latency_routing.py            | 136 ------------------
 .../router_strategy/test_lowest_latency.py    | 127 ++++++++++++++++
 3 files changed, 149 insertions(+), 168 deletions(-)

diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py
index bb6877a032b..805d4ff9080 100644
--- a/litellm/router_strategy/lowest_latency.py
+++ b/litellm/router_strategy/lowest_latency.py
@@ -32,6 +32,12 @@ def _average_latency(samples: Sequence[float]) -> float:
     return sum(samples) / len(samples)
 
 
+def _ttft_seconds(elapsed: timedelta | float) -> float:
+    if isinstance(elapsed, timedelta):
+        return elapsed.total_seconds()
+    return float(elapsed)
+
+
 class LowestLatencyLoggingHandler(CustomLogger):
     test_flag: bool = False
     logged_success: int = 0
@@ -86,14 +92,13 @@ class LowestLatencyLoggingHandler(CustomLogger):
                     # breaks JSON serialization when the router cache syncs to
                     # Redis (issue #33169)
                     response_ms = response_ms.total_seconds()
-                time_to_first_token_response_time = None
+                time_to_first_token: float | None = None
 
                 if kwargs.get("stream", None) is not None and kwargs["stream"] is True:
                     # only log ttft for streaming request
-                    time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time
+                    time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time)
 
                 final_value: float = response_ms
-                time_to_first_token: float | None = None
                 total_tokens = 0
 
                 if isinstance(response_obj, ModelResponse):
@@ -111,13 +116,6 @@ class LowestLatencyLoggingHandler(CustomLogger):
                         else:
                             final_value = response_seconds
 
-                        if time_to_first_token_response_time is not None:
-                            if isinstance(time_to_first_token_response_time, timedelta):
-                                ttft_seconds = time_to_first_token_response_time.total_seconds()
-                            else:
-                                ttft_seconds = time_to_first_token_response_time
-                            time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens)
-
                 # ------------
                 # Update usage
                 # ------------
@@ -138,14 +136,14 @@ class LowestLatencyLoggingHandler(CustomLogger):
                 ## Time to first token
                 if time_to_first_token is not None:
                     if (
-                        len(request_count_dict[id].get("time_to_first_token", []))
+                        len(request_count_dict[id].get("time_to_first_token_seconds", []))
                         < self.routing_args.max_latency_list_size
                     ):
-                        request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token)
+                        request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token)
                     else:
-                        request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][
-                            1:
-                        ] + [time_to_first_token]
+                        request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][
+                            "time_to_first_token_seconds"
+                        ][1:] + [time_to_first_token]
 
                 if precise_minute not in request_count_dict[id]:
                     request_count_dict[id][precise_minute] = {}
@@ -252,7 +250,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
                     {model_group}_map: {
                         id: {
                             "latency": [..]
-                            "time_to_first_token": [..]
+                            "time_to_first_token_seconds": [..]
                             f"{date:hour:minute}" : {"tpm": 34, "rpm": 3}
                         }
                     }
@@ -273,14 +271,13 @@ class LowestLatencyLoggingHandler(CustomLogger):
                     # breaks JSON serialization when the router cache syncs to
                     # Redis (issue #33169)
                     response_ms = response_ms.total_seconds()
-                time_to_first_token_response_time = None
+                time_to_first_token: float | None = None
                 if kwargs.get("stream", None) is not None and kwargs["stream"] is True:
                     # only log ttft for streaming request
-                    time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time
+                    time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time)
 
                 final_value: float = response_ms
                 total_tokens = 0
-                time_to_first_token: float | None = None
 
                 if isinstance(response_obj, ModelResponse):
                     _usage: Final = getattr(response_obj, "usage", None)
@@ -296,13 +293,6 @@ class LowestLatencyLoggingHandler(CustomLogger):
                             final_value = float(normalized_value)
                         else:
                             final_value = response_seconds
-
-                        if time_to_first_token_response_time is not None:
-                            if isinstance(time_to_first_token_response_time, timedelta):
-                                ttft_seconds = time_to_first_token_response_time.total_seconds()
-                            else:
-                                ttft_seconds = time_to_first_token_response_time
-                            time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens)
                 # ------------
                 # Update usage
                 # ------------
@@ -328,14 +318,14 @@ class LowestLatencyLoggingHandler(CustomLogger):
                 ## Time to first token
                 if time_to_first_token is not None:
                     if (
-                        len(request_count_dict[id].get("time_to_first_token", []))
+                        len(request_count_dict[id].get("time_to_first_token_seconds", []))
                         < self.routing_args.max_latency_list_size
                     ):
-                        request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token)
+                        request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token)
                     else:
-                        request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][
-                            1:
-                        ] + [time_to_first_token]
+                        request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][
+                            "time_to_first_token_seconds"
+                        ][1:] + [time_to_first_token]
 
                 if precise_minute not in request_count_dict[id]:
                     request_count_dict[id][precise_minute] = {}
@@ -433,7 +423,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
                 or float("inf")
             )
             item_latency = item_map.get("latency", [])
-            item_ttft_latency = item_map.get("time_to_first_token", [])
+            item_ttft_latency = item_map.get("time_to_first_token_seconds", [])
             item_rpm = item_map.get(precise_minute, {}).get("rpm", 0)
             item_tpm = item_map.get(precise_minute, {}).get("tpm", 0)
 
diff --git a/tests/local_testing/test_lowest_latency_routing.py b/tests/local_testing/test_lowest_latency_routing.py
index 598b1dbcaf9..aba4500199a 100644
--- a/tests/local_testing/test_lowest_latency_routing.py
+++ b/tests/local_testing/test_lowest_latency_routing.py
@@ -1077,73 +1077,6 @@ async def test_latency_list_trimming_discards_oldest_entry_async():
         ), f"Oldest latency {oldest_latency} should have been discarded"
 
 
-def test_ttft_list_trimming_discards_oldest_entry():
-    """
-    The time_to_first_token list trims the oldest entry when full, matching
-    the behavior of the latency list.
-    """
-    max_size = 3
-    test_cache = DualCache()
-    lowest_latency_logger = LowestLatencyLoggingHandler(
-        router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
-    )
-
-    model_group = "gpt-3.5-turbo"
-    deployment_id = "test-deployment"
-
-    ttft_values = []
-    for i in range(max_size + 1):
-        start_time = time.time()
-        expected_ttft = float(i + 1) * 0.1  # 0.1, 0.2, 0.3, 0.4
-        completion_start_time = start_time + expected_ttft
-        end_time = start_time + float(i + 1)
-        ttft_values.append(expected_ttft)
-
-        kwargs = {
-            "litellm_params": {
-                "metadata": {
-                    "model_group": model_group,
-                    "deployment": "azure/gpt-4.1-mini",
-                },
-                "model_info": {"id": deployment_id},
-            },
-            "stream": True,
-            "completion_start_time": completion_start_time,
-        }
-        # TTFT is only recorded when response_obj is a ModelResponse.
-        response_obj = litellm.ModelResponse(
-            usage=litellm.Usage(completion_tokens=1, total_tokens=1)
-        )
-
-        lowest_latency_logger.log_success_event(
-            response_obj=response_obj,
-            kwargs=kwargs,
-            start_time=start_time,
-            end_time=end_time,
-        )
-
-    latency_key = f"{model_group}_map"
-    cached_data = test_cache.get_cache(key=latency_key)
-    ttft_list = cached_data[deployment_id].get("time_to_first_token", [])
-
-    assert (
-        len(ttft_list) == max_size
-    ), f"Expected {max_size} entries, got {len(ttft_list)}"
-
-    newest_ttft = ttft_values[-1]
-    oldest_ttft = ttft_values[0]
-    tolerance = 0.05
-
-    assert (
-        abs(ttft_list[-1] - newest_ttft) < tolerance
-    ), f"Newest TTFT {newest_ttft} should be at end of list"
-
-    for ttft in ttft_list:
-        assert (
-            abs(ttft - oldest_ttft) > tolerance
-        ), f"Oldest TTFT {oldest_ttft} should have been discarded"
-
-
 @pytest.mark.asyncio
 async def test_timeout_penalty_discards_oldest_entry():
     """
@@ -1269,72 +1202,3 @@ def test_list_order_preserved_after_multiple_trims():
         assert (
             abs(latency_list[i] - expected) < tolerance
         ), f"At index {i}, expected ~{expected}, got {latency_list[i]}"
-
-
-@pytest.mark.asyncio
-async def test_ttft_list_trimming_discards_oldest_entry_async():
-    """
-    Async counterpart: the time_to_first_token list trims the oldest entry
-    when full. Exercises the async_log_success_event TTFT path, which only
-    runs when response_obj is a ModelResponse and the call is marked as
-    streaming with a completion_start_time.
-    """
-    max_size = 3
-    test_cache = DualCache()
-    lowest_latency_logger = LowestLatencyLoggingHandler(
-        router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
-    )
-
-    model_group = "gpt-3.5-turbo"
-    deployment_id = "test-deployment"
-
-    ttft_values = []
-    for i in range(max_size + 1):
-        start_time = time.time()
-        expected_ttft = float(i + 1) * 0.1  # 0.1, 0.2, 0.3, 0.4
-        completion_start_time = start_time + expected_ttft
-        end_time = start_time + float(i + 1)
-        ttft_values.append(expected_ttft)
-
-        kwargs = {
-            "litellm_params": {
-                "metadata": {
-                    "model_group": model_group,
-                    "deployment": "azure/gpt-4.1-mini",
-                },
-                "model_info": {"id": deployment_id},
-            },
-            "stream": True,
-            "completion_start_time": completion_start_time,
-        }
-        response_obj = litellm.ModelResponse(
-            usage=litellm.Usage(completion_tokens=1, total_tokens=1)
-        )
-
-        await lowest_latency_logger.async_log_success_event(
-            response_obj=response_obj,
-            kwargs=kwargs,
-            start_time=start_time,
-            end_time=end_time,
-        )
-
-    latency_key = f"{model_group}_map"
-    cached_data = await test_cache.async_get_cache(key=latency_key)
-    ttft_list = cached_data[deployment_id].get("time_to_first_token", [])
-
-    assert (
-        len(ttft_list) == max_size
-    ), f"Expected {max_size} entries, got {len(ttft_list)}"
-
-    newest_ttft = ttft_values[-1]
-    oldest_ttft = ttft_values[0]
-    tolerance = 0.05
-
-    assert (
-        abs(ttft_list[-1] - newest_ttft) < tolerance
-    ), f"Newest TTFT {newest_ttft} should be at end of list"
-
-    for ttft in ttft_list:
-        assert (
-            abs(ttft - oldest_ttft) > tolerance
-        ), f"Oldest TTFT {oldest_ttft} should have been discarded"
diff --git a/tests/test_litellm/router_strategy/test_lowest_latency.py b/tests/test_litellm/router_strategy/test_lowest_latency.py
index eb02459be68..812d7bbff32 100644
--- a/tests/test_litellm/router_strategy/test_lowest_latency.py
+++ b/tests/test_litellm/router_strategy/test_lowest_latency.py
@@ -165,6 +165,133 @@ def test_sync_chat_zero_completion_tokens_falls_back_to_seconds():
     json.dumps({"latency": latencies})
 
 
+MODEL_GROUP = "gpt-4o-mini"
+FAST_TTFT_ID = "fast-ttft-short-output"
+SLOW_TTFT_ID = "slow-ttft-long-output"
+STREAMING_DEPLOYMENTS = [
+    {"model_info": {"id": FAST_TTFT_ID}, "litellm_params": {}},
+    {"model_info": {"id": SLOW_TTFT_ID}, "litellm_params": {}},
+]
+
+
+def _streaming_kwargs(deployment_id: str, start_time: datetime, ttft_seconds: float):
+    return {
+        "litellm_params": {
+            "metadata": {"model_group": MODEL_GROUP},
+            "model_info": {"id": deployment_id},
+        },
+        "stream": True,
+        "completion_start_time": start_time + timedelta(seconds=ttft_seconds),
+    }
+
+
+def _recorded_ttft(cache: DualCache, deployment_id: str):
+    cached = cache.get_cache(key=f"{MODEL_GROUP}_map") or {}
+    return cached.get(deployment_id, {}).get("time_to_first_token_seconds", [])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"])
+async def test_streaming_ttft_ranking_ignores_completion_length(sync_mode: bool):
+    """Deployment A: TTFT 1s, 50 completion tokens. Deployment B: TTFT 3s, 500
+    completion tokens. Dividing TTFT by completion tokens made B look faster
+    (3/500 = 0.006 beats 1/50 = 0.02); actual TTFT must win."""
+    cache = DualCache()
+    handler = LowestLatencyLoggingHandler(router_cache=cache)
+    start_time = datetime(2026, 1, 1, 12, 0, 0)
+    end_time = start_time + timedelta(seconds=10)
+
+    samples = (
+        (FAST_TTFT_ID, 1.0, 50),
+        (SLOW_TTFT_ID, 3.0, 500),
+    )
+    for deployment_id, ttft, completion_tokens in samples:
+        kwargs = _streaming_kwargs(deployment_id, start_time, ttft)
+        response_obj = _chat_response(completion_tokens=completion_tokens)
+        if sync_mode:
+            handler.log_success_event(
+                response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=end_time
+            )
+        else:
+            await handler.async_log_success_event(
+                response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=end_time
+            )
+
+    assert _recorded_ttft(cache, FAST_TTFT_ID) == [pytest.approx(1.0)]
+    assert _recorded_ttft(cache, SLOW_TTFT_ID) == [pytest.approx(3.0)]
+
+    request_kwargs = {"stream": True, "metadata": {}}
+    if sync_mode:
+        picked = handler.get_available_deployments(
+            model_group=MODEL_GROUP, healthy_deployments=STREAMING_DEPLOYMENTS, request_kwargs=request_kwargs
+        )
+    else:
+        picked = await handler.async_get_available_deployments(
+            model_group=MODEL_GROUP, healthy_deployments=STREAMING_DEPLOYMENTS, request_kwargs=request_kwargs
+        )
+
+    assert picked is not None
+    assert picked["model_info"]["id"] == FAST_TTFT_ID
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"])
+async def test_ttft_window_keeps_newest_samples_when_full(sync_mode: bool):
+    """Float timestamps, as the SDK passes them. Once max_latency_list_size
+    samples exist the oldest TTFT is dropped so the window slides."""
+    max_size = 3
+    cache = DualCache()
+    handler = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"max_latency_list_size": max_size})
+    start_time = 1_700_000_000.0
+    ttfts = (0.1, 0.2, 0.3, 0.4)
+
+    for ttft in ttfts:
+        kwargs = {
+            "litellm_params": {
+                "metadata": {"model_group": MODEL_GROUP},
+                "model_info": {"id": FAST_TTFT_ID},
+            },
+            "stream": True,
+            "completion_start_time": start_time + ttft,
+        }
+        response_obj = _chat_response(completion_tokens=1)
+        if sync_mode:
+            handler.log_success_event(
+                response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=start_time + 1.0
+            )
+        else:
+            await handler.async_log_success_event(
+                response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=start_time + 1.0
+            )
+
+    assert _recorded_ttft(cache, FAST_TTFT_ID) == [pytest.approx(ttft) for ttft in ttfts[-max_size:]]
+
+
+@pytest.mark.asyncio
+async def test_streaming_routing_ignores_per_token_ttft_samples_from_older_workers():
+    """Workers on the previous release share the Redis map and keep writing
+    seconds-per-token under the old "time_to_first_token" key during a rolling
+    deploy. Those samples favor SLOW; routing must only read the seconds key."""
+    cache = DualCache()
+    handler = LowestLatencyLoggingHandler(router_cache=cache)
+    cache.set_cache(
+        key=f"{MODEL_GROUP}_map",
+        value={
+            FAST_TTFT_ID: {"time_to_first_token": [0.02], "time_to_first_token_seconds": [1.0]},
+            SLOW_TTFT_ID: {"time_to_first_token": [0.006], "time_to_first_token_seconds": [3.0]},
+        },
+    )
+
+    picked = await handler.async_get_available_deployments(
+        model_group=MODEL_GROUP,
+        healthy_deployments=STREAMING_DEPLOYMENTS,
+        request_kwargs={"stream": True, "metadata": {}},
+    )
+
+    assert picked is not None
+    assert picked["model_info"]["id"] == FAST_TTFT_ID
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "cached_entry",

From 978f93ce9cf3e2748494a1c8b67b48488930b06b Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 8 Sep 2026 10:05:32 -0700
Subject: [PATCH 162/164] fix(tool_permission): log expected skip and deny
 events below WARNING (#40208)

The no-tools skip path in async_pre_call_hook now logs at DEBUG and the
denied-by-rule messages in async_pre_call_hook and _evaluate_tool_calls
log at INFO. The malformed tool arguments warning is unchanged. Adds
regression tests pinning each level

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../guardrail_hooks/tool_permission.py        |  6 +-
 .../guardrail_hooks/test_tool_permission.py   | 75 +++++++++++++++++++
 2 files changed, 78 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py
index a8b33109900..e20f0b320b9 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py
@@ -547,7 +547,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
 
         for _tool_call, is_allowed, _rule_id, message in checked:
             if not is_allowed and message is not None:
-                verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message)
+                verbose_proxy_logger.info("Tool Permission Guardrail: %s", message)
                 if self.on_disallowed_action == "block":
                     raise GuardrailRaisedException(
                         guardrail_name=self.guardrail_name, message=message, blocked_content=True
@@ -809,7 +809,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
 
         new_tools: Final = self._collect_request_tools(data)
         if not new_tools:
-            verbose_proxy_logger.warning(
+            verbose_proxy_logger.debug(
                 "Tool Permission Guardrail: not running guardrail. No tools or functions in data"
             )
             return data
@@ -820,7 +820,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
             is_allowed, _, message = self._check_tool_permission(tool_name, tool_type)
 
             if not is_allowed and message is not None:
-                verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message)
+                verbose_proxy_logger.info("Tool Permission Guardrail: %s", message)
                 if self.on_disallowed_action == "block":
                     raise HTTPException(
                         status_code=400,
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py
index 0dbd4591ac9..87a1b84acc5 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py
@@ -3,6 +3,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics)
 """
 
 import json
+import logging
 import re
 from unittest.mock import patch
 
@@ -520,6 +521,80 @@ class TestToolPermissionGuardrail:
                 )
         assert excinfo.value.status_code == 400
 
+    @pytest.mark.asyncio
+    async def test_async_pre_call_hook_without_tools_logs_skip_at_debug(self, caplog):
+        data = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}
+
+        with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"):
+            with patch.object(self.guardrail, "should_run_guardrail", return_value=True):
+                result = await self.guardrail.async_pre_call_hook(
+                    user_api_key_dict=UserAPIKeyAuth(),
+                    cache=DualCache(default_in_memory_ttl=1),
+                    data=data,
+                    call_type="completion",
+                )
+
+        assert result is data
+        skip_levels = [r.levelno for r in caplog.records if "No tools or functions in data" in r.getMessage()]
+        assert skip_levels == [logging.DEBUG]
+        assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == []
+
+    @pytest.mark.asyncio
+    async def test_async_pre_call_hook_denied_tool_logs_at_info(self, caplog):
+        data = {"tools": [{"type": "function", "function": {"name": "Read"}}]}
+
+        with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"):
+            with patch.object(self.guardrail, "should_run_guardrail", return_value=True):
+                with pytest.raises(HTTPException):
+                    await self.guardrail.async_pre_call_hook(
+                        user_api_key_dict=UserAPIKeyAuth(),
+                        cache=DualCache(default_in_memory_ttl=1),
+                        data=data,
+                        call_type="completion",
+                    )
+
+        denied_levels = [
+            r.levelno
+            for r in caplog.records
+            if r.getMessage() == "Tool Permission Guardrail: Tool 'Read' denied by rule 'deny_read'"
+        ]
+        assert denied_levels == [logging.INFO]
+        assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == []
+
+    @pytest.mark.asyncio
+    async def test_async_post_call_success_hook_denied_tool_logs_at_info(self, caplog):
+        tool_call = {"function": {"name": "Read", "arguments": "{}"}, "type": "function"}
+        response = ModelResponse(choices=[Choices(message={"tool_calls": [tool_call]})])
+
+        with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"):
+            with patch.object(self.guardrail, "should_run_guardrail", return_value=True):
+                with pytest.raises(GuardrailRaisedException):
+                    await self.guardrail.async_post_call_success_hook(
+                        data={"guardrails": ["test-tool-permission"]},
+                        user_api_key_dict=UserAPIKeyAuth(),
+                        response=response,
+                    )
+
+        denied_levels = [
+            r.levelno
+            for r in caplog.records
+            if r.getMessage() == "Tool Permission Guardrail: Tool 'Read' denied by rule 'deny_read'"
+        ]
+        assert denied_levels == [logging.INFO]
+        assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == []
+
+    def test_parse_tool_call_arguments_malformed_json_logs_warning(self, caplog):
+        tool_call = ChatCompletionMessageToolCall(function={"name": "Bash", "arguments": "{not json"}, id="call_1")
+
+        with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"):
+            parsed, error = self.guardrail._parse_tool_call_arguments(tool_call)
+
+        assert parsed is None
+        assert error == "arguments could not be parsed"
+        warning_messages = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING]
+        assert len(warning_messages) == 1
+        assert warning_messages[0].startswith("Tool Permission Guardrail: Failed to decode arguments for tool Bash")
+
     @pytest.mark.asyncio
     async def test_async_pre_call_hook_blocks_legacy_functions(self):
         data = {

From 99824533ff7a99c07e0c768f61fcd5b5204a7794 Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 8 Sep 2026 10:06:53 -0700
Subject: [PATCH 163/164] fix(proxy): kill the whole prisma process group when
 a boot migration command times out (#39509)

* fix(proxy): kill the whole prisma process group when a boot migration command times out

* fix(proxy): keep boot alive when the installed proxy extras has no prisma runner
---
 litellm/proxy/db/check_migration.py           | 39 ++++++---
 litellm/proxy/db/prisma_client.py             | 22 ++++--
 tests/test_litellm/proxy/db/conftest.py       | 79 +++++++++++++++++++
 .../proxy/db/test_check_migration.py          | 32 ++++++++
 .../proxy/db/test_prisma_client.py            | 59 +++++++++-----
 5 files changed, 197 insertions(+), 34 deletions(-)

diff --git a/litellm/proxy/db/check_migration.py b/litellm/proxy/db/check_migration.py
index b7a07d4eeea..6e2a06c96e1 100644
--- a/litellm/proxy/db/check_migration.py
+++ b/litellm/proxy/db/check_migration.py
@@ -46,17 +46,32 @@ def extract_sql_commands(diff_output: str) -> list[str]:
 
 def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]:
     """Checks for differences between current database and Prisma schema.
+
+    Never raises: a diff that cannot be produced, because the runner is missing,
+    because the command failed, or because it outlived its budget, is reported as
+    "no diff" so boot continues.
+
     Returns:
         A tuple containing:
         - A boolean indicating if differences were found (True) or not (False).
-        - A string with the diff output or error message.
-    Raises:
-        subprocess.CalledProcessError: If the Prisma command fails.
-        Exception: For any other errors during execution.
+        - The SQL commands that would close the diff, empty when there is none.
     """
-    verbose_logger.debug("Checking for Prisma schema diff...")
     try:
-        result: Final = subprocess.run(
+        from litellm_proxy_extras.prisma_toolchain import (
+            PRISMA_COMMAND_TIMEOUT_ENV_VAR,
+            prisma_command_timeout,
+            run_prisma,
+        )
+    except ImportError as e:
+        print(  # noqa: T201  # boot-time operator output, same channel as this helper's other messages
+            f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}"
+        )
+        return False, []
+
+    verbose_logger.debug("Checking for Prisma schema diff...")
+    timeout: Final = prisma_command_timeout()
+    try:
+        result: Final = run_prisma(
             [
                 "prisma",
                 "migrate",
@@ -67,12 +82,10 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]:
                 "./schema.prisma",
                 "--script",
             ],
-            capture_output=True,
-            text=True,
-            check=True,
+            timeout=timeout,
+            env=os.environ.copy(),
         )
 
-        # return True, "Migration diff generated successfully."
         sql_commands: Final = extract_sql_commands(result.stdout)
 
         if sql_commands:
@@ -83,6 +96,12 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]:
             return True, sql_commands
         else:
             return False, []
+    except subprocess.TimeoutExpired:
+        print(  # noqa: T201  # boot-time operator output, same channel as this helper's other messages
+            f"Timed out after {timeout}s generating the migration diff. "
+            f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer."
+        )
+        return False, []
     except subprocess.CalledProcessError as e:
         error_message: Final = f"Failed to generate migration diff. Error: {e.stderr}"
         print(error_message)  # noqa: T201
diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py
index 2190ae55fd2..9ea2432f2f5 100644
--- a/litellm/proxy/db/prisma_client.py
+++ b/litellm/proxy/db/prisma_client.py
@@ -937,9 +937,17 @@ class PrismaManager:
                         use_v2_resolver=use_v2_resolver,
                     )
                 else:
+                    try:
+                        from litellm_proxy_extras.prisma_toolchain import (
+                            prisma_command_timeout,
+                            run_prisma,
+                        )
+                    except ImportError as e:
+                        verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e)
+                        return False
+
                     PrismaManager._raise_if_partitioned_spend_logs()
-                    # Use prisma db push with increased timeout
-                    subprocess.run(
+                    run_prisma(
                         [
                             "prisma",
                             "db",
@@ -947,13 +955,15 @@ class PrismaManager:
                             "--accept-data-loss",
                             "--skip-generate",
                         ],
-                        timeout=60,
-                        check=True,
+                        timeout=prisma_command_timeout(),
+                        env=os.environ.copy(),
+                        stdout=None,
+                        stderr=None,
                     )
                     PrismaManager._apply_replica_identity_full_if_requested()
                     return True
-            except subprocess.TimeoutExpired:
-                verbose_proxy_logger.warning("Attempt %s timed out", attempt + 1)
+            except subprocess.TimeoutExpired as e:
+                verbose_proxy_logger.warning("Attempt %s timed out after %.0fs", attempt + 1, e.timeout)
                 time.sleep(random.randrange(5, 15))
             except subprocess.CalledProcessError as e:
                 attempts_left = 3 - attempt
diff --git a/tests/test_litellm/proxy/db/conftest.py b/tests/test_litellm/proxy/db/conftest.py
index 6f67b91ac1d..03d7ea81257 100644
--- a/tests/test_litellm/proxy/db/conftest.py
+++ b/tests/test_litellm/proxy/db/conftest.py
@@ -1,5 +1,11 @@
+import json
 import os
+import signal
+import sys
+import time
 from collections.abc import Generator
+from dataclasses import dataclass
+from pathlib import Path
 from typing import Optional
 
 import pytest
@@ -75,3 +81,76 @@ def reset_entra_token_provider_cache() -> Generator[None, None, None]:
 def unset_database_url(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setenv("DATABASE_URL", "about-to-be-unset")
     monkeypatch.delenv("DATABASE_URL")
+
+
+FAKE_PRISMA_CLI = """#!{python}
+import json
+import os
+import pathlib
+import subprocess
+import sys
+import time
+
+calls_file = pathlib.Path(os.environ["FAKE_PRISMA_CALLS"])
+earlier_calls = calls_file.read_text().splitlines() if calls_file.exists() else []
+with calls_file.open("a") as log:
+    print(json.dumps(sys.argv[1:]), file=log)
+if not earlier_calls and os.environ.get("FAKE_PRISMA_HANG_FIRST"):
+    grandchild = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(600)"])
+    pathlib.Path(os.environ["FAKE_PRISMA_GRANDCHILD_PIDFILE"]).write_text(str(grandchild.pid))
+    time.sleep(600)
+sys.exit(0)
+"""
+
+
+@dataclass(frozen=True, slots=True)
+class FakePrismaCli:
+    """A stand-in `prisma` on PATH, recording every invocation.
+
+    With FAKE_PRISMA_HANG_FIRST set it hangs on its first call from a process tree
+    of its own, the way the real CLI wraps Node around a Rust schema engine, so a
+    timeout that kills only the direct child leaves the rest of that tree running.
+    """
+
+    calls_file: Path
+    grandchild_pidfile: Path
+
+    @property
+    def calls(self) -> list[list[str]]:
+        if not self.calls_file.exists():
+            return []
+        return [json.loads(line) for line in self.calls_file.read_text().splitlines()]
+
+    def grandchild_is_gone(self, within_seconds: float) -> bool:
+        deadline = time.monotonic() + within_seconds
+        while time.monotonic() < deadline:
+            try:
+                os.kill(int(self.grandchild_pidfile.read_text()), 0)
+            except ProcessLookupError:
+                return True
+            time.sleep(0.05)
+        return False
+
+
+@pytest.fixture
+def fake_prisma_cli(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generator[FakePrismaCli, None, None]:
+    bin_dir = tmp_path / "fakebin"
+    bin_dir.mkdir()
+    script = bin_dir / "prisma"
+    script.write_text(FAKE_PRISMA_CLI.format(python=sys.executable))
+    script.chmod(0o755)
+    cli = FakePrismaCli(
+        calls_file=tmp_path / "calls.jsonl",
+        grandchild_pidfile=tmp_path / "grandchild.pid",
+    )
+    monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}")
+    monkeypatch.setenv("FAKE_PRISMA_CALLS", str(cli.calls_file))
+    monkeypatch.setenv("FAKE_PRISMA_GRANDCHILD_PIDFILE", str(cli.grandchild_pidfile))
+    monkeypatch.setenv("LITELLM_PRISMA_COMMAND_TIMEOUT", "1")
+    monkeypatch.delenv("FAKE_PRISMA_HANG_FIRST", raising=False)
+    yield cli
+    if cli.grandchild_pidfile.exists():
+        try:
+            os.kill(int(cli.grandchild_pidfile.read_text()), signal.SIGKILL)
+        except ProcessLookupError:
+            pass
diff --git a/tests/test_litellm/proxy/db/test_check_migration.py b/tests/test_litellm/proxy/db/test_check_migration.py
index 9e2f6a1089c..74a6841cb23 100644
--- a/tests/test_litellm/proxy/db/test_check_migration.py
+++ b/tests/test_litellm/proxy/db/test_check_migration.py
@@ -37,3 +37,35 @@ def test_check_migration_out_of_sync(mocker):
     check_migration.verbose_logger.exception.assert_called_once()
     actual_message = check_migration.verbose_logger.exception.call_args[0][0]
     assert "prisma schema out of sync with db" in actual_message
+
+
+@pytest.mark.timeout(30)
+def test_migrate_diff_stops_at_its_budget_and_takes_its_process_tree_with_it(fake_prisma_cli, monkeypatch):
+    """
+    `prisma migrate diff` ran unbounded, so a database that never answers hung boot
+    before uvicorn ever started, and interrupting the proxy orphaned the schema engine.
+    """
+    from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper
+
+    monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1")
+
+    assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, [])
+    assert fake_prisma_cli.calls == [
+        ["migrate", "diff", "--from-url", "postgresql://u:p@localhost:9/x",
+         "--to-schema-datamodel", "./schema.prisma", "--script"]
+    ]
+    assert fake_prisma_cli.grandchild_is_gone(within_seconds=5)
+
+
+def test_migrate_diff_without_the_prisma_runner_skips_instead_of_crashing_boot(monkeypatch):
+    """
+    Boot calls this helper directly, so an ImportError here takes the proxy down before
+    uvicorn starts. An install without the runner must lose the diagnostic, not the proxy.
+    """
+    import sys
+
+    from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper
+
+    monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None)
+
+    assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, [])
diff --git a/tests/test_litellm/proxy/db/test_prisma_client.py b/tests/test_litellm/proxy/db/test_prisma_client.py
index f0983d6bf62..7f43e483557 100644
--- a/tests/test_litellm/proxy/db/test_prisma_client.py
+++ b/tests/test_litellm/proxy/db/test_prisma_client.py
@@ -10,7 +10,7 @@ from fastapi.testclient import TestClient
 
 
 
-from litellm.proxy.db.prisma_client import PrismaWrapper, should_update_prisma_schema
+from litellm.proxy.db.prisma_client import PrismaManager, PrismaWrapper, should_update_prisma_schema
 
 
 @pytest.fixture(autouse=True)
@@ -193,7 +193,10 @@ async def test_recreate_prisma_client_recovers_from_disconnected_client(
     mock_new_prisma.connect.assert_awaited_once()
 
 
-def test_db_push_applies_replica_identity_full_when_requested(monkeypatch):
+DB_PUSH_ARGV = ["db", "push", "--accept-data-loss", "--skip-generate"]
+
+
+def test_db_push_applies_replica_identity_full_when_requested(monkeypatch, fake_prisma_cli, unset_database_url):
     """`prisma db push` bypasses litellm-proxy-extras, so it needs its own call
     into the opt-in REPLICA IDENTITY FULL step."""
     from litellm.proxy.db.prisma_client import PrismaManager
@@ -208,14 +211,13 @@ def test_db_push_applies_replica_identity_full_when_requested(monkeypatch):
         staticmethod(lambda: applied.append(True)),
     )
 
-    with patch("litellm.proxy.db.prisma_client.subprocess.run") as mock_run:
-        assert PrismaManager.setup_database(use_migrate=False) is True
+    assert PrismaManager.setup_database(use_migrate=False) is True
 
-    assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"]
+    assert fake_prisma_cli.calls == [DB_PUSH_ARGV]
     assert applied == [True]
 
 
-def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch):
+def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch, fake_prisma_cli, unset_database_url):
     """A doc-partitioned LiteLLM_SpendLogs makes `prisma db push` rewrite the
     primary key back to ("request_id"), which Postgres rejects; the guard must
     fail fast with guidance instead of running the push."""
@@ -228,29 +230,23 @@ def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch):
     monkeypatch.setattr(
         ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: True)
     )
-    with patch(  # test-quality-ok: subprocess.run is the external prisma CLI boundary, asserted never reached
-        "litellm.proxy.db.prisma_client.subprocess.run"
-    ) as mock_run:
-        with pytest.raises(RuntimeError) as err:
-            PrismaManager.setup_database(use_migrate=False)
+    with pytest.raises(RuntimeError) as err:
+        PrismaManager.setup_database(use_migrate=False)
 
     assert str(err.value) == PARTITIONED_SPEND_LOGS_PUSH_ERROR
-    mock_run.assert_not_called()
+    assert fake_prisma_cli.calls == []
 
 
-def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch):
+def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch, fake_prisma_cli, unset_database_url):
     from litellm.proxy.db.prisma_client import PrismaManager
     from litellm_proxy_extras.utils import ProxyExtrasDBManager
 
     monkeypatch.setattr(
         ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: False)
     )
-    with patch(  # test-quality-ok: subprocess.run is the external prisma CLI boundary, not SDK logic
-        "litellm.proxy.db.prisma_client.subprocess.run"
-    ) as mock_run:
-        assert PrismaManager.setup_database(use_migrate=False) is True
+    assert PrismaManager.setup_database(use_migrate=False) is True
 
-    assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"]
+    assert fake_prisma_cli.calls == [DB_PUSH_ARGV]
 
 
 def _entra_jwt(expires_in_seconds: int) -> str:
@@ -377,3 +373,30 @@ def test_minting_without_the_database_env_vars_names_them(azure_env, monkeypatch
 
     with pytest.raises(RuntimeError, match="DATABASE_HOST"):
         wrapper.get_rds_iam_token()
+
+
+@pytest.mark.timeout(45)
+def test_db_push_timeout_takes_its_process_tree_with_it(fake_prisma_cli, unset_database_url, monkeypatch):
+    """
+    A timed-out `db push` used to leave Node and the schema engine writing the schema,
+    so the next attempt pushed into a database the abandoned one was still mutating.
+    """
+    monkeypatch.delenv("LITELLM_SET_REPLICA_IDENTITY_FULL", raising=False)
+    monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1")
+
+    assert PrismaManager.setup_database(use_migrate=False) is True
+    assert fake_prisma_cli.calls == [DB_PUSH_ARGV, DB_PUSH_ARGV]
+    assert fake_prisma_cli.grandchild_is_gone(within_seconds=5)
+
+
+def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashing_boot(
+    fake_prisma_cli, unset_database_url, monkeypatch
+):
+    """
+    An ImportError out of setup_database escapes the caller's RuntimeError handler and
+    kills boot, bypassing the operator's enforce_prisma_migration_check choice.
+    """
+    monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None)
+
+    assert PrismaManager.setup_database(use_migrate=False) is False
+    assert fake_prisma_cli.calls == []

From a85c3152ca59b5b56d1168e2a384e29a4fb0b861 Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 8 Sep 2026 10:08:24 -0700
Subject: [PATCH 164/164] fix(proxy): load db credentials in the model
 reconcile so a worker never serves a model before its credential (#39876)

* fix(proxy): load db credentials inside the model reconcile so a worker never serves a model before its credential

* fix(proxy): load db credentials in the model read-through so a request miss never adds a model before its credential

* fix(proxy): read credentials from the writer db before the router update and look a credential up once

* test(proxy): assert the credential is loaded when db models reach the router instead of the call order
---
 .../litellm_core_utils/credential_accessor.py |  15 ++-
 .../common_utils/registry_read_through.py     |   1 +
 litellm/proxy/proxy_server.py                 |  31 +----
 litellm/utils.py                              |  18 ++-
 .../common_utils/test_config_sync_pubsub.py   |   7 +-
 .../test_registry_read_through.py             |  45 +++++++
 .../test_scheduled_job_stagger.py             |   2 +-
 .../proxy/proxy_server/test_proxy_config.py   | 123 ++++++++++++++++++
 tests/test_litellm/proxy/test_proxy_server.py |  68 +++++-----
 tests/test_litellm/test_utils.py              |  44 +++++++
 10 files changed, 278 insertions(+), 76 deletions(-)

diff --git a/litellm/litellm_core_utils/credential_accessor.py b/litellm/litellm_core_utils/credential_accessor.py
index 7071750b970..7b4e8240b69 100644
--- a/litellm/litellm_core_utils/credential_accessor.py
+++ b/litellm/litellm_core_utils/credential_accessor.py
@@ -7,16 +7,19 @@ from litellm.types.utils import CredentialItem
 
 
 class CredentialAccessor:
+    @staticmethod
+    def find_credential(credential_name: str) -> CredentialItem | None:
+        return next(
+            (credential for credential in litellm.credential_list if credential.credential_name == credential_name),
+            None,
+        )
+
     @staticmethod
     def get_credential_values(credential_name: str) -> dict:
         """Safe accessor for credentials."""
 
-        if not litellm.credential_list:
-            return {}
-        for credential in litellm.credential_list:
-            if credential.credential_name == credential_name:
-                return credential.credential_values.copy()
-        return {}
+        credential: Final = CredentialAccessor.find_credential(credential_name)
+        return {} if credential is None else credential.credential_values.copy()
 
     @staticmethod
     def upsert_credentials(credentials: list[CredentialItem]):
diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py
index 460b348e188..63da5d15207 100644
--- a/litellm/proxy/common_utils/registry_read_through.py
+++ b/litellm/proxy/common_utils/registry_read_through.py
@@ -125,6 +125,7 @@ async def _resync_model_deployments(model_name: str) -> bool:
         )
         return proxy_server.llm_router is not None
     async with proxy_server.MODEL_RECONCILE_LOCK:
+        await proxy_server.proxy_config.get_credentials(prisma_client=prisma_client)
         proxy_server.proxy_config._add_deployment(db_models=rows)
         proxy_server.llm_model_list = router.get_model_list()
     return True
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 32b6b841af7..1d6baf5fded 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -7109,11 +7109,10 @@ class ProxyConfig:
                 ],
             )
 
-            # Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set)
-            if self._should_load_db_object(object_type="models"):
-                new_models: Final = await self._get_models_from_db(prisma_client=prisma_client)
-
-                # update llm router
+            load_models: Final = self._should_load_db_object(object_type="models")
+            new_models: Final = await self._get_models_from_db(prisma_client=prisma_client) if load_models else None
+            await self.get_credentials(prisma_client=prisma_client)
+            if load_models:
                 still_desired_ids = await self._update_llm_router(
                     new_models=new_models, proxy_logging_obj=proxy_logging_obj
                 )
@@ -7153,12 +7152,9 @@ class ProxyConfig:
         async def _resync_config_from_db() -> None:
             await self.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
 
-        async def _resync_credentials_from_db() -> None:
-            await self.get_credentials(prisma_client=prisma_client)
-
         subscriber: Final = ConfigSyncSubscriber(
             redis_cache=redis_cache,
-            resync_callbacks=(_resync_config_from_db, _resync_credentials_from_db),
+            resync_callbacks=(_resync_config_from_db,),
         )
         self.config_sync_subscriber = subscriber
         subscriber.start()
@@ -8013,7 +8009,7 @@ class ProxyConfig:
 
     async def get_credentials(self, prisma_client: PrismaClient):
         try:
-            credentials = await CredentialsRepository(prisma_client).find_all()
+            credentials = await CredentialsRepository(WriterPinnedClient(prisma_client.db)).find_all()
             credentials = [self.decrypt_credentials(cred) for cred in credentials]
             await self.delete_credentials(credentials)  # delete credentials that are not in the all-up list
             CredentialAccessor.upsert_credentials(credentials)  # upsert credentials that are in the all-up list
@@ -9597,19 +9593,6 @@ class ProxyStartupEvent:
         )
 
         if store_model_in_db is True:
-            ### GET STORED CREDENTIALS ###
-            scheduler.add_job(
-                proxy_config.get_credentials,
-                "interval",
-                seconds=config_reload_interval_seconds,
-                # REMOVED jitter parameter - major cause of memory leak
-                args=[prisma_client],
-                id="get_credentials_job",
-                replace_existing=True,
-                misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
-            )
-            await proxy_config.get_credentials(prisma_client=prisma_client)
-
             # MEMORY LEAK FIX: Increase interval from 10s to 30s minimum
             # Frequent polling was causing excessive memory allocations
             scheduler.add_job(
@@ -9623,7 +9606,7 @@ class ProxyStartupEvent:
                 misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
             )
 
-            # this will load all existing models on proxy startup
+            # this will load all existing credentials and models on proxy startup
             await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
 
             proxy_config.start_config_sync_subscriber(
diff --git a/litellm/utils.py b/litellm/utils.py
index d0e11bc9551..141ec323776 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -672,11 +672,19 @@ def load_credentials_from_list(kwargs: dict):
     CredentialAccessor: Final = getattr(sys.modules[__name__], "CredentialAccessor")
 
     credential_name: Final = kwargs.get("litellm_credential_name")
-    if credential_name and litellm.credential_list:
-        credential_accessor: Final[Mapping[str, object]] = CredentialAccessor.get_credential_values(credential_name)
-        for key, value in credential_accessor.items():
-            if key not in kwargs:
-                kwargs[key] = value
+    if not credential_name:
+        return
+    credential: Final = CredentialAccessor.find_credential(credential_name)
+    if credential is None:
+        verbose_logger.warning(
+            "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it",
+            credential_name,
+            len(litellm.credential_list),
+        )
+        return
+    for key, value in credential.credential_values.items():
+        if key not in kwargs:
+            kwargs[key] = value
 
 
 def get_dynamic_callbacks(
diff --git a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py
index 3bb3f75b9b8..2a0ebf9492a 100644
--- a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py
+++ b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py
@@ -824,7 +824,7 @@ class _StopFailingSubscriber(ConfigSyncSubscriber):
         raise RuntimeError("stop failed")
 
 
-async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() -> None:
+async def test_proxy_config_subscriber_resyncs_deployments_only() -> None:
     from litellm.proxy.proxy_server import ProxyConfig
 
     cache = _FakeRedisCache(_ScriptedPubSubRedisClient([_QueuePubSub()]))
@@ -852,10 +852,7 @@ async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() ->
         await callback()
     await config.stop_config_sync_subscriber()
 
-    assert calls == [
-        ("add_deployment", prisma_client, proxy_logging_obj),
-        ("get_credentials", prisma_client, None),
-    ]
+    assert calls == [("add_deployment", prisma_client, proxy_logging_obj)]
     assert config.config_sync_subscriber is None
     assert subscriber._task is None
 
diff --git a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py
index f0fbdea4e85..6f7c20166c5 100644
--- a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py
+++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py
@@ -393,6 +393,51 @@ async def test_resync_model_deployments_mutates_router_under_model_reconcile_loc
     assert not proxy_server.MODEL_RECONCILE_LOCK.locked()
 
 
+@pytest.mark.asyncio
+async def test_resync_model_deployments_loads_db_credentials_before_reconciling_models(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
+    from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments
+    from litellm.types.utils import CredentialItem
+
+    rows: Final = [MagicMock()]
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=rows)
+    router: Final = MagicMock()
+    router.get_model_list.return_value = []
+    installed: Final = MagicMock()
+
+    async def load_credentials_from_db(prisma_client: object) -> None:
+        CredentialAccessor.upsert_credentials(
+            [
+                CredentialItem(
+                    credential_name="openai-cred",
+                    credential_values={"api_key": "sk-from-db"},
+                    credential_info={},
+                )
+            ]
+        )
+
+    def install_models(db_models: object) -> None:
+        installed(db_models=db_models, credential=CredentialAccessor.get_credential_values("openai-cred"))
+
+    monkeypatch.setattr(litellm, "credential_list", [])
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+    monkeypatch.setattr(proxy_server, "llm_router", router)
+    monkeypatch.setattr(proxy_server, "llm_model_list", None)
+    monkeypatch.setattr(proxy_server.proxy_config, "get_credentials", load_credentials_from_db)
+    monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", install_models)
+
+    assert await _resync_model_deployments("model-created-on-a-sibling-replica") is True
+    installed.assert_called_once_with(db_models=rows, credential={"api_key": "sk-from-db"})
+
+
 @pytest.mark.asyncio
 async def test_resync_model_deployments_respects_supported_db_objects(monkeypatch):
     from unittest.mock import AsyncMock, MagicMock
diff --git a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py
index ca4d62737b6..54e0aa74a25 100644
--- a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py
+++ b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py
@@ -23,7 +23,7 @@ from litellm.proxy.common_utils.scheduled_job_stagger import (
 )
 
 OPERATOR_CRON_JOB_ID = "spend_log_cleanup_job"
-SHARED_INTERVAL_JOB_IDS = ("periodic_reload_job", "get_credentials_job", "add_deployment_job")
+SHARED_INTERVAL_JOB_IDS = ("periodic_reload_job", "add_deployment_job")
 
 
 async def _noop() -> None: ...
diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
index 770cec1834e..ae4ec4086bb 100644
--- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
+++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
@@ -2939,6 +2939,129 @@ async def test_ProxyConfig_add_deployment_applies_db_router_settings(monkeypatch
     fake_router.update_settings.assert_called_once_with(routing_strategy="latency-based-routing")
 
 
+def _stub_add_deployment_collaborators(
+    monkeypatch: pytest.MonkeyPatch, pc: ProxyConfig, fake_prisma: MagicMock
+) -> None:
+    from litellm.proxy import proxy_server
+
+    fake_router = MagicMock()
+    fake_router.get_model_list = MagicMock(return_value=[])
+
+    async def fake_get_config(*args: object, **kwargs: object) -> dict[str, object]:
+        return {}
+
+    monkeypatch.setattr(litellm, "credential_list", [])
+    monkeypatch.setattr(pc, "get_config", fake_get_config)
+    monkeypatch.setattr(pc, "_init_non_llm_objects_in_db", AsyncMock())
+    monkeypatch.setattr(proxy_server, "prefetch_config_params", AsyncMock())
+    monkeypatch.setattr(proxy_server, "get_config_param", AsyncMock(return_value=None))
+    monkeypatch.setattr(proxy_server, "llm_router", fake_router)
+    monkeypatch.setattr(proxy_server, "master_key", "sk-master")
+    monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
+    monkeypatch.setattr(proxy_server, "proxy_config", pc)
+    monkeypatch.delenv("LITELLM_SALT_KEY", raising=False)
+
+
+def _encrypted_credential_row(credential_name: str, api_key: str) -> dict[str, object]:
+    from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
+
+    return {
+        "credential_name": credential_name,
+        "credential_values": {"api_key": encrypt_value_helper(api_key, new_encryption_key="sk-master")},
+        "credential_info": {"custom_llm_provider": "openai"},
+    }
+
+
+def _fake_prisma_with_encrypted_credential(credential_name: str, api_key: str) -> MagicMock:
+    fake_prisma = MagicMock()
+    fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(
+        return_value=[_encrypted_credential_row(credential_name, api_key)]
+    )
+    return fake_prisma
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig_add_deployment_loads_db_credentials_before_reconciling_models(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
+    from litellm.proxy import proxy_server
+    from litellm.utils import load_credentials_from_list
+
+    pc = ProxyConfig()
+    fake_prisma = MagicMock()
+    fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(return_value=[])
+    _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma)
+    monkeypatch.setattr(proxy_server, "general_settings", {})
+    installed = MagicMock()
+
+    async def read_models_while_a_credential_lands(prisma_client: object) -> list[MagicMock]:
+        fake_prisma.db.litellm_credentialstable.find_many.return_value = [
+            _encrypted_credential_row("openai-cred", "sk-from-db")
+        ]
+        return [MagicMock()]
+
+    async def install_models(new_models: object, proxy_logging_obj: object) -> None:
+        installed(credential=CredentialAccessor.get_credential_values("openai-cred"))
+
+    monkeypatch.setattr(pc, "_get_models_from_db", read_models_while_a_credential_lands)
+    monkeypatch.setattr(pc, "_update_llm_router", install_models)
+
+    await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock())
+
+    installed.assert_called_once_with(credential={"api_key": "sk-from-db"})
+    assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"}
+    request_kwargs = {"litellm_credential_name": "openai-cred"}
+    load_credentials_from_list(request_kwargs)
+    assert request_kwargs == {"litellm_credential_name": "openai-cred", "api_key": "sk-from-db"}
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig_add_deployment_loads_db_credentials_even_when_models_are_not_db_objects(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
+    from litellm.proxy import proxy_server
+
+    pc = ProxyConfig()
+    fake_prisma = _fake_prisma_with_encrypted_credential("openai-cred", "sk-from-db")
+    _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma)
+    monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["mcp"]})
+    models_fetch = AsyncMock(return_value=[])
+    monkeypatch.setattr(pc, "_get_models_from_db", models_fetch)
+
+    await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock())
+
+    models_fetch.assert_not_awaited()
+    assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"}
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig_get_credentials_reads_from_writer_not_replica(monkeypatch: pytest.MonkeyPatch) -> None:
+    from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
+    from litellm.proxy.db.prisma_client import PrismaWrapper
+    from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
+
+    pc = ProxyConfig()
+    writer_inner = MagicMock(name="writer_prisma")
+    reader_inner = MagicMock(name="reader_prisma")
+    writer_inner.litellm_credentialstable.find_many = AsyncMock(
+        return_value=[_encrypted_credential_row("openai-cred", "sk-from-writer")]
+    )
+    reader_inner.litellm_credentialstable.find_many = AsyncMock(return_value=[])
+    fake_prisma = MagicMock()
+    fake_prisma.db = RoutingPrismaWrapper(
+        writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False),
+        reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False),
+    )
+    _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma)
+
+    await pc.get_credentials(prisma_client=fake_prisma)
+
+    assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-writer"}
+    reader_inner.litellm_credentialstable.find_many.assert_not_awaited()
+
+
 # ---------------------------------------------------------------------------
 # ProxyConfig._add_general_settings_from_db_config
 # ---------------------------------------------------------------------------
diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py
index b7bb58378d4..937e4f15741 100644
--- a/tests/test_litellm/proxy/test_proxy_server.py
+++ b/tests/test_litellm/proxy/test_proxy_server.py
@@ -822,16 +822,16 @@ def _mock_scheduled_proxy_config() -> MagicMock:
 
 
 @pytest.mark.asyncio
-async def test_initialize_scheduled_jobs_credentials(monkeypatch):
-    """
-    Test that get_credentials is only called when store_model_in_db is True
-    """
+async def test_initialize_scheduled_jobs_loads_credentials_only_through_add_deployment(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
     monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
     monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
+    from apscheduler.schedulers.asyncio import AsyncIOScheduler
+
     from litellm.proxy.proxy_server import ProxyStartupEvent
     from litellm.proxy.utils import ProxyLogging
 
-    # Mock dependencies
     mock_prisma_client = MagicMock()
     mock_proxy_logging = MagicMock(spec=ProxyLogging)
     mock_proxy_logging.slack_alerting_instance = MagicMock()
@@ -841,25 +841,6 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch):
     with (
         patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
         patch("litellm.proxy.proxy_server.store_model_in_db", False),
-    ):  # set store_model_in_db to False
-        # Test when store_model_in_db is False
-        await ProxyStartupEvent.initialize_scheduled_background_jobs(
-            general_settings={},
-            prisma_client=mock_prisma_client,
-            proxy_budget_rescheduler_min_time=1,
-            proxy_budget_rescheduler_max_time=2,
-            proxy_batch_write_at=5,
-            proxy_logging_obj=mock_proxy_logging,
-        )
-
-        # Verify get_credentials was not called
-        mock_proxy_config.get_credentials.assert_not_called()
-
-    # Now test with store_model_in_db = True
-    with (
-        patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
-        patch("litellm.proxy.proxy_server.store_model_in_db", True),
-        patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True),
     ):
         await ProxyStartupEvent.initialize_scheduled_background_jobs(
             general_settings={},
@@ -870,12 +851,31 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch):
             proxy_logging_obj=mock_proxy_logging,
         )
 
-        # Verify get_credentials was called both directly and scheduled
-        assert mock_proxy_config.get_credentials.call_count == 1  # Direct call
+        mock_proxy_config.get_credentials.assert_not_called()
+        mock_proxy_config.add_deployment.assert_not_called()
 
-        # Verify a scheduled job was added for get_credentials
-        mock_scheduler_calls = [call[0] for call in mock_proxy_config.get_credentials.mock_calls]
-        assert len(mock_scheduler_calls) > 0
+    scheduler = AsyncIOScheduler()
+    try:
+        with (
+            patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
+            patch("litellm.proxy.proxy_server.store_model_in_db", True),
+            patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler),
+        ):
+            await ProxyStartupEvent.initialize_scheduled_background_jobs(
+                general_settings={},
+                prisma_client=mock_prisma_client,
+                proxy_budget_rescheduler_min_time=1,
+                proxy_budget_rescheduler_max_time=2,
+                proxy_batch_write_at=5,
+                proxy_logging_obj=mock_proxy_logging,
+            )
+
+        assert scheduler.get_job("get_credentials_job") is None
+        assert scheduler.get_job("add_deployment_job") is not None
+        mock_proxy_config.get_credentials.assert_not_called()
+        assert mock_proxy_config.add_deployment.call_count == 1
+    finally:
+        scheduler.shutdown(wait=False)
 
 
 @pytest.mark.asyncio
@@ -924,7 +924,7 @@ async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypat
 @pytest.mark.asyncio
 async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch):
     """
-    The DB config-reload jobs (add_deployment, get_credentials) that keep multi-pod
+    The DB config-reload job (add_deployment) that keeps multi-pod
     deployments in sync must be scheduled at the configured
     proxy_config_reload_interval_seconds, not a hardcoded value.
     """
@@ -967,7 +967,7 @@ async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(
         if "id" in job_call.kwargs
     }
     assert scheduled_seconds["add_deployment_job"] == configured_interval
-    assert scheduled_seconds["get_credentials_job"] == configured_interval
+    assert "get_credentials_job" not in scheduled_seconds
 
 
 @pytest.mark.asyncio
@@ -1011,7 +1011,7 @@ async def test_initialize_scheduled_jobs_rejects_non_positive_config_reload_inte
         if "id" in job_call.kwargs
     }
     assert scheduled_seconds["add_deployment_job"] == 30
-    assert scheduled_seconds["get_credentials_job"] == 30
+    assert "get_credentials_job" not in scheduled_seconds
 
 
 @pytest.mark.asyncio
@@ -7446,10 +7446,8 @@ async def test_store_model_in_db_db_override_when_config_false():
         # store_model_in_db should now be True (overridden by DB)
         assert ps.store_model_in_db is True
 
-        # add_deployment and get_credentials should have been called
-        # since store_model_in_db is now True
         assert mock_proxy_config.add_deployment.call_count == 1
-        assert mock_proxy_config.get_credentials.call_count == 1
+        mock_proxy_config.get_credentials.assert_not_called()
 
 
 @pytest.mark.asyncio
diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py
index 8a56a84ade7..4ce68b71079 100644
--- a/tests/test_litellm/test_utils.py
+++ b/tests/test_litellm/test_utils.py
@@ -6092,3 +6092,47 @@ class TestFinalOptionalParamsLineRedaction:
 
         assert "'max_tokens': 17" in printed
         assert "'temperature': 0.25" in printed
+
+
+def _credential_warnings(caplog: pytest.LogCaptureFixture) -> list[str]:
+    return [record.getMessage() for record in caplog.records if "litellm_credential_name=" in record.getMessage()]
+
+
+def test_load_credentials_from_list_warns_when_the_named_credential_is_not_loaded(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    from litellm.utils import load_credentials_from_list
+
+    monkeypatch.setattr(litellm, "credential_list", [])
+    request_kwargs = {"litellm_credential_name": "openai-cred", "model": "openai/gpt-5.4-mini"}
+    with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
+        load_credentials_from_list(request_kwargs)
+
+    assert request_kwargs == {"litellm_credential_name": "openai-cred", "model": "openai/gpt-5.4-mini"}
+    assert _credential_warnings(caplog) == [
+        "litellm_credential_name=openai-cred matched none of the 0 loaded credentials; the request runs without it"
+    ]
+
+
+def test_load_credentials_from_list_fills_kwargs_from_the_loaded_credential_without_warning(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    from litellm.types.utils import CredentialItem
+    from litellm.utils import load_credentials_from_list
+
+    loaded = CredentialItem(
+        credential_name="openai-cred",
+        credential_values={"api_key": "sk-from-db", "api_base": "https://credential.example"},
+        credential_info={},
+    )
+    monkeypatch.setattr(litellm, "credential_list", [loaded])
+    request_kwargs = {"litellm_credential_name": "openai-cred", "api_base": "https://request.example"}
+    with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
+        load_credentials_from_list(request_kwargs)
+
+    assert request_kwargs == {
+        "litellm_credential_name": "openai-cred",
+        "api_base": "https://request.example",
+        "api_key": "sk-from-db",
+    }
+    assert _credential_warnings(caplog) == []