From 429ad06972b7263435a551a1a5f12f3f985afb53 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 29 Aug 2026 14:23:57 -0700
Subject: [PATCH 01/43] fix(guardrails): write structured_messages rewrites
 back into /v1/responses input

Message-rewriting guardrails such as Headroom return their rewrite in
structured_messages and leave texts untouched. The responses guardrail
translation only mapped texts back, so compression never reached the
upstream request on /v1/responses while the retrieve tool still got
injected. Convert the returned messages back to Responses input (plus
instructions) the way the chat and Anthropic handlers already do, and
keep developer messages as input_text in the chat-to-responses bridge.

Resolves LIT-6494
---
 .../transformation.py                         |  2 +-
 .../guardrail_translation/handler.py          | 47 +++++++--
 ...responses_transformation_transformation.py | 13 +++
 ...test_openai_responses_guardrail_handler.py | 96 +++++++++++++++++++
 .../guardrail_hooks/test_headroom.py          | 34 +++++++
 5 files changed, 181 insertions(+), 11 deletions(-)

diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py
index 85fb0bc8dc6..63c66a4c2c5 100644
--- a/litellm/completion_extras/litellm_responses_transformation/transformation.py
+++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py
@@ -931,7 +931,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode)
 
     def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]:
-        if role == "user" or role == "system" or role == "tool":
+        if role in ("user", "system", "developer", "tool"):
             return {"type": "input_text", "text": content}
         else:
             return {"type": "output_text", "text": content}
diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py
index 7c5d8ac99ad..4320edb8414 100644
--- a/litellm/llms/openai/responses/guardrail_translation/handler.py
+++ b/litellm/llms/openai/responses/guardrail_translation/handler.py
@@ -38,6 +38,7 @@ from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
 from litellm.completion_extras.litellm_responses_transformation.transformation import (
+    LiteLLMResponsesTransformationHandler,
     OpenAiResponsesToChatCompletionStreamIterator,
 )
 from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
@@ -149,8 +150,15 @@ class OpenAIResponsesHandler(BaseTranslation):
                 input_type="request",
                 logging_obj=litellm_logging_obj,
             )
-            guardrailed_texts = guardrailed_inputs.get("texts", [])
-            data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data
+            guardrailed_structured_messages = guardrailed_inputs.get("structured_messages")
+            if (
+                guardrailed_structured_messages is not None
+                and guardrailed_structured_messages is not structured_messages
+            ):
+                self._write_back_structured_messages(data, guardrailed_structured_messages)
+            else:
+                guardrailed_texts = guardrailed_inputs.get("texts", [])
+                data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data
             self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools"))
             verbose_proxy_logger.debug("OpenAI Responses API: Processed string input")
             return data
@@ -198,24 +206,43 @@ class OpenAIResponsesHandler(BaseTranslation):
                 logging_obj=litellm_logging_obj,
             )
 
-            guardrailed_texts = guardrailed_inputs.get("texts", [])
             self._apply_guardrailed_tools_to_data(
                 data,
                 original_tools_list,
                 guardrailed_inputs.get("tools"),
             )
 
-            # Step 3: Map guardrail responses back to original input structure
-            await self._apply_guardrail_responses_to_input(
-                messages=input_data,
-                responses=guardrailed_texts,
-                task_mappings=task_mappings,
-            )
+            guardrailed_structured_messages = guardrailed_inputs.get("structured_messages")
+            if (
+                guardrailed_structured_messages is not None
+                and guardrailed_structured_messages is not structured_messages
+            ):
+                self._write_back_structured_messages(data, guardrailed_structured_messages)
+            else:
+                # Step 3: Map guardrail responses back to original input structure
+                await self._apply_guardrail_responses_to_input(
+                    messages=input_data,
+                    responses=guardrailed_inputs.get("texts", []),
+                    task_mappings=task_mappings,
+                )
 
-        verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", input_data)
+        verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input"))
 
         return data
 
+    @staticmethod
+    def _write_back_structured_messages(data: dict, structured_messages: Sequence[AllMessageValues]) -> None:
+        input_items, instructions = (
+            LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(
+                list(structured_messages)
+            )
+        )
+        data["input"] = input_items
+        if instructions is None:
+            data.pop("instructions", None)
+            return
+        data["instructions"] = instructions
+
     def extract_request_tool_names(self, data: dict) -> list[str]:
         """Extract tool names from Responses API request (tools[].name for function
         and custom, tools[].server_label for mcp)."""
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 21b60d7a216..6590718878d 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
@@ -1295,6 +1295,19 @@ def test_text_plus_tool_calls_sequence():
 # =============================================================================
 
 
+def test_developer_message_content_uses_input_text():
+    handler = LiteLLMResponsesTransformationHandler()
+
+    input_items, instructions = handler.convert_chat_completion_messages_to_responses_api(
+        [{"role": "developer", "content": "Always answer in French."}]
+    )
+
+    assert instructions is None
+    assert input_items == [
+        {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "Always answer in French."}]}
+    ]
+
+
 def test_tool_message_output_uses_input_text_not_output_text():
     """
     Test that tool message content uses input_text type, not output_text.
diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
index 447175b09a6..63f8babf760 100644
--- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
+++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
@@ -1229,3 +1229,99 @@ class TestOpenAIResponsesHandlerToolInjection:
         names = [t.get("name") for t in result["tools"]]
         assert "get_weather" in names
         assert "injected_tool" in names
+
+
+COMPRESSED_MARKER = "[compressed document; retrieve the full text with hash=b573993006976af767214fac]"
+
+
+class StructuredRewriteGuardrail(CustomGuardrail):
+    """Guardrail that rewrites whole messages via structured_messages and leaves
+    texts untouched, the way message-compressing guardrails do."""
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        messages = list(inputs.get("structured_messages") or [])
+        first_user = next(i for i, m in enumerate(messages) if m.get("role") == "user")
+        rewritten = [
+            {**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages)
+        ]
+        return {**inputs, "structured_messages": rewritten}
+
+
+def _texts(item: dict) -> list[str]:
+    content = item.get("content")
+    if isinstance(content, str):
+        return [content]
+    return [part["text"] for part in content]
+
+
+class TestStructuredMessagesWriteBack:
+    """A guardrail's structured_messages rewrite must land in the Responses request,
+    not only the per-text mapping the chat handler shares with it."""
+
+    @pytest.mark.asyncio
+    async def test_list_input_gets_rewritten_messages_and_keeps_instructions(self):
+        handler = OpenAIResponsesHandler()
+        data = {
+            "model": "gpt-5.6",
+            "instructions": "Answer from the memo only.",
+            "input": [
+                {"role": "user", "content": "memo " * 400},
+                {"role": "assistant", "content": "Understood."},
+                {"role": "user", "content": "What is the codename?"},
+            ],
+        }
+
+        result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
+
+        assert result["instructions"] == "Answer from the memo only."
+        user_items = [item for item in result["input"] if item.get("role") == "user"]
+        assert [_texts(item) for item in user_items] == [[COMPRESSED_MARKER], ["What is the codename?"]]
+        assert not any(item.get("role") == "system" for item in result["input"])
+        assert _texts(next(item for item in result["input"] if item.get("role") == "assistant")) == ["Understood."]
+
+    @pytest.mark.asyncio
+    async def test_string_input_becomes_rewritten_message_list(self):
+        handler = OpenAIResponsesHandler()
+        data = {"model": "gpt-5.6", "input": "memo " * 400}
+
+        result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
+
+        assert [_texts(item) for item in result["input"]] == [[COMPRESSED_MARKER]]
+        assert "instructions" not in result
+
+    @pytest.mark.asyncio
+    async def test_developer_item_survives_write_back_as_input_text(self):
+        handler = OpenAIResponsesHandler()
+        data = {
+            "model": "gpt-5.6",
+            "input": [
+                {"role": "developer", "content": "Always answer in French."},
+                {"role": "user", "content": "memo " * 400},
+                {"role": "user", "content": "What is the codename?"},
+            ],
+        }
+
+        result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
+
+        developer = next(item for item in result["input"] if item.get("role") == "developer")
+        assert developer["content"] == [{"type": "input_text", "text": "Always answer in French."}]
+
+    @pytest.mark.asyncio
+    async def test_same_inputs_object_back_keeps_the_text_mapping(self):
+        handler = OpenAIResponsesHandler()
+        original_input = [
+            {"role": "user", "content": "Hello"},
+            {"role": "user", "content": [{"type": "input_text", "text": "Again"}]},
+        ]
+        data = {"model": "gpt-5.6", "input": original_input}
+
+        result = await handler.process_input_messages(data, MockGuardrail())
+
+        assert result["input"] is original_input
+        assert [_texts(item) for item in result["input"]] == [["Hello [GUARDRAILED]"], ["Again [GUARDRAILED]"]]
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
index 1fbc975e40a..37b9852d3e1 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
@@ -954,6 +954,40 @@ async def test_passthrough_handler_does_not_log_headroom_as_run(
     assert "headroom" not in _applied_guardrails(data)
 
 
+@pytest.mark.asyncio
+async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstream(
+    guardrail: HeadroomGuardrail,
+):
+    """Regression for LIT-6494: on /v1/responses the compressed messages must be
+    written back into `input`, not only the retrieve tool into `tools`, or the
+    model keeps reading the full document and never calls headroom_retrieve."""
+    from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler
+
+    data = {
+        "model": "gpt-5.6",
+        "instructions": ORIGINAL_MESSAGES[0]["content"],
+        "input": [{"role": m["role"], "content": m["content"]} for m in ORIGINAL_MESSAGES[1:]],
+        "tools": [{"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}}],
+    }
+
+    with patch.object(
+        guardrail.async_handler,
+        "post",
+        new_callable=AsyncMock,
+        return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH),
+    ):
+        result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail)
+
+    assert result["instructions"] == ORIGINAL_MESSAGES[0]["content"]
+    assert [item["content"][0]["text"] for item in result["input"]] == [
+        COMPRESSED_MESSAGES_WITH_HASH[0]["content"],
+        ORIGINAL_MESSAGES[2]["content"],
+        ORIGINAL_MESSAGES[3]["content"],
+    ]
+    assert "A" * 5000 not in json.dumps(result["input"])
+    assert [tool["name"] for tool in result["tools"]] == ["get_weather", HEADROOM_RETRIEVE_TOOL_NAME]
+
+
 @pytest.mark.asyncio
 async def test_apply_guardrail_http_error_raises():
     guardrail = _make_guardrail()

From ef96af51211318ea9bbc9d49947c75b71a3f00b5 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 29 Aug 2026 14:34:42 -0700
Subject: [PATCH 02/43] fix(headroom): resolve CCR retrieval on streaming
 /v1/responses

Streaming /v1/responses requests that carried Headroom's retrieve tool were
sent upstream as streams, so the model's headroom_retrieve function_call was
streamed straight back to a client that never declared the tool and the
retrieval never resolved. Chat completions already avoid this by converting
the request to non-stream in the pre-call deployment hook, letting the
agentic loop resolve the retrieve, and fake-streaming the final answer.

The hook now converts responses call types too, the responses handler wraps
the resolved result as a fake stream whenever any interception converted the
stream (shared converted_stream_requested helper instead of per-integration
key checks), and the follow-up request filter drops every non-code-interpreter
interception key through is_interception_internal_key.

Resolves LIT-6481
---
 litellm/llms/custom_httpx/llm_http_handler.py |  11 +-
 .../guardrail_hooks/headroom/headroom.py      |   5 +-
 litellm/types/integrations/custom_logger.py   |   8 ++
 .../guardrail_hooks/test_headroom.py          | 117 +++++++++++++++++-
 4 files changed, 132 insertions(+), 9 deletions(-)

diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index 573ba85416f..52a91a1d06d 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -97,9 +97,12 @@ from litellm.types.containers.main import (
 )
 from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig
 from litellm.types.integrations.custom_logger import (
+    NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
     AgenticLoopPlan,
     AgenticLoopRequestPatch,
     AgenticLoopSafetyError,
+    converted_stream_requested,
+    is_interception_internal_key,
 )
 from litellm.types.llms.anthropic_messages.anthropic_response import (
     AnthropicMessagesResponse,
@@ -2932,10 +2935,7 @@ class BaseLLMHTTPHandler:
         )
 
         result: Final = final_response if final_response is not None else initial_response
-        interception_converted_stream: Final = litellm_params.get(
-            "_code_interpreter_interception_converted_stream"
-        ) or litellm_params.get("_websearch_interception_converted_stream")
-        if interception_converted_stream and not litellm_params.get("_agentic_loop_depth"):
+        if converted_stream_requested(litellm_params) and not litellm_params.get("_agentic_loop_depth"):
             return self._wrap_responses_response_as_fake_stream(
                 result=result,
                 model=model,
@@ -5399,8 +5399,7 @@ class BaseLLMHTTPHandler:
         kwargs_for_followup: Final = {
             k: v
             for k, v in kwargs.items()
-            if not k.startswith("_websearch_interception")
-            and not k.startswith("_compression_interception")
+            if not is_interception_internal_key(k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES)
             and k != "_code_interpreter_interception_converted_stream"
             and k not in internal_keys
             and k not in optional_params
diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
index e2d2fffb2df..08753bdaed9 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
@@ -49,6 +49,9 @@ if TYPE_CHECKING:
     from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
 
 BYPASS_HEADER: Final = "x-headroom-bypass"
+_STREAM_CONVERTIBLE_CALL_TYPES: Final = frozenset(
+    (CallTypes.completion, CallTypes.acompletion, CallTypes.responses, CallTypes.aresponses)
+)
 HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve"
 _HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})")
 _HASH_CACHE_TTL_SECONDS: Final = 15 * 60
@@ -724,7 +727,7 @@ class HeadroomGuardrail(CustomGuardrail):
     ) -> dict[str, Any] | None:  # mutable-ok: overrides CustomLogger hook whose contract is a plain dict
         base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type)
         effective: Final = base_result if base_result is not None else kwargs
-        if call_type not in (CallTypes.completion, CallTypes.acompletion):
+        if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES:
             return base_result
         if not effective.get("stream"):
             return base_result
diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py
index 9a714e1724e..5de58a20242 100644
--- a/litellm/types/integrations/custom_logger.py
+++ b/litellm/types/integrations/custom_logger.py
@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from typing import Any, Final
 
 from pydantic import BaseModel, Field
@@ -29,6 +30,13 @@ def is_interception_internal_key(
     return any(key.startswith(prefix) for prefix in prefixes)
 
 
+CONVERTED_STREAM_KEYS: Final = frozenset(f"{prefix}_converted_stream" for prefix in INTERCEPTION_INTERNAL_PREFIXES)
+
+
+def converted_stream_requested(params: Mapping[str, object]) -> bool:
+    return any(bool(params.get(key)) for key in CONVERTED_STREAM_KEYS)
+
+
 class AgenticLoopSafetyError(ValueError):
     """
     Raised when an agentic-loop safety rail refuses a rerun.
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
index 37b9852d3e1..e1c4cc0922a 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
@@ -1984,6 +1984,58 @@ def _openai_text_payload(content: str) -> dict:
     return _openai_completion_payload({"role": "assistant", "content": content}, "stop")
 
 
+def _responses_retrieve_tool_definition() -> dict:
+    return {"type": "function", **_retrieve_tool_definition()["function"]}
+
+
+def _openai_responses_payload(output_item: dict) -> dict:
+    return {
+        "id": "resp_ccr",
+        "object": "response",
+        "created_at": 1700000000,
+        "status": "completed",
+        "model": "gpt-4o",
+        "output": [output_item],
+        "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
+        "parallel_tool_calls": True,
+        "tool_choice": "auto",
+        "tools": [],
+        "error": None,
+        "incomplete_details": None,
+        "instructions": None,
+        "metadata": {},
+        "temperature": 1.0,
+        "top_p": 1.0,
+        "text": {"format": {"type": "text"}},
+        "truncation": "disabled",
+    }
+
+
+def _openai_responses_retrieve_call_payload() -> dict:
+    return _openai_responses_payload(
+        {
+            "type": "function_call",
+            "id": "fc_ccr",
+            "call_id": "call_ccr",
+            "name": HEADROOM_RETRIEVE_TOOL_NAME,
+            "arguments": json.dumps({"hash": CCR_HASH}),
+            "status": "completed",
+        }
+    )
+
+
+def _openai_responses_text_payload(text: str) -> dict:
+    return _openai_responses_payload(
+        {
+            "type": "message",
+            "id": "msg_ccr",
+            "role": "assistant",
+            "status": "completed",
+            "content": [{"type": "output_text", "text": text, "annotations": []}],
+        }
+    )
+
+
 @pytest.mark.parametrize(
     "call_type, stream, tools, expect_conversion",
     [
@@ -1992,12 +2044,14 @@ def _openai_text_payload(content: str) -> dict:
         (CallTypes.acompletion, False, [_retrieve_tool_definition()], False),
         (CallTypes.acompletion, True, [{"type": "function", "function": {"name": "get_weather"}}], False),
         (CallTypes.acompletion, True, None, False),
-        (CallTypes.aresponses, True, [_retrieve_tool_definition()], False),
+        (CallTypes.aresponses, True, [_retrieve_tool_definition()], True),
+        (CallTypes.responses, True, [_responses_retrieve_tool_definition()], True),
+        (CallTypes.aresponses, False, [_retrieve_tool_definition()], False),
         (CallTypes.anthropic_messages, True, [_retrieve_tool_definition()], False),
     ],
 )
 @pytest.mark.asyncio
-async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions(
+async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions_and_responses(
     guardrail: HeadroomGuardrail,
     call_type: CallTypes,
     stream: bool,
@@ -2128,6 +2182,65 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end(
     assert not any(key.startswith("_headroom_interception") for key in followup_body)
 
 
+@pytest.mark.asyncio
+async def test_streaming_responses_resolves_ccr_retrieval_end_to_end(
+    guardrail: HeadroomGuardrail,
+    respx_mock: respx.MockRouter,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Regression test for LIT-6481: streaming /v1/responses must resolve the
+    retrieve tool call server-side exactly like streaming /chat/completions does,
+    instead of streaming a headroom_retrieve function_call to the client."""
+    original_content = "the full uncompressed document"
+    final_answer = "the document says hello"
+    guardrail._issued_hashes_by_call_id["ccr-call-id"] = (
+        frozenset({CCR_HASH}),
+        time.monotonic() + 999,
+    )
+
+    monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
+    monkeypatch.setattr(litellm, "callbacks", [guardrail])
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    upstream = respx_mock.post("https://api.openai.com/v1/responses").mock(
+        side_effect=[
+            httpx.Response(200, json=_openai_responses_retrieve_call_payload()),
+            httpx.Response(200, json=_openai_responses_text_payload(final_answer)),
+        ]
+    )
+
+    with patch.object(
+        guardrail.async_handler,
+        "get",
+        new_callable=AsyncMock,
+        return_value=_make_retrieve_response(original_content),
+    ) as mock_get:
+        response = await litellm.aresponses(
+            model="openai/gpt-4o",
+            input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}],
+            tools=[_responses_retrieve_tool_definition()],
+            stream=True,
+            litellm_call_id="ccr-call-id",
+        )
+        events = [event async for event in response]
+
+    streamed_text = "".join(
+        getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta"
+    )
+    assert streamed_text == final_answer
+    assert not any("function_call" in str(getattr(event, "type", "")) for event in events)
+    assert not any(
+        getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events
+    )
+    mock_get.assert_called_once()
+    assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0])
+
+    assert len(upstream.calls) == 2
+    followup_body = json.loads(upstream.calls[1].request.content)
+    assert not followup_body.get("stream")
+    assert original_content in json.dumps(followup_body["input"])
+    assert not any(key.startswith("_headroom_interception") for key in followup_body)
+
+
 # ---------------------------------------------------------------------------
 # LIT-5018: the turn the model is being asked to act on is never compressed.
 #

From 201ada99828bd7e62de465cd4ceed97002af99a6 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 29 Aug 2026 14:41:08 -0700
Subject: [PATCH 03/43] fix(headroom): check the responses converted-stream
 flag on the agentic kwargs dict

---
 litellm/llms/custom_httpx/llm_http_handler.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index 52a91a1d06d..ad25a44ab7f 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -2921,6 +2921,7 @@ class BaseLLMHTTPHandler:
             logging_obj=logging_obj,
         )
 
+        agentic_kwargs: Final = dict(litellm_params)
         final_response: Final = await self._call_agentic_completion_hooks(
             response=initial_response,
             model=model,
@@ -2930,12 +2931,12 @@ class BaseLLMHTTPHandler:
             logging_obj=logging_obj,
             stream=False,
             custom_llm_provider=custom_llm_provider,
-            kwargs=dict(litellm_params),
+            kwargs=agentic_kwargs,
             api_surface="responses",
         )
 
         result: Final = final_response if final_response is not None else initial_response
-        if converted_stream_requested(litellm_params) and not litellm_params.get("_agentic_loop_depth"):
+        if converted_stream_requested(agentic_kwargs) and not agentic_kwargs.get("_agentic_loop_depth"):
             return self._wrap_responses_response_as_fake_stream(
                 result=result,
                 model=model,

From 1695b7f7b1931bcc59c50975c8ccdbb460ad7c04 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 29 Aug 2026 15:02:20 -0700
Subject: [PATCH 04/43] fix(headroom): fake-stream converted sync /v1/responses
 calls too

The sync response_api_handler agentic branch returned the completed ResponsesAPIResponse for a request the Headroom guardrail had converted from streaming, so litellm.responses(stream=True) handed callers a non-iterable object. Wrap it in the same fake stream the async path uses
---
 litellm/llms/custom_httpx/llm_http_handler.py | 14 ++++-
 .../guardrail_hooks/test_headroom.py          | 52 +++++++++++++++++++
 2 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index ad25a44ab7f..2ec5377e217 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -2744,6 +2744,7 @@ class BaseLLMHTTPHandler:
         )
 
         if self._has_agentic_completion_hook(logging_obj):
+            agentic_kwargs: Final = dict(litellm_params)
             final_response: Final = run_async_function(
                 self._call_agentic_completion_hooks,
                 response=initial_response,
@@ -2754,10 +2755,19 @@ class BaseLLMHTTPHandler:
                 logging_obj=logging_obj,
                 stream=False,
                 custom_llm_provider=custom_llm_provider,
-                kwargs=dict(litellm_params),
+                kwargs=agentic_kwargs,
                 api_surface="responses",
             )
-            return final_response if final_response is not None else initial_response
+            result: Final = final_response if final_response is not None else initial_response
+            if converted_stream_requested(agentic_kwargs) and not agentic_kwargs.get("_agentic_loop_depth"):
+                return self._wrap_responses_response_as_fake_stream(
+                    result=result,
+                    model=model,
+                    responses_api_provider_config=responses_api_provider_config,
+                    logging_obj=logging_obj,
+                    custom_llm_provider=custom_llm_provider,
+                )
+            return result
 
         return initial_response
 
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
index e1c4cc0922a..c297f593617 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
@@ -2241,6 +2241,58 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end(
     assert not any(key.startswith("_headroom_interception") for key in followup_body)
 
 
+def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end(
+    guardrail: HeadroomGuardrail,
+    respx_mock: respx.MockRouter,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """The synchronous responses() path converts the stream the same way, so it
+    must hand back a stream iterator with the resolved answer rather than the
+    completed response object."""
+    original_content = "the full uncompressed document"
+    final_answer = "the document says hello"
+    guardrail._issued_hashes_by_call_id["ccr-call-id"] = (
+        frozenset({CCR_HASH}),
+        time.monotonic() + 999,
+    )
+
+    monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setattr(litellm, "callbacks", [guardrail])
+    upstream = respx_mock.post("https://api.openai.com/v1/responses").mock(
+        side_effect=[
+            httpx.Response(200, json=_openai_responses_retrieve_call_payload()),
+            httpx.Response(200, json=_openai_responses_text_payload(final_answer)),
+        ]
+    )
+
+    with patch.object(
+        guardrail.async_handler,
+        "get",
+        new_callable=AsyncMock,
+        return_value=_make_retrieve_response(original_content),
+    ) as mock_get:
+        response = litellm.responses(
+            model="openai/gpt-4o",
+            input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}],
+            tools=[_responses_retrieve_tool_definition()],
+            stream=True,
+            litellm_call_id="ccr-call-id",
+        )
+        events = list(response)
+
+    streamed_text = "".join(
+        getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta"
+    )
+    assert streamed_text == final_answer
+    assert not any(
+        getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events
+    )
+    mock_get.assert_called_once()
+    assert len(upstream.calls) == 2
+    assert not json.loads(upstream.calls[1].request.content).get("stream")
+
+
 # ---------------------------------------------------------------------------
 # LIT-5018: the turn the model is being asked to act on is never compressed.
 #

From 6bd3699d4338193b23f8e937061cbab7c4af2e49 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 29 Aug 2026 16:39:14 -0700
Subject: [PATCH 05/43] fix(responses): keep guardrailed input items and bridge
 stream usage intact

- _write_back_structured_messages now patches only the rewritten rows back
  into the original input items, so reasoning items (encrypted_content),
  function_call ids, and web_search_call items survive a guardrail rewrite
  verbatim; rewrites that cannot be row-mapped fall back to the previous
  full conversion
- the responses bridge stream snapshot restores usage hidden in
  _hidden_params when stream_options is unset, so converted fake streams
  report real input_tokens instead of 0
---
 litellm/llms/custom_httpx/llm_http_handler.py |   4 +-
 .../guardrail_translation/handler.py          | 135 +++++++++++++++++-
 .../streaming_iterator.py                     |   6 +
 ...test_openai_responses_guardrail_handler.py | 129 ++++++++++++++++-
 .../guardrail_hooks/test_headroom.py          |   2 +-
 .../test_streaming_iterator_transformation.py |  15 ++
 6 files changed, 279 insertions(+), 12 deletions(-)

diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index 2ec5377e217..fcb988705eb 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -2744,7 +2744,7 @@ class BaseLLMHTTPHandler:
         )
 
         if self._has_agentic_completion_hook(logging_obj):
-            agentic_kwargs: Final = dict(litellm_params)
+            agentic_kwargs: Final = dict(litellm_params)  # mutable-ok: agentic hooks mutate kwargs in place
             final_response: Final = run_async_function(
                 self._call_agentic_completion_hooks,
                 response=initial_response,
@@ -2931,7 +2931,7 @@ class BaseLLMHTTPHandler:
             logging_obj=logging_obj,
         )
 
-        agentic_kwargs: Final = dict(litellm_params)
+        agentic_kwargs: Final = dict(litellm_params)  # mutable-ok: agentic hooks mutate kwargs in place
         final_response: Final = await self._call_agentic_completion_hooks(
             response=initial_response,
             model=model,
diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py
index 4320edb8414..8c014feeded 100644
--- a/litellm/llms/openai/responses/guardrail_translation/handler.py
+++ b/litellm/llms/openai/responses/guardrail_translation/handler.py
@@ -28,7 +28,8 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
     - text: str
 """
 
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Union, cast
 
 from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
@@ -50,6 +51,7 @@ from litellm.types.llms.openai import (
     ChatCompletionToolCallChunk,
     ChatCompletionToolParam,
     OpenAIMcpServerTool,
+    ResponsesAPIOptionalRequestParams,
     ResponsesAPIStreamEvents,
 )
 from litellm.types.responses.main import (
@@ -81,6 +83,119 @@ class ResponsesStreamChunk(TypedDict, total=False):
     text: ReadOnly[str]
 
 
+_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
+    {"function_call_output": "output", "message": "content"}
+)
+
+_EMPTY_RESPONSES_REQUEST: Final[ResponsesAPIOptionalRequestParams] = {}
+
+
+def _item_rewrite_field(item: Mapping[str, object]) -> str | None:
+    item_type: Final = item.get("type")
+    if item_type is None:
+        return "content" if "content" in item else None
+    if not isinstance(item_type, str):
+        return None
+    return _PATCHABLE_ITEM_FIELDS.get(item_type)
+
+
+def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None:
+    field: Final = _item_rewrite_field(item)
+    if field is None or not isinstance(rewritten, Mapping):
+        return None
+    rewritten_content: Final = rewritten.get("content")
+    if isinstance(item.get(field), str) and isinstance(rewritten_content, str):
+        return {**item, field: rewritten_content}  # mutable-ok: request input items must stay JSON-plain dicts
+    rewritten_row: Final = cast("AllMessageValues", rewritten)  # cast-ok: guardrails hand back chat-shaped rows
+    converted_items, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(
+        [rewritten_row]  # mutable-ok: converter signature takes a list
+    )
+    if len(converted_items) != 1 or not isinstance(converted_items[0], Mapping):
+        return None
+    first_converted: Final = cast("Mapping[str, object]", converted_items[0])  # cast-ok: isinstance-checked above
+    converted_value: Final = first_converted.get(field)
+    if converted_value is None:
+        return None
+    return {**item, field: converted_value}  # mutable-ok: request input items must stay JSON-plain dicts
+
+
+def _input_item_provenance(
+    raw_input: Sequence[object],
+    expected_messages: Sequence[object],
+) -> tuple[Mapping[int, int], frozenset[int]] | None:
+    if not all(isinstance(item, Mapping) for item in raw_input):
+        return None
+    prefixes: Final = tuple(
+        LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
+            input=cast("ResponseInputParam", raw_input[:count]),  # cast-ok: items checked as Mappings above
+            responses_api_request=_EMPTY_RESPONSES_REQUEST,
+        )
+        for count in range(len(raw_input) + 1)
+    )
+    if tuple(prefixes[-1]) != tuple(expected_messages):
+        return None
+    item_for_message: Final = MappingProxyType(
+        {
+            message_index: item_index
+            for item_index in range(len(raw_input))
+            for message_index in range(len(prefixes[item_index]), len(prefixes[item_index + 1]))
+        }
+    )
+    tainted: Final = frozenset(
+        message_index
+        for item_index in range(len(raw_input))
+        for message_index in range(len(prefixes[item_index]))
+        if prefixes[item_index + 1][message_index] != prefixes[item_index][message_index]
+    )
+    return item_for_message, tainted
+
+
+def _patch_rewritten_rows_into_input(
+    data: dict,
+    original_messages: Sequence[object],
+    structured_messages: Sequence[object],
+) -> bool:
+    raw_input: Final = data.get("input")
+    if not isinstance(raw_input, list) or len(original_messages) != len(structured_messages):
+        return False
+    offset: Final = 1 if data.get("instructions") else 0
+    provenance: Final = _input_item_provenance(raw_input, tuple(original_messages)[offset:])
+    if provenance is None:
+        return False
+    item_for_message, tainted = provenance
+    changed: Final = tuple(
+        (index, rewritten)
+        for index, (original, rewritten) in enumerate(zip(original_messages, structured_messages))
+        if original != rewritten
+    )
+    instruction_rewrites: Final = tuple(rewritten for index, rewritten in changed if index < offset)
+    rewritten_instructions: Final = (
+        instruction_rewrites[0].get("content")
+        if instruction_rewrites and isinstance(instruction_rewrites[0], Mapping)
+        else None
+    )
+    if instruction_rewrites and not isinstance(rewritten_instructions, str):
+        return False
+    body_changes: Final = tuple((index - offset, rewritten) for index, rewritten in changed if index >= offset)
+    if any(message_index in tainted or message_index not in item_for_message for message_index, _ in body_changes):
+        return False
+    replacements: Final = MappingProxyType(
+        {
+            item_for_message[message_index]: _rewritten_input_item(
+                cast("Mapping[str, object]", raw_input[item_for_message[message_index]]),  # cast-ok: checked Mappings
+                rewritten,
+            )
+            for message_index, rewritten in body_changes
+        }
+    )
+    if len(replacements) != len(body_changes) or any(item is None for item in replacements.values()):
+        return False
+    data["input"] = [replacements.get(index, item) for index, item in enumerate(raw_input)]  # mutable-ok: JSON body
+    if isinstance(rewritten_instructions, str):
+        data["instructions"] = rewritten_instructions
+    return True
+
+
 class OpenAIResponsesHandler(BaseTranslation):
     """
     Handler for processing OpenAI Responses API with guardrails.
@@ -155,9 +270,9 @@ class OpenAIResponsesHandler(BaseTranslation):
                 guardrailed_structured_messages is not None
                 and guardrailed_structured_messages is not structured_messages
             ):
-                self._write_back_structured_messages(data, guardrailed_structured_messages)
+                self._write_back_structured_messages(data, structured_messages or (), guardrailed_structured_messages)
             else:
-                guardrailed_texts = guardrailed_inputs.get("texts", [])
+                guardrailed_texts = guardrailed_inputs.get("texts") or ()
                 data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data
             self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools"))
             verbose_proxy_logger.debug("OpenAI Responses API: Processed string input")
@@ -217,12 +332,12 @@ class OpenAIResponsesHandler(BaseTranslation):
                 guardrailed_structured_messages is not None
                 and guardrailed_structured_messages is not structured_messages
             ):
-                self._write_back_structured_messages(data, guardrailed_structured_messages)
+                self._write_back_structured_messages(data, structured_messages or (), guardrailed_structured_messages)
             else:
                 # Step 3: Map guardrail responses back to original input structure
                 await self._apply_guardrail_responses_to_input(
                     messages=input_data,
-                    responses=guardrailed_inputs.get("texts", []),
+                    responses=guardrailed_inputs.get("texts", []),  # mutable-ok: callee signature takes a list
                     task_mappings=task_mappings,
                 )
 
@@ -231,10 +346,16 @@ class OpenAIResponsesHandler(BaseTranslation):
         return data
 
     @staticmethod
-    def _write_back_structured_messages(data: dict, structured_messages: Sequence[AllMessageValues]) -> None:
+    def _write_back_structured_messages(
+        data: dict,
+        original_messages: Sequence[object],
+        structured_messages: Sequence[AllMessageValues],
+    ) -> None:
+        if _patch_rewritten_rows_into_input(data, original_messages, structured_messages):
+            return
         input_items, instructions = (
             LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(
-                list(structured_messages)
+                list(structured_messages)  # mutable-ok: converter signature takes a list
             )
         )
         data["input"] = input_items
diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py
index 8b1eeb30306..f32599e8098 100644
--- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py
+++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py
@@ -557,6 +557,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
         hidden_params: Final = getattr(chunk, "_hidden_params", None)
         if hidden_params is not None:
             chunk_dict["_hidden_params"] = dict(hidden_params) if isinstance(hidden_params, dict) else hidden_params
+        if (
+            chunk_dict.get("usage") is None
+            and isinstance(hidden_params, dict)
+            and hidden_params.get("usage") is not None
+        ):
+            chunk_dict["usage"] = hidden_params["usage"]
         return chunk_dict
 
     def create_reasoning_summary_text_done_event(
diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
index 63f8babf760..87449a4b1e8 100644
--- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
+++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
@@ -1253,6 +1253,43 @@ class StructuredRewriteGuardrail(CustomGuardrail):
         return {**inputs, "structured_messages": rewritten}
 
 
+class ToolOutputRewriteGuardrail(CustomGuardrail):
+    """Guardrail that compresses the first tool-result row, the way Headroom does."""
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        messages = list(inputs.get("structured_messages") or [])
+        first_tool = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "tool")
+        rewritten = [
+            {**m, "content": COMPRESSED_MARKER} if i == first_tool else m for i, m in enumerate(messages)
+        ]
+        return {**inputs, "structured_messages": rewritten}
+
+
+class DroppingRewriteGuardrail(CustomGuardrail):
+    """Guardrail that rewrites the first user row and drops the last row, so the
+    rewrite can only land through the full-conversion fallback."""
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        messages = list(inputs.get("structured_messages") or [])
+        first_user = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "user")
+        rewritten = [
+            {**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages)
+        ]
+        return {**inputs, "structured_messages": rewritten[:-1]}
+
+
 def _texts(item: dict) -> list[str]:
     content = item.get("content")
     if isinstance(content, str):
@@ -1296,7 +1333,93 @@ class TestStructuredMessagesWriteBack:
         assert "instructions" not in result
 
     @pytest.mark.asyncio
-    async def test_developer_item_survives_write_back_as_input_text(self):
+    async def test_developer_item_preserved_verbatim_by_row_patch(self):
+        handler = OpenAIResponsesHandler()
+        developer_item = {"role": "developer", "content": "Always answer in French."}
+        data = {
+            "model": "gpt-5.6",
+            "input": [
+                developer_item,
+                {"role": "user", "content": "memo " * 400},
+                {"role": "user", "content": "What is the codename?"},
+            ],
+        }
+
+        result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
+
+        assert result["input"][0] is developer_item
+        assert developer_item["content"] == "Always answer in French."
+        assert _texts(result["input"][1]) == [COMPRESSED_MARKER]
+        assert _texts(result["input"][2]) == ["What is the codename?"]
+
+    @pytest.mark.asyncio
+    async def test_reasoning_and_function_call_items_survive_tool_output_compression(self):
+        handler = OpenAIResponsesHandler()
+        reasoning_item = {
+            "id": "rs_123",
+            "type": "reasoning",
+            "summary": [],
+            "encrypted_content": "gAAAAA-signed-reasoning",
+        }
+        function_call_item = {
+            "id": "fc_123",
+            "type": "function_call",
+            "call_id": "call_abc",
+            "name": "read_document",
+            "arguments": '{"path": "memo.txt"}',
+            "status": "completed",
+        }
+        data = {
+            "model": "gpt-5.6",
+            "instructions": "Answer from the memo only.",
+            "input": [
+                reasoning_item,
+                function_call_item,
+                {"type": "function_call_output", "call_id": "call_abc", "output": "memo " * 400},
+                {"role": "user", "content": "What is the codename?"},
+            ],
+        }
+
+        result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail())
+
+        assert result["instructions"] == "Answer from the memo only."
+        assert result["input"][0] is reasoning_item
+        assert reasoning_item["encrypted_content"] == "gAAAAA-signed-reasoning"
+        assert result["input"][1] is function_call_item
+        assert function_call_item["id"] == "fc_123"
+        assert result["input"][2] == {
+            "type": "function_call_output",
+            "call_id": "call_abc",
+            "output": COMPRESSED_MARKER,
+        }
+        assert result["input"][3] == {"role": "user", "content": "What is the codename?"}
+
+    @pytest.mark.asyncio
+    async def test_web_search_call_item_preserved_verbatim(self):
+        handler = OpenAIResponsesHandler()
+        web_search_item = {
+            "id": "ws_123",
+            "type": "web_search_call",
+            "status": "completed",
+            "action": {"type": "search", "query": "codename memo"},
+        }
+        data = {
+            "model": "gpt-5.6",
+            "input": [
+                web_search_item,
+                {"role": "user", "content": "memo " * 400},
+                {"role": "user", "content": "What is the codename?"},
+            ],
+        }
+
+        result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
+
+        assert result["input"][0] is web_search_item
+        assert _texts(result["input"][1]) == [COMPRESSED_MARKER]
+        assert _texts(result["input"][2]) == ["What is the codename?"]
+
+    @pytest.mark.asyncio
+    async def test_row_count_change_falls_back_to_full_conversion(self):
         handler = OpenAIResponsesHandler()
         data = {
             "model": "gpt-5.6",
@@ -1307,10 +1430,12 @@ class TestStructuredMessagesWriteBack:
             ],
         }
 
-        result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
+        result = await handler.process_input_messages(data, DroppingRewriteGuardrail())
 
+        assert len(result["input"]) == 2
         developer = next(item for item in result["input"] if item.get("role") == "developer")
         assert developer["content"] == [{"type": "input_text", "text": "Always answer in French."}]
+        assert _texts(next(item for item in result["input"] if item.get("role") == "user")) == [COMPRESSED_MARKER]
 
     @pytest.mark.asyncio
     async def test_same_inputs_object_back_keeps_the_text_mapping(self):
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
index c297f593617..22c4f46e3fb 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
@@ -979,7 +979,7 @@ async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstre
         result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail)
 
     assert result["instructions"] == ORIGINAL_MESSAGES[0]["content"]
-    assert [item["content"][0]["text"] for item in result["input"]] == [
+    assert [item["content"] for item in result["input"]] == [
         COMPRESSED_MESSAGES_WITH_HASH[0]["content"],
         ORIGINAL_MESSAGES[2]["content"],
         ORIGINAL_MESSAGES[3]["content"],
diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
index 823f656ddc5..3ffd44b2fbe 100644
--- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
+++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
@@ -24,6 +24,7 @@ from litellm.types.utils import (
     ModelResponse,
     ModelResponseStream,
     StreamingChoices,
+    Usage,
 )
 
 CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256"
@@ -523,3 +524,17 @@ async def test_streaming_response_id_falls_back_when_upstream_yields_nothing():
     assert response_ids
     assert len(set(response_ids)) == 1
     assert response_ids[0].startswith("resp_")
+
+
+def test_completed_event_restores_usage_hidden_by_stream_options_none():
+    final_chunk = _chunk("", finish_reason="stop")
+    final_chunk._hidden_params = {"usage": Usage(prompt_tokens=117, completion_tokens=5, total_tokens=122)}
+    iterator = _build_iterator([_chunk("the document says hello"), final_chunk])
+
+    events = list(iterator)
+
+    completed = next(
+        event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
+    )
+    assert completed.response.usage.input_tokens == 117
+    assert completed.response.usage.output_tokens == 5

From 98a9a7e525590d9af15d26ea3db0257d14b42d43 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 29 Aug 2026 21:12:41 -0700
Subject: [PATCH 06/43] fix(streaming): carry hidden usage on the async
 fake-stream final chunk

The sync __next__ exhaustion branch stores calculate_total_usage() in the
final chunk's _hidden_params when stream_options is None, but the async
__anext__ sibling branch never did. Converted (fake) streams, like the ones
the Headroom guardrail produces by flipping streaming /v1/responses calls to
non-streaming, are consumed async, so their real usage never reached the
completion-to-responses bridge and it token-counted from scratch, reporting
input_tokens=0. Mirror the sync branch's hidden-usage block into the async
exhaustion branch and add a regression test that async-iterates a
CustomStreamWrapper over a MockResponseIterator and asserts the final chunk
carries the mock response's usage.
---
 .../litellm_core_utils/streaming_handler.py   |  3 ++
 .../test_streaming_handler.py                 | 37 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index 1e0b778d244..d9641a6306b 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -2336,6 +2336,9 @@ class CustomStreamWrapper:
         else:
             self.sent_last_chunk = True
             processed_chunk: Final = self.finish_reason_handler()
+            if self.stream_options is None:
+                usage: Final = calculate_total_usage(chunks=self.chunks)
+                processed_chunk._hidden_params["usage"] = usage  # pyright: ignore[reportPrivateUsage]  # sync parity
             # see sync __next__'s sibling branch: deliberately do NOT restore
             # here - this chunk is still this call's own data, and restoring
             # before returning it would corrupt the caller's own log
diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
index 7f54fbfb4c2..7caf401ce38 100644
--- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
+++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
@@ -4692,3 +4692,40 @@ async def test_async_stream_assembled_response_keeps_vertex_traffic_type(logging
     assembled = litellm.stream_chunk_builder(chunks=received, messages=[{"role": "user", "content": "hi"}])
     assert assembled is not None
     assert assembled._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND_FLEX"
+
+
+@pytest.mark.asyncio
+async def test_async_fake_stream_final_chunk_carries_hidden_usage(logging_obj: Logging):
+    from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
+    from litellm.types.utils import ModelResponse
+
+    model_response = ModelResponse(
+        id="chatcmpl-fake-stream",
+        model="my-random-model",
+        choices=[
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "hello world"},
+                "finish_reason": "stop",
+            }
+        ],
+    )
+    model_response.usage = Usage(prompt_tokens=1234, completion_tokens=7, total_tokens=1241)
+
+    wrapper = CustomStreamWrapper(
+        completion_stream=MockResponseIterator(model_response=model_response),
+        model="my-random-model",
+        custom_llm_provider="anthropic",
+        logging_obj=logging_obj,
+    )
+
+    final_chunk = None
+    async for chunk in wrapper:
+        final_chunk = chunk
+
+    assert final_chunk is not None
+    hidden_usage = final_chunk._hidden_params.get("usage")
+    assert hidden_usage is not None
+    assert hidden_usage.prompt_tokens == 1234
+    assert hidden_usage.completion_tokens == 7
+    assert hidden_usage.total_tokens == 1241

From 574010a2cec368ba2fcffbc6fd8717161b33a3b1 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 29 Aug 2026 22:04:26 -0700
Subject: [PATCH 07/43] fix(responses): make guardrail input provenance O(n)
 and guard non-list structured_messages

_input_item_provenance converted every input prefix, so an n-item request paid
for n+1 full conversions. It now converts each item once, glues consecutive
function_call items (plus their trailing-assistant context) into units so the
transform's tool_call merging is reproduced inside the unit conversion, and
verifies the unit concatenation against one full conversion, bailing to the
full-conversion fallback on any mismatch. Messages from multi-item units are
tainted, which keeps parallel tool calls patchable exactly like the old prefix
pass while unpredicted merges fall back safely.

A guardrail handing back a non-list structured_messages payload (the
HiddenLayer v2 evaluation dict) previously fell through the length-mismatch
fallback and 500ed converting the dict's keys as messages. The write-back is
now skipped for non-list payloads, restoring the previous no-write-back
behavior on the Responses surface.

Also refreshes the compresr texts-mirror docstring, which still claimed the
Responses translation cannot round-trip structured_messages.
---
 .../guardrail_translation/handler.py          |  75 ++++-
 .../guardrail_hooks/compresr/compresr.py      |   7 +-
 ...test_openai_responses_guardrail_handler.py | 296 ++++++++++++++++++
 3 files changed, 365 insertions(+), 13 deletions(-)

diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py
index 8c014feeded..bd58fb4cdbc 100644
--- a/litellm/llms/openai/responses/guardrail_translation/handler.py
+++ b/litellm/llms/openai/responses/guardrail_translation/handler.py
@@ -29,6 +29,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
 """
 
 from collections.abc import Mapping, Sequence
+from itertools import accumulate
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Union, cast
 
@@ -119,33 +120,85 @@ def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapp
     return {**item, field: converted_value}  # mutable-ok: request input items must stay JSON-plain dicts
 
 
+def _is_function_call_item(item: object) -> bool:
+    return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call")
+
+
+def _last_message_role(messages: Sequence[object]) -> str | None:
+    if not messages:
+        return None
+    last: Final = messages[-1]
+    role: Final = last.get("role") if isinstance(last, Mapping) else getattr(last, "role", None)
+    return role if isinstance(role, str) else None
+
+
+def _provenance_unit_bounds(
+    raw_input: Sequence[object],
+    solo_conversions: Sequence[Sequence[object]],
+) -> tuple[tuple[int, int], ...]:
+    trailing_roles: Final = tuple(
+        accumulate(
+            (_last_message_role(messages) for messages in solo_conversions),
+            lambda previous, current: current if current is not None else previous,
+        )
+    )
+    start_indexes: Final = tuple(
+        index
+        for index in range(len(raw_input))
+        if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant")
+    )
+    return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input))))
+
+
 def _input_item_provenance(
     raw_input: Sequence[object],
     expected_messages: Sequence[object],
 ) -> tuple[Mapping[int, int], frozenset[int]] | None:
     if not all(isinstance(item, Mapping) for item in raw_input):
         return None
-    prefixes: Final = tuple(
+    solo_conversions: Final = tuple(
         LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
-            input=cast("ResponseInputParam", raw_input[:count]),  # cast-ok: items checked as Mappings above
+            input=cast("ResponseInputParam", [item]),  # cast-ok: items checked as Mappings above
             responses_api_request=_EMPTY_RESPONSES_REQUEST,
         )
-        for count in range(len(raw_input) + 1)
+        for item in raw_input
     )
-    if tuple(prefixes[-1]) != tuple(expected_messages):
+    full_conversion: Final = tuple(
+        LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
+            input=cast("ResponseInputParam", list(raw_input)),  # cast-ok: items checked as Mappings above
+            responses_api_request=_EMPTY_RESPONSES_REQUEST,
+        )
+    )
+    if full_conversion != tuple(expected_messages):
         return None
+    units: Final = _provenance_unit_bounds(raw_input, solo_conversions)
+    unit_messages: Final = tuple(
+        tuple(solo_conversions[start])
+        if end - start == 1
+        else tuple(
+            LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
+                input=cast("ResponseInputParam", list(raw_input[start:end])),  # cast-ok: checked as Mappings above
+                responses_api_request=_EMPTY_RESPONSES_REQUEST,
+            )
+        )
+        for start, end in units
+    )
+    if tuple(message for messages in unit_messages for message in messages) != full_conversion:
+        return None
+    boundaries: Final = tuple(accumulate((len(messages) for messages in unit_messages), initial=0))
     item_for_message: Final = MappingProxyType(
         {
-            message_index: item_index
-            for item_index in range(len(raw_input))
-            for message_index in range(len(prefixes[item_index]), len(prefixes[item_index + 1]))
+            message_index: start
+            for unit_index, (start, end) in enumerate(units)
+            if end - start == 1
+            for message_index in range(boundaries[unit_index], boundaries[unit_index + 1])
         }
     )
     tainted: Final = frozenset(
         message_index
-        for item_index in range(len(raw_input))
-        for message_index in range(len(prefixes[item_index]))
-        if prefixes[item_index + 1][message_index] != prefixes[item_index][message_index]
+        for unit_index, (start, end) in enumerate(units)
+        if end - start > 1
+        for message_index in range(boundaries[unit_index], boundaries[unit_index + 1])
     )
     return item_for_message, tainted
 
@@ -351,6 +404,8 @@ class OpenAIResponsesHandler(BaseTranslation):
         original_messages: Sequence[object],
         structured_messages: Sequence[AllMessageValues],
     ) -> None:
+        if not isinstance(structured_messages, list):
+            return
         if _patch_rewritten_rows_into_input(data, original_messages, structured_messages):
             return
         input_items, instructions = (
diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py
index 5c14d03f50e..84b328a9c60 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py
@@ -916,9 +916,10 @@ class CompresrGuardrail(CustomGuardrail):
     def _mirror_texts_channel(input_texts: object, applied: _CompressionResult) -> list[object] | None:
         """Compressed content mirrored into the Responses `texts` channel.
 
-        The chat/Anthropic handlers round-trip ``structured_messages``; the
-        Responses translation cannot rebuild its input from chat messages and
-        instead writes back through ``texts``. This matches by value, so a
+        The chat/Anthropic/Responses handlers round-trip
+        ``structured_messages``; translations without that round-trip write
+        back through ``texts``, so the compressed content is mirrored there
+        too. This matches by value, so a
         replacement is applied only when it is unambiguous: one compression per
         text, and every occurrence in ``texts`` accounted for by a compressed
         target. Anything else is left uncompressed rather than risk a wrong or
diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
index 87449a4b1e8..86c023dd7a4 100644
--- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
+++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
@@ -1450,3 +1450,299 @@ class TestStructuredMessagesWriteBack:
 
         assert result["input"] is original_input
         assert [_texts(item) for item in result["input"]] == [["Hello [GUARDRAILED]"], ["Again [GUARDRAILED]"]]
+
+
+class AllToolOutputsRewriteGuardrail(CustomGuardrail):
+    """Guardrail that compresses every tool-result row."""
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        messages = list(inputs.get("structured_messages") or [])
+        rewritten = [
+            {**m, "content": COMPRESSED_MARKER} if isinstance(m, dict) and m.get("role") == "tool" else m
+            for m in messages
+        ]
+        return {**inputs, "structured_messages": rewritten}
+
+
+class AssistantRewriteGuardrail(CustomGuardrail):
+    """Guardrail that rewrites the first assistant row's content."""
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        messages = list(inputs.get("structured_messages") or [])
+        first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "assistant")
+        rewritten = [{**m, "content": COMPRESSED_MARKER} if i == first else m for i, m in enumerate(messages)]
+        return {**inputs, "structured_messages": rewritten}
+
+
+class DictStructuredMessagesGuardrail(CustomGuardrail):
+    """Guardrail that hands back a raw evaluation dict instead of a message list,
+    the way HiddenLayer v2 does."""
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        return {**inputs, "structured_messages": {"evaluation": "allowed", "messages": []}}
+
+
+def _parallel_tool_call_input() -> list:
+    return [
+        {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"},
+        {"id": "fc_2", "type": "function_call", "call_id": "call_2", "name": "read_b", "arguments": "{}"},
+        {"type": "function_call_output", "call_id": "call_1", "output": "memo " * 400},
+        {"type": "function_call_output", "call_id": "call_2", "output": "note " * 400},
+        {"role": "user", "content": "What is the codename?"},
+    ]
+
+
+class TestProvenancePatching:
+    """The O(n) provenance pass must keep patching rewritten rows in place for the
+    shapes real agent loops produce, and fall back safely everywhere else."""
+
+    @pytest.mark.asyncio
+    async def test_parallel_tool_call_outputs_both_patched(self):
+        handler = OpenAIResponsesHandler()
+        raw_input = _parallel_tool_call_input()
+        fc_1, fc_2 = raw_input[0], raw_input[1]
+        data = {"model": "gpt-5.6", "input": raw_input}
+
+        result = await handler.process_input_messages(data, AllToolOutputsRewriteGuardrail())
+
+        assert result["input"][0] is fc_1
+        assert result["input"][1] is fc_2
+        assert result["input"][2] == {"type": "function_call_output", "call_id": "call_1", "output": COMPRESSED_MARKER}
+        assert result["input"][3] == {"type": "function_call_output", "call_id": "call_2", "output": COMPRESSED_MARKER}
+        assert result["input"][4] == {"role": "user", "content": "What is the codename?"}
+
+    @pytest.mark.asyncio
+    async def test_assistant_turn_with_tool_call_keeps_items_verbatim(self):
+        handler = OpenAIResponsesHandler()
+        assistant_item = {"role": "assistant", "content": "Let me read the memo."}
+        function_call_item = {
+            "id": "fc_9",
+            "type": "function_call",
+            "call_id": "call_9",
+            "name": "read_document",
+            "arguments": '{"path": "memo.txt"}',
+        }
+        data = {
+            "model": "gpt-5.6",
+            "input": [
+                assistant_item,
+                function_call_item,
+                {"type": "function_call_output", "call_id": "call_9", "output": "memo " * 400},
+                {"role": "user", "content": "What is the codename?"},
+            ],
+        }
+
+        result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail())
+
+        assert result["input"][0] is assistant_item
+        assert result["input"][1] is function_call_item
+        assert result["input"][2] == {"type": "function_call_output", "call_id": "call_9", "output": COMPRESSED_MARKER}
+
+    @pytest.mark.asyncio
+    async def test_rewrite_of_merged_tool_call_message_falls_back(self):
+        handler = OpenAIResponsesHandler()
+        raw_input = _parallel_tool_call_input()
+        data = {"model": "gpt-5.6", "input": raw_input}
+
+        result = await handler.process_input_messages(data, AssistantRewriteGuardrail())
+
+        assert not any(item is original for item in result["input"] for original in raw_input)
+        assistant_items = [item for item in result["input"] if item.get("role") == "assistant"]
+        assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]]
+
+    @pytest.mark.asyncio
+    async def test_rewrite_of_lone_function_call_message_falls_back(self):
+        handler = OpenAIResponsesHandler()
+        data = {
+            "model": "gpt-5.6",
+            "input": [
+                {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"},
+                {"type": "function_call_output", "call_id": "call_1", "output": "memo memo"},
+                {"role": "user", "content": "What is the codename?"},
+            ],
+        }
+
+        raw_input = data["input"]
+        result = await handler.process_input_messages(data, AssistantRewriteGuardrail())
+
+        assert not any(item is original for item in result["input"] for original in raw_input)
+        assistant_items = [item for item in result["input"] if item.get("role") == "assistant"]
+        assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]]
+
+    def test_provenance_bails_on_non_mapping_item(self):
+        from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance
+
+        assert _input_item_provenance(["not a mapping"], []) is None
+
+    def test_provenance_bails_when_expected_messages_disagree(self):
+        from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance
+
+        assert _input_item_provenance([{"role": "user", "content": "hi"}], [{"role": "user", "content": "bye"}]) is None
+
+    def test_provenance_bails_on_unpredicted_merge(self):
+        from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance
+        from litellm.responses.litellm_completion_transformation.transformation import (
+            LiteLLMCompletionResponsesConfig,
+        )
+
+        raw_input = [
+            {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"},
+            {"role": "assistant", "content": "Reading the memo now."},
+        ]
+        expected = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
+            input=raw_input, responses_api_request={}
+        )
+        assert len(expected) == 1
+        assert _input_item_provenance(raw_input, expected) is None
+
+    def test_provenance_maps_and_taints_parallel_tool_calls(self):
+        from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance
+        from litellm.responses.litellm_completion_transformation.transformation import (
+            LiteLLMCompletionResponsesConfig,
+        )
+
+        raw_input = _parallel_tool_call_input()
+        expected = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
+            input=raw_input, responses_api_request={}
+        )
+        provenance = _input_item_provenance(raw_input, expected)
+        assert provenance is not None
+        item_for_message, tainted = provenance
+        assert tainted == {0}
+        assert dict(item_for_message) == {1: 2, 2: 3, 3: 4}
+
+
+class TestDictStructuredMessagesGuard:
+    """A guardrail handing back a non-list structured_messages payload must not
+    blow up the request; the write-back is skipped instead."""
+
+    @pytest.mark.asyncio
+    async def test_list_input_survives_dict_structured_messages(self):
+        handler = OpenAIResponsesHandler()
+        original_input = [{"role": "user", "content": "Hello"}]
+        data = {"model": "gpt-5.6", "input": original_input}
+
+        result = await handler.process_input_messages(data, DictStructuredMessagesGuardrail())
+
+        assert result["input"] is original_input
+        assert result["input"] == [{"role": "user", "content": "Hello"}]
+
+    @pytest.mark.asyncio
+    async def test_string_input_survives_dict_structured_messages(self):
+        handler = OpenAIResponsesHandler()
+        data = {"model": "gpt-5.6", "input": "Hello there"}
+
+        result = await handler.process_input_messages(data, DictStructuredMessagesGuardrail())
+
+        assert result["input"] == "Hello there"
+
+
+class SystemRewriteGuardrail(CustomGuardrail):
+    """Guardrail that rewrites the system row, the way prompt-hardening guardrails do."""
+
+    def __init__(self, rewritten_content: Any = COMPRESSED_MARKER):
+        super().__init__()
+        self.rewritten_content = rewritten_content
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        messages = list(inputs.get("structured_messages") or [])
+        first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "system")
+        rewritten = [
+            {**m, "content": self.rewritten_content} if i == first else m for i, m in enumerate(messages)
+        ]
+        return {**inputs, "structured_messages": rewritten}
+
+
+class TestPatchEdgeBranches:
+    @pytest.mark.asyncio
+    async def test_multimodal_user_item_rewritten_through_conversion(self):
+        handler = OpenAIResponsesHandler()
+        data = {
+            "model": "gpt-5.6",
+            "input": [
+                {"role": "user", "content": [{"type": "input_text", "text": "memo " * 400}]},
+                {"role": "user", "content": "What is the codename?"},
+            ],
+        }
+
+        result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
+
+        assert _texts(result["input"][0]) == [COMPRESSED_MARKER]
+        assert result["input"][1] == {"role": "user", "content": "What is the codename?"}
+
+    @pytest.mark.asyncio
+    async def test_instructions_rewrite_lands_in_instructions_field(self):
+        handler = OpenAIResponsesHandler()
+        user_item = {"role": "user", "content": "What is the codename?"}
+        data = {
+            "model": "gpt-5.6",
+            "instructions": "Answer from the memo only.",
+            "input": [user_item],
+        }
+
+        result = await handler.process_input_messages(data, SystemRewriteGuardrail())
+
+        assert result["instructions"] == COMPRESSED_MARKER
+        assert result["input"][0] is user_item
+
+    @pytest.mark.asyncio
+    async def test_non_string_instructions_rewrite_falls_back(self):
+        handler = OpenAIResponsesHandler()
+        user_item = {"role": "user", "content": "What is the codename?"}
+        data = {
+            "model": "gpt-5.6",
+            "instructions": "Answer from the memo only.",
+            "input": [user_item],
+        }
+
+        result = await handler.process_input_messages(
+            data, SystemRewriteGuardrail(rewritten_content=[{"type": "text", "text": COMPRESSED_MARKER}])
+        )
+
+        assert result["input"][0] is not user_item
+
+    @pytest.mark.asyncio
+    async def test_unpredicted_merge_falls_back_through_patch(self):
+        handler = OpenAIResponsesHandler()
+        raw_input = [
+            {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"},
+            {"role": "assistant", "content": "Reading the memo now."},
+            {"type": "function_call_output", "call_id": "call_1", "output": "memo memo"},
+            {"role": "user", "content": "memo " * 400},
+        ]
+        data = {"model": "gpt-5.6", "input": raw_input}
+
+        result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
+
+        assert not any(item is original for item in result["input"] for original in raw_input)
+        user_items = [item for item in result["input"] if item.get("role") == "user"]
+        assert _texts(user_items[0]) == [COMPRESSED_MARKER]
+
+    def test_item_rewrite_field_ignores_non_string_type(self):
+        from litellm.llms.openai.responses.guardrail_translation.handler import _item_rewrite_field
+
+        assert _item_rewrite_field({"type": 123, "content": "hello"}) is None

From 3004b12e900dc521e68f3f4e02e2e6bb53d50dd8 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 30 Aug 2026 12:40:40 -0700
Subject: [PATCH 08/43] refactor(responses): unify guardrail input processing
 and drop mutable request params

The staging merge tightened the ruff-strict and type-discipline budgets, so the
two `data: dict` parameters the write-back helpers introduced (LIT001) and the
17-branch `process_input_messages` (C901) no longer fit.

Fold the duplicated string/list guardrail flow into one path: a pure
`_extract_guardrail_inputs` builds the guardrail payload, the write-back
helpers become pure functions returning `_RequestFields` (patched input items
plus the resulting instructions value), and the request dict is only mutated
in `process_input_messages` itself. `_apply_guardrail_responses_to_input`
takes Sequence views since it only reads. A non-list `structured_messages`
payload now falls through to the plain texts write-back, matching the
pre-write-back behavior for guardrails that never touch structured messages.

Handler file deltas vs the merge base: LIT001 57 -> 53, LIT002 42 -> 40,
LIT010 28 -> 18, C901 3 -> 3.
---
 .../guardrail_translation/handler.py          | 258 ++++++++----------
 1 file changed, 121 insertions(+), 137 deletions(-)

diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py
index bd58fb4cdbc..76a755799d0 100644
--- a/litellm/llms/openai/responses/guardrail_translation/handler.py
+++ b/litellm/llms/openai/responses/guardrail_translation/handler.py
@@ -31,7 +31,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
 from collections.abc import Mapping, Sequence
 from itertools import accumulate
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Union, cast
+from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast
 
 from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
 from openai.types.responses.tool_param import FunctionToolParam
@@ -203,18 +203,28 @@ def _input_item_provenance(
     return item_for_message, tainted
 
 
-def _patch_rewritten_rows_into_input(
-    data: dict,
+class _RequestFields(NamedTuple):
+    input: tuple[object, ...]
+    instructions: str | None
+
+
+class _ExtractedInputs(NamedTuple):
+    inputs: GenericGuardrailAPIInputs
+    task_mappings: tuple[tuple[int, int | None], ...]
+
+
+def _patched_request_fields(
+    raw_input: object,
+    instructions: object,
     original_messages: Sequence[object],
     structured_messages: Sequence[object],
-) -> bool:
-    raw_input: Final = data.get("input")
+) -> _RequestFields | None:
     if not isinstance(raw_input, list) or len(original_messages) != len(structured_messages):
-        return False
-    offset: Final = 1 if data.get("instructions") else 0
+        return None
+    offset: Final = 1 if instructions else 0
     provenance: Final = _input_item_provenance(raw_input, tuple(original_messages)[offset:])
     if provenance is None:
-        return False
+        return None
     item_for_message, tainted = provenance
     changed: Final = tuple(
         (index, rewritten)
@@ -225,13 +235,14 @@ def _patch_rewritten_rows_into_input(
     rewritten_instructions: Final = (
         instruction_rewrites[0].get("content")
         if instruction_rewrites and isinstance(instruction_rewrites[0], Mapping)
-        else None
+        else instructions
     )
-    if instruction_rewrites and not isinstance(rewritten_instructions, str):
-        return False
+    instructions_value: Final = rewritten_instructions if isinstance(rewritten_instructions, str) else None
+    if rewritten_instructions is not None and instructions_value is None:
+        return None
     body_changes: Final = tuple((index - offset, rewritten) for index, rewritten in changed if index >= offset)
     if any(message_index in tainted or message_index not in item_for_message for message_index, _ in body_changes):
-        return False
+        return None
     replacements: Final = MappingProxyType(
         {
             item_for_message[message_index]: _rewritten_input_item(
@@ -242,11 +253,28 @@ def _patch_rewritten_rows_into_input(
         }
     )
     if len(replacements) != len(body_changes) or any(item is None for item in replacements.values()):
-        return False
-    data["input"] = [replacements.get(index, item) for index, item in enumerate(raw_input)]  # mutable-ok: JSON body
-    if isinstance(rewritten_instructions, str):
-        data["instructions"] = rewritten_instructions
-    return True
+        return None
+    return _RequestFields(
+        input=tuple(replacements.get(index, item) for index, item in enumerate(raw_input)),
+        instructions=instructions_value,
+    )
+
+
+def _written_back_request_fields(
+    raw_input: object,
+    instructions: object,
+    original_messages: Sequence[object],
+    structured_messages: Sequence[AllMessageValues],
+) -> _RequestFields | None:
+    if not isinstance(structured_messages, list):
+        return None
+    patched: Final = _patched_request_fields(raw_input, instructions, original_messages, structured_messages)
+    if patched is not None:
+        return patched
+    input_items, converted_instructions = (
+        LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(structured_messages)
+    )
+    return _RequestFields(input=tuple(input_items), instructions=converted_instructions)
 
 
 class OpenAIResponsesHandler(BaseTranslation):
@@ -288,136 +316,92 @@ class OpenAIResponsesHandler(BaseTranslation):
         Handles both string input and list of message objects.
         """
         input_data: Final[str | ResponseInputParam | None] = data.get("input")
-        tools_to_check: Final[list[ChatCompletionToolParam]] = []
-        if input_data is None:
+        if not isinstance(input_data, (str, list)):
             return data
-
         structured_messages: Final = self.get_structured_messages(data)
-
-        # Handle simple string input
-        if isinstance(input_data, str):
-            inputs = GenericGuardrailAPIInputs(texts=[input_data])
-            original_tools: list[dict[str, object]] = []
-
-            # Extract and transform tools if present
-            if "tools" in data and data["tools"]:
-                original_tools = list(data["tools"])
-                self._extract_and_transform_tools(data["tools"], tools_to_check)
-                if tools_to_check:
-                    inputs["tools"] = tools_to_check
-            if structured_messages:
-                inputs["structured_messages"] = structured_messages
-            # Include model information if available
-            model = data.get("model")
-            if model:
-                inputs["model"] = model
-
-            guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
-                inputs=inputs,
-                request_data=data,
-                input_type="request",
-                logging_obj=litellm_logging_obj,
-            )
-            guardrailed_structured_messages = guardrailed_inputs.get("structured_messages")
-            if (
-                guardrailed_structured_messages is not None
-                and guardrailed_structured_messages is not structured_messages
-            ):
-                self._write_back_structured_messages(data, structured_messages or (), guardrailed_structured_messages)
+        extracted: Final = self._extract_guardrail_inputs(data, input_data)
+        if not extracted.inputs.get("texts"):
+            return data
+        if structured_messages:
+            extracted.inputs["structured_messages"] = structured_messages
+        original_tools: Final[list[dict[str, object]]] = list(data.get("tools") or [])
+        guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
+            inputs=extracted.inputs,
+            request_data=data,
+            input_type="request",
+            logging_obj=litellm_logging_obj,
+        )
+        self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools"))
+        written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs)
+        if written_back is not None:
+            data["input"] = list(written_back.input)  # mutable-ok: JSON body
+            if written_back.instructions is None:
+                data.pop("instructions", None)
             else:
-                guardrailed_texts = guardrailed_inputs.get("texts") or ()
-                data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data
-            self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools"))
-            verbose_proxy_logger.debug("OpenAI Responses API: Processed string input")
-            return data
-
-        # Handle list input (ResponseInputParam)
-        if not isinstance(input_data, list):
-            return data
+                data["instructions"] = written_back.instructions
+        elif isinstance(input_data, str):
+            guardrailed_texts: Final = guardrailed_inputs.get("texts") or ()
+            data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data
+        else:
+            await self._apply_guardrail_responses_to_input(
+                messages=input_data,
+                responses=guardrailed_inputs.get("texts") or (),
+                task_mappings=extracted.task_mappings,
+            )
+        verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input"))
+        return data
 
+    def _extract_guardrail_inputs(
+        self,
+        data: Mapping[str, object],
+        input_data: "str | ResponseInputParam",
+    ) -> _ExtractedInputs:
         texts_to_check: Final[list[str]] = []
         images_to_check: Final[list[str]] = []
         task_mappings: Final[list[tuple[int, int | None]]] = []
-        original_tools_list: Final[list[dict[str, object]]] = list(data.get("tools") or [])
-
-        # Step 1: Extract all text content, images, and tools
-        for msg_idx, message in enumerate(input_data):
-            self._extract_input_text_and_images(
-                message=message,
-                msg_idx=msg_idx,
-                texts_to_check=texts_to_check,
-                images_to_check=images_to_check,
-                task_mappings=task_mappings,
-            )
-
-        # Extract and transform tools if present
-        if "tools" in data and data["tools"]:
-            self._extract_and_transform_tools(data["tools"], tools_to_check)
-
-        # Step 2: Apply guardrail to all texts in batch
-        if texts_to_check:
-            inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
-            if images_to_check:
-                inputs["images"] = images_to_check
-            if tools_to_check:
-                inputs["tools"] = tools_to_check
-            if structured_messages:
-                inputs["structured_messages"] = structured_messages
-            # Include model information if available
-            model = data.get("model")
-            if model:
-                inputs["model"] = model
-            guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
-                inputs=inputs,
-                request_data=data,
-                input_type="request",
-                logging_obj=litellm_logging_obj,
-            )
-
-            self._apply_guardrailed_tools_to_data(
-                data,
-                original_tools_list,
-                guardrailed_inputs.get("tools"),
-            )
-
-            guardrailed_structured_messages = guardrailed_inputs.get("structured_messages")
-            if (
-                guardrailed_structured_messages is not None
-                and guardrailed_structured_messages is not structured_messages
-            ):
-                self._write_back_structured_messages(data, structured_messages or (), guardrailed_structured_messages)
-            else:
-                # Step 3: Map guardrail responses back to original input structure
-                await self._apply_guardrail_responses_to_input(
-                    messages=input_data,
-                    responses=guardrailed_inputs.get("texts", []),  # mutable-ok: callee signature takes a list
+        tools_to_check: Final[list[ChatCompletionToolParam]] = []
+        if isinstance(input_data, str):
+            texts_to_check.append(input_data)
+        else:
+            for msg_idx, message in enumerate(input_data):
+                self._extract_input_text_and_images(
+                    message=message,
+                    msg_idx=msg_idx,
+                    texts_to_check=texts_to_check,
+                    images_to_check=images_to_check,
                     task_mappings=task_mappings,
                 )
-
-        verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input"))
-
-        return data
+        tools: Final = data.get("tools")
+        if tools:
+            self._extract_and_transform_tools(
+                cast("list[FunctionToolParam | OpenAIMcpServerTool]", tools),  # cast-ok: request body tools
+                tools_to_check,
+            )
+        inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
+        if images_to_check:
+            inputs["images"] = images_to_check
+        if tools_to_check:
+            inputs["tools"] = tools_to_check
+        model: Final = data.get("model")
+        if isinstance(model, str):
+            inputs["model"] = model
+        return _ExtractedInputs(inputs=inputs, task_mappings=tuple(task_mappings))
 
     @staticmethod
-    def _write_back_structured_messages(
-        data: dict,
-        original_messages: Sequence[object],
-        structured_messages: Sequence[AllMessageValues],
-    ) -> None:
-        if not isinstance(structured_messages, list):
-            return
-        if _patch_rewritten_rows_into_input(data, original_messages, structured_messages):
-            return
-        input_items, instructions = (
-            LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(
-                list(structured_messages)  # mutable-ok: converter signature takes a list
-            )
+    def _written_back_request_fields(
+        data: Mapping[str, object],
+        structured_messages: Sequence[AllMessageValues] | None,
+        guardrailed_inputs: GenericGuardrailAPIInputs,
+    ) -> _RequestFields | None:
+        guardrailed: Final = guardrailed_inputs.get("structured_messages")
+        if guardrailed is None or guardrailed is structured_messages:
+            return None
+        return _written_back_request_fields(
+            data.get("input"),
+            data.get("instructions"),
+            structured_messages or (),
+            guardrailed,
         )
-        data["input"] = input_items
-        if instructions is None:
-            data.pop("instructions", None)
-            return
-        data["instructions"] = instructions
 
     def extract_request_tool_names(self, data: dict) -> list[str]:
         """Extract tool names from Responses API request (tools[].name for function
@@ -543,8 +527,8 @@ class OpenAIResponsesHandler(BaseTranslation):
     async def _apply_guardrail_responses_to_input(
         self,
         messages: Any,  # Can be List[Dict[str, Any]] or ResponseInputParam
-        responses: list[str],
-        task_mappings: list[tuple[int, int | None]],
+        responses: Sequence[str],
+        task_mappings: Sequence[tuple[int, int | None]],
     ) -> None:
         """
         Apply guardrail responses back to input messages.

From 58c5223ab4b34d6bf013b7bafb673691e9d23d2a Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 30 Aug 2026 12:59:58 -0700
Subject: [PATCH 09/43] refactor(responses): rename write-back helper out of
 the method's name

The recursion detector in code-quality reads the staticmethod
_written_back_request_fields calling the module-level function of the
same name as a recursive call. Renaming the module-level helper to
_patch_or_convert_request_fields removes the shadowing and describes
what it does: patch changed rows in place, else fall back to full
conversion.
---
 .../llms/openai/responses/guardrail_translation/handler.py    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py
index e7a00a15a23..c676eaa7711 100644
--- a/litellm/llms/openai/responses/guardrail_translation/handler.py
+++ b/litellm/llms/openai/responses/guardrail_translation/handler.py
@@ -264,7 +264,7 @@ def _patched_request_fields(
     )
 
 
-def _written_back_request_fields(
+def _patch_or_convert_request_fields(
     raw_input: object,
     instructions: object,
     original_messages: Sequence[object],
@@ -408,7 +408,7 @@ class OpenAIResponsesHandler(BaseTranslation):
         guardrailed: Final = guardrailed_inputs.get("structured_messages")
         if guardrailed is None or guardrailed is structured_messages:
             return None
-        return _written_back_request_fields(
+        return _patch_or_convert_request_fields(
             data.get("input"),
             data.get("instructions"),
             structured_messages or (),

From 1072de94de8e9ce6de0e1e6197caa9ab207abbf6 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Mon, 31 Aug 2026 10:00:55 -0700
Subject: [PATCH 10/43] fix(azure_ai): only reclassify as azure when api_base
 is a classic Azure OpenAI endpoint

---
 litellm/llms/azure_ai/chat/transformation.py  | 28 +++++++------
 .../chat/test_azure_ai_transformation.py      | 40 +++++++++++++++++++
 .../test_gpt_5_5_model_metadata.py            |  3 +-
 3 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py
index 7fe9d3dec52..8797f4dc400 100644
--- a/litellm/llms/azure_ai/chat/transformation.py
+++ b/litellm/llms/azure_ai/chat/transformation.py
@@ -207,20 +207,22 @@ class AzureAIStudioConfig(OpenAIConfig):
                 message["content"] = texts
         return stripped_messages
 
-    def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool:
-        try:
-            if "/" in model:
-                model = model.split("/", 1)[1]
-            if (
-                model in litellm.open_ai_chat_completion_models
-                or model in litellm.open_ai_text_completion_models
-                or model in litellm.open_ai_embedding_models
-            ):
-                return True
-
-        except Exception:
+    def _is_foundry_model_inference_base(self, api_base: str) -> bool:
+        parsed: Final = urlparse(api_base)
+        host: Final = parsed.hostname
+        if host is None or not host.endswith(".services.ai.azure.com"):
             return False
-        return False
+        return "/openai/deployments" not in parsed.path
+
+    def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool:
+        if api_base is None or self._is_foundry_model_inference_base(api_base):
+            return False
+        stripped_model: Final = model.split("/", 1)[1] if "/" in model else model
+        return (
+            stripped_model in litellm.open_ai_chat_completion_models
+            or stripped_model in litellm.open_ai_text_completion_models
+            or stripped_model in litellm.open_ai_embedding_models
+        )
 
     def _get_openai_compatible_provider_info(
         self,
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 11a727c9635..33fbb4e8fc7 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
@@ -31,6 +31,46 @@ async def test_get_openai_compatible_provider_info():
     assert custom_llm_provider == "azure"
 
 
+@pytest.mark.parametrize(
+    "model, api_base, expected_provider",
+    [
+        ("azure_ai/gpt-4o", "https://my-resource.services.ai.azure.com", "azure_ai"),
+        ("azure_ai/gpt-4o", "https://my-resource.services.ai.azure.com/models", "azure_ai"),
+        ("azure_ai/gpt-5.4-nano", "https://my-resource.services.ai.azure.com", "azure_ai"),
+        ("azure_ai/gpt-4o", "https://my-resource.openai.azure.com", "azure"),
+        (
+            "azure_ai/gpt-4o",
+            "https://my-resource.services.ai.azure.com/openai/deployments/gpt-4o/chat/completions"
+            "?api-version=2024-08-01-preview",
+            "azure",
+        ),
+        ("azure_ai/mistral-large-latest", "https://my-resource.services.ai.azure.com", "azure_ai"),
+        ("azure_ai/mistral-large-latest", "https://my-resource.openai.azure.com", "azure_ai"),
+    ],
+)
+def test_foundry_base_keeps_azure_ai_provider(model: str, api_base: str, expected_provider: str):
+    """Regression for #38276: a Foundry .services.ai.azure.com base must not be reclassified as azure."""
+    config = AzureAIStudioConfig()
+    (
+        _,
+        _,
+        custom_llm_provider,
+    ) = config._get_openai_compatible_provider_info(
+        model=model,
+        api_base=api_base,
+        api_key="my-key",
+        custom_llm_provider="azure_ai",
+    )
+    assert custom_llm_provider == expected_provider
+
+
+def test_is_azure_openai_model_without_api_base_keeps_azure_ai():
+    """Metadata lookups (get_model_info, supports_* checks) carry no api_base and must not flip the provider."""
+    config = AzureAIStudioConfig()
+    assert config._is_azure_openai_model(model="azure_ai/gpt-4o", api_base=None) is False
+    assert config._is_azure_openai_model(model="azure_ai/gpt-4o", api_base="https://my-res.openai.azure.com") is True
+
+
 def test_azure_ai_validate_environment():
     config = AzureAIStudioConfig()
     headers = config.validate_environment(
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 1c12a48ed9d..a60fa9466e6 100644
--- a/tests/test_litellm/test_gpt_5_5_model_metadata.py
+++ b/tests/test_litellm/test_gpt_5_5_model_metadata.py
@@ -47,8 +47,7 @@ def test_azure_ai_gpt_5_5_model_info(model):
 
     routed_model, provider, _, _ = get_llm_provider(model=model)
     assert routed_model == model.split("/", 1)[1]
-    # azure_ai/* models resolve under the azure provider in get_llm_provider
-    assert provider == "azure"
+    assert provider == "azure_ai"
 
 
 def test_azure_ai_gpt_5_5_backup_matches_main():

From 604f1fde5014d3a966e1a2ee814e28f83bf97549 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Mon, 31 Aug 2026 11:46:33 -0700
Subject: [PATCH 11/43] fix(azure_ai): route Foundry embeddings to the /models
 inference route

---
 litellm/llms/azure_ai/chat/transformation.py  |  7 +-
 litellm/llms/azure_ai/common_utils.py         |  9 +++
 litellm/llms/azure_ai/embed/handler.py        | 19 ++++-
 .../embed/test_azure_ai_embed_handler.py      | 69 +++++++++++++++++++
 4 files changed, 96 insertions(+), 8 deletions(-)
 create mode 100644 tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py

diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py
index 8797f4dc400..f2d405e9a17 100644
--- a/litellm/llms/azure_ai/chat/transformation.py
+++ b/litellm/llms/azure_ai/chat/transformation.py
@@ -15,6 +15,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
     filter_value_from_dict,
 )
 from litellm.llms.azure.common_utils import BaseAzureLLM
+from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base
 from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
 from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error
 from litellm.llms.openai.openai import OpenAIConfig
@@ -208,11 +209,7 @@ class AzureAIStudioConfig(OpenAIConfig):
         return stripped_messages
 
     def _is_foundry_model_inference_base(self, api_base: str) -> bool:
-        parsed: Final = urlparse(api_base)
-        host: Final = parsed.hostname
-        if host is None or not host.endswith(".services.ai.azure.com"):
-            return False
-        return "/openai/deployments" not in parsed.path
+        return is_foundry_model_inference_base(api_base)
 
     def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool:
         if api_base is None or self._is_foundry_model_inference_base(api_base):
diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py
index 26a90157455..aa34bab5b2e 100644
--- a/litellm/llms/azure_ai/common_utils.py
+++ b/litellm/llms/azure_ai/common_utils.py
@@ -1,5 +1,6 @@
 from collections.abc import Mapping
 from typing import Final, Literal
+from urllib.parse import urlparse
 
 import litellm
 from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
@@ -10,6 +11,14 @@ from litellm.types.router import GenericLiteLLMParams
 AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"]
 
 
+def is_foundry_model_inference_base(api_base: str) -> bool:
+    parsed: Final = urlparse(api_base)
+    host: Final = parsed.hostname
+    if host is None or not host.endswith(".services.ai.azure.com"):
+        return False
+    return "/openai/deployments" not in parsed.path
+
+
 def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None:
     """
     Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment.
diff --git a/litellm/llms/azure_ai/embed/handler.py b/litellm/llms/azure_ai/embed/handler.py
index 65c3997c099..c65edbf56e6 100644
--- a/litellm/llms/azure_ai/embed/handler.py
+++ b/litellm/llms/azure_ai/embed/handler.py
@@ -1,8 +1,10 @@
 from typing import Final
+from urllib.parse import urlsplit, urlunsplit
 
 from openai import OpenAI
 
 import litellm
+from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     HTTPHandler,
@@ -16,6 +18,16 @@ from litellm.utils import convert_to_model_response_object
 from .cohere_transformation import AzureAICohereConfig
 
 
+def _foundry_models_route_base(api_base: str | None) -> str | None:
+    if api_base is None or not is_foundry_model_inference_base(api_base):
+        return api_base
+    parts: Final = urlsplit(api_base)
+    path: Final = parts.path.rstrip("/")
+    if path.endswith("/models"):
+        return api_base
+    return urlunsplit((parts.scheme, parts.netloc, f"{path}/models", parts.query, parts.fragment))
+
+
 class AzureAIEmbedding(OpenAIChatCompletion):
     def _process_response(
         self,
@@ -214,6 +226,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
 
         assemble result in-order, and return
         """
+        resolved_api_base: Final = _foundry_models_route_base(api_base)
         if aembedding is True:
             return self.async_embedding(
                 model,
@@ -223,7 +236,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
                 model_response,
                 optional_params,
                 api_key,
-                api_base,
+                resolved_api_base,
                 client,
             )
 
@@ -245,7 +258,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
                 model_response=model_response,
                 optional_params=optional_params,
                 api_key=api_key,
-                api_base=api_base,
+                api_base=resolved_api_base,
                 client=client,
             )
 
@@ -262,7 +275,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
                 model_response,
                 optional_params,
                 api_key,
-                api_base,
+                resolved_api_base,
                 client=(client if client is not None and isinstance(client, OpenAI) else None),
                 aembedding=aembedding,
                 shared_session=shared_session,
diff --git a/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py b/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py
new file mode 100644
index 00000000000..0401629171d
--- /dev/null
+++ b/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py
@@ -0,0 +1,69 @@
+import httpx
+import pytest
+import respx
+
+from litellm import embedding
+from litellm.llms.azure_ai.embed.handler import _foundry_models_route_base
+
+EMBEDDING_PAYLOAD = {
+    "object": "list",
+    "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}],
+    "model": "text-embedding-3-small",
+    "usage": {"prompt_tokens": 2, "total_tokens": 2},
+}
+
+
+@pytest.mark.parametrize(
+    ("api_base", "expected"),
+    [
+        (
+            "https://my-foundry.services.ai.azure.com",
+            "https://my-foundry.services.ai.azure.com/models",
+        ),
+        (
+            "https://my-foundry.services.ai.azure.com/",
+            "https://my-foundry.services.ai.azure.com/models",
+        ),
+        (
+            "https://my-foundry.services.ai.azure.com?api-version=2024-05-01-preview",
+            "https://my-foundry.services.ai.azure.com/models?api-version=2024-05-01-preview",
+        ),
+        (
+            "https://my-foundry.services.ai.azure.com/models",
+            "https://my-foundry.services.ai.azure.com/models",
+        ),
+        (
+            "https://my-foundry.services.ai.azure.com/openai/deployments/text-embedding-3-small",
+            "https://my-foundry.services.ai.azure.com/openai/deployments/text-embedding-3-small",
+        ),
+        (
+            "https://my-resource.openai.azure.com",
+            "https://my-resource.openai.azure.com",
+        ),
+        (
+            "https://Mistral-serverless.eastus2.models.ai.azure.com",
+            "https://Mistral-serverless.eastus2.models.ai.azure.com",
+        ),
+        (None, None),
+    ],
+)
+def test_foundry_models_route_base(api_base, expected):
+    assert _foundry_models_route_base(api_base) == expected
+
+
+@respx.mock
+def test_azure_ai_embedding_calls_foundry_models_route():
+    route = respx.post("https://my-foundry.services.ai.azure.com/models/embeddings").mock(
+        return_value=httpx.Response(200, json=EMBEDDING_PAYLOAD)
+    )
+
+    response = embedding(
+        model="azure_ai/text-embedding-3-small",
+        input=["hello world"],
+        api_base="https://my-foundry.services.ai.azure.com",
+        api_key="fake-key",
+    )
+
+    assert route.called
+    assert response.data is not None
+    assert len(response.data) == 1

From c493fc855c85041e331f77fc53ad39950089e8f2 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Tue, 1 Sep 2026 15:46:14 -0700
Subject: [PATCH 12/43] fix(ui): show MCP servers and agents inherited from
 access groups on the team overview

The team Overview and Settings tabs fed only object_permission into the Object
Permissions card, so a team whose access group grants MCP servers or agents read
"MCP Servers 0" and "Agents 0" while the Models card next to it already listed the
inherited models. /team/info has returned access_group_mcp_server_ids and
access_group_agent_ids for a while, nothing in the dashboard read them.

ObjectPermissionsView now accepts the inherited ids and MCPServerPermissions /
AgentPermissions merge them into their lists with an Inherited tag, deduped
against direct grants, so an admin can tell a group grant from a direct one.

Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL
---
 .../components/object_permissions_view.tsx    | 12 ++++-
 .../permissions/AgentPermissions.test.tsx     | 50 +++++++++++++++++++
 .../permissions/AgentPermissions.tsx          | 34 ++++++++++---
 .../permissions/MCPServerPermissions.test.tsx | 43 ++++++++++++++++
 .../permissions/MCPServerPermissions.tsx      | 34 ++++++++++---
 .../src/components/team/TeamInfo.test.tsx     | 28 +++++++++++
 .../src/components/team/TeamInfo.tsx          | 10 +++-
 7 files changed, 194 insertions(+), 17 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx

diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx
index 011791fa18d..7d1c4ce4ac9 100644
--- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx
+++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx
@@ -6,6 +6,8 @@ import type { ObjectPermission } from "./object_permission_types";
 
 interface ObjectPermissionsViewProps {
   objectPermission?: ObjectPermission | null;
+  inheritedMcpServerIds?: string[];
+  inheritedAgentIds?: string[];
   variant?: "card" | "inline";
   className?: string;
   accessToken?: string | null;
@@ -13,6 +15,8 @@ interface ObjectPermissionsViewProps {
 
 export function ObjectPermissionsView({
   objectPermission,
+  inheritedMcpServerIds = [],
+  inheritedAgentIds = [],
   variant = "card",
   className = "",
   accessToken,
@@ -34,9 +38,15 @@ export function ObjectPermissionsView({
         mcpAccessGroups={mcpAccessGroups}
         mcpToolPermissions={mcpToolPermissions}
         mcpToolsets={mcpToolsets}
+        inheritedMcpServers={inheritedMcpServerIds}
+        accessToken={accessToken}
+      />
+      <AgentPermissions
+        agents={agents}
+        agentAccessGroups={agentAccessGroups}
+        inheritedAgents={inheritedAgentIds}
         accessToken={accessToken}
       />
-      <AgentPermissions agents={agents} agentAccessGroups={agentAccessGroups} accessToken={accessToken} />
       <div className="min-w-0 rounded-md border border-border p-4">
         <p className="text-sm font-medium text-foreground">Search tools</p>
         {searchTools.length === 0 ? (
diff --git a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx
new file mode 100644
index 00000000000..14fe6177ac5
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx
@@ -0,0 +1,50 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { render, screen } from "@testing-library/react";
+import AgentPermissions from "./AgentPermissions";
+import * as networking from "../networking";
+
+vi.mock("../networking");
+
+describe("AgentPermissions", () => {
+  const accessToken = "test-token";
+  const agentId = "90337622-756e-4f25-98f0-01fc8174aa24";
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("lists agents inherited from access groups with an Inherited tag and counts them", async () => {
+    vi.mocked(networking.getAgentsList).mockResolvedValue({
+      agents: [{ agent_id: agentId, agent_name: "support_agent" }],
+    });
+
+    render(<AgentPermissions agents={[]} inheritedAgents={[agentId]} accessToken={accessToken} />);
+
+    expect(await screen.findByText(/support_agent/)).toBeInTheDocument();
+    expect(screen.getByText("Inherited")).toBeInTheDocument();
+    expect(screen.getByText("1")).toBeInTheDocument();
+    expect(screen.queryByText("No agents or access groups configured")).not.toBeInTheDocument();
+    expect(networking.getAgentsList).toHaveBeenCalledWith(accessToken);
+  });
+
+  it("does not double-list an agent that is both granted directly and inherited", async () => {
+    vi.mocked(networking.getAgentsList).mockResolvedValue({
+      agents: [{ agent_id: agentId, agent_name: "support_agent" }],
+    });
+
+    render(<AgentPermissions agents={[agentId]} inheritedAgents={[agentId]} accessToken={accessToken} />);
+
+    expect(await screen.findByText(/support_agent/)).toBeInTheDocument();
+    expect(screen.getAllByText(/support_agent/)).toHaveLength(1);
+    expect(screen.queryByText("Inherited")).not.toBeInTheDocument();
+    expect(screen.getByText("1")).toBeInTheDocument();
+  });
+
+  it("shows the empty state when nothing is granted directly or inherited", () => {
+    render(<AgentPermissions agents={[]} inheritedAgents={[]} accessToken={accessToken} />);
+
+    expect(screen.getByText("No agents or access groups configured")).toBeInTheDocument();
+    expect(screen.getByText("0")).toBeInTheDocument();
+    expect(networking.getAgentsList).not.toHaveBeenCalled();
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx
index d84bd1834f8..21bacc23ce1 100644
--- a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx
+++ b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx
@@ -14,16 +14,26 @@ interface Agent {
 interface AgentPermissionsProps {
   agents: string[];
   agentAccessGroups?: string[];
+  inheritedAgents?: string[];
   accessToken?: string | null;
 }
 
-export function AgentPermissions({ agents, agentAccessGroups = [], accessToken }: AgentPermissionsProps) {
+const INHERITED_AGENT_TOOLTIP = "Granted through one of the team's access groups";
+
+export function AgentPermissions({
+  agents,
+  agentAccessGroups = [],
+  inheritedAgents = [],
+  accessToken,
+}: AgentPermissionsProps) {
   const [agentDetails, setAgentDetails] = useState<Agent[]>([]);
+  const inheritedOnlyAgents = inheritedAgents.filter((agent) => !agents.includes(agent));
+  const agentIdCount = agents.length + inheritedOnlyAgents.length;
 
   // Fetch agent details when component mounts
   useEffect(() => {
     const fetchAgentDetails = async () => {
-      if (accessToken && agents.length > 0) {
+      if (accessToken && agentIdCount > 0) {
         try {
           const response = await getAgentsList(accessToken);
           if (response && response.agents && Array.isArray(response.agents)) {
@@ -35,7 +45,7 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken }
       }
     };
     fetchAgentDetails();
-  }, [accessToken, agents.length]);
+  }, [accessToken, agentIdCount]);
 
   // Function to get display name for agent
   const getAgentDisplayName = (agentId: string) => {
@@ -47,10 +57,11 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken }
     return agentId;
   };
 
-  // Merge agents and access groups into one list
+  // Merge agents, inherited agents and access groups into one list
   const mergedItems = [
-    ...agents.map((agent) => ({ type: "agent", value: agent })),
-    ...agentAccessGroups.map((group) => ({ type: "accessGroup", value: group })),
+    ...agents.map((agent) => ({ type: "agent", value: agent, inherited: false })),
+    ...inheritedOnlyAgents.map((agent) => ({ type: "agent", value: agent, inherited: true })),
+    ...agentAccessGroups.map((group) => ({ type: "accessGroup", value: group, inherited: false })),
   ];
   const totalCount = mergedItems.length;
 
@@ -76,8 +87,17 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken }
                           <span className="text-sm font-medium text-foreground truncate">
                             {getAgentDisplayName(item.value)}
                           </span>
+                          {item.inherited && (
+                            <span className="ml-1 px-1.5 py-0.5 text-[9px] font-semibold text-purple-600 bg-purple-50 border border-purple-200 rounded-sm uppercase tracking-wide shrink-0 dark:text-purple-300 dark:bg-purple-950 dark:border-purple-800">
+                              Inherited
+                            </span>
+                          )}
                         </TooltipTrigger>
-                        <TooltipContent>{`Full ID: ${item.value}`}</TooltipContent>
+                        <TooltipContent>
+                          {item.inherited
+                            ? `${INHERITED_AGENT_TOOLTIP}. Full ID: ${item.value}`
+                            : `Full ID: ${item.value}`}
+                        </TooltipContent>
                       </Tooltip>
                     </TooltipProvider>
                   ) : (
diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx
index c2df945f367..fe7e3a7b971 100644
--- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx
+++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx
@@ -406,4 +406,47 @@ describe("MCPServerPermissions", () => {
     );
     await waitFor(() => expect(screen.getByText("Blocked")).toHaveAttribute("data-variant", "destructive"));
   });
+
+  it("lists servers inherited from access groups with an Inherited tag and counts them", async () => {
+    vi.mocked(networking.fetchMCPServers).mockResolvedValue([
+      { server_id: mockServerId1, server_name: mockServerName1, alias: mockServerName1 },
+    ]);
+
+    render(
+      <MCPServerPermissions
+        mcpServers={[]}
+        mcpAccessGroups={[]}
+        mcpToolPermissions={{}}
+        inheritedMcpServers={[mockServerId1]}
+        accessToken={mockAccessToken}
+      />,
+    );
+
+    expect(await screen.findByText(/DW_MCP/)).toBeInTheDocument();
+    expect(screen.getByText("Inherited")).toBeInTheDocument();
+    expect(screen.getByText("1")).toBeInTheDocument();
+    expect(screen.queryByText("No MCP servers, access groups, or toolsets configured")).not.toBeInTheDocument();
+    expect(networking.fetchMCPServers).toHaveBeenCalledWith(mockAccessToken);
+  });
+
+  it("does not double-list a server that is both granted directly and inherited", async () => {
+    vi.mocked(networking.fetchMCPServers).mockResolvedValue([
+      { server_id: mockServerId2, server_name: mockServerName2, alias: mockServerName2 },
+    ]);
+
+    render(
+      <MCPServerPermissions
+        mcpServers={[mockServerId2]}
+        mcpAccessGroups={[]}
+        mcpToolPermissions={{}}
+        inheritedMcpServers={[mockServerId2]}
+        accessToken={mockAccessToken}
+      />,
+    );
+
+    expect(await screen.findByText(/Test Server/)).toBeInTheDocument();
+    expect(screen.getAllByText(/Test Server/)).toHaveLength(1);
+    expect(screen.queryByText("Inherited")).not.toBeInTheDocument();
+    expect(screen.getByText("1")).toBeInTheDocument();
+  });
 });
diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx
index f11bc6b9627..742409f3f8d 100644
--- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx
+++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx
@@ -11,14 +11,18 @@ interface MCPServerPermissionsProps {
   mcpAccessGroups?: string[];
   mcpToolPermissions?: Record<string, string[]>;
   mcpToolsets?: string[];
+  inheritedMcpServers?: string[];
   accessToken?: string | null;
 }
 
+const INHERITED_MCP_SERVER_TOOLTIP = "Granted through one of the team's access groups";
+
 export function MCPServerPermissions({
   mcpServers,
   mcpAccessGroups = [],
   mcpToolPermissions = {},
   mcpToolsets = [],
+  inheritedMcpServers = [],
   accessToken,
 }: MCPServerPermissionsProps) {
   const [mcpServerDetails, setMCPServerDetails] = useState<MCPServer[]>([]);
@@ -50,10 +54,16 @@ export function MCPServerPermissions({
     });
   };
 
+  const directServerIds = mcpServers.filter(
+    (server) => server !== NO_MCP_SERVERS_SENTINEL && server !== ALL_PROXY_MCP_SERVERS_SENTINEL,
+  );
+  const inheritedOnlyServerIds = inheritedMcpServers.filter((server) => !mcpServers.includes(server));
+  const serverIdCount = directServerIds.length + inheritedOnlyServerIds.length;
+
   // Fetch MCP server details when component mounts
   useEffect(() => {
     const fetchMCPServerDetails = async () => {
-      if (accessToken && mcpServers.length > 0) {
+      if (accessToken && serverIdCount > 0) {
         try {
           const response = await fetchMCPServers(accessToken);
           if (response && Array.isArray(response)) {
@@ -67,7 +77,7 @@ export function MCPServerPermissions({
       }
     };
     fetchMCPServerDetails();
-  }, [accessToken, mcpServers.length]);
+  }, [accessToken, serverIdCount]);
 
   // Fetch toolset details
   useEffect(() => {
@@ -98,12 +108,11 @@ export function MCPServerPermissions({
   const blocksAllMcpServers = mcpServers.includes(NO_MCP_SERVERS_SENTINEL);
   const grantsAllProxyMcpServers = mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL);
 
-  // Merge servers and access groups into one list
+  // Merge servers, inherited servers and access groups into one list
   const mergedItems = [
-    ...mcpServers
-      .filter((server) => server !== NO_MCP_SERVERS_SENTINEL && server !== ALL_PROXY_MCP_SERVERS_SENTINEL)
-      .map((server) => ({ type: "server", value: server })),
-    ...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group })),
+    ...directServerIds.map((server) => ({ type: "server", value: server, inherited: false })),
+    ...inheritedOnlyServerIds.map((server) => ({ type: "server", value: server, inherited: true })),
+    ...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group, inherited: false })),
   ];
   const totalCount = mergedItems.length + mcpToolsets.length;
 
@@ -152,8 +161,17 @@ export function MCPServerPermissions({
                           <span className="text-sm font-medium text-foreground truncate">
                             {getMCPServerDisplayName(item.value)}
                           </span>
+                          {item.inherited && (
+                            <span className="ml-1 px-1.5 py-0.5 text-[9px] font-semibold text-info bg-info/10 border border-info/20 rounded-sm uppercase tracking-wide shrink-0">
+                              Inherited
+                            </span>
+                          )}
                         </TooltipTrigger>
-                        <TooltipContent>{`Full ID: ${item.value}`}</TooltipContent>
+                        <TooltipContent>
+                          {item.inherited
+                            ? `${INHERITED_MCP_SERVER_TOOLTIP}. Full ID: ${item.value}`
+                            : `Full ID: ${item.value}`}
+                        </TooltipContent>
                       </Tooltip>
                     ) : (
                       <div className="inline-flex items-center gap-2 min-w-0">
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
index ca1e0413dcb..40a3d2ada1a 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
@@ -38,6 +38,9 @@ vi.mock("@/components/networking", () => ({
   organizationInfoCall: vi.fn(),
   getRouterSettingsCall: vi.fn().mockResolvedValue({ fields: [] }),
   getPassThroughEndpointsCall: vi.fn(),
+  fetchMCPServers: vi.fn().mockResolvedValue([]),
+  fetchMCPToolsets: vi.fn().mockResolvedValue([]),
+  getAgentsList: vi.fn().mockResolvedValue({ agents: [] }),
 }));
 
 const can = vi.fn();
@@ -302,6 +305,31 @@ describe("TeamInfoView", () => {
       );
     });
 
+    it("shows MCP servers and agents inherited from access groups in the Object Permissions card", async () => {
+      vi.mocked(networking.fetchMCPServers).mockResolvedValue([
+        { server_id: "mcp-github-1234", server_name: "github", alias: "github" },
+      ]);
+      vi.mocked(networking.getAgentsList).mockResolvedValue({
+        agents: [{ agent_id: "agent-support-5678", agent_name: "support_agent" }],
+      });
+      vi.mocked(networking.teamInfoCall).mockResolvedValue(
+        createMockTeamData({
+          object_permission: null,
+          access_group_ids: ["ag-1"],
+          access_group_mcp_server_ids: ["mcp-github-1234"],
+          access_group_agent_ids: ["agent-support-5678"],
+        }),
+      );
+
+      renderWithProviders(<TeamInfoView {...defaultProps} />);
+
+      expect(await screen.findByText(/github \(mcp\.\.\.1234\)/)).toBeInTheDocument();
+      expect(await screen.findByText(/support_agent \(age\.\.\.5678\)/)).toBeInTheDocument();
+      expect(screen.getAllByText("Inherited")).toHaveLength(2);
+      expect(screen.queryByText("No MCP servers, access groups, or toolsets configured")).not.toBeInTheDocument();
+      expect(screen.queryByText("No agents or access groups configured")).not.toBeInTheDocument();
+    });
+
     it("keeps the all-proxy-models badge non-clickable", async () => {
       vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["all-proxy-models"] }));
 
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
index 44f0d420fd6..7ebef0691e6 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
@@ -1033,7 +1033,13 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
             </div>
           </Card>
 
-          <ObjectPermissionsView objectPermission={info.object_permission} variant="card" accessToken={accessToken} />
+          <ObjectPermissionsView
+            objectPermission={info.object_permission}
+            inheritedMcpServerIds={info.access_group_mcp_server_ids}
+            inheritedAgentIds={info.access_group_agent_ids}
+            variant="card"
+            accessToken={accessToken}
+          />
 
           <Card className="block p-6">
             <GuardrailSettingsView
@@ -1883,6 +1889,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
 
               <ObjectPermissionsView
                 objectPermission={info.object_permission}
+                inheritedMcpServerIds={info.access_group_mcp_server_ids}
+                inheritedAgentIds={info.access_group_agent_ids}
                 variant="inline"
                 className="pt-4 border-t border-border"
                 accessToken={accessToken}

From 944d018c72cae0930fd8c13c61ebe702ffca85d3 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Tue, 1 Sep 2026 16:25:15 -0700
Subject: [PATCH 13/43] test(ui): query the screen in models page tests and
 drop restating comments

The merge of #38872 into staging kept the destructured render queries
in the models-and-endpoints page test, which pushes
testing-library/prefer-screen-queries to 21 against a budget of 18 and
fails frontend-lint for every PR on top of it. Two comments that only
restated the list merge below them are gone as well.

Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL
---
 .../app/(dashboard)/models-and-endpoints/page.test.tsx | 10 +++++-----
 .../src/components/permissions/AgentPermissions.tsx    |  1 -
 .../components/permissions/MCPServerPermissions.tsx    |  1 -
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx
index 1199b66621f..84a05113177 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx
@@ -111,17 +111,17 @@ describe("ModelsAndEndpointsPage", () => {
   // POST /model/new 403s a proxy_admin_viewer, so the form's tab must not render for one.
   it("hides the Add Model tab for a view-only admin session", () => {
     mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN);
-    const { getByRole, queryByRole } = renderPage();
-    expect(queryByRole("tab", { name: "Add Model" })).not.toBeInTheDocument();
-    expect(getByRole("tab", { name: "All Models" })).toBeInTheDocument();
+    renderPage();
+    expect(screen.queryByRole("tab", { name: "Add Model" })).not.toBeInTheDocument();
+    expect(screen.getByRole("tab", { name: "All Models" })).toBeInTheDocument();
   });
 
   // Read parity: the Auto-Routers list stays reachable for a view-only admin; only the
   // create affordance inside it is withheld, which AutoRoutersTabPanel decides.
   it("keeps the Auto-Routers tab for a view-only admin session", () => {
     mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN);
-    const { getByRole } = renderPage();
-    expect(getByRole("tab", { name: /Auto-Routers/ })).toBeInTheDocument();
+    renderPage();
+    expect(screen.getByRole("tab", { name: /Auto-Routers/ })).toBeInTheDocument();
   });
 
   // Auto-routers are excluded from the All Models table, so this tab is their home: the only
diff --git a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx
index 21bacc23ce1..11f6c414922 100644
--- a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx
+++ b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx
@@ -57,7 +57,6 @@ export function AgentPermissions({
     return agentId;
   };
 
-  // Merge agents, inherited agents and access groups into one list
   const mergedItems = [
     ...agents.map((agent) => ({ type: "agent", value: agent, inherited: false })),
     ...inheritedOnlyAgents.map((agent) => ({ type: "agent", value: agent, inherited: true })),
diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx
index 742409f3f8d..52199dab3c7 100644
--- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx
+++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx
@@ -108,7 +108,6 @@ export function MCPServerPermissions({
   const blocksAllMcpServers = mcpServers.includes(NO_MCP_SERVERS_SENTINEL);
   const grantsAllProxyMcpServers = mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL);
 
-  // Merge servers, inherited servers and access groups into one list
   const mergedItems = [
     ...directServerIds.map((server) => ({ type: "server", value: server, inherited: false })),
     ...inheritedOnlyServerIds.map((server) => ({ type: "server", value: server, inherited: true })),

From 7b942fd983af2b033c2767f783777ae86df6cec8 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 2 Sep 2026 11:05:38 -0700
Subject: [PATCH 14/43] fix(azure_ai): route audio and realtime calls on
 Foundry hosts through the Azure OpenAI handlers

---
 litellm/constants.py                         |  1 +
 litellm/main.py                              | 10 +++--
 litellm/realtime_api/main.py                 |  3 +-
 tests/test_litellm/realtime_api/test_main.py | 37 ++++++++++++++++++
 tests/test_litellm/test_main.py              | 40 ++++++++++++++++++++
 5 files changed, 87 insertions(+), 4 deletions(-)

diff --git a/litellm/constants.py b/litellm/constants.py
index 1bd977dd9a9..041a83f53e6 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -7,6 +7,7 @@ from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_ran
 
 DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
 AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
+AZURE_OPENAI_AUDIO_PROVIDERS: Final = frozenset({"azure", "azure_ai"})
 ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5))
 ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000
 DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
diff --git a/litellm/main.py b/litellm/main.py
index c4c5bbefc4f..f6b6453ab45 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -60,6 +60,7 @@ if TYPE_CHECKING:
     from litellm.types.utils import TokenCountResponse
 
 from litellm.constants import (
+    AZURE_OPENAI_AUDIO_PROVIDERS,
     DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
     DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
 )
@@ -7770,7 +7771,7 @@ def transcription(
         provider=LlmProviders(custom_llm_provider),
     )
 
-    if custom_llm_provider == "azure" and provider_config is None:
+    if custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS and provider_config is None:
         # azure configs
         api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
 
@@ -8057,7 +8058,10 @@ def speech(
         custom_llm_provider=custom_llm_provider,
     )
     response: HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent] | None = None
-    if custom_llm_provider == "openai" or custom_llm_provider in litellm.openai_compatible_providers:
+    if custom_llm_provider == "openai" or (
+        custom_llm_provider in litellm.openai_compatible_providers
+        and custom_llm_provider not in AZURE_OPENAI_AUDIO_PROVIDERS
+    ):
         if voice is None or not (isinstance(voice, str)):
             raise litellm.BadRequestError(
                 message="'voice' is required to be passed as a string for OpenAI TTS",
@@ -8111,7 +8115,7 @@ def speech(
             aspeech=aspeech,
             shared_session=shared_session,
         )
-    elif custom_llm_provider == "azure":
+    elif custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS:
         # Check if this is Azure Speech Service (Cognitive Services TTS)
         if model.startswith("speech/"):
             from litellm.llms.azure.text_to_speech.transformation import (
diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py
index d4b9f4e8cce..3862aec445f 100644
--- a/litellm/realtime_api/main.py
+++ b/litellm/realtime_api/main.py
@@ -8,6 +8,7 @@ from typing import Any, Final, Literal, cast
 
 import litellm
 from litellm.constants import (
+    AZURE_OPENAI_AUDIO_PROVIDERS,
     REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS,
     REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES,
     request_timeout,
@@ -400,7 +401,7 @@ async def _arealtime(
             litellm_metadata=_build_litellm_metadata(kwargs),
             query_params=query_params,
         )
-    elif _custom_llm_provider == "azure":
+    elif _custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS:
         api_base = dynamic_api_base or litellm_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
         # set API KEY
         api_key = dynamic_api_key or litellm.api_key or litellm.openai_key or get_secret_str("AZURE_API_KEY")
diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py
index a3dd5688ad1..761e87ac764 100644
--- a/tests/test_litellm/realtime_api/test_main.py
+++ b/tests/test_litellm/realtime_api/test_main.py
@@ -1,6 +1,7 @@
 import asyncio
 import time
 from types import TracebackType
+from typing import Final
 from unittest.mock import MagicMock, patch
 
 
@@ -294,3 +295,39 @@ async def test_azure_health_check_honors_deployment_realtime_protocol():
             model_params={"realtime_protocol": "GA"},
         )
     assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview"
+
+
+class _ConnectThatStopsAfterCapturingTheUrl:
+    url: str | None = None
+
+    def __call__(self, url: str, **kwargs: object) -> "_ConnectThatStopsAfterCapturingTheUrl":
+        self.url = url
+        return self
+
+    async def __aenter__(self) -> None:
+        raise RuntimeError("backend url captured, nothing to bridge")
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        return None
+
+
+@pytest.mark.asyncio
+async def test_arealtime_azure_ai_on_a_foundry_host_connects_to_the_azure_openai_realtime_route():
+    connect: Final = _ConnectThatStopsAfterCapturingTheUrl()
+    with patch("websockets.connect", connect):
+        await realtime_main._arealtime.__wrapped__(
+            model="azure_ai/gpt-realtime-mini",
+            websocket=MagicMock(),
+            api_base="https://my-project.services.ai.azure.com",
+            api_key="fake-key",
+            litellm_logging_obj=FakeLogging(),
+        )
+    assert connect.url == (
+        "wss://my-project.services.ai.azure.com/openai/realtime"
+        "?api-version=2024-10-01-preview&deployment=gpt-realtime-mini"
+    )
diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py
index 8cf878d05d9..30975033b7d 100644
--- a/tests/test_litellm/test_main.py
+++ b/tests/test_litellm/test_main.py
@@ -3181,3 +3181,43 @@ 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
+
+
+FOUNDRY_HOST: Final = "https://my-project.services.ai.azure.com"
+
+
+def test_azure_ai_transcription_on_a_foundry_host_uses_the_azure_openai_deployment_route(
+    respx_mock: respx.MockRouter,
+):
+    route: Final = respx_mock.post(
+        url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/whisper-1/audio/transcriptions\?api-version=.+"
+    ).mock(return_value=httpx.Response(200, json={"text": "hello"}))
+
+    response: Final = litellm.transcription(
+        model="azure_ai/whisper-1",
+        file=("tone.wav", b"RIFF\x00\x00\x00\x00WAVE", "audio/wav"),
+        api_base=FOUNDRY_HOST,
+        api_key="fake-key",
+    )
+
+    assert route.called
+    assert response.text == "hello"
+
+
+def test_azure_ai_speech_on_a_foundry_host_uses_the_azure_openai_deployment_route(
+    respx_mock: respx.MockRouter,
+):
+    route: Final = respx_mock.post(
+        url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/tts-1/audio/speech\?api-version=.+"
+    ).mock(return_value=httpx.Response(200, content=b"mp3-bytes"))
+
+    response: Final = litellm.speech(
+        model="azure_ai/tts-1",
+        input="hello",
+        voice="alloy",
+        api_base=FOUNDRY_HOST,
+        api_key="fake-key",
+    )
+
+    assert route.called
+    assert response.content == b"mp3-bytes"

From de643c028fce8d7634191550968eae7c4c991f04 Mon Sep 17 00:00:00 2001
From: yassin <yassin@berri.ai>
Date: Thu, 3 Sep 2026 00:25:50 +0000
Subject: [PATCH 15/43] style: format handler after merge

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../llms/openai/responses/guardrail_translation/handler.py    | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py
index 0c19ecf909e..59ca601d8be 100644
--- a/litellm/llms/openai/responses/guardrail_translation/handler.py
+++ b/litellm/llms/openai/responses/guardrail_translation/handler.py
@@ -374,7 +374,9 @@ class OpenAIResponsesHandler(BaseTranslation):
             input_type="request",
             logging_obj=litellm_logging_obj,
         )
-        self._apply_guardrailed_tools_to_data(data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools"))
+        self._apply_guardrailed_tools_to_data(
+            data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools")
+        )
         written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs)
         if written_back is not None:
             data["input"] = list(written_back.input)  # mutable-ok: JSON body

From 45f44fea6019e3e9afd75e7157992899f0846bd0 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 2 Sep 2026 17:53:39 -0700
Subject: [PATCH 16/43] fix(headroom): leave background responses requests
 uncompressed and unconverted

---
 .../guardrail_hooks/headroom/headroom.py      |  6 ++-
 .../guardrail_hooks/test_headroom.py          | 43 +++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
index 5fbf98d0aae..9d993384461 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
@@ -728,6 +728,10 @@ class HeadroomGuardrail(CustomGuardrail):
             verbose_proxy_logger.debug("Headroom: %s header set; skipping compression", BYPASS_HEADER)
             return inputs
 
+        if request_data.get("background"):
+            verbose_proxy_logger.debug("Headroom: background request; skipping compression")
+            return inputs
+
         structured_messages: Final = inputs.get("structured_messages")
         if not _is_object_list(structured_messages) or not structured_messages:
             return inputs
@@ -831,7 +835,7 @@ class HeadroomGuardrail(CustomGuardrail):
         effective: Final = base_result if base_result is not None else kwargs
         if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES:
             return base_result
-        if not effective.get("stream"):
+        if not effective.get("stream") or effective.get("background"):
             return base_result
         if not has_headroom_retrieve_tool(effective.get("tools")):
             return base_result
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
index a940e4c3394..400eaf8ab3f 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
@@ -193,6 +193,33 @@ async def test_apply_guardrail_compresses_and_returns_structured_messages(
     assert "headroom" in _applied_guardrails(request_data)
 
 
+@pytest.mark.asyncio
+async def test_apply_guardrail_leaves_background_requests_uncompressed(
+    guardrail: HeadroomGuardrail,
+):
+    inputs = GenericGuardrailAPIInputs(
+        texts=["A" * 5000],
+        structured_messages=ORIGINAL_MESSAGES,
+    )
+    request_data = {"model": "gpt-4o", "background": True}
+
+    with patch.object(
+        guardrail.async_handler,
+        "post",
+        new_callable=AsyncMock,
+        return_value=_make_compress_response(COMPRESSED_MESSAGES),
+    ) as post:
+        result = await guardrail.apply_guardrail(
+            inputs=inputs,
+            request_data=request_data,
+            input_type="request",
+        )
+
+    assert result is inputs
+    post.assert_not_awaited()
+    assert _recorded_guardrail_entries(request_data) == []
+
+
 def _recorded_guardrail_response(request_data: dict) -> dict:
     entries = request_data["metadata"]["standard_logging_guardrail_information"]
     assert len(entries) == 1
@@ -2074,6 +2101,22 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple
     assert kwargs["stream"] is True
 
 
+@pytest.mark.asyncio
+async def test_pre_call_deployment_hook_leaves_background_streams_alone(guardrail: HeadroomGuardrail):
+    kwargs = {
+        "model": "gpt-4o",
+        "stream": True,
+        "background": True,
+        "tools": [_responses_retrieve_tool_definition()],
+    }
+
+    result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.aresponses)
+
+    assert result is kwargs
+    assert HEADROOM_CONVERTED_STREAM_KEY not in kwargs
+    assert kwargs["stream"] is True
+
+
 @pytest.mark.asyncio
 async def test_pre_call_deployment_hook_still_compresses_for_deployment_level_configs(
     guardrail: HeadroomGuardrail,

From 7a5e4b9ba4f6d7d61122dd32b7e9c923cbd18bcf Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 10:04:18 -0700
Subject: [PATCH 17/43] fix(openai): bridge gpt-5.4+ tool calls to
 /v1/responses on every api.openai.com host

The auto-bridge that moves gpt-5.4+ requests carrying function tools and no
reasoning_effort onto /v1/responses only fired when the resolved api_base was
the literal https://api.openai.com/v1, so a deployment pointed at an OpenAI
PrivateLink hostname (<region>.privatelink.api.openai.com) or a port-qualified
or trailing-slash default stayed on Chat Completions and got OpenAI's 400 back.
Gate on the resolved URL's hostname instead: api.openai.com or any subdomain of
it bridges, every other custom base still stays on chat
---
 litellm/main.py                 | 26 +++++++----
 tests/test_litellm/test_main.py | 76 +++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+), 9 deletions(-)

diff --git a/litellm/main.py b/litellm/main.py
index 0bca4a7350e..6f8310a9d8a 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -26,6 +26,7 @@ from copy import deepcopy
 from functools import partial
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args
+from urllib.parse import urlsplit
 
 from litellm._logging import _redact_string
 from litellm._uuid import uuid
@@ -984,6 +985,12 @@ def mock_completion(
 
 
 _OPENAI_DEFAULT_API_BASE: Final = "https://api.openai.com/v1"
+_OPENAI_API_HOST: Final = "api.openai.com"
+
+
+def _is_openai_backed_api_base(api_base: str) -> bool:
+    hostname: Final = urlsplit(api_base).hostname
+    return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}"))
 
 
 def _resolve_openai_api_base(api_base: str | None) -> str:
@@ -1053,7 +1060,7 @@ def responses_api_bridge_check(
     #   natively by Chat Completions with reasoning on, so custom-only requests stay on
     #   chat and keep their native custom tool_call response shape.
     # - The UNSET-effort arm only fires against endpoints known to enforce that
-    #   constraint (the default OpenAI endpoint, or Azure OpenAI where api_base is
+    #   constraint (any api.openai.com host, or Azure OpenAI where api_base is
     #   always set): chat-only OpenAI-compatible backends registered under the openai
     #   provider with a custom api_base and gpt-5.4+ model names serve tools without
     #   reasoning fine and have no /responses route, so they keep pre-existing
@@ -1068,14 +1075,15 @@ def responses_api_bridge_check(
         reasoning_active = reasoning_effort.get("effort") != "none" or reasoning_effort.get("summary") is not None
     else:
         reasoning_active = reasoning_effort != "none"
-    # The reasoning+tools constraint is enforced only by the real OpenAI endpoint (and Azure OpenAI).
-    # Resolve the effective base arg>global>env>default exactly as the chat handler does, so a custom
-    # base set via litellm.api_base or OPENAI_BASE_URL/OPENAI_API_BASE isn't misread as the default and
-    # bridged to a /responses route it lacks. A whitespace-only base collapses to the default too.
-    resolved_api_base: Final = _resolve_openai_api_base(api_base)
-    on_constraint_enforcing_endpoint: Final = custom_llm_provider == "azure" or resolved_api_base.strip() in (
-        "",
-        _OPENAI_DEFAULT_API_BASE,
+    # The reasoning+tools constraint is enforced by the real OpenAI backend behind any api.openai.com
+    # host (the default URL or a PrivateLink hostname such as <region>.privatelink.api.openai.com) and
+    # by Azure OpenAI. Resolve the effective base arg>global>env>default exactly as the chat handler
+    # does, so a custom base set via litellm.api_base or OPENAI_BASE_URL/OPENAI_API_BASE isn't misread
+    # as the default and bridged to a /responses route it lacks. A whitespace-only base collapses to
+    # the default too.
+    resolved_api_base: Final = _resolve_openai_api_base(api_base).strip()
+    on_constraint_enforcing_endpoint: Final = (
+        custom_llm_provider == "azure" or resolved_api_base == "" or _is_openai_backed_api_base(resolved_api_base)
     )
     if (
         custom_llm_provider in ("openai", "azure")
diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py
index 3c8bf142835..4cba5048e5d 100644
--- a/tests/test_litellm/test_main.py
+++ b/tests/test_litellm/test_main.py
@@ -1133,6 +1133,82 @@ def test_responses_api_bridge_check_custom_api_base_via_env_with_unset_effort_st
     assert model_info.get("mode") != "responses"
 
 
+@pytest.mark.parametrize(
+    "api_base",
+    [
+        "https://southcentralus.privatelink.api.openai.com/v1",
+        "https://privatelink.corp.api.openai.com/v1",
+        "https://api.openai.com:443/v1",
+        "https://api.openai.com/v1/",
+        "HTTPS://API.OPENAI.COM/v1",
+    ],
+)
+def test_responses_api_bridge_check_openai_backed_custom_api_base_with_unset_effort_routes_to_responses(api_base):
+    """
+    A custom api_base whose host is api.openai.com or a subdomain of it (a PrivateLink hostname, a
+    port-qualified or trailing-slash default) still reaches the real OpenAI backend, which rejects
+    function tools with reasoning on Chat Completions, so the unset-effort arm must bridge exactly as
+    it does for the literal default URL. Regression guard for GH #39353.
+    """
+    from litellm.main import responses_api_bridge_check
+
+    model_info, model = responses_api_bridge_check(
+        model="gpt-5.6",
+        custom_llm_provider="openai",
+        tools=[{"type": "function", "function": {"name": "get_capital"}}],
+        reasoning_effort=None,
+        api_base=api_base,
+        )
+
+    assert model == "gpt-5.6"
+    assert model_info.get("mode") == "responses"
+
+
+@pytest.mark.parametrize(
+    "api_base",
+    [
+        "https://api.openai.com.evil.example/v1",
+        "https://notapi.openai.com/v1",
+        "https://gateway.example/v1?upstream=api.openai.com",
+        "https://openai.internal.example/api.openai.com/v1",
+    ],
+)
+def test_responses_api_bridge_check_lookalike_custom_api_base_with_unset_effort_stays_chat(api_base):
+    """Only the host decides: api.openai.com appearing elsewhere in the URL is still a foreign backend."""
+    from litellm.main import responses_api_bridge_check
+
+    model_info, model = responses_api_bridge_check(
+        model="gpt-5.6",
+        custom_llm_provider="openai",
+        tools=[{"type": "function", "function": {"name": "get_capital"}}],
+        reasoning_effort=None,
+        api_base=api_base,
+        )
+
+    assert model == "gpt-5.6"
+    assert model_info.get("mode") != "responses"
+
+
+def test_responses_api_bridge_check_privatelink_api_base_via_env_with_unset_effort_routes_to_responses(monkeypatch):
+    """A PrivateLink base set through OPENAI_BASE_URL resolves the way the chat handler's does and still bridges."""
+    import litellm
+    from litellm.main import responses_api_bridge_check
+
+    monkeypatch.setattr(litellm, "api_base", None)
+    monkeypatch.delenv("OPENAI_API_BASE", raising=False)
+    monkeypatch.setenv("OPENAI_BASE_URL", "https://southcentralus.privatelink.api.openai.com/v1")
+    model_info, model = responses_api_bridge_check(
+        model="gpt-5.6",
+        custom_llm_provider="openai",
+        tools=[{"type": "function", "function": {"name": "get_capital"}}],
+        reasoning_effort=None,
+        api_base=None,
+        )
+
+    assert model == "gpt-5.6"
+    assert model_info.get("mode") == "responses"
+
+
 def test_responses_api_bridge_check_custom_api_base_with_explicit_effort_still_routes():
     """Explicit reasoning_effort keeps its pre-existing bridging behavior on any api_base."""
     from litellm.main import responses_api_bridge_check

From 2fd6e190517b7ce050c603489556efdb9870137a Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Thu, 3 Sep 2026 10:20:10 -0700
Subject: [PATCH 18/43] fix(ui): name the granting access group on hover
 instead of an Inherited tag

`/team/info` access_group_details now carries mcp_server_ids and agent_ids per
group next to models, so the dashboard can say which group granted a server or
agent. The Object Permissions rows drop the Inherited badge and the row tooltip
reads "Granted via access group <name>. Full ID: <id>", listing every group
when more than one grants the same id and falling back to "an access group"
when the proxy did not say.

Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL
---
 litellm/proxy/_types.py                       |  2 +
 .../management_endpoints/team_endpoints.py    |  2 +
 .../test_team_endpoints.py                    |  6 +-
 .../components/object_permissions_view.tsx    | 13 ++--
 .../permissions/AgentPermissions.test.tsx     | 35 ++++++++---
 .../permissions/AgentPermissions.tsx          | 24 +++-----
 .../permissions/MCPServerPermissions.test.tsx | 22 ++++---
 .../permissions/MCPServerPermissions.tsx      | 30 ++++-----
 .../permissions/inheritedGrants.test.ts       | 61 +++++++++++++++++++
 .../components/permissions/inheritedGrants.ts | 26 ++++++++
 .../src/components/team/TeamInfo.test.tsx     | 26 ++++++--
 .../src/components/team/TeamInfo.tsx          | 20 ++++--
 .../src/components/team/teamModelAccess.ts    |  4 +-
 13 files changed, 204 insertions(+), 67 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/permissions/inheritedGrants.test.ts
 create mode 100644 ui/litellm-dashboard/src/components/permissions/inheritedGrants.ts

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index e0a2097919b..57f6f31d4d9 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -4172,6 +4172,8 @@ class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase):
     access_group_id: str
     access_group_name: str
     models: tuple[str, ...]
+    mcp_server_ids: tuple[str, ...] = ()
+    agent_ids: tuple[str, ...] = ()
 
 
 class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index 714cf252e69..72cdc75c29e 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -4318,6 +4318,8 @@ async def _resolve_team_access_group_resources(
                     access_group_id=group.access_group_id,
                     access_group_name=group.access_group_name,
                     models=tuple(group.access_model_names or ()),
+                    mcp_server_ids=tuple(group.access_mcp_server_ids or ()),
+                    agent_ids=tuple(group.access_agent_ids or ()),
                 )
                 for group in resolved_groups
             ),
diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
index 30b2ab86b9a..5bb7ceb8fb1 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
@@ -9897,11 +9897,11 @@ class TestResolveTeamAccessGroupResources:
         assert resolved.access_group_mcp_server_ids == ["mcp-1"]
         assert resolved.access_group_agent_ids == ["agent-1"]
         assert [
-            (d.access_group_id, d.access_group_name, d.models)
+            (d.access_group_id, d.access_group_name, d.models, d.mcp_server_ids, d.agent_ids)
             for d in (resolved.access_group_details or [])
         ] == [
-            ("ag-1", "shared-models", ("gpt-4", "claude-3")),
-            ("ag-2", "extra-models", ("claude-3", "gemini")),
+            ("ag-1", "shared-models", ("gpt-4", "claude-3"), ("mcp-1",), ()),
+            ("ag-2", "extra-models", ("claude-3", "gemini"), (), ("agent-1",)),
         ]
 
     @pytest.mark.asyncio
diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx
index 7d1c4ce4ac9..e90645a9e9f 100644
--- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx
+++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx
@@ -3,11 +3,12 @@ import VectorStorePermissions from "./permissions/VectorStorePermissions";
 import MCPServerPermissions from "./permissions/MCPServerPermissions";
 import AgentPermissions from "./permissions/AgentPermissions";
 import type { ObjectPermission } from "./object_permission_types";
+import type { InheritedGrant } from "./permissions/inheritedGrants";
 
 interface ObjectPermissionsViewProps {
   objectPermission?: ObjectPermission | null;
-  inheritedMcpServerIds?: string[];
-  inheritedAgentIds?: string[];
+  inheritedMcpServers?: InheritedGrant[];
+  inheritedAgents?: InheritedGrant[];
   variant?: "card" | "inline";
   className?: string;
   accessToken?: string | null;
@@ -15,8 +16,8 @@ interface ObjectPermissionsViewProps {
 
 export function ObjectPermissionsView({
   objectPermission,
-  inheritedMcpServerIds = [],
-  inheritedAgentIds = [],
+  inheritedMcpServers = [],
+  inheritedAgents = [],
   variant = "card",
   className = "",
   accessToken,
@@ -38,13 +39,13 @@ export function ObjectPermissionsView({
         mcpAccessGroups={mcpAccessGroups}
         mcpToolPermissions={mcpToolPermissions}
         mcpToolsets={mcpToolsets}
-        inheritedMcpServers={inheritedMcpServerIds}
+        inheritedMcpServers={inheritedMcpServers}
         accessToken={accessToken}
       />
       <AgentPermissions
         agents={agents}
         agentAccessGroups={agentAccessGroups}
-        inheritedAgents={inheritedAgentIds}
+        inheritedAgents={inheritedAgents}
         accessToken={accessToken}
       />
       <div className="min-w-0 rounded-md border border-border p-4">
diff --git a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx
index 14fe6177ac5..01c7b60b8a5 100644
--- a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx
+++ b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx
@@ -1,5 +1,6 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
 import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
 import AgentPermissions from "./AgentPermissions";
 import * as networking from "../networking";
 
@@ -13,31 +14,51 @@ describe("AgentPermissions", () => {
     vi.clearAllMocks();
   });
 
-  it("lists agents inherited from access groups with an Inherited tag and counts them", async () => {
+  it("lists agents inherited from access groups, counts them, and names the groups on hover", async () => {
+    const user = userEvent.setup();
     vi.mocked(networking.getAgentsList).mockResolvedValue({
       agents: [{ agent_id: agentId, agent_name: "support_agent" }],
     });
 
-    render(<AgentPermissions agents={[]} inheritedAgents={[agentId]} accessToken={accessToken} />);
+    render(
+      <AgentPermissions
+        agents={[]}
+        inheritedAgents={[{ id: agentId, accessGroupNames: ["platform-tools", "support"] }]}
+        accessToken={accessToken}
+      />,
+    );
 
-    expect(await screen.findByText(/support_agent/)).toBeInTheDocument();
-    expect(screen.getByText("Inherited")).toBeInTheDocument();
+    const row = await screen.findByText(/support_agent/);
     expect(screen.getByText("1")).toBeInTheDocument();
     expect(screen.queryByText("No agents or access groups configured")).not.toBeInTheDocument();
     expect(networking.getAgentsList).toHaveBeenCalledWith(accessToken);
+
+    await user.hover(row);
+    expect(
+      await screen.findByText(`Granted via access groups platform-tools, support. Full ID: ${agentId}`),
+    ).toBeInTheDocument();
   });
 
   it("does not double-list an agent that is both granted directly and inherited", async () => {
+    const user = userEvent.setup();
     vi.mocked(networking.getAgentsList).mockResolvedValue({
       agents: [{ agent_id: agentId, agent_name: "support_agent" }],
     });
 
-    render(<AgentPermissions agents={[agentId]} inheritedAgents={[agentId]} accessToken={accessToken} />);
+    render(
+      <AgentPermissions
+        agents={[agentId]}
+        inheritedAgents={[{ id: agentId, accessGroupNames: ["support"] }]}
+        accessToken={accessToken}
+      />,
+    );
 
-    expect(await screen.findByText(/support_agent/)).toBeInTheDocument();
+    const row = await screen.findByText(/support_agent/);
     expect(screen.getAllByText(/support_agent/)).toHaveLength(1);
-    expect(screen.queryByText("Inherited")).not.toBeInTheDocument();
     expect(screen.getByText("1")).toBeInTheDocument();
+
+    await user.hover(row);
+    expect(await screen.findByText(`Full ID: ${agentId}`)).toBeInTheDocument();
   });
 
   it("shows the empty state when nothing is granted directly or inherited", () => {
diff --git a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx
index 11f6c414922..d1ca25c975b 100644
--- a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx
+++ b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx
@@ -3,6 +3,7 @@ import { UserGroupIcon } from "@heroicons/react/outline";
 import { Badge } from "@/components/ui/badge";
 import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
 import { getAgentsList } from "../networking";
+import { InheritedGrant, inheritedGrantTooltip } from "./inheritedGrants";
 
 interface Agent {
   agent_id: string;
@@ -14,12 +15,10 @@ interface Agent {
 interface AgentPermissionsProps {
   agents: string[];
   agentAccessGroups?: string[];
-  inheritedAgents?: string[];
+  inheritedAgents?: InheritedGrant[];
   accessToken?: string | null;
 }
 
-const INHERITED_AGENT_TOOLTIP = "Granted through one of the team's access groups";
-
 export function AgentPermissions({
   agents,
   agentAccessGroups = [],
@@ -27,7 +26,7 @@ export function AgentPermissions({
   accessToken,
 }: AgentPermissionsProps) {
   const [agentDetails, setAgentDetails] = useState<Agent[]>([]);
-  const inheritedOnlyAgents = inheritedAgents.filter((agent) => !agents.includes(agent));
+  const inheritedOnlyAgents = inheritedAgents.filter((grant) => !agents.includes(grant.id));
   const agentIdCount = agents.length + inheritedOnlyAgents.length;
 
   // Fetch agent details when component mounts
@@ -58,9 +57,9 @@ export function AgentPermissions({
   };
 
   const mergedItems = [
-    ...agents.map((agent) => ({ type: "agent", value: agent, inherited: false })),
-    ...inheritedOnlyAgents.map((agent) => ({ type: "agent", value: agent, inherited: true })),
-    ...agentAccessGroups.map((group) => ({ type: "accessGroup", value: group, inherited: false })),
+    ...agents.map((agent) => ({ type: "agent", value: agent, tooltip: `Full ID: ${agent}` })),
+    ...inheritedOnlyAgents.map((grant) => ({ type: "agent", value: grant.id, tooltip: inheritedGrantTooltip(grant) })),
+    ...agentAccessGroups.map((group) => ({ type: "accessGroup", value: group, tooltip: "" })),
   ];
   const totalCount = mergedItems.length;
 
@@ -86,17 +85,8 @@ export function AgentPermissions({
                           <span className="text-sm font-medium text-foreground truncate">
                             {getAgentDisplayName(item.value)}
                           </span>
-                          {item.inherited && (
-                            <span className="ml-1 px-1.5 py-0.5 text-[9px] font-semibold text-purple-600 bg-purple-50 border border-purple-200 rounded-sm uppercase tracking-wide shrink-0 dark:text-purple-300 dark:bg-purple-950 dark:border-purple-800">
-                              Inherited
-                            </span>
-                          )}
                         </TooltipTrigger>
-                        <TooltipContent>
-                          {item.inherited
-                            ? `${INHERITED_AGENT_TOOLTIP}. Full ID: ${item.value}`
-                            : `Full ID: ${item.value}`}
-                        </TooltipContent>
+                        <TooltipContent>{item.tooltip}</TooltipContent>
                       </Tooltip>
                     </TooltipProvider>
                   ) : (
diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx
index fe7e3a7b971..78ede9865a4 100644
--- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx
+++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx
@@ -407,7 +407,8 @@ describe("MCPServerPermissions", () => {
     await waitFor(() => expect(screen.getByText("Blocked")).toHaveAttribute("data-variant", "destructive"));
   });
 
-  it("lists servers inherited from access groups with an Inherited tag and counts them", async () => {
+  it("lists servers inherited from access groups, counts them, and names the group on hover", async () => {
+    const user = userEvent.setup();
     vi.mocked(networking.fetchMCPServers).mockResolvedValue([
       { server_id: mockServerId1, server_name: mockServerName1, alias: mockServerName1 },
     ]);
@@ -417,19 +418,24 @@ describe("MCPServerPermissions", () => {
         mcpServers={[]}
         mcpAccessGroups={[]}
         mcpToolPermissions={{}}
-        inheritedMcpServers={[mockServerId1]}
+        inheritedMcpServers={[{ id: mockServerId1, accessGroupNames: ["platform-tools"] }]}
         accessToken={mockAccessToken}
       />,
     );
 
-    expect(await screen.findByText(/DW_MCP/)).toBeInTheDocument();
-    expect(screen.getByText("Inherited")).toBeInTheDocument();
+    const row = await screen.findByText(/DW_MCP/);
     expect(screen.getByText("1")).toBeInTheDocument();
     expect(screen.queryByText("No MCP servers, access groups, or toolsets configured")).not.toBeInTheDocument();
     expect(networking.fetchMCPServers).toHaveBeenCalledWith(mockAccessToken);
+
+    await user.hover(row);
+    expect(
+      await screen.findByText(`Granted via access group platform-tools. Full ID: ${mockServerId1}`),
+    ).toBeInTheDocument();
   });
 
   it("does not double-list a server that is both granted directly and inherited", async () => {
+    const user = userEvent.setup();
     vi.mocked(networking.fetchMCPServers).mockResolvedValue([
       { server_id: mockServerId2, server_name: mockServerName2, alias: mockServerName2 },
     ]);
@@ -439,14 +445,16 @@ describe("MCPServerPermissions", () => {
         mcpServers={[mockServerId2]}
         mcpAccessGroups={[]}
         mcpToolPermissions={{}}
-        inheritedMcpServers={[mockServerId2]}
+        inheritedMcpServers={[{ id: mockServerId2, accessGroupNames: ["platform-tools"] }]}
         accessToken={mockAccessToken}
       />,
     );
 
-    expect(await screen.findByText(/Test Server/)).toBeInTheDocument();
+    const row = await screen.findByText(/Test Server/);
     expect(screen.getAllByText(/Test Server/)).toHaveLength(1);
-    expect(screen.queryByText("Inherited")).not.toBeInTheDocument();
     expect(screen.getByText("1")).toBeInTheDocument();
+
+    await user.hover(row);
+    expect(await screen.findByText(`Full ID: ${mockServerId2}`)).toBeInTheDocument();
   });
 });
diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx
index 52199dab3c7..65f210addfa 100644
--- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx
+++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx
@@ -5,18 +5,17 @@ import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip
 import { fetchMCPServers, fetchMCPToolsets } from "../networking";
 import { MCPServer, MCPToolset } from "../mcp_tools/types";
 import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants";
+import { InheritedGrant, inheritedGrantTooltip } from "./inheritedGrants";
 
 interface MCPServerPermissionsProps {
   mcpServers: string[];
   mcpAccessGroups?: string[];
   mcpToolPermissions?: Record<string, string[]>;
   mcpToolsets?: string[];
-  inheritedMcpServers?: string[];
+  inheritedMcpServers?: InheritedGrant[];
   accessToken?: string | null;
 }
 
-const INHERITED_MCP_SERVER_TOOLTIP = "Granted through one of the team's access groups";
-
 export function MCPServerPermissions({
   mcpServers,
   mcpAccessGroups = [],
@@ -57,8 +56,8 @@ export function MCPServerPermissions({
   const directServerIds = mcpServers.filter(
     (server) => server !== NO_MCP_SERVERS_SENTINEL && server !== ALL_PROXY_MCP_SERVERS_SENTINEL,
   );
-  const inheritedOnlyServerIds = inheritedMcpServers.filter((server) => !mcpServers.includes(server));
-  const serverIdCount = directServerIds.length + inheritedOnlyServerIds.length;
+  const inheritedOnlyServers = inheritedMcpServers.filter((grant) => !mcpServers.includes(grant.id));
+  const serverIdCount = directServerIds.length + inheritedOnlyServers.length;
 
   // Fetch MCP server details when component mounts
   useEffect(() => {
@@ -109,9 +108,13 @@ export function MCPServerPermissions({
   const grantsAllProxyMcpServers = mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL);
 
   const mergedItems = [
-    ...directServerIds.map((server) => ({ type: "server", value: server, inherited: false })),
-    ...inheritedOnlyServerIds.map((server) => ({ type: "server", value: server, inherited: true })),
-    ...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group, inherited: false })),
+    ...directServerIds.map((server) => ({ type: "server", value: server, tooltip: `Full ID: ${server}` })),
+    ...inheritedOnlyServers.map((grant) => ({
+      type: "server",
+      value: grant.id,
+      tooltip: inheritedGrantTooltip(grant),
+    })),
+    ...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group, tooltip: "" })),
   ];
   const totalCount = mergedItems.length + mcpToolsets.length;
 
@@ -160,17 +163,8 @@ export function MCPServerPermissions({
                           <span className="text-sm font-medium text-foreground truncate">
                             {getMCPServerDisplayName(item.value)}
                           </span>
-                          {item.inherited && (
-                            <span className="ml-1 px-1.5 py-0.5 text-[9px] font-semibold text-info bg-info/10 border border-info/20 rounded-sm uppercase tracking-wide shrink-0">
-                              Inherited
-                            </span>
-                          )}
                         </TooltipTrigger>
-                        <TooltipContent>
-                          {item.inherited
-                            ? `${INHERITED_MCP_SERVER_TOOLTIP}. Full ID: ${item.value}`
-                            : `Full ID: ${item.value}`}
-                        </TooltipContent>
+                        <TooltipContent>{item.tooltip}</TooltipContent>
                       </Tooltip>
                     ) : (
                       <div className="inline-flex items-center gap-2 min-w-0">
diff --git a/ui/litellm-dashboard/src/components/permissions/inheritedGrants.test.ts b/ui/litellm-dashboard/src/components/permissions/inheritedGrants.test.ts
new file mode 100644
index 00000000000..766f8f23f0b
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/permissions/inheritedGrants.test.ts
@@ -0,0 +1,61 @@
+import { describe, expect, it } from "vitest";
+import { computeInheritedGrants, inheritedGrantTooltip } from "./inheritedGrants";
+import { TeamAccessGroupModelGrant } from "../team/teamModelAccess";
+
+const GRANTS: TeamAccessGroupModelGrant[] = [
+  { access_group_id: "ag-1", access_group_name: "platform-tools", models: [], mcp_server_ids: ["mcp-1", "mcp-2"] },
+  {
+    access_group_id: "ag-2",
+    access_group_name: "support",
+    models: [],
+    mcp_server_ids: ["mcp-2"],
+    agent_ids: ["agent-1"],
+  },
+];
+
+describe("computeInheritedGrants", () => {
+  it("attributes each id to every group that grants it, in group order", () => {
+    expect(computeInheritedGrants(["mcp-1", "mcp-2"], GRANTS, (g) => g.mcp_server_ids)).toEqual([
+      { id: "mcp-1", accessGroupNames: ["platform-tools"] },
+      { id: "mcp-2", accessGroupNames: ["platform-tools", "support"] },
+    ]);
+  });
+
+  it("keeps ids the flat list carries but no group detail explains, with no group names", () => {
+    expect(computeInheritedGrants(["agent-1", "agent-legacy"], GRANTS, (g) => g.agent_ids)).toEqual([
+      { id: "agent-1", accessGroupNames: ["support"] },
+      { id: "agent-legacy", accessGroupNames: [] },
+    ]);
+  });
+
+  it("falls back to the group details when the flat list is missing, without duplicates", () => {
+    expect(computeInheritedGrants(undefined, GRANTS, (g) => g.mcp_server_ids).map((g) => g.id)).toEqual([
+      "mcp-1",
+      "mcp-2",
+    ]);
+  });
+
+  it("returns nothing when neither source has ids", () => {
+    expect(computeInheritedGrants(undefined, undefined, (g) => g.agent_ids)).toEqual([]);
+  });
+});
+
+describe("inheritedGrantTooltip", () => {
+  it("names a single group", () => {
+    expect(inheritedGrantTooltip({ id: "mcp-1", accessGroupNames: ["platform-tools"] })).toBe(
+      "Granted via access group platform-tools. Full ID: mcp-1",
+    );
+  });
+
+  it("lists several groups", () => {
+    expect(inheritedGrantTooltip({ id: "mcp-2", accessGroupNames: ["platform-tools", "support"] })).toBe(
+      "Granted via access groups platform-tools, support. Full ID: mcp-2",
+    );
+  });
+
+  it("stays generic when the proxy did not say which group granted it", () => {
+    expect(inheritedGrantTooltip({ id: "agent-legacy", accessGroupNames: [] })).toBe(
+      "Granted via an access group. Full ID: agent-legacy",
+    );
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/permissions/inheritedGrants.ts b/ui/litellm-dashboard/src/components/permissions/inheritedGrants.ts
new file mode 100644
index 00000000000..fb78de5ed77
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/permissions/inheritedGrants.ts
@@ -0,0 +1,26 @@
+import { describeGroups, TeamAccessGroupModelGrant } from "../team/teamModelAccess";
+
+export interface InheritedGrant {
+  id: string;
+  accessGroupNames: string[];
+}
+
+export function computeInheritedGrants(
+  ids: string[] | undefined,
+  grants: TeamAccessGroupModelGrant[] | undefined,
+  idsOf: (grant: TeamAccessGroupModelGrant) => string[] | undefined,
+): InheritedGrant[] {
+  const known = grants ?? [];
+  const allIds = [...new Set([...(ids ?? []), ...known.flatMap((grant) => idsOf(grant) ?? [])])];
+  return allIds.map((id) => ({
+    id,
+    accessGroupNames: known
+      .filter((grant) => (idsOf(grant) ?? []).includes(id))
+      .map((grant) => grant.access_group_name),
+  }));
+}
+
+export const inheritedGrantTooltip = (grant: InheritedGrant): string => {
+  const source = grant.accessGroupNames.length > 0 ? describeGroups(grant.accessGroupNames) : "an access group";
+  return `Granted via ${source}. Full ID: ${grant.id}`;
+};
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
index 40a3d2ada1a..ae78ac06a9c 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
@@ -305,7 +305,8 @@ describe("TeamInfoView", () => {
       );
     });
 
-    it("shows MCP servers and agents inherited from access groups in the Object Permissions card", async () => {
+    it("shows MCP servers and agents inherited from access groups in the Object Permissions card, naming the group on hover", async () => {
+      const user = userEvent.setup();
       vi.mocked(networking.fetchMCPServers).mockResolvedValue([
         { server_id: "mcp-github-1234", server_name: "github", alias: "github" },
       ]);
@@ -318,16 +319,33 @@ describe("TeamInfoView", () => {
           access_group_ids: ["ag-1"],
           access_group_mcp_server_ids: ["mcp-github-1234"],
           access_group_agent_ids: ["agent-support-5678"],
+          access_group_details: [
+            {
+              access_group_id: "ag-1",
+              access_group_name: "platform-tools",
+              models: [],
+              mcp_server_ids: ["mcp-github-1234"],
+              agent_ids: ["agent-support-5678"],
+            },
+          ],
         }),
       );
 
       renderWithProviders(<TeamInfoView {...defaultProps} />);
 
-      expect(await screen.findByText(/github \(mcp\.\.\.1234\)/)).toBeInTheDocument();
-      expect(await screen.findByText(/support_agent \(age\.\.\.5678\)/)).toBeInTheDocument();
-      expect(screen.getAllByText("Inherited")).toHaveLength(2);
+      const serverRow = await screen.findByText(/github \(mcp\.\.\.1234\)/);
+      const agentRow = await screen.findByText(/support_agent \(age\.\.\.5678\)/);
       expect(screen.queryByText("No MCP servers, access groups, or toolsets configured")).not.toBeInTheDocument();
       expect(screen.queryByText("No agents or access groups configured")).not.toBeInTheDocument();
+
+      await user.hover(serverRow);
+      expect(
+        await screen.findByText("Granted via access group platform-tools. Full ID: mcp-github-1234"),
+      ).toBeInTheDocument();
+      await user.hover(agentRow);
+      expect(
+        await screen.findByText("Granted via access group platform-tools. Full ID: agent-support-5678"),
+      ).toBeInTheDocument();
     });
 
     it("keeps the all-proxy-models badge non-clickable", async () => {
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
index 7ebef0691e6..3f6d6a96972 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
@@ -58,6 +58,7 @@ import {
   TeamModelBadge,
   TeamModelBadgeKind,
 } from "./teamModelAccess";
+import { computeInheritedGrants } from "../permissions/inheritedGrants";
 import MetadataKeyValueFields, {
   metadataObjectToPairs,
   metadataPairsSchema,
@@ -936,6 +937,17 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
 
   const { team_info: info } = teamData;
 
+  const inheritedMcpServers = computeInheritedGrants(
+    info.access_group_mcp_server_ids,
+    info.access_group_details,
+    (grant) => grant.mcp_server_ids,
+  );
+  const inheritedAgents = computeInheritedGrants(
+    info.access_group_agent_ids,
+    info.access_group_details,
+    (grant) => grant.agent_ids,
+  );
+
   const initialKillSwitchOn = info.metadata?.disable_global_guardrails === true;
 
   const allGuardrails: GuardrailListItem[] = guardrailsData?.guardrails ?? [];
@@ -1035,8 +1047,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
 
           <ObjectPermissionsView
             objectPermission={info.object_permission}
-            inheritedMcpServerIds={info.access_group_mcp_server_ids}
-            inheritedAgentIds={info.access_group_agent_ids}
+            inheritedMcpServers={inheritedMcpServers}
+            inheritedAgents={inheritedAgents}
             variant="card"
             accessToken={accessToken}
           />
@@ -1889,8 +1901,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
 
               <ObjectPermissionsView
                 objectPermission={info.object_permission}
-                inheritedMcpServerIds={info.access_group_mcp_server_ids}
-                inheritedAgentIds={info.access_group_agent_ids}
+                inheritedMcpServers={inheritedMcpServers}
+                inheritedAgents={inheritedAgents}
                 variant="inline"
                 className="pt-4 border-t border-border"
                 accessToken={accessToken}
diff --git a/ui/litellm-dashboard/src/components/team/teamModelAccess.ts b/ui/litellm-dashboard/src/components/team/teamModelAccess.ts
index 91ddf0f4045..bc8022592bd 100644
--- a/ui/litellm-dashboard/src/components/team/teamModelAccess.ts
+++ b/ui/litellm-dashboard/src/components/team/teamModelAccess.ts
@@ -5,6 +5,8 @@ export interface TeamAccessGroupModelGrant {
   access_group_id: string;
   access_group_name: string;
   models: string[];
+  mcp_server_ids?: string[];
+  agent_ids?: string[];
 }
 
 export type TeamModelBadgeKind = "all-proxy" | "no-default" | "direct" | "access-group";
@@ -19,7 +21,7 @@ export function normalizeTeamModelSelection(models: string[] | undefined): strin
   return models && models.length > 0 ? models : [NO_DEFAULT_MODELS];
 }
 
-const describeGroups = (names: string[]): string =>
+export const describeGroups = (names: string[]): string =>
   names.length > 1 ? `access groups ${names.join(", ")}` : `access group ${names[0]}`;
 
 export function computeTeamModelBadges(

From b9e030ddd662cb7abb98b8fa5139efb14d57cb6c Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 10:38:25 -0700
Subject: [PATCH 19/43] fix(cost): apply off_peak_pricing in the dashscope cost
 calculator

---
 .../litellm_core_utils/llm_cost_calc/utils.py |   4 +-
 litellm/llms/dashscope/cost_calculator.py     | 120 ++++++++-------
 .../test_dashscope_cost_calculator.py         | 138 +++++++++++++++++-
 3 files changed, 209 insertions(+), 53 deletions(-)

diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py
index b34c416cd40..21587af73aa 100644
--- a/litellm/litellm_core_utils/llm_cost_calc/utils.py
+++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py
@@ -428,7 +428,7 @@ def _coerce_off_peak_rate(value: object, default: float) -> float:
     return default
 
 
-def _apply_off_peak_pricing(
+def apply_off_peak_pricing(
     model_info: ModelInfo,
     current_time: datetime | None,
     prompt_base_cost: float,
@@ -462,7 +462,7 @@ def _apply_off_peak_to_base_costs(
     has no field for them.
     """
     prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs
-    off_peak_prompt, off_peak_completion, off_peak_cache_read = _apply_off_peak_pricing(
+    off_peak_prompt, off_peak_completion, off_peak_cache_read = apply_off_peak_pricing(
         model_info, current_time, prompt, completion, cache_read
     )
     return (off_peak_prompt, off_peak_completion, cache_creation, cache_creation_above_1hr, off_peak_cache_read)
diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py
index dd5bee1fe8b..d8eb1f9f8d7 100644
--- a/litellm/llms/dashscope/cost_calculator.py
+++ b/litellm/llms/dashscope/cost_calculator.py
@@ -7,11 +7,13 @@ cached, cache-creation, output, reasoning) is billed at that one tier's rate.
 See https://help.aliyun.com/zh/model-studio/billing-for-model-studio
 """
 
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
+from datetime import datetime
 from typing import Final
 
 from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate
 from litellm.litellm_core_utils.llm_cost_calc.utils import (
+    apply_off_peak_pricing,
     parse_completion_tokens_details,
     parse_prompt_tokens_details,
 )
@@ -32,6 +34,19 @@ class TokenBreakdown:
         return self.text_tokens + self.cached_tokens + self.cache_creation_tokens
 
 
+@dataclass(frozen=True, slots=True)
+class TokenRates:
+    input_rate: float
+    cache_read_rate: float
+    cache_creation_rate: float
+    output_rate: float
+    reasoning_rate: float | None
+
+    @property
+    def billed_reasoning_rate(self) -> float:
+        return self.output_rate if self.reasoning_rate is None else self.reasoning_rate
+
+
 def _extract_token_breakdown(usage: Usage) -> TokenBreakdown:
     prompt_details: Final = parse_prompt_tokens_details(usage)
     cached_tokens: Final = prompt_details["cache_hit_tokens"]
@@ -57,69 +72,75 @@ def _flat_rate(model_info: ModelInfo, cost_key: str, fallback_cost_key: str) ->
     return float(value)
 
 
-def _calculate_prompt_cost(
-    breakdown: TokenBreakdown,
-    model_info: ModelInfo,
-    tier: dict | None,
-) -> float:
-    if tier is not None:
-        return (
-            (breakdown.text_tokens * tier_rate(tier, "input_cost_per_token"))
-            + (breakdown.cached_tokens * tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token"))
-            + (
-                breakdown.cache_creation_tokens
-                * tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token")
-            )
-        )
-
-    input_cost: Final = float(model_info.get("input_cost_per_token") or 0.0)
-    cache_read_cost: Final = _flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token")
-    cache_creation_cost: Final = _flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token")
-
-    return (
-        (breakdown.text_tokens * input_cost)
-        + (breakdown.cached_tokens * cache_read_cost)
-        + (breakdown.cache_creation_tokens * cache_creation_cost)
+def _flat_rates(model_info: ModelInfo) -> TokenRates:
+    reasoning_rate: Final = model_info.get("output_cost_per_reasoning_token")
+    return TokenRates(
+        input_rate=float(model_info.get("input_cost_per_token") or 0.0),
+        cache_read_rate=_flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token"),
+        cache_creation_rate=_flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token"),
+        output_rate=float(model_info.get("output_cost_per_token") or 0.0),
+        reasoning_rate=None if reasoning_rate is None else float(reasoning_rate),
     )
 
 
-def _calculate_completion_cost(
-    breakdown: TokenBreakdown,
-    model_info: ModelInfo,
-    tier: dict | None,
-) -> float:
+def _tier_rates(model_info: ModelInfo, tier: dict) -> TokenRates:
     # A tier that declares output rates keeps the request on them, all-or-nothing. A tier table
     # spelling out only input rates would serve every completion for free, so there the model's
     # own output rates stand in
-    tier_declares_output: Final = tier is not None and "output_cost_per_token" in tier
-    output_cost: Final = (
-        tier_rate(tier, "output_cost_per_token")
-        if tier_declares_output
-        else float(model_info.get("output_cost_per_token") or 0.0)
-    )
-    tier_declares_reasoning: Final = tier is not None and "output_cost_per_reasoning_token" in tier
-    model_reasoning_rate: Final = None if tier_declares_output else model_info.get("output_cost_per_reasoning_token")
-    reasoning_cost: Final = (
-        tier_rate(tier, "output_cost_per_reasoning_token", "output_cost_per_token")
-        if tier_declares_reasoning
-        else float(model_reasoning_rate)
-        if model_reasoning_rate is not None
-        else output_cost
+    flat_rates: Final = _flat_rates(model_info)
+    tier_declares_output: Final = "output_cost_per_token" in tier
+    tier_declares_reasoning: Final = "output_cost_per_reasoning_token" in tier
+    return TokenRates(
+        input_rate=tier_rate(tier, "input_cost_per_token"),
+        cache_read_rate=tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token"),
+        cache_creation_rate=tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token"),
+        output_rate=tier_rate(tier, "output_cost_per_token") if tier_declares_output else flat_rates.output_rate,
+        reasoning_rate=(
+            tier_rate(tier, "output_cost_per_reasoning_token")
+            if tier_declares_reasoning
+            else None
+            if tier_declares_output
+            else flat_rates.reasoning_rate
+        ),
     )
 
-    return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost)
+
+def _off_peak_rates(model_info: ModelInfo, current_time: datetime | None, rates: TokenRates) -> TokenRates:
+    input_rate, output_rate, cache_read_rate = apply_off_peak_pricing(
+        model_info, current_time, rates.input_rate, rates.output_rate, rates.cache_read_rate
+    )
+    return replace(rates, input_rate=input_rate, output_rate=output_rate, cache_read_rate=cache_read_rate)
 
 
-def cost_per_token(model: str, usage: Usage, custom_llm_provider: str = "dashscope") -> tuple[float, float]:
+def _bill(breakdown: TokenBreakdown, rates: TokenRates) -> tuple[float, float]:
+    prompt_cost: Final = (
+        (breakdown.text_tokens * rates.input_rate)
+        + (breakdown.cached_tokens * rates.cache_read_rate)
+        + (breakdown.cache_creation_tokens * rates.cache_creation_rate)
+    )
+    completion_cost: Final = (breakdown.completion_tokens * rates.output_rate) + (
+        breakdown.reasoning_tokens * rates.billed_reasoning_rate
+    )
+    return prompt_cost, completion_cost
+
+
+def cost_per_token(
+    model: str,
+    usage: Usage,
+    custom_llm_provider: str = "dashscope",
+    current_time: datetime | None = None,
+) -> tuple[float, float]:
     """
     Calculate cost per token for Dashscope models.
 
-    Supports both tiered and flat pricing with cached and reasoning tokens.
+    Supports both tiered and flat pricing with cached and reasoning tokens, and swaps in the
+    model's off_peak_pricing rates while one of its windows is open.
 
     Args:
         model: Model name without provider prefix
         usage: LiteLLM Usage block
         custom_llm_provider: The provider id the request resolved to; dashscope or one of its brand aliases
+        current_time: The moment the request is billed at; defaults to now, UTC
 
     Returns:
         Tuple[float, float] - (prompt_cost_in_usd, completion_cost_in_usd)
@@ -133,8 +154,7 @@ def cost_per_token(model: str, usage: Usage, custom_llm_provider: str = "dashsco
         if tiered_pricing
         else None
     )
+    standard_rates: Final = _flat_rates(model_info) if tier is None else _tier_rates(model_info, tier)
+    rates: Final = _off_peak_rates(model_info, current_time, standard_rates)
 
-    prompt_cost: Final = _calculate_prompt_cost(breakdown=breakdown, model_info=model_info, tier=tier)
-    completion_cost: Final = _calculate_completion_cost(breakdown=breakdown, model_info=model_info, tier=tier)
-
-    return prompt_cost, completion_cost
+    return _bill(breakdown, rates)
diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py
index 8dc4620dd1b..b6281834f24 100644
--- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py
+++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py
@@ -10,11 +10,11 @@ Tests the cost calculation for Dashscope models including:
 
 import math
 import os
+from datetime import datetime, timezone
 
 import pytest
 
 # Add the project root to Python path
-
 import litellm
 from litellm.llms.dashscope.cost_calculator import (
     cost_per_token as dashscope_cost_per_token,
@@ -526,3 +526,139 @@ class TestDashscopeCostCalculator:
 
         assert prompt_cost == 0.0
         assert math.isclose(completion_cost, 500 * 1.6e-06, rel_tol=1e-10)
+
+    OFF_PEAK_WINDOW = "14:00-00:00"
+    INSIDE_WINDOW = datetime(2026, 9, 3, 17, 25, tzinfo=timezone.utc)
+    OUTSIDE_WINDOW = datetime(2026, 9, 3, 9, 0, tzinfo=timezone.utc)
+
+    def _register_off_peak_flat_model(self, model_key: str, off_peak_pricing: dict) -> None:
+        litellm.model_cost[model_key] = {
+            "litellm_provider": "dashscope",
+            "mode": "chat",
+            "input_cost_per_token": 2.4e-06,
+            "output_cost_per_token": 4.8e-06,
+            "cache_read_input_token_cost": 2e-07,
+            "cache_creation_input_token_cost": 3e-06,
+            "off_peak_pricing": off_peak_pricing,
+        }
+
+    def test_dashscope_off_peak_window_swaps_in_the_off_peak_rates(self):
+        """
+        Regression (LIT-6782): a deployment configured with off_peak_pricing kept billing the
+        standard dashscope rates inside its window, while the same block on a deepseek
+        deployment billed the off-peak rates.
+        """
+        self._register_off_peak_flat_model(
+            "dashscope/deepseek-off-peak-test",
+            {
+                "hours_utc": self.OFF_PEAK_WINDOW,
+                "input_cost_per_token": 1.2e-06,
+                "output_cost_per_token": 2.4e-06,
+                "cache_read_input_token_cost": 1e-07,
+            },
+        )
+        usage = Usage(
+            prompt_tokens=1000,
+            completion_tokens=200,
+            prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300, cache_creation_tokens=100),
+        )
+
+        prompt_cost, completion_cost = dashscope_cost_per_token(
+            model="deepseek-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW
+        )
+
+        assert math.isclose(prompt_cost, (600 * 1.2e-06) + (300 * 1e-07) + (100 * 3e-06), rel_tol=1e-10)
+        assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10)
+
+        peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token(
+            model="deepseek-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW
+        )
+
+        assert math.isclose(peak_prompt_cost, (600 * 2.4e-06) + (300 * 2e-07) + (100 * 3e-06), rel_tol=1e-10)
+        assert math.isclose(peak_completion_cost, 200 * 4.8e-06, rel_tol=1e-10)
+
+    def test_dashscope_off_peak_window_overrides_the_selected_tier(self):
+        """An open off-peak window bills the whole request at the flat off-peak rates, whichever tier
+        the input volume selected."""
+        self._register_tiered_model(
+            "dashscope/qwen-tiered-off-peak-test",
+            [
+                {"range": [0, 1000], "input_cost_per_token": 4e-07, "output_cost_per_token": 1.6e-06},
+                {"range": [1000, 2000], "input_cost_per_token": 8e-07, "output_cost_per_token": 3.2e-06},
+            ],
+        )
+        litellm.model_cost["dashscope/qwen-tiered-off-peak-test"]["off_peak_pricing"] = {
+            "hours_utc": self.OFF_PEAK_WINDOW,
+            "input_cost_per_token": 1e-07,
+            "output_cost_per_token": 4e-07,
+        }
+        usage = Usage(prompt_tokens=1500, completion_tokens=300)
+
+        prompt_cost, completion_cost = dashscope_cost_per_token(
+            model="qwen-tiered-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW
+        )
+
+        assert math.isclose(prompt_cost, 1500 * 1e-07, rel_tol=1e-10)
+        assert math.isclose(completion_cost, 300 * 4e-07, rel_tol=1e-10)
+
+        peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token(
+            model="qwen-tiered-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW
+        )
+
+        assert math.isclose(peak_prompt_cost, 1500 * 8e-07, rel_tol=1e-10)
+        assert math.isclose(peak_completion_cost, 300 * 3.2e-06, rel_tol=1e-10)
+
+    def test_dashscope_off_peak_rates_left_unset_keep_the_standard_rates(self):
+        """A block that only overrides the input rate leaves output and cache reads on the standard
+        rates, and an explicit reasoning rate is never swapped out."""
+        self._register_off_peak_flat_model(
+            "dashscope/qwen-partial-off-peak-test",
+            {"hours_utc": self.OFF_PEAK_WINDOW, "input_cost_per_token": 1.2e-06},
+        )
+        litellm.model_cost["dashscope/qwen-partial-off-peak-test"]["output_cost_per_reasoning_token"] = 9e-06
+        usage = Usage(
+            prompt_tokens=1000,
+            completion_tokens=200,
+            prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300),
+            completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50),
+        )
+
+        prompt_cost, completion_cost = dashscope_cost_per_token(
+            model="qwen-partial-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW
+        )
+
+        assert math.isclose(prompt_cost, (700 * 1.2e-06) + (300 * 2e-07), rel_tol=1e-10)
+        assert math.isclose(completion_cost, (150 * 4.8e-06) + (50 * 9e-06), rel_tol=1e-10)
+
+    def test_dashscope_off_peak_output_rate_covers_reasoning_without_a_dedicated_rate(self):
+        """Reasoning tokens on a model with no dedicated reasoning rate follow the off-peak output
+        rate, the same way they follow the standard output rate outside the window."""
+        self._register_off_peak_flat_model(
+            "dashscope/qwen-reasoning-off-peak-test",
+            {"hours_utc": self.OFF_PEAK_WINDOW, "output_cost_per_token": 2.4e-06},
+        )
+        usage = Usage(
+            prompt_tokens=100,
+            completion_tokens=200,
+            completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50),
+        )
+
+        _, completion_cost = dashscope_cost_per_token(
+            model="qwen-reasoning-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW
+        )
+
+        assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10)
+
+    def test_dashscope_off_peak_defaults_to_the_current_time(self):
+        """The proxy's cost dispatch passes no clock, so an all-day window has to apply on the
+        default current time."""
+        self._register_off_peak_flat_model(
+            "dashscope/qwen-all-day-off-peak-test",
+            {"hours_utc": "00:00-00:00", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 2.4e-06},
+        )
+        usage = Usage(prompt_tokens=1000, completion_tokens=200)
+
+        prompt_cost, completion_cost = dashscope_cost_per_token(model="qwen-all-day-off-peak-test", usage=usage)
+
+        assert math.isclose(prompt_cost, 1000 * 1.2e-06, rel_tol=1e-10)
+        assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10)

From 68ffa1db230ee2a03e851f567ce03ec48ce83e9b Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 10:43:41 -0700
Subject: [PATCH 20/43] fix(router): pin JWT-authenticated callers by user id
 in deployment_affinity

---
 .../deployment_affinity_check.py              |  36 +--
 .../test_deployment_affinity_check.py         | 207 ++++++++++++++++++
 2 files changed, 225 insertions(+), 18 deletions(-)

diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py
index 39d3e25aacb..ea450864604 100644
--- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py
+++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py
@@ -82,6 +82,7 @@ class DeploymentAffinityCheck(CustomLogger):
     """
 
     CACHE_KEY_PREFIX = "deployment_affinity:v1"
+    USER_ID_AFFINITY_PREFIX: Final = "user_id:"
 
     def __init__(
         self,
@@ -253,15 +254,6 @@ class DeploymentAffinityCheck(CustomLogger):
         hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped"
         return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}"
 
-    @staticmethod
-    def _get_user_key_from_metadata_dict(metadata: dict) -> str | None:
-        # NOTE: affinity is keyed on the *API key hash* provided by the proxy (not the
-        # OpenAI `user` parameter, which is an end-user identifier).
-        user_key: Final = metadata.get("user_api_key_hash")
-        if user_key is None:
-            return None
-        return str(user_key)
-
     @staticmethod
     def _get_session_id_from_metadata_dict(metadata: dict) -> str | None:
         session_id: Final = metadata.get("session_id")
@@ -285,22 +277,30 @@ class DeploymentAffinityCheck(CustomLogger):
         return metadata_dicts
 
     @staticmethod
-    def _get_user_key_from_request_kwargs(request_kwargs: dict) -> str | None:
+    def _first_metadata_value(metadata_dicts: Sequence[dict], key: str) -> str | None:
+        value: Final = next((metadata[key] for metadata in metadata_dicts if metadata.get(key) is not None), None)
+        return None if value is None else str(value)
+
+    @classmethod
+    def _get_user_key_from_request_kwargs(cls, request_kwargs: dict) -> str | None:
         """
         Extract a stable affinity key from request kwargs.
 
-        Source (proxy): `metadata.user_api_key_hash`
+        Source (proxy): `metadata.user_api_key_hash` for virtual-key callers. JWT-authenticated
+        callers carry no key hash, so their `metadata.user_api_key_user_id` stands in for it,
+        namespaced under `USER_ID_AFFINITY_PREFIX` so a user id can never alias a key hash.
 
         Note: the OpenAI `user` parameter is an end-user identifier and is intentionally
         not used for deployment affinity.
         """
-        # Check metadata dicts (Proxy usage)
-        for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs):
-            user_key = DeploymentAffinityCheck._get_user_key_from_metadata_dict(metadata=metadata)
-            if user_key is not None:
-                return user_key
-
-        return None
+        metadata_dicts: Final = cls._iter_metadata_dicts(request_kwargs)
+        user_api_key_hash: Final = cls._first_metadata_value(metadata_dicts, "user_api_key_hash")
+        if user_api_key_hash is not None:
+            return user_api_key_hash
+        user_id: Final = cls._first_metadata_value(metadata_dicts, "user_api_key_user_id")
+        if user_id is None:
+            return None
+        return f"{cls.USER_ID_AFFINITY_PREFIX}{user_id}"
 
     @staticmethod
     def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None:
diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py
index 60433921de6..1852d641f0a 100644
--- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py
+++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py
@@ -998,3 +998,210 @@ async def test_model_group_affinity_config_overrides_global():
     )
     # All deployments returned (user-key affinity disabled for this group)
     assert len(filtered) == 2
+
+
+def _jwt_metadata(user_id: str) -> dict:
+    return {"user_api_key_hash": None, "user_api_key_user_id": user_id}
+
+
+def _two_deployments(model_group: str) -> list[dict]:
+    return [
+        {
+            "model_name": model_group,
+            "litellm_params": {"model": "openai/gpt-5.4-mini"},
+            "model_info": {"id": "openai-deployment-a"},
+        },
+        {
+            "model_name": model_group,
+            "litellm_params": {"model": "openai/gpt-5.4-mini"},
+            "model_info": {"id": "openai-deployment-b"},
+        },
+    ]
+
+
+@pytest.mark.asyncio
+async def test_async_jwt_user_affinity_routes_to_same_deployment():
+    """
+    JWT-authenticated proxy requests carry no `user_api_key_hash`, only `user_api_key_user_id`.
+    They must still pin to one deployment per user.
+    """
+    model_group = "gpt-5.4-mini"
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": model_group,
+                "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock-api-key-a"},
+                "model_info": {"id": "openai-deployment-a"},
+            },
+            {
+                "model_name": model_group,
+                "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock-api-key-b"},
+                "model_info": {"id": "openai-deployment-b"},
+            },
+        ],
+        optional_pre_call_checks=["deployment_affinity"],
+    )
+
+    choice_calls = {"count": 0}
+
+    def deterministic_choice(seq):
+        choice_calls["count"] += 1
+        if choice_calls["count"] == 1:
+            return seq[0]
+        return seq[1] if len(seq) > 1 else seq[0]
+
+    with patch(  # test-quality-ok: simple-shuffle has no injectable RNG; forcing the other pick is what proves the pin overrides the strategy
+        "litellm.router_strategy.simple_shuffle.random.choice",
+        side_effect=deterministic_choice,
+    ):
+        first_response = await router.acompletion(
+            model=model_group,
+            messages=[{"role": "user", "content": "Reply with the single word ok"}],
+            mock_response="ok",
+            metadata=_jwt_metadata("jwt-user-alice"),
+        )
+        second_response = await router.acompletion(
+            model=model_group,
+            messages=[{"role": "user", "content": "Reply with the single word ok"}],
+            mock_response="ok",
+            metadata=_jwt_metadata("jwt-user-alice"),
+        )
+
+    first_model_id = first_response._hidden_params["model_id"]
+    assert first_model_id in ("openai-deployment-a", "openai-deployment-b")
+    assert second_response._hidden_params["model_id"] == first_model_id
+
+
+@pytest.mark.asyncio
+async def test_proxy_jwt_auth_metadata_pins_per_user():
+    """
+    The metadata the proxy stamps for a JWT caller (`UserAPIKeyAuth(api_key=None, user_id=<sub>)`)
+    must claim a pin and be read back by the filter, and another JWT user must not inherit it.
+    """
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
+
+    model_group = "gpt-5.4-mini"
+    healthy_deployments = _two_deployments(model_group)
+    callback = DeploymentAffinityCheck(
+        cache=DualCache(),
+        ttl_seconds=60,
+        enable_user_key_affinity=True,
+        enable_responses_api_affinity=False,
+    )
+
+    def proxy_request(user_id: str) -> dict:
+        return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
+            data={"model": model_group, "messages": [{"role": "user", "content": "hi"}], "metadata": {}},
+            user_api_key_dict=UserAPIKeyAuth(api_key=None, user_id=user_id),
+            _metadata_variable_name="metadata",
+        )
+
+    alice_request = proxy_request("jwt-user-alice")
+    assert alice_request["metadata"]["user_api_key_hash"] is None
+
+    await callback.async_pre_call_deployment_hook(
+        kwargs={
+            **alice_request,
+            "metadata": {**alice_request["metadata"], "deployment_model_name": model_group},
+            "model_info": {"id": "openai-deployment-b"},
+        },
+        call_type=None,
+    )
+
+    alice_pinned = await callback.async_filter_deployments(
+        model=model_group,
+        healthy_deployments=healthy_deployments,
+        messages=None,
+        request_kwargs=alice_request,
+        parent_otel_span=None,
+    )
+    assert [deployment["model_info"]["id"] for deployment in alice_pinned] == ["openai-deployment-b"]
+
+    bob_filtered = await callback.async_filter_deployments(
+        model=model_group,
+        healthy_deployments=healthy_deployments,
+        messages=None,
+        request_kwargs=proxy_request("jwt-user-bob"),
+        parent_otel_span=None,
+    )
+    assert bob_filtered == healthy_deployments
+
+
+@pytest.mark.asyncio
+async def test_jwt_user_id_never_reads_a_virtual_key_pin():
+    """
+    A JWT user id that happens to equal a virtual key's 64-hex hash must not read that key's pin.
+    """
+    model_group = "gpt-5.4-mini"
+    healthy_deployments = _two_deployments(model_group)
+    callback = DeploymentAffinityCheck(
+        cache=DualCache(),
+        ttl_seconds=60,
+        enable_user_key_affinity=True,
+        enable_responses_api_affinity=False,
+    )
+    key_hash = "a" * 64
+
+    await callback.async_pre_call_deployment_hook(
+        kwargs={
+            "metadata": {"user_api_key_hash": key_hash, "deployment_model_name": model_group},
+            "model_info": {"id": "openai-deployment-b"},
+        },
+        call_type=None,
+    )
+
+    key_pinned = await callback.async_filter_deployments(
+        model=model_group,
+        healthy_deployments=healthy_deployments,
+        messages=None,
+        request_kwargs={"metadata": {"user_api_key_hash": key_hash}},
+        parent_otel_span=None,
+    )
+    assert [deployment["model_info"]["id"] for deployment in key_pinned] == ["openai-deployment-b"]
+
+    lookalike_jwt_user = await callback.async_filter_deployments(
+        model=model_group,
+        healthy_deployments=healthy_deployments,
+        messages=None,
+        request_kwargs={"metadata": _jwt_metadata(key_hash)},
+        parent_otel_span=None,
+    )
+    assert lookalike_jwt_user == healthy_deployments
+
+
+@pytest.mark.asyncio
+async def test_virtual_key_hash_wins_over_user_id_for_affinity():
+    """
+    A virtual-key caller with a user id pins on the key hash, so two keys owned by one user
+    keep independent pins.
+    """
+    model_group = "gpt-5.4-mini"
+    healthy_deployments = _two_deployments(model_group)
+    callback = DeploymentAffinityCheck(
+        cache=DualCache(),
+        ttl_seconds=60,
+        enable_user_key_affinity=True,
+        enable_responses_api_affinity=False,
+    )
+
+    await callback.async_pre_call_deployment_hook(
+        kwargs={
+            "metadata": {
+                "user_api_key_hash": "key-one",
+                "user_api_key_user_id": "shared-user",
+                "deployment_model_name": model_group,
+            },
+            "model_info": {"id": "openai-deployment-b"},
+        },
+        call_type=None,
+    )
+
+    other_key_same_user = await callback.async_filter_deployments(
+        model=model_group,
+        healthy_deployments=healthy_deployments,
+        messages=None,
+        request_kwargs={"metadata": {"user_api_key_hash": "key-two", "user_api_key_user_id": "shared-user"}},
+        parent_otel_span=None,
+    )
+    assert other_key_same_user == healthy_deployments

From 1b42b81f4e048db9403c967365f15165f07a0dd4 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 10:50:21 -0700
Subject: [PATCH 21/43] fix(router): log the hashed affinity key so JWT callers
 stay distinguishable

---
 .../pre_call_checks/deployment_affinity_check.py            | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py
index ea450864604..6f3ea8eb78a 100644
--- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py
+++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py
@@ -533,9 +533,9 @@ class DeploymentAffinityCheck(CustomLogger):
             return typed_healthy_deployments
 
         verbose_router_logger.debug(
-            "DeploymentAffinityCheck: api-key affinity hit -> deployment=%s user_key=%s",
+            "DeploymentAffinityCheck: caller affinity hit -> deployment=%s user_key=%s",
             model_id,
-            self._shorten_for_logs(user_key),
+            self._shorten_for_logs(self._hash_user_key(user_key)),
         )
         return [deployment]
 
@@ -626,7 +626,7 @@ class DeploymentAffinityCheck(CustomLogger):
                         deployment_model_name,
                         model_id,
                         self.ttl_seconds,
-                        self._shorten_for_logs(user_key),
+                        self._shorten_for_logs(self._hash_user_key(user_key)),
                     )
                 else:
                     verbose_router_logger.debug(

From aaef5d219aa3ed1ba262302705df8805570a6a45 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 11:07:29 -0700
Subject: [PATCH 22/43] fix(proxy): drop anthropic-beta on the Vertex
 passthrough count-tokens route

---
 .../llm_passthrough_endpoints.py              | 12 ++-
 .../test_vertex_passthrough_load_balancing.py | 88 +++++++++++++++++++
 2 files changed, 99 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
index b48b8d81494..29f216fd450 100644
--- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
@@ -1730,6 +1730,16 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict:
     return headers
 
 
+def _is_vertex_anthropic_count_tokens_route(endpoint: str) -> bool:
+    return endpoint.rsplit("/", 1)[-1].split(":", 1)[0] == "count-tokens"
+
+
+def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str]) -> Mapping[str, str]:
+    if not _is_vertex_anthropic_count_tokens_route(endpoint):
+        return headers
+    return MappingProxyType({name: value for name, value in headers.items() if name.lower() != "anthropic-beta"})
+
+
 def get_vertex_pass_through_handler(
     call_type: Literal["discovery", "aiplatform"],  # noqa: UP037  # ruff reports quoted Literal values here
 ) -> BaseVertexAIPassThroughHandler:
@@ -2128,7 +2138,7 @@ async def _base_vertex_proxy_route(
     endpoint_func: Final = create_pass_through_route(
         endpoint=endpoint,
         target=target,
-        custom_headers=headers,
+        custom_headers=_upstream_headers_for_vertex_route(endpoint, headers),
         is_streaming_request=is_streaming_request,
     )  # dynamically construct pass-through endpoint based on incoming path
 
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py
index 961479c0393..6735f2a3780 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py
@@ -5,6 +5,7 @@ import pytest
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
     _base_vertex_proxy_route,
+    _upstream_headers_for_vertex_route,
 )
 from litellm.types.router import DeploymentTypedDict
 
@@ -348,6 +349,93 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
         assert headers_passed_through is False
 
 
+VERTEX_ANTHROPIC_MODELS_PREFIX = "v1/projects/test-project/locations/global/publishers/anthropic/models/"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("model_segment", "expects_anthropic_beta"),
+    [
+        ("count-tokens:rawPredict", False),
+        ("claude-sonnet-4-6:streamRawPredict", True),
+    ],
+)
+async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
+    model_segment: str, expects_anthropic_beta: bool
+):
+    with (
+        patch(  # test-quality-ok: the route reads this proxy global at call time, nothing injects it
+            "litellm.proxy.proxy_server.llm_router", None
+        ),
+        patch(  # test-quality-ok: the route reads this proxy global at call time, nothing injects it
+            "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
+        ) as mock_pt_router,
+        patch(  # test-quality-ok: the route offers no injection point for its header preparation
+            "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
+            new_callable=AsyncMock,
+        ) as mock_prep_headers,
+        patch(  # test-quality-ok: the upstream call is captured here, the route offers no injection point
+            "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
+        ) as mock_create_route,
+        patch(  # test-quality-ok: the route calls auth directly rather than through Depends
+            "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
+            new_callable=AsyncMock,
+        ) as mock_auth,
+        patch(  # test-quality-ok: the route reads the request body for this, a MagicMock request has none
+            "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn",
+            new_callable=AsyncMock,
+            return_value=False,
+        ),
+    ):
+        mock_pt_router.get_vertex_credentials.return_value = MagicMock()
+        mock_prep_headers.return_value = (
+            {
+                "anthropic-beta": "tool-search-tool-2025-10-19,web-search-2025-03-05",
+                "content-type": "application/json",
+                "Authorization": "Bearer vertex-access-token",
+            },
+            "https://aiplatform.googleapis.com",
+            False,
+            "test-project",
+            "global",
+        )
+        mock_create_route.return_value = AsyncMock()
+        mock_auth.return_value = UserAPIKeyAuth(api_key="sk-litellm-secret-key")
+
+        await _base_vertex_proxy_route(
+            endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}",
+            request=MagicMock(),
+            fastapi_response=MagicMock(),
+            get_vertex_pass_through_handler=MagicMock(),
+        )
+
+        upstream_headers = mock_create_route.call_args.kwargs["custom_headers"]
+        assert ("anthropic-beta" in upstream_headers) is expects_anthropic_beta
+        assert upstream_headers["Authorization"] == "Bearer vertex-access-token"
+        assert upstream_headers["content-type"] == "application/json"
+
+
+def test_upstream_headers_for_vertex_route_filters_anthropic_beta_by_route():
+    headers = {
+        "Anthropic-Beta": "effort-2025-11-24",
+        "content-type": "application/json",
+        "Authorization": "Bearer vertex-access-token",
+    }
+
+    count_tokens_headers = _upstream_headers_for_vertex_route(
+        f"{VERTEX_ANTHROPIC_MODELS_PREFIX}count-tokens:rawPredict", headers
+    )
+    model_headers = _upstream_headers_for_vertex_route(
+        f"{VERTEX_ANTHROPIC_MODELS_PREFIX}claude-sonnet-4-6:rawPredict", headers
+    )
+
+    assert dict(count_tokens_headers) == {
+        "content-type": "application/json",
+        "Authorization": "Bearer vertex-access-token",
+    }
+    assert dict(model_headers) == headers
+
+
 @pytest.mark.asyncio
 async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
     """

From b3325750ae4d5b462c7162084c5f74261f458c42 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Thu, 3 Sep 2026 11:09:44 -0700
Subject: [PATCH 23/43] fix(ui): aggregate session token usage in the logs
 table

The logs table already rolled up cost per session but the Tokens column only showed the representative call's usage. The per-session aggregate query now also sums prompt, completion and total tokens, and the Tokens cell switches to those sums for multi-call sessions the same way the Cost cell does.

Claude-Session: https://claude.ai/code/session_01CNasFqyjnLN3Rqman25vde
---
 .../spend_management_endpoints.py             | 21 ++++-
 .../test_spend_management_endpoints.py        | 85 +++++++++++++++++++
 .../RequestLogsTableColumns.test.tsx          | 43 +++++++++-
 .../view_logs/RequestLogsTableColumns.tsx     | 17 ++--
 .../src/components/view_logs/columns.tsx      |  3 +
 5 files changed, 160 insertions(+), 9 deletions(-)

diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py
index 295faaa980a..2a50d5170f0 100644
--- a/litellm/proxy/spend_tracking/spend_management_endpoints.py
+++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py
@@ -173,6 +173,9 @@ class _SessionSpendRow(TypedDict):
     session_cache_hit_count: ReadOnly[int]
     session_llm_count: ReadOnly[int]
     session_agent_count: ReadOnly[int]
+    session_total_prompt_tokens: ReadOnly[int]
+    session_total_completion_tokens: ReadOnly[int]
+    session_total_tokens: ReadOnly[int]
     session_models: ReadOnly[Sequence[str]]
 
 
@@ -188,6 +191,9 @@ class _SessionSpendStats(NamedTuple):
     session_cache_hit_count: int
     session_llm_count: int
     session_agent_count: int
+    session_total_prompt_tokens: int
+    session_total_completion_tokens: int
+    session_total_tokens: int
     session_models: Sequence[str]
     session_models_truncated: bool
 
@@ -4287,8 +4293,8 @@ async def _build_ui_spend_logs_response(
     Build the paginated response for the UI spend-logs endpoint.
 
     When ``enrich_session_counts`` is ``True`` (the default for the v1/UI
-    endpoint), each row is enriched with ``session_total_count`` plus spend
-    and call-type aggregates so the frontend knows which sessions are
+    endpoint), each row is enriched with ``session_total_count`` plus spend,
+    token and call-type aggregates so the frontend knows which sessions are
     expandable (multi-call sessions).  One ``GROUP BY (session_id, api_key)``
     query serves every referenced session, keyed per api key so two callers
     reusing a session id never see each other's totals.  Rows without a
@@ -4356,7 +4362,10 @@ async def _build_ui_spend_logs_response(
                            COUNT(*) FILTER (
                                WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL}
                            )::int AS session_llm_count,
-                           COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count
+                           COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count,
+                           COALESCE(SUM(prompt_tokens), 0)::bigint AS session_total_prompt_tokens,
+                           COALESCE(SUM(completion_tokens), 0)::bigint AS session_total_completion_tokens,
+                           COALESCE(SUM(total_tokens), 0)::bigint AS session_total_tokens
                     FROM "LiteLLM_SpendLogs"
                     WHERE session_id = ANY($1::text[])
                       AND api_key = ANY($2::text[])
@@ -4389,6 +4398,9 @@ async def _build_ui_spend_logs_response(
                     session_cache_hit_count=int(row.get("session_cache_hit_count") or 0),
                     session_llm_count=int(row.get("session_llm_count") or 0),
                     session_agent_count=int(row.get("session_agent_count") or 0),
+                    session_total_prompt_tokens=int(row.get("session_total_prompt_tokens") or 0),
+                    session_total_completion_tokens=int(row.get("session_total_completion_tokens") or 0),
+                    session_total_tokens=int(row.get("session_total_tokens") or 0),
                     session_models=models[:_SESSION_MODELS_LIMIT],
                     session_models_truncated=len(models) > _SESSION_MODELS_LIMIT,
                 )
@@ -4418,6 +4430,9 @@ async def _build_ui_spend_logs_response(
                 row_dict["session_cache_hit_count"] = session_stats.session_cache_hit_count
                 row_dict["session_llm_count"] = session_stats.session_llm_count
                 row_dict["session_agent_count"] = session_stats.session_agent_count
+                row_dict["session_total_prompt_tokens"] = session_stats.session_total_prompt_tokens
+                row_dict["session_total_completion_tokens"] = session_stats.session_total_completion_tokens
+                row_dict["session_total_tokens"] = session_stats.session_total_tokens
                 row_dict["session_models"] = session_stats.session_models
                 row_dict["session_models_truncated"] = session_stats.session_models_truncated
             enriched.append(row_dict)
diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
index 60a32102946..f4cd8814bc1 100644
--- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
+++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
@@ -4234,6 +4234,91 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_spend():
     assert call_args[2] == [api_key]
 
 
+@pytest.mark.asyncio
+async def test_build_ui_spend_logs_response_sums_multi_round_session_tokens():
+    """
+    Regression test for LIT-4929: the logs table showed the summed session cost but
+    only the last call's token usage.  Every row of a multi-round session must carry
+    the session-wide prompt, completion and total token sums from the aggregate
+    query, while rows outside a session carry none of them.
+    """
+    from litellm.proxy.spend_tracking.spend_management_endpoints import (
+        _build_ui_spend_logs_response,
+    )
+
+    session_id = "sess-multi-round-tokens"
+    api_key = "hashed-key-xyz"
+    dict_rows = [
+        {
+            "request_id": "req-1",
+            "session_id": session_id,
+            "call_type": "completion",
+            "api_key": api_key,
+            "total_tokens": 10,
+            "prompt_tokens": 7,
+            "completion_tokens": 3,
+        },
+        {
+            "request_id": "req-2",
+            "session_id": session_id,
+            "call_type": "completion",
+            "api_key": api_key,
+            "total_tokens": 50,
+            "prompt_tokens": 35,
+            "completion_tokens": 15,
+        },
+        {
+            "request_id": "req-3",
+            "session_id": None,
+            "call_type": "completion",
+            "api_key": api_key,
+            "total_tokens": 5,
+            "prompt_tokens": 4,
+            "completion_tokens": 1,
+        },
+    ]
+
+    mock_prisma = MagicMock()
+    mock_prisma.db.query_raw = AsyncMock(
+        return_value=[
+            {
+                "session_id": session_id,
+                "api_key": api_key,
+                "session_total_count": 2,
+                "session_total_spend": 0.06,
+                "mcp_tool_call_count": 0,
+                "mcp_tool_call_spend": 0.0,
+                "session_total_prompt_tokens": 42,
+                "session_total_completion_tokens": 18,
+                "session_total_tokens": 60,
+            }
+        ]
+    )
+
+    result = await _build_ui_spend_logs_response(
+        prisma_client=mock_prisma,
+        data=dict_rows,
+        total_records=3,
+        page=1,
+        page_size=50,
+        total_pages=1,
+        enrich_session_counts=True,
+    )
+
+    rows = result["data"]
+    session_rows = rows[:2]
+    assert [row["session_total_tokens"] for row in session_rows] == [60, 60]
+    assert [row["session_total_prompt_tokens"] for row in session_rows] == [42, 42]
+    assert [row["session_total_completion_tokens"] for row in session_rows] == [18, 18]
+    assert [(row["total_tokens"], row["prompt_tokens"], row["completion_tokens"]) for row in session_rows] == [
+        (10, 7, 3),
+        (50, 35, 15),
+    ]
+
+    token_keys = ("session_total_tokens", "session_total_prompt_tokens", "session_total_completion_tokens")
+    assert all(key not in rows[2] for key in token_keys)
+
+
 @pytest.mark.asyncio
 async def test_build_ui_spend_logs_response_session_cache_hit_count():
     """
diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx
index e3bacc0908a..d2c84173fd1 100644
--- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx
@@ -1,4 +1,4 @@
-import { render, screen } from "@testing-library/react";
+import { render, screen, within } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { describe, expect, it, vi } from "vitest";
 
@@ -75,6 +75,47 @@ describe("Cost column", () => {
   });
 });
 
+describe("Tokens column", () => {
+  it("shows the summed session token usage, not the representative call's tokens, for a multi-round session", () => {
+    renderRows([
+      logEntry({
+        request_id: "req-session-tokens",
+        total_tokens: 10,
+        prompt_tokens: 7,
+        completion_tokens: 3,
+        session_id: "sess-1",
+        session_total_count: 3,
+        session_total_tokens: 60,
+        session_total_prompt_tokens: 42,
+        session_total_completion_tokens: 18,
+      }),
+    ]);
+
+    const tokensCell = screen.getByText("60").closest("td")!;
+    expect(within(tokensCell).getByText("(42+18)")).toBeInTheDocument();
+    expect(within(tokensCell).getByText("session total")).toBeInTheDocument();
+    expect(screen.queryByText("10")).not.toBeInTheDocument();
+    expect(screen.queryByText("(7+3)")).not.toBeInTheDocument();
+  });
+
+  it("falls back to the call's own tokens with no session label when the backend sent no session token sums", () => {
+    renderRows([
+      logEntry({
+        request_id: "req-no-token-aggregate",
+        total_tokens: 10,
+        prompt_tokens: 7,
+        completion_tokens: 3,
+        session_id: "sess-2",
+        session_total_count: 3,
+      }),
+    ]);
+
+    const tokensCell = screen.getByText("10").closest("td")!;
+    expect(within(tokensCell).getByText("(7+3)")).toBeInTheDocument();
+    expect(within(tokensCell).queryByText("session total")).not.toBeInTheDocument();
+  });
+});
+
 describe("Type column", () => {
   it("shows the conversation badge and composition even when an MCP call represents the conversation", async () => {
     const user = userEvent.setup();
diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx
index 8db0b106851..9d4dc4f7898 100644
--- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx
@@ -263,13 +263,20 @@ export const getRequestLogsTableColumns = ({
     meta: { numeric: true },
     cell: ({ row }) => {
       const log = row.original;
+      const showSessionTotal = (log.session_total_count || 1) > 1 && log.session_total_tokens != null;
+      const total = showSessionTotal ? log.session_total_tokens : log.total_tokens;
+      const prompt = showSessionTotal ? log.session_total_prompt_tokens : log.prompt_tokens;
+      const completion = showSessionTotal ? log.session_total_completion_tokens : log.completion_tokens;
       return (
-        <span className="text-sm">
-          {String(log.total_tokens || "0")}
-          <span className="text-muted-foreground text-xs ml-1">
-            ({String(log.prompt_tokens || "0")}+{String(log.completion_tokens || "0")})
+        <div className="flex flex-col items-end">
+          <span className="text-sm">
+            {String(total || "0")}
+            <span className="text-muted-foreground text-xs ml-1">
+              ({String(prompt || "0")}+{String(completion || "0")})
+            </span>
           </span>
-        </span>
+          {showSessionTotal && <span className="text-[10px] text-muted-foreground">session total</span>}
+        </div>
       );
     },
   },
diff --git a/ui/litellm-dashboard/src/components/view_logs/columns.tsx b/ui/litellm-dashboard/src/components/view_logs/columns.tsx
index 21e09faf454..b2e29c3a0c1 100644
--- a/ui/litellm-dashboard/src/components/view_logs/columns.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/columns.tsx
@@ -42,6 +42,9 @@ export type LogEntry = {
   request_duration_ms?: number;
   session_total_count?: number;
   session_total_spend?: number;
+  session_total_tokens?: number;
+  session_total_prompt_tokens?: number;
+  session_total_completion_tokens?: number;
   session_cache_hit_count?: number;
   mcp_tool_call_count?: number;
   mcp_tool_call_spend?: number;

From 3b13a5fda968620d699d5f8cc15d7485f118bc5e Mon Sep 17 00:00:00 2001
From: ryan <ryan@berri.ai>
Date: Thu, 3 Sep 2026 18:31:43 +0000
Subject: [PATCH 24/43] test(ui): query the tokens cell by role instead of
 walking the DOM

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../RequestLogsTableColumns.test.tsx          | 56 +++++++++----------
 1 file changed, 25 insertions(+), 31 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx
index d2c84173fd1..3c9e6543c1c 100644
--- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx
@@ -1,4 +1,4 @@
-import { render, screen, within } from "@testing-library/react";
+import { render, screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { describe, expect, it, vi } from "vitest";
 
@@ -76,43 +76,37 @@ describe("Cost column", () => {
 });
 
 describe("Tokens column", () => {
-  it("shows the summed session token usage, not the representative call's tokens, for a multi-round session", () => {
-    renderRows([
-      logEntry({
-        request_id: "req-session-tokens",
-        total_tokens: 10,
-        prompt_tokens: 7,
-        completion_tokens: 3,
-        session_id: "sess-1",
-        session_total_count: 3,
-        session_total_tokens: 60,
-        session_total_prompt_tokens: 42,
-        session_total_completion_tokens: 18,
-      }),
-    ]);
+  const sessionRow: Partial<LogEntry> = {
+    request_id: "req-session-tokens",
+    total_tokens: 10,
+    prompt_tokens: 7,
+    completion_tokens: 3,
+    session_id: "sess-1",
+    session_total_count: 3,
+  };
 
-    const tokensCell = screen.getByText("60").closest("td")!;
-    expect(within(tokensCell).getByText("(42+18)")).toBeInTheDocument();
-    expect(within(tokensCell).getByText("session total")).toBeInTheDocument();
+  it("shows the summed session token usage, not the representative call's tokens, for a multi-round session", () => {
+    const aggregatedRow: Partial<LogEntry> = {
+      ...sessionRow,
+      session_total_tokens: 60,
+      session_total_prompt_tokens: 42,
+      session_total_completion_tokens: 18,
+    };
+    renderRows([logEntry(aggregatedRow)]);
+
+    const tokensCell = screen.getByRole("cell", { name: /\(42\+18\)/ });
+    expect(tokensCell).toHaveTextContent("60");
+    expect(tokensCell).toHaveTextContent("session total");
     expect(screen.queryByText("10")).not.toBeInTheDocument();
     expect(screen.queryByText("(7+3)")).not.toBeInTheDocument();
   });
 
   it("falls back to the call's own tokens with no session label when the backend sent no session token sums", () => {
-    renderRows([
-      logEntry({
-        request_id: "req-no-token-aggregate",
-        total_tokens: 10,
-        prompt_tokens: 7,
-        completion_tokens: 3,
-        session_id: "sess-2",
-        session_total_count: 3,
-      }),
-    ]);
+    renderRows([logEntry(sessionRow)]);
 
-    const tokensCell = screen.getByText("10").closest("td")!;
-    expect(within(tokensCell).getByText("(7+3)")).toBeInTheDocument();
-    expect(within(tokensCell).queryByText("session total")).not.toBeInTheDocument();
+    const tokensCell = screen.getByRole("cell", { name: /\(7\+3\)/ });
+    expect(tokensCell).toHaveTextContent("10");
+    expect(tokensCell).not.toHaveTextContent("session total");
   });
 });
 

From 4d2ffe2e8e8dd1769a8fd2f1a3ebc857a4a3701c Mon Sep 17 00:00:00 2001
From: ryan <ryan@berri.ai>
Date: Thu, 3 Sep 2026 18:38:22 +0000
Subject: [PATCH 25/43] test(ui): hoist inherited-grant fixture out of the
 inline createMockTeamData arg

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../src/components/team/TeamInfo.test.tsx     | 32 +++++++++----------
 1 file changed, 15 insertions(+), 17 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
index ae78ac06a9c..a9c1077e96b 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
@@ -313,23 +313,21 @@ describe("TeamInfoView", () => {
       vi.mocked(networking.getAgentsList).mockResolvedValue({
         agents: [{ agent_id: "agent-support-5678", agent_name: "support_agent" }],
       });
-      vi.mocked(networking.teamInfoCall).mockResolvedValue(
-        createMockTeamData({
-          object_permission: null,
-          access_group_ids: ["ag-1"],
-          access_group_mcp_server_ids: ["mcp-github-1234"],
-          access_group_agent_ids: ["agent-support-5678"],
-          access_group_details: [
-            {
-              access_group_id: "ag-1",
-              access_group_name: "platform-tools",
-              models: [],
-              mcp_server_ids: ["mcp-github-1234"],
-              agent_ids: ["agent-support-5678"],
-            },
-          ],
-        }),
-      );
+      const platformToolsGroup = {
+        access_group_id: "ag-1",
+        access_group_name: "platform-tools",
+        models: [],
+        mcp_server_ids: ["mcp-github-1234"],
+        agent_ids: ["agent-support-5678"],
+      };
+      const inheritedGrants = {
+        object_permission: null,
+        access_group_ids: ["ag-1"],
+        access_group_mcp_server_ids: ["mcp-github-1234"],
+        access_group_agent_ids: ["agent-support-5678"],
+        access_group_details: [platformToolsGroup],
+      };
+      vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData(inheritedGrants));
 
       renderWithProviders(<TeamInfoView {...defaultProps} />);
 

From f87b9097eaffdde471df332c323a1a59be9c2014 Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 11:50:51 -0700
Subject: [PATCH 26/43] test(bedrock): drop EOL cohere.command-r-plus-v1:0 from
 local_testing

Bedrock retired cohere.command-r-plus-v1:0 on 2026-08-19 and lists no
Cohere command chat model anymore, so the three local_testing cases that
pinned it fail with a 404 end-of-life error on every pipeline. Drop the
case from test_completion_bedrock_httpx_models and move the
parallel-streaming Bedrock entry to mistral.mistral-7b-instruct-v0:2,
which still takes the invoke route and is ACTIVE in the CI account.
---
 tests/local_testing/test_completion.py | 1 -
 tests/local_testing/test_streaming.py  | 3 +--
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py
index ef8d6c55148..f8f23ea015a 100644
--- a/tests/local_testing/test_completion.py
+++ b/tests/local_testing/test_completion.py
@@ -2879,7 +2879,6 @@ def response_format_tests(response: litellm.ModelResponse):
     "model",
     [
         "bedrock/mistral.mistral-large-2407-v1:0",
-        "bedrock/cohere.command-r-plus-v1:0",
         "us.anthropic.claude-sonnet-4-5-20250929-v1:0",
         "mistral.mistral-7b-instruct-v0:2",
         "meta.llama3-8b-instruct-v1:0",
diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py
index 07d693af447..bf39d3155b7 100644
--- a/tests/local_testing/test_streaming.py
+++ b/tests/local_testing/test_streaming.py
@@ -1168,7 +1168,6 @@ async def test_completion_replicate_llama3_streaming(sync_mode):
     "model, region",
     [
         # ["bedrock/ai21.jamba-instruct-v1:0", "us-east-1"],
-        # ["bedrock/cohere.command-r-plus-v1:0", None],
         ["us.anthropic.claude-sonnet-4-5-20250929-v1:0", None],
         # ["mistral.mistral-7b-instruct-v0:2", None],
         # ["meta.llama3-8b-instruct-v1:0", None],
@@ -1271,7 +1270,7 @@ def test_bedrock_claude_3_streaming():
     "model",
     [
         "claude-haiku-4-5-20251001",
-        "cohere.command-r-plus-v1:0",  # bedrock
+        "bedrock/mistral.mistral-7b-instruct-v0:2",
         "gpt-3.5-turbo",
     ],
 )

From 9c7c7a05ac92aebf1694ba087501f4b7c0c1416d Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 11:58:25 -0700
Subject: [PATCH 27/43] test(router): type the deployment affinity JWT test
 helpers

---
 .../test_deployment_affinity_check.py         | 24 ++++++++++---------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py
index 1852d641f0a..b5651062098 100644
--- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py
+++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py
@@ -1,11 +1,12 @@
 import asyncio
+import itertools
+import json
+from collections.abc import Sequence
+from typing import Final
 from unittest.mock import AsyncMock, patch
 
 import pytest
 
-
-import json
-
 import litellm
 from litellm.caching.dual_cache import DualCache
 from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
@@ -102,11 +103,10 @@ async def test_async_user_key_affinity_routes_to_same_deployment():
 
     # Deterministic routing: first selection uses seq[0], second selection attempts seq[1]
     # unless the list has been filtered to length=1 by deployment affinity.
-    choice_calls = {"count": 0}
+    choice_calls: Final = itertools.count(1)
 
-    def deterministic_choice(seq):
-        choice_calls["count"] += 1
-        if choice_calls["count"] == 1:
+    def deterministic_choice(seq: Sequence[dict[str, object]]) -> dict[str, object]:
+        if next(choice_calls) == 1:
             return seq[0]
         return seq[1] if len(seq) > 1 else seq[0]
 
@@ -1000,7 +1000,7 @@ async def test_model_group_affinity_config_overrides_global():
     assert len(filtered) == 2
 
 
-def _jwt_metadata(user_id: str) -> dict:
+def _jwt_metadata(user_id: str) -> dict[str, str | None]:
     return {"user_api_key_hash": None, "user_api_key_user_id": user_id}
 
 
@@ -1090,7 +1090,7 @@ async def test_proxy_jwt_auth_metadata_pins_per_user():
         enable_responses_api_affinity=False,
     )
 
-    def proxy_request(user_id: str) -> dict:
+    def proxy_request(user_id: str) -> dict[str, object]:
         return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
             data={"model": model_group, "messages": [{"role": "user", "content": "hi"}], "metadata": {}},
             user_api_key_dict=UserAPIKeyAuth(api_key=None, user_id=user_id),
@@ -1098,12 +1098,14 @@ async def test_proxy_jwt_auth_metadata_pins_per_user():
         )
 
     alice_request = proxy_request("jwt-user-alice")
-    assert alice_request["metadata"]["user_api_key_hash"] is None
+    alice_metadata = alice_request["metadata"]
+    assert isinstance(alice_metadata, dict)
+    assert alice_metadata["user_api_key_hash"] is None
 
     await callback.async_pre_call_deployment_hook(
         kwargs={
             **alice_request,
-            "metadata": {**alice_request["metadata"], "deployment_model_name": model_group},
+            "metadata": {**alice_metadata, "deployment_model_name": model_group},
             "model_info": {"id": "openai-deployment-b"},
         },
         call_type=None,

From 92122086ecc9271640d284fe4e1db32bdfc205ff Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Thu, 3 Sep 2026 19:12:11 +0000
Subject: [PATCH 28/43] fix: stop a cleared Organization field from failing key
 creation (#39316)

* fix: stop a cleared Organization field from failing key creation

Clearing the Organization combobox in the Create Key modal left organization_id set to an empty string, so /key/generate looked up an organization named "" and failed with "Organization doesn't exist in db. Organization=".

OrganizationDropdown now emits null on clear, and GenerateKeyRequest normalizes an empty organization_id or project_id to None the same way it already does for team_id.

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

* test: drop customer-specific docstring from key request normalization test

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

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: yassin <yassin@berri.ai>
---
 litellm/proxy/_types.py                              |  2 +-
 .../test_key_management_endpoints.py                 | 12 ++++++++++++
 .../common_components/OrganizationDropdown.test.tsx  | 11 +++++++++++
 .../common_components/OrganizationDropdown.tsx       |  4 ++--
 .../src/components/organisms/create_key_button.tsx   |  6 +++---
 .../src/components/templates/key_edit_view.tsx       |  6 +++---
 6 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 0aea72be1e2..98464d3a127 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -1216,7 +1216,7 @@ class GenerateKeyRequest(KeyRequestBase):
     organization_id: str | None = None
     project_id: str | None = None
 
-    @field_validator("team_id", "organization_id", mode="before")
+    @field_validator("team_id", "organization_id", "project_id", mode="before")
     @classmethod
     def treat_cleared_id_as_unset(cls, v: object) -> object:
         if v == "":
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 68704f476b5..7954a4693cc 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
@@ -17505,6 +17505,18 @@ def test_generate_key_request_blank_team_id_is_personal():
     assert GenerateKeyRequest(team_id="team-1").team_id == "team-1"
 
 
+def test_generate_key_request_blank_organization_and_project_id_are_unset():
+    from litellm.proxy._types import RegenerateKeyRequest
+
+    cleared = GenerateKeyRequest(organization_id="", project_id="")
+    assert cleared.organization_id is None
+    assert cleared.project_id is None
+    assert "organization_id" not in cleared.model_dump(exclude_none=True)
+    assert RegenerateKeyRequest(organization_id="").organization_id is None
+    assert GenerateKeyRequest(organization_id="org-1", project_id="proj-1").organization_id == "org-1"
+    assert GenerateKeyRequest(organization_id="org-1", project_id="proj-1").project_id == "proj-1"
+
+
 def test_key_request_blank_organization_id_is_unset():
     from litellm.proxy._types import RegenerateKeyRequest, UpdateKeyRequest
 
diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx
index e0b2b897b36..524ec6a715e 100644
--- a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx
+++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx
@@ -58,6 +58,17 @@ describe("OrganizationDropdown", () => {
     expect(onChange.mock.calls[0][0]).toBe("org-1");
   });
 
+  it("emits null, never the empty string, when the selection is cleared", async () => {
+    const onChange = vi.fn();
+    const user = userEvent.setup();
+    render(<OrganizationDropdown organizations={MOCK_ORGS} value="org-1" onChange={onChange} />);
+
+    await user.click(screen.getByRole("button", { name: "Clear" }));
+
+    expect(onChange).toHaveBeenCalledTimes(1);
+    expect(onChange).toHaveBeenCalledWith(null);
+  });
+
   it("should filter options by organization id", async () => {
     const user = userEvent.setup();
     render(<OrganizationDropdown organizations={MOCK_ORGS} />);
diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx
index 8da35ecd02e..663028b2d92 100644
--- a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx
+++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx
@@ -5,7 +5,7 @@ import { Organization } from "../networking";
 interface OrganizationDropdownProps {
   organizations?: Organization[] | null;
   value?: string;
-  onChange?: (value: string) => void;
+  onChange?: (value: string | null) => void;
   disabled?: boolean;
   loading?: boolean;
   style?: React.CSSProperties;
@@ -32,7 +32,7 @@ const OrganizationDropdown: React.FC<OrganizationDropdownProps> = ({
           sublabel: org.organization_id,
         }))}
         value={value}
-        onValueChange={(organizationId) => onChange?.(organizationId)}
+        onValueChange={(organizationId) => onChange?.(organizationId || null)}
         placeholder={placeholder}
         emptyText={loading ? "Loading organizations…" : "No organizations found"}
         disabled={disabled}
diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
index 5749541dcea..ee3a88acba0 100644
--- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
+++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
@@ -587,9 +587,9 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
     }
   };
 
-  const changeOrganization = (write: FieldWrite) => (orgId: string) => {
-    write(orgId || undefined);
-    setSelectedOrganizationId(orgId || null);
+  const changeOrganization = (write: FieldWrite) => (orgId: string | null) => {
+    write(orgId ?? undefined);
+    setSelectedOrganizationId(orgId);
     // Clear team and project when org changes
     setSelectedCreateKeyTeam(null);
     setSelectedProjectId(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 3e772fd0e9b..1330be2788b 100644
--- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
+++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
@@ -303,9 +303,9 @@ export function KeyEditView({
     }
   };
 
-  const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | undefined) => {
-    setField(orgId || null);
-    setSelectedOrganizationId(orgId || null);
+  const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => {
+    setField(orgId);
+    setSelectedOrganizationId(orgId);
     form.setValue("team_id", undefined);
   };
 

From 51d821ae45aef7fa95e222ebcb55486210d08543 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 12:21:54 -0700
Subject: [PATCH 29/43] fix(cost): bill bedrock_mantle web search at $12 per 1k
 queries using Bedrock's reported count

---
 .../llm_cost_calc/tool_call_cost_tracking.py  |  34 +++--
 ...odel_prices_and_context_window_backup.json |  25 ++++
 litellm/types/llms/openai.py                  |  13 ++
 model_prices_and_context_window.json          |  25 ++++
 .../test_tool_call_cost_tracking.py           | 118 ++++++++++++++++++
 5 files changed, 207 insertions(+), 8 deletions(-)

diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py
index 5504756ceb8..bf99035a6b1 100644
--- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py
+++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py
@@ -5,6 +5,8 @@ Helper utilities for tracking the cost of built-in tools.
 from collections.abc import Mapping
 from typing import Final, Literal
 
+from pydantic import ValidationError
+
 import litellm
 from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
 from litellm.litellm_core_utils.llm_cost_calc.utils import (
@@ -13,6 +15,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
 from litellm.types.llms.openai import (
     FileSearchTool,
     ResponsesAPIResponse,
+    ResponsesToolUsage,
     WebSearchOptions,
 )
 from litellm.types.utils import (
@@ -32,6 +35,17 @@ def _output_item_type(output_item: object) -> str | None:
     return item_type if isinstance(item_type, str) else None
 
 
+def _reported_web_search_requests(response_object: ResponsesAPIResponse) -> int | None:
+    tool_usage: Final = getattr(response_object, "tool_usage", None)
+    if tool_usage is None:
+        return None
+    try:
+        web_search: Final = ResponsesToolUsage.model_validate(tool_usage).web_search
+    except ValidationError:
+        return None
+    return None if web_search is None else web_search.num_requests
+
+
 def _usage_reports_server_side_web_search_calls(usage: Usage) -> bool:
     details: Final = getattr(usage, "server_side_tool_usage_details", None)
     if not isinstance(details, Mapping):
@@ -182,15 +196,19 @@ class StandardBuiltInToolCostTracking:
 
         Providers that report a request count in usage (gemini, anthropic, xai, vertex) are handled by
         get_cost_for_web_search_request and never reach here. This path prices per call, so it must count
-        the web_search_call items. Chat-completions responses only expose url_citation annotations with no
-        count, so they floor to a single billable search.
+        the web_search_call items, unless the response reports the billable count itself
+        (Bedrock's tool_usage.web_search.num_requests, which excludes open_page fetches). Chat-completions
+        responses only expose url_citation annotations with no count, so they floor to a single billable search.
         """
-        if isinstance(response_object, ResponsesAPIResponse):
-            count = sum(
-                1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call"
-            )
-            return max(count, 1)
-        return 1
+        if not isinstance(response_object, ResponsesAPIResponse):
+            return 1
+        reported: Final = _reported_web_search_requests(response_object)
+        if reported is not None:
+            return reported
+        count: Final = sum(
+            1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call"
+        )
+        return max(count, 1)
 
     @staticmethod
     def _handle_file_search_cost(
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 9b358a1cedd..ffac34ea37e 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -52911,6 +52911,11 @@
         "cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
         "output_cost_per_token": 3.3e-05,
         "output_cost_per_token_above_272k_tokens": 4.95e-05,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.012,
+            "search_context_size_low": 0.012,
+            "search_context_size_medium": 0.012
+        },
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
         "max_output_tokens": 128000,
@@ -52945,6 +52950,11 @@
         "cache_read_input_token_cost_above_272k_tokens": 4.4e-07,
         "output_cost_per_token": 1.32e-05,
         "output_cost_per_token_above_272k_tokens": 1.98e-05,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.012,
+            "search_context_size_low": 0.012,
+            "search_context_size_medium": 0.012
+        },
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
         "max_output_tokens": 128000,
@@ -53007,6 +53017,11 @@
         "cache_read_input_token_cost_above_272k_tokens": 4.4e-08,
         "output_cost_per_token": 1.32e-06,
         "output_cost_per_token_above_272k_tokens": 1.98e-06,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.012,
+            "search_context_size_low": 0.012,
+            "search_context_size_medium": 0.012
+        },
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
         "max_output_tokens": 128000,
@@ -53195,6 +53210,11 @@
         "cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
         "output_cost_per_token": 3.3e-05,
         "output_cost_per_token_above_272k_tokens": 4.95e-05,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.012,
+            "search_context_size_low": 0.012,
+            "search_context_size_medium": 0.012
+        },
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
         "max_output_tokens": 128000,
@@ -53226,6 +53246,11 @@
         "cache_read_input_token_cost_above_272k_tokens": 5.5e-07,
         "output_cost_per_token": 1.65e-05,
         "output_cost_per_token_above_272k_tokens": 2.475e-05,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.012,
+            "search_context_size_low": 0.012,
+            "search_context_size_medium": 0.012
+        },
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
         "max_output_tokens": 128000,
diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py
index 32d88da0085..b33ff954c35 100644
--- a/litellm/types/llms/openai.py
+++ b/litellm/types/llms/openai.py
@@ -66,6 +66,7 @@ from pydantic import (
     ConfigDict,
     Discriminator,
     Field,
+    NonNegativeInt,
     PrivateAttr,
     SerializerFunctionWrapHandler,
     field_serializer,
@@ -1321,6 +1322,18 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject):
     model_config = {"extra": "allow"}
 
 
+class WebSearchToolUsage(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    num_requests: NonNegativeInt
+
+
+class ResponsesToolUsage(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    web_search: WebSearchToolUsage | None = None
+
+
 ResponsesAPIStatus = Literal["completed", "failed", "in_progress", "cancelled", "queued", "incomplete"]
 """
 The status of the response generation.
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 9b358a1cedd..ffac34ea37e 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -52911,6 +52911,11 @@
         "cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
         "output_cost_per_token": 3.3e-05,
         "output_cost_per_token_above_272k_tokens": 4.95e-05,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.012,
+            "search_context_size_low": 0.012,
+            "search_context_size_medium": 0.012
+        },
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
         "max_output_tokens": 128000,
@@ -52945,6 +52950,11 @@
         "cache_read_input_token_cost_above_272k_tokens": 4.4e-07,
         "output_cost_per_token": 1.32e-05,
         "output_cost_per_token_above_272k_tokens": 1.98e-05,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.012,
+            "search_context_size_low": 0.012,
+            "search_context_size_medium": 0.012
+        },
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
         "max_output_tokens": 128000,
@@ -53007,6 +53017,11 @@
         "cache_read_input_token_cost_above_272k_tokens": 4.4e-08,
         "output_cost_per_token": 1.32e-06,
         "output_cost_per_token_above_272k_tokens": 1.98e-06,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.012,
+            "search_context_size_low": 0.012,
+            "search_context_size_medium": 0.012
+        },
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
         "max_output_tokens": 128000,
@@ -53195,6 +53210,11 @@
         "cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
         "output_cost_per_token": 3.3e-05,
         "output_cost_per_token_above_272k_tokens": 4.95e-05,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.012,
+            "search_context_size_low": 0.012,
+            "search_context_size_medium": 0.012
+        },
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
         "max_output_tokens": 128000,
@@ -53226,6 +53246,11 @@
         "cache_read_input_token_cost_above_272k_tokens": 5.5e-07,
         "output_cost_per_token": 1.65e-05,
         "output_cost_per_token_above_272k_tokens": 2.475e-05,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.012,
+            "search_context_size_low": 0.012,
+            "search_context_size_medium": 0.012
+        },
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
         "max_output_tokens": 128000,
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
index fd795ffcc96..0d75301a57a 100644
--- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
+++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
@@ -928,3 +928,121 @@ def test_web_search_gate_reads_server_side_tool_usage_details_without_citations(
         standard_built_in_tools_params=None,
     )
     assert cost == 3 * _DEFAULT_WEB_SEARCH_COST_PER_CALL
+
+
+_BEDROCK_MANTLE_WEB_SEARCH_MODELS = (
+    "bedrock_mantle/openai.gpt-5.6-sol",
+    "bedrock_mantle/openai.gpt-5.6-terra",
+    "bedrock_mantle/openai.gpt-5.6-luna",
+    "bedrock_mantle/openai.gpt-5.5",
+    "bedrock_mantle/openai.gpt-5.4",
+)
+
+_BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012
+
+
+def _bedrock_mantle_responses_with_web_search(model, actions, tool_usage=None):
+    from litellm.types.llms.openai import ResponsesAPIResponse
+
+    payload = {
+        "id": "resp_1",
+        "created_at": 1756900000,
+        "model": model.split("/", 1)[-1],
+        "object": "response",
+        "status": "completed",
+        "output": [
+            {"type": "web_search_call", "id": f"ws_{i}", "status": "completed", "action": action}
+            for i, action in enumerate(actions)
+        ],
+    }
+    return ResponsesAPIResponse.model_validate(
+        payload if tool_usage is None else {**payload, "tool_usage": tool_usage}
+    )
+
+
+def _bedrock_mantle_web_search_cost(model, response):
+    from litellm.types.utils import Usage
+
+    return StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
+        model=model,
+        response_object=response,
+        usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
+        custom_llm_provider="bedrock_mantle",
+        standard_built_in_tools_params=None,
+    )
+
+
+@pytest.mark.parametrize("model", _BEDROCK_MANTLE_WEB_SEARCH_MODELS)
+def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model):
+    """
+    Regression for LIT-6870: the bedrock_mantle GPT ids forward the web_search tool but carried no
+    search_context_cost_per_query, so every Bedrock web search (billed at $12 per 1,000 queries)
+    was costed at $0. Two reported queries must bill 2 x $0.012, whichever model prefix shape the
+    cost path resolves the deployment under.
+    """
+    pricing = litellm.get_model_info(model)["search_context_cost_per_query"]
+    assert pricing == {
+        "search_context_size_low": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
+        "search_context_size_medium": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
+        "search_context_size_high": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
+    }
+
+    response = _bedrock_mantle_responses_with_web_search(
+        model,
+        actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}],
+        tool_usage={"web_search": {"num_requests": 2}},
+    )
+    for cost_model in (model, model.split("/", 1)[1]):
+        cost = _bedrock_mantle_web_search_cost(cost_model, response)
+        assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), (
+            f"{cost_model} must bill 2 x ${_BEDROCK_MANTLE_WEB_SEARCH_RATE} for 2 web searches, got ${cost}"
+        )
+
+
+@pytest.mark.parametrize("num_requests", [1, 0])
+def test_web_search_call_count_prefers_provider_reported_num_requests(local_model_cost_map, num_requests):
+    """
+    Regression for LIT-6870: Bedrock bills one query per search and reports the billable count as
+    tool_usage.web_search.num_requests, while its open_page fetches share the web_search_call item
+    type. A search plus an open_page must bill the reported count, never the two items.
+    """
+    model = "bedrock_mantle/openai.gpt-5.6-sol"
+    response = _bedrock_mantle_responses_with_web_search(
+        model,
+        actions=[
+            {"type": "search", "query": "litellm"},
+            {"type": "open_page", "url": "https://docs.litellm.ai/"},
+        ],
+        tool_usage={"web_search": {"num_requests": num_requests}},
+    )
+
+    cost = _bedrock_mantle_web_search_cost(model, response)
+
+    assert cost == pytest.approx(num_requests * _BEDROCK_MANTLE_WEB_SEARCH_RATE), (
+        f"{num_requests} reported web search requests must bill {num_requests} x "
+        f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}"
+    )
+
+
+@pytest.mark.parametrize(
+    "tool_usage",
+    [None, {}, {"web_search": None}, {"web_search": {"num_requests": "many"}}, {"web_search": {"num_requests": -1}}],
+)
+def test_web_search_call_count_falls_back_to_items_without_reported_count(local_model_cost_map, tool_usage):
+    """
+    Without a usable reported count (no tool_usage, no web_search block, or a malformed one) the
+    per-call path must keep counting web_search_call items instead of raising or billing zero.
+    """
+    model = "bedrock_mantle/openai.gpt-5.6-sol"
+    response = _bedrock_mantle_responses_with_web_search(
+        model,
+        actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}],
+        tool_usage=tool_usage,
+    )
+
+    cost = _bedrock_mantle_web_search_cost(model, response)
+
+    assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), (
+        f"2 web_search_call items with tool_usage={tool_usage!r} must bill 2 x "
+        f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}"
+    )

From 339da4183d998b1ff2db54bc0468fa994a620095 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 12:34:54 -0700
Subject: [PATCH 30/43] test(cost): type the web search cost helpers and cover
 OpenAI-shaped tool_usage

---
 .../test_tool_call_cost_tracking.py           | 59 ++++++++++---------
 1 file changed, 32 insertions(+), 27 deletions(-)

diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
index 0d75301a57a..6cc3dcceebc 100644
--- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
+++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
@@ -1,4 +1,5 @@
 import os
+from collections.abc import Mapping, Sequence
 
 import pytest
 
@@ -6,7 +7,7 @@ import litellm
 from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
     StandardBuiltInToolCostTracking,
 )
-from litellm.types.llms.openai import FileSearchTool, WebSearchOptions
+from litellm.types.llms.openai import FileSearchTool, ResponsesAPIResponse, WebSearchOptions
 from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams
 
 
@@ -941,9 +942,9 @@ _BEDROCK_MANTLE_WEB_SEARCH_MODELS = (
 _BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012
 
 
-def _bedrock_mantle_responses_with_web_search(model, actions, tool_usage=None):
-    from litellm.types.llms.openai import ResponsesAPIResponse
-
+def _responses_with_web_search(
+    model: str, actions: Sequence[Mapping[str, str]], tool_usage: Mapping[str, object] | None = None
+) -> ResponsesAPIResponse:
     payload = {
         "id": "resp_1",
         "created_at": 1756900000,
@@ -960,26 +961,21 @@ def _bedrock_mantle_responses_with_web_search(model, actions, tool_usage=None):
     )
 
 
-def _bedrock_mantle_web_search_cost(model, response):
+def _web_search_cost(model: str, response: ResponsesAPIResponse, custom_llm_provider: str) -> float:
     from litellm.types.utils import Usage
 
     return StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
         model=model,
         response_object=response,
         usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
-        custom_llm_provider="bedrock_mantle",
+        custom_llm_provider=custom_llm_provider,
         standard_built_in_tools_params=None,
     )
 
 
 @pytest.mark.parametrize("model", _BEDROCK_MANTLE_WEB_SEARCH_MODELS)
 def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model):
-    """
-    Regression for LIT-6870: the bedrock_mantle GPT ids forward the web_search tool but carried no
-    search_context_cost_per_query, so every Bedrock web search (billed at $12 per 1,000 queries)
-    was costed at $0. Two reported queries must bill 2 x $0.012, whichever model prefix shape the
-    cost path resolves the deployment under.
-    """
+    """Two Bedrock-reported web searches bill 2 x $0.012 under the prefixed and the bare model id alike."""
     pricing = litellm.get_model_info(model)["search_context_cost_per_query"]
     assert pricing == {
         "search_context_size_low": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
@@ -987,13 +983,13 @@ def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model)
         "search_context_size_high": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
     }
 
-    response = _bedrock_mantle_responses_with_web_search(
+    response = _responses_with_web_search(
         model,
         actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}],
         tool_usage={"web_search": {"num_requests": 2}},
     )
     for cost_model in (model, model.split("/", 1)[1]):
-        cost = _bedrock_mantle_web_search_cost(cost_model, response)
+        cost = _web_search_cost(cost_model, response, "bedrock_mantle")
         assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), (
             f"{cost_model} must bill 2 x ${_BEDROCK_MANTLE_WEB_SEARCH_RATE} for 2 web searches, got ${cost}"
         )
@@ -1001,13 +997,9 @@ def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model)
 
 @pytest.mark.parametrize("num_requests", [1, 0])
 def test_web_search_call_count_prefers_provider_reported_num_requests(local_model_cost_map, num_requests):
-    """
-    Regression for LIT-6870: Bedrock bills one query per search and reports the billable count as
-    tool_usage.web_search.num_requests, while its open_page fetches share the web_search_call item
-    type. A search plus an open_page must bill the reported count, never the two items.
-    """
+    """A search plus an open_page fetch bills tool_usage.web_search.num_requests, never the two items."""
     model = "bedrock_mantle/openai.gpt-5.6-sol"
-    response = _bedrock_mantle_responses_with_web_search(
+    response = _responses_with_web_search(
         model,
         actions=[
             {"type": "search", "query": "litellm"},
@@ -1016,7 +1008,7 @@ def test_web_search_call_count_prefers_provider_reported_num_requests(local_mode
         tool_usage={"web_search": {"num_requests": num_requests}},
     )
 
-    cost = _bedrock_mantle_web_search_cost(model, response)
+    cost = _web_search_cost(model, response, "bedrock_mantle")
 
     assert cost == pytest.approx(num_requests * _BEDROCK_MANTLE_WEB_SEARCH_RATE), (
         f"{num_requests} reported web search requests must bill {num_requests} x "
@@ -1029,20 +1021,33 @@ def test_web_search_call_count_prefers_provider_reported_num_requests(local_mode
     [None, {}, {"web_search": None}, {"web_search": {"num_requests": "many"}}, {"web_search": {"num_requests": -1}}],
 )
 def test_web_search_call_count_falls_back_to_items_without_reported_count(local_model_cost_map, tool_usage):
-    """
-    Without a usable reported count (no tool_usage, no web_search block, or a malformed one) the
-    per-call path must keep counting web_search_call items instead of raising or billing zero.
-    """
+    """Without a usable reported count the per-call path keeps counting web_search_call items."""
     model = "bedrock_mantle/openai.gpt-5.6-sol"
-    response = _bedrock_mantle_responses_with_web_search(
+    response = _responses_with_web_search(
         model,
         actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}],
         tool_usage=tool_usage,
     )
 
-    cost = _bedrock_mantle_web_search_cost(model, response)
+    cost = _web_search_cost(model, response, "bedrock_mantle")
 
     assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), (
         f"2 web_search_call items with tool_usage={tool_usage!r} must bill 2 x "
         f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}"
     )
+
+
+def test_web_search_call_count_reads_reported_count_beside_other_tool_usage_entries(local_model_cost_map):
+    """OpenAI reports web_search.num_requests next to other tool entries, which must not disable the reported count."""
+    response = _responses_with_web_search(
+        "gpt-5.6",
+        actions=[{"type": "search", "query": "S&P 500 close"}, {"type": "open_page", "url": "https://example.com/"}],
+        tool_usage={
+            "image_gen": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
+            "web_search": {"num_requests": 1},
+        },
+    )
+
+    cost = _web_search_cost("gpt-5.6", response, "openai")
+
+    assert cost == pytest.approx(0.01), f"1 reported OpenAI web search must bill 1 x $0.01, not the 2 items, got ${cost}"

From 1e75668a259ee9b6c7d77abd17ef51df70250b1c Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 12:37:17 -0700
Subject: [PATCH 31/43] fix(openai): default stream usage on PrivateLink and
 regional api.openai.com hosts

---
 litellm/llms/openai/common_utils.py           |  9 ++++
 litellm/llms/openai/openai.py                 |  8 ++-
 tests/test_litellm/llms/openai/test_openai.py | 52 +++++++++++++++++++
 .../llms/openai/test_openai_common_utils.py   | 21 +++++++-
 4 files changed, 84 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_litellm/llms/openai/test_openai.py

diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py
index bcd4ea43243..2db6d78a218 100644
--- a/litellm/llms/openai/common_utils.py
+++ b/litellm/llms/openai/common_utils.py
@@ -11,6 +11,7 @@ import time
 import uuid
 from collections.abc import AsyncIterator, Iterator, Mapping
 from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Optional
+from urllib.parse import urlsplit
 
 import httpx
 import openai
@@ -43,6 +44,14 @@ _OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(OpenAI)
 _AZURE_OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(AzureOpenAI)
 
 
+_OPENAI_API_HOST: Final[str] = "api.openai.com"
+
+
+def is_openai_backed_api_base(api_base: str) -> bool:
+    hostname: Final = urlsplit(api_base).hostname
+    return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}"))
+
+
 class OpenAIError(BaseLLMException):
     def __init__(
         self,
diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py
index 1cfc6e06ee9..edc8d64d9c2 100644
--- a/litellm/llms/openai/openai.py
+++ b/litellm/llms/openai/openai.py
@@ -2,7 +2,6 @@ import time
 import types
 from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
-from urllib.parse import urlparse
 
 import httpx
 
@@ -55,6 +54,7 @@ from .common_utils import (
     OpenAIError,
     build_output_token_limit_response,
     drop_params_from_unprocessable_entity_error,
+    is_openai_backed_api_base,
     is_output_token_limit_error,
 )
 from .workload_identity import resolve_openai_workload_identity_config
@@ -1190,10 +1190,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         """
         if stream_options is not None:
             return {"stream_options": stream_options}
-        else:
-            # by default litellm will include usage for openai endpoints
-            if api_base is None or urlparse(api_base).hostname == "api.openai.com":
-                return {"stream_options": {"include_usage": True}}
+        if api_base is None or is_openai_backed_api_base(api_base):
+            return {"stream_options": {"include_usage": True}}
         return {}
 
     # Embedding
diff --git a/tests/test_litellm/llms/openai/test_openai.py b/tests/test_litellm/llms/openai/test_openai.py
new file mode 100644
index 00000000000..136b837f191
--- /dev/null
+++ b/tests/test_litellm/llms/openai/test_openai.py
@@ -0,0 +1,52 @@
+import pytest
+
+from litellm.llms.openai.openai import OpenAIChatCompletion
+
+
+@pytest.mark.parametrize(
+    "api_base",
+    [
+        None,
+        "https://api.openai.com/v1",
+        "https://api.openai.com:443/v1",
+        "https://southcentralus.privatelink.api.openai.com/v1",
+        "https://eu.api.openai.com/v1",
+        "https://us.api.openai.com/v1",
+        "HTTPS://API.OPENAI.COM/v1/",
+    ],
+)
+def test_get_stream_options_defaults_include_usage_on_every_openai_backed_host(api_base):
+    """
+    PrivateLink and regional hostnames reach the real OpenAI backend, so a stream with no caller
+    stream_options must ask for the usage chunk exactly as the default base does. Regression guard
+    for LIT-6875: spend for those deployments fell back to local token counting.
+    """
+    assert OpenAIChatCompletion().get_stream_options(stream_options=None, api_base=api_base) == {
+        "stream_options": {"include_usage": True}
+    }
+
+
+@pytest.mark.parametrize(
+    "api_base",
+    [
+        "https://my-gateway.example/v1",
+        "https://api.openai.com.evil.example/v1",
+        "https://notapi.openai.com/v1",
+        "https://gateway.example/v1?upstream=api.openai.com",
+        "https://openai.internal.example/api.openai.com/v1",
+    ],
+)
+def test_get_stream_options_leaves_foreign_hosts_without_a_usage_default(api_base):
+    """Only the host decides: an OpenAI-compatible backend elsewhere may not support stream_options at all."""
+    assert OpenAIChatCompletion().get_stream_options(stream_options=None, api_base=api_base) == {}
+
+
+@pytest.mark.parametrize(
+    "api_base",
+    ["https://southcentralus.privatelink.api.openai.com/v1", "https://my-gateway.example/v1"],
+)
+def test_get_stream_options_passes_caller_stream_options_through_on_any_host(api_base):
+    caller_options = {"include_usage": False}
+    assert OpenAIChatCompletion().get_stream_options(stream_options=caller_options, api_base=api_base) == {
+        "stream_options": caller_options
+    }
diff --git a/tests/test_litellm/llms/openai/test_openai_common_utils.py b/tests/test_litellm/llms/openai/test_openai_common_utils.py
index 3ae29e411e8..d3c21c5bd5a 100644
--- a/tests/test_litellm/llms/openai/test_openai_common_utils.py
+++ b/tests/test_litellm/llms/openai/test_openai_common_utils.py
@@ -7,7 +7,7 @@ import pytest
 
 import litellm
 from litellm.litellm_core_utils.token_counter import token_counter
-from litellm.llms.openai.common_utils import BaseOpenAILLM
+from litellm.llms.openai.common_utils import BaseOpenAILLM, is_openai_backed_api_base
 
 # Test parameters for different API functions
 API_FUNCTION_PARAMS = [
@@ -392,3 +392,22 @@ async def test_async_genuine_bad_request_still_raises(provider, stream):
 
     with pytest.raises(litellm.BadRequestError):
         await _call_and_drain()
+
+
+@pytest.mark.parametrize(
+    ("api_base", "expected"),
+    [
+        ("https://api.openai.com/v1", True),
+        ("https://api.openai.com:443/v1/", True),
+        ("https://southcentralus.privatelink.api.openai.com/v1", True),
+        ("https://eu.api.openai.com/v1", True),
+        ("HTTPS://API.OPENAI.COM/v1", True),
+        ("https://my-gateway.example/v1", False),
+        ("https://api.openai.com.evil.example/v1", False),
+        ("https://notapi.openai.com/v1", False),
+        ("https://gateway.example/v1?upstream=api.openai.com", False),
+        ("not a url", False),
+    ],
+)
+def test_is_openai_backed_api_base_decides_by_hostname_only(api_base, expected):
+    assert is_openai_backed_api_base(api_base) is expected

From 425e3069b93ec005e3186f5d32fa21ef70effe17 Mon Sep 17 00:00:00 2001
From: moe-berri <moe@berri.ai>
Date: Thu, 3 Sep 2026 12:40:22 -0700
Subject: [PATCH 32/43] fix(proxy): expose configured model mode

---
 litellm/proxy/utils.py                       |  3 ++
 litellm/router.py                            | 11 ++++++
 litellm/types/proxy/model_listing.py         |  4 +-
 tests/test_litellm/proxy/test_proxy_utils.py | 41 ++++++++++++++++++++
 4 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index cab2bd6d9db..24a65457928 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -7531,6 +7531,9 @@ def create_model_info_response(
             max_input_tokens = configured_input
         if configured_output is not None:
             max_output_tokens = configured_output
+        configured_mode: Final = llm_router.get_configured_mode(model_id)
+        if isinstance(configured_mode, str):
+            base["mode"] = configured_mode
 
     if max_input_tokens is not None:
         base["max_input_tokens"] = max_input_tokens
diff --git a/litellm/router.py b/litellm/router.py
index 2b8b342d253..cfb080a24a0 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -9981,6 +9981,17 @@ class Router:
             coerce_token_limit(model_info.get("max_output_tokens")),
         )
 
+    def get_configured_mode(self, model_name: str) -> "str | None":
+        """Return the mode explicitly configured for a concrete deployment."""
+        deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name)
+        if deployment is None:
+            return None
+
+        mode: Final = deployment.model_info.get("mode")
+        if isinstance(mode, str) and mode.strip():
+            return mode
+        return None
+
     def get_configured_display_name(self, model_name: str) -> "str | None":
         """
         Return the display_name explicitly configured in a concrete deployment's
diff --git a/litellm/types/proxy/model_listing.py b/litellm/types/proxy/model_listing.py
index b59c0f2cf19..24cfa85eee4 100644
--- a/litellm/types/proxy/model_listing.py
+++ b/litellm/types/proxy/model_listing.py
@@ -11,8 +11,8 @@ class ModelInfoMetadata(TypedDict):
 
 class ModelInfoResponse(TypedDict):
     """OpenAI-compatible model object. `mode`, `max_input_tokens`, and
-    `max_output_tokens` are attached when the cost map knows them; `metadata`
-    is present only when the endpoint is called with include_metadata=true.
+    `max_output_tokens` are attached when the cost map or deployment config
+    knows them; `metadata` is present only with include_metadata=true.
     """
 
     id: str
diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py
index dcaad968663..f16c6c937d0 100644
--- a/tests/test_litellm/proxy/test_proxy_utils.py
+++ b/tests/test_litellm/proxy/test_proxy_utils.py
@@ -942,6 +942,47 @@ def test_create_model_info_response_uses_deployment_limits_when_not_in_cost_map(
     assert response["max_output_tokens"] == 8000
 
 
+def test_create_model_info_response_uses_deployment_mode_for_auto_router():
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "claude-sonnet",
+                "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"},
+            },
+            {
+                "model_name": "claude-auto",
+                "litellm_params": {
+                    "model": "auto_router/complexity_router",
+                    "complexity_router_config": {
+                        "tiers": {
+                            "SIMPLE": "claude-sonnet",
+                            "MEDIUM": "claude-sonnet",
+                            "COMPLEX": "claude-sonnet",
+                        }
+                    },
+                    "complexity_router_default_model": "claude-sonnet",
+                },
+                "model_info": {
+                    "mode": "chat",
+                    "max_input_tokens": 1_000_000,
+                    "max_output_tokens": 128_000,
+                },
+            },
+        ]
+    )
+
+    response = create_model_info_response(
+        model_id="claude-auto",
+        provider="openai",
+        llm_router=router,
+        get_model_info=_raise_unmapped,
+    )
+
+    assert response["mode"] == "chat"
+    assert response["max_input_tokens"] == 1_000_000
+    assert response["max_output_tokens"] == 128_000
+
+
 def test_create_model_info_response_deployment_limits_override_cost_map():
     router = MagicMock()
     router.get_configured_token_limits.return_value = (200000, None)

From 897fba08c8ddb6dba99288b040ae5c7cad8a5757 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 11:08:03 -0700
Subject: [PATCH 33/43] feat(models): add gpt-6-astra pricing and metadata

Adds the OpenAI gpt-6-astra entry to both price files with standard, flex, priority (fast mode), batch, and above-272K long-context rates, and regression tests covering each tier and the batch rates.
---
 ...odel_prices_and_context_window_backup.json | 68 +++++++++++++++++++
 model_prices_and_context_window.json          | 68 +++++++++++++++++++
 .../llm_cost_calc/test_llm_cost_calc_utils.py | 48 +++++++++++++
 tests/test_litellm/test_cost_calculator.py    | 14 ++++
 ...penai_service_tier_long_context_pricing.py |  7 ++
 5 files changed, 205 insertions(+)

diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 9b358a1cedd..04a67200ced 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -29277,6 +29277,74 @@
         "supports_xhigh_reasoning_effort": true,
         "supports_minimal_reasoning_effort": true
     },
+    "gpt-6-astra": {
+        "cache_creation_input_token_cost": 1.25e-05,
+        "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05,
+        "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05,
+        "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05,
+        "cache_creation_input_token_cost_flex": 6.25e-06,
+        "cache_creation_input_token_cost_priority": 2.5e-05,
+        "cache_read_input_token_cost": 1e-06,
+        "cache_read_input_token_cost_above_272k_tokens": 2e-06,
+        "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06,
+        "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06,
+        "cache_read_input_token_cost_flex": 5e-07,
+        "cache_read_input_token_cost_priority": 2e-06,
+        "input_cost_per_token": 1e-05,
+        "input_cost_per_token_above_272k_tokens": 2e-05,
+        "input_cost_per_token_above_272k_tokens_flex": 1e-05,
+        "input_cost_per_token_above_272k_tokens_priority": 4e-05,
+        "input_cost_per_token_batches": 5e-06,
+        "input_cost_per_token_flex": 5e-06,
+        "input_cost_per_token_priority": 2e-05,
+        "litellm_provider": "openai",
+        "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,
+        "output_cost_per_token_above_272k_tokens_flex": 3.75e-05,
+        "output_cost_per_token_above_272k_tokens_priority": 0.00015,
+        "output_cost_per_token_batches": 2.5e-05,
+        "output_cost_per_token_flex": 2.5e-05,
+        "output_cost_per_token_priority": 0.0001,
+        "regional_processing_uplift_multiplier_eu": 1.1,
+        "regional_processing_uplift_multiplier_us": 1.1,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.01,
+            "search_context_size_low": 0.01,
+            "search_context_size_medium": 0.01
+        },
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/batch",
+            "/v1/responses"
+        ],
+        "supported_modalities": [
+            "text",
+            "image"
+        ],
+        "supported_output_modalities": [
+            "text"
+        ],
+        "supports_computer_use": true,
+        "supports_function_calling": 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_cache_breakpoint": 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
+    },
     "gpt-5.6": {
         "cache_creation_input_token_cost": 5e-06,
         "cache_creation_input_token_cost_above_272k_tokens": 1e-05,
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 9b358a1cedd..04a67200ced 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -29277,6 +29277,74 @@
         "supports_xhigh_reasoning_effort": true,
         "supports_minimal_reasoning_effort": true
     },
+    "gpt-6-astra": {
+        "cache_creation_input_token_cost": 1.25e-05,
+        "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05,
+        "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05,
+        "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05,
+        "cache_creation_input_token_cost_flex": 6.25e-06,
+        "cache_creation_input_token_cost_priority": 2.5e-05,
+        "cache_read_input_token_cost": 1e-06,
+        "cache_read_input_token_cost_above_272k_tokens": 2e-06,
+        "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06,
+        "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06,
+        "cache_read_input_token_cost_flex": 5e-07,
+        "cache_read_input_token_cost_priority": 2e-06,
+        "input_cost_per_token": 1e-05,
+        "input_cost_per_token_above_272k_tokens": 2e-05,
+        "input_cost_per_token_above_272k_tokens_flex": 1e-05,
+        "input_cost_per_token_above_272k_tokens_priority": 4e-05,
+        "input_cost_per_token_batches": 5e-06,
+        "input_cost_per_token_flex": 5e-06,
+        "input_cost_per_token_priority": 2e-05,
+        "litellm_provider": "openai",
+        "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,
+        "output_cost_per_token_above_272k_tokens_flex": 3.75e-05,
+        "output_cost_per_token_above_272k_tokens_priority": 0.00015,
+        "output_cost_per_token_batches": 2.5e-05,
+        "output_cost_per_token_flex": 2.5e-05,
+        "output_cost_per_token_priority": 0.0001,
+        "regional_processing_uplift_multiplier_eu": 1.1,
+        "regional_processing_uplift_multiplier_us": 1.1,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.01,
+            "search_context_size_low": 0.01,
+            "search_context_size_medium": 0.01
+        },
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/batch",
+            "/v1/responses"
+        ],
+        "supported_modalities": [
+            "text",
+            "image"
+        ],
+        "supported_output_modalities": [
+            "text"
+        ],
+        "supports_computer_use": true,
+        "supports_function_calling": 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_cache_breakpoint": 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
+    },
     "gpt-5.6": {
         "cache_creation_input_token_cost": 5e-06,
         "cache_creation_input_token_cost_above_272k_tokens": 1e-05,
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 b7f0ca1efe1..0b6832d4bef 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
@@ -1662,6 +1662,54 @@ def test_generic_cost_per_token_gpt56_cyber(
     assert completion_cost == pytest.approx(completion_tokens * output_rate)
 
 
+@pytest.mark.parametrize(
+    "service_tier,tier_multiplier",
+    [(None, 1.0), ("flex", 0.5), ("priority", 2.0), ("fast", 2.0)],
+)
+@pytest.mark.parametrize(
+    "prompt_tokens,input_side_multiplier,output_multiplier",
+    [(100000, 1.0, 1.0), (300000, 2.0, 1.5)],
+)
+def test_generic_cost_per_token_gpt_6_astra_price_sheet(
+    _local_model_cost_map,
+    service_tier,
+    tier_multiplier,
+    prompt_tokens,
+    input_side_multiplier,
+    output_multiplier,
+):
+    """gpt-6-astra launch price sheet: $10 input, $1 cache read, $12.50 cache write, $50 output per 1M tokens.
+
+    Above 272K prompt tokens the input-side rates double and the output rate is 1.5x on the whole
+    request. Flex is half the applicable rate and fast mode, billed as priority, is double it.
+    """
+    cached_tokens = 50000
+    cache_write_tokens = 40000
+    text_tokens = prompt_tokens - cached_tokens - cache_write_tokens
+    completion_tokens = 1000
+    usage = Usage(
+        prompt_tokens=prompt_tokens,
+        completion_tokens=completion_tokens,
+        total_tokens=prompt_tokens + completion_tokens,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            cached_tokens=cached_tokens, cache_write_tokens=cache_write_tokens
+        ),
+    )
+
+    prompt_cost, completion_cost = generic_cost_per_token(
+        model="gpt-6-astra",
+        usage=usage,
+        custom_llm_provider="openai",
+        service_tier=service_tier,
+    )
+
+    input_side = tier_multiplier * input_side_multiplier
+    assert prompt_cost == pytest.approx(
+        input_side * (text_tokens * 1e-5 + cached_tokens * 1e-6 + cache_write_tokens * 1.25e-5)
+    )
+    assert completion_cost == pytest.approx(tier_multiplier * output_multiplier * completion_tokens * 5e-5)
+
+
 @pytest.mark.parametrize(
     "model,input_cost,output_cost,cache_read_cost",
     [
diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index 7c2174018e8..6dc3b2790c9 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -4473,3 +4473,17 @@ def test_explicit_pricing_precedes_private_provider_response_model(
     )
 
     assert selected == expected
+
+
+def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_model_cost_map):
+    """gpt-6-astra batch pricing is 50% off the standard $10 input and $50 output rates per 1M tokens."""
+    from litellm.cost_calculator import batch_cost_calculator
+
+    usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
+
+    prompt_cost, completion_cost = batch_cost_calculator(
+        usage=usage, model="gpt-6-astra", custom_llm_provider="openai"
+    )
+
+    assert prompt_cost == pytest.approx(1000 * 5e-6)
+    assert completion_cost == pytest.approx(500 * 2.5e-5)
diff --git a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py
index c0860a5b55f..70e9c2720b8 100644
--- a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py
+++ b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py
@@ -52,6 +52,12 @@ PRIORITY_LONG_CONTEXT = {
         "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08,
         "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06,
     },
+    "gpt-6-astra": {
+        "input_cost_per_token_above_272k_tokens_priority": 4e-05,
+        "output_cost_per_token_above_272k_tokens_priority": 0.00015,
+        "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06,
+        "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05,
+    },
 }
 
 EXPECTED = {**FLEX_LONG_CONTEXT, **PRIORITY_LONG_CONTEXT}
@@ -114,6 +120,7 @@ TIERED_COST_CASES = [
     ("gpt-5.6-sol", "priority", 1.6e-05, 6e-05),
     ("gpt-5.6-terra", "priority", 8e-06, 3.6e-05),
     ("gpt-5.6-luna", "priority", 8e-07, 3.6e-06),
+    ("gpt-6-astra", "priority", 4e-05, 0.00015),
 ]
 
 

From 4991d0bf3e58fc1022d97ce8baf1a517c656d5a2 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 12:47:13 -0700
Subject: [PATCH 34/43] fix(models): match gpt-6-astra reasoning effort levels
 to OpenAI docs

OpenAI documents low, medium, high, xhigh, and max for gpt-6-astra, with no none level, so the entry stops advertising none and starts advertising max.
---
 .../model_prices_and_context_window_backup.json |  3 ++-
 model_prices_and_context_window.json            |  3 ++-
 .../test_reasoning_effort_capability.py         | 17 +++++++++++++++++
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 04a67200ced..385cf22d06c 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -29330,9 +29330,10 @@
         ],
         "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_none_reasoning_effort": false,
         "supports_parallel_function_calling": true,
         "supports_pdf_input": true,
         "supports_prompt_cache_breakpoint": true,
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 04a67200ced..385cf22d06c 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -29330,9 +29330,10 @@
         ],
         "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_none_reasoning_effort": false,
         "supports_parallel_function_calling": true,
         "supports_pdf_input": true,
         "supports_prompt_cache_breakpoint": 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 a0dbf3b6637..7b2e45ab3ed 100644
--- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py
+++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py
@@ -371,3 +371,20 @@ class TestKimiK3AdvertisesItsDocumentedLevels:
             "low",
             "high",
         )
+
+
+class TestGpt6AstraAdvertisesItsDocumentedLevels:
+    def test_the_entry_advertises_low_through_max_without_none(self, local_model_cost_map):
+        """OpenAI documents low, medium, high, xhigh and max for gpt-6-astra. Unlike gpt-5.6-sol it
+        does not take none, so a group must not offer none and must offer max."""
+        from litellm.utils import _get_model_info_helper
+
+        model_info = dict(_get_model_info_helper(model="gpt-6-astra", custom_llm_provider="openai"))
+
+        assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == (
+            "low",
+            "medium",
+            "high",
+            "xhigh",
+            "max",
+        )

From f40f14ae39b4452cb55f117ab3b06949bb3135a7 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 13:09:46 -0700
Subject: [PATCH 35/43] fix(tests): fold the local price map into the provider
 model sets

CI unit shards load the price map from main at import, so a model that only
exists on the branch never reaches open_ai_chat_completion_models and
cost_per_token cannot infer its provider. Refresh the sets after swapping in
the local map so the tier pricing cases resolve gpt-6-astra before merge
---
 .../test_openai_service_tier_long_context_pricing.py             | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py
index 70e9c2720b8..bdb2dc26813 100644
--- a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py
+++ b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py
@@ -69,6 +69,7 @@ NO_PUBLISHED_PRIORITY_LONG_CONTEXT = ("gpt-5.4", "gpt-5.5")
 def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
     monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
+    litellm.add_known_models()
 
 
 @lru_cache(maxsize=2)

From 1f20b381151d1c01d38edb56743dee4b771fa25a Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Thu, 3 Sep 2026 13:17:48 -0700
Subject: [PATCH 36/43] fix(vector_stores): only list vector stores the caller
 was granted (#39612)

* fix(vector_stores): only list vector stores the caller was granted

/vector_store/list returned every managed vector store with no team_id to any
key, and let a dashboard session see stores created from the dashboard because
every session shares the litellm-dashboard team id. Non-admin listings now show
a store only when the key or one of the caller's real teams is allowlisted for
it via object_permission.vector_stores, or the team owns it

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

* fix(vector_stores): keep a dashboard session key's own grants when the user has no teams

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>
---
 .../management_endpoints.py                   |  14 ++-
 litellm/proxy/vector_store_endpoints/utils.py |  78 +++++++++++-
 .../test_vector_store_access_control.py       | 116 +++++++++++++++++-
 3 files changed, 192 insertions(+), 16 deletions(-)

diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py
index c928398a87f..fe4732c6492 100644
--- a/litellm/proxy/vector_store_endpoints/management_endpoints.py
+++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py
@@ -30,7 +30,10 @@ from litellm.proxy._types import (
 )
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
-from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store
+from litellm.proxy.vector_store_endpoints.utils import (
+    can_user_access_vector_store,
+    filter_listable_vector_stores,
+)
 from litellm.repositories.prisma_protocols import TableActions
 from litellm.repositories.table_repositories import ManagedVectorStoresRepository
 from litellm.types.vector_stores import (
@@ -390,11 +393,10 @@ async def list_vector_stores(
 
         # Filter vector stores based on access control
         accessible_vector_stores: Final = []
-        for vs in vector_store_map.values():
-            if await _check_vector_store_access(vs, user_api_key_dict):
-                redacted = LiteLLM_ManagedVectorStore(**vs)
-                redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params"))
-                accessible_vector_stores.append(redacted)
+        for vs in await filter_listable_vector_stores(vector_store_map.values(), user_api_key_dict):
+            redacted = LiteLLM_ManagedVectorStore(**vs)
+            redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params"))
+            accessible_vector_stores.append(redacted)
 
         total_count: Final = len(accessible_vector_stores)
         total_pages: Final = (total_count + page_size - 1) // page_size
diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py
index 93f1510bf22..6e94a5a88ac 100644
--- a/litellm/proxy/vector_store_endpoints/utils.py
+++ b/litellm/proxy/vector_store_endpoints/utils.py
@@ -1,11 +1,17 @@
 import json
 import re
+from collections.abc import Iterable
+from types import MappingProxyType
 from typing import Any, Final, Literal
 
 from fastapi import HTTPException, Request
 
 import litellm
 from litellm._logging import verbose_proxy_logger
+from litellm.proxy._experimental.mcp_server.ui_session_utils import (
+    is_ui_session_credential,
+    resolve_ui_session_team_ids,
+)
 from litellm.proxy._types import (
     LiteLLM_ObjectPermissionTable,
     LitellmUserRoles,
@@ -160,10 +166,16 @@ async def can_user_access_vector_store(
     if _is_proxy_admin(user_api_key_dict):
         return True
 
-    vector_store_team_id: Final = vector_store.get("team_id")
-    if vector_store_team_id is None:
+    if vector_store.get("team_id") is None:
         return True
 
+    return await _is_vector_store_granted(vector_store, user_api_key_dict)
+
+
+async def _is_vector_store_granted(
+    vector_store: LiteLLM_ManagedVectorStore,
+    user_api_key_dict: UserAPIKeyAuth,
+) -> bool:
     vector_store_id: Final = vector_store.get("vector_store_id") or ""
 
     key_object_permission = user_api_key_dict.object_permission
@@ -178,12 +190,70 @@ async def can_user_access_vector_store(
     if _object_permission_allows_vector_store(team_object_permission, vector_store_id):
         return True
 
-    if user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store_team_id:
-        return True
+    return user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store.get("team_id")
 
+
+async def _team_auth_context(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth:
+    from litellm.proxy.auth.auth_checks import get_team_object
+    from litellm.proxy.proxy_server import (
+        prisma_client,
+        proxy_logging_obj,
+        user_api_key_cache,
+    )
+
+    team: Final = await get_team_object(
+        team_id=team_id,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        parent_otel_span=user_api_key_dict.parent_otel_span,
+        proxy_logging_obj=proxy_logging_obj,
+    )
+    return user_api_key_dict.model_copy(
+        update=MappingProxyType(
+            {
+                "team_id": team_id,
+                "team_object_permission": team.object_permission,
+                "team_object_permission_id": team.object_permission_id,
+            }
+        )
+    )
+
+
+async def _vector_store_listing_auth_contexts(
+    user_api_key_dict: UserAPIKeyAuth,
+) -> tuple[UserAPIKeyAuth, ...]:
+    if not is_ui_session_credential(user_api_key_dict):
+        return (user_api_key_dict,)
+    session_key_context: Final = user_api_key_dict.model_copy(
+        update=MappingProxyType({"team_id": None, "team_object_permission": None, "team_object_permission_id": None})
+    )
+    team_ids: Final = await resolve_ui_session_team_ids(user_api_key_dict)
+    team_contexts: Final = tuple([await _team_auth_context(team_id, user_api_key_dict) for team_id in team_ids])
+    return (session_key_context, *team_contexts)
+
+
+async def _is_vector_store_granted_to_any(
+    vector_store: LiteLLM_ManagedVectorStore,
+    auth_contexts: tuple[UserAPIKeyAuth, ...],
+) -> bool:
+    for auth_context in auth_contexts:
+        if await _is_vector_store_granted(vector_store, auth_context):
+            return True
     return False
 
 
+async def filter_listable_vector_stores(
+    vector_stores: Iterable[LiteLLM_ManagedVectorStore],
+    user_api_key_dict: UserAPIKeyAuth,
+) -> tuple[LiteLLM_ManagedVectorStore, ...]:
+    """Non-admins only see stores their key, one of their teams' object_permission, or team ownership grants."""
+    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)
+    return tuple([vs for vs in vector_stores if await _is_vector_store_granted_to_any(vs, auth_contexts)])
+
+
 async def get_litellm_managed_vector_store(
     vector_store_id: str,
 ) -> LiteLLM_ManagedVectorStore | None:
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 7d72121456a..93049b21460 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
@@ -120,9 +120,7 @@ async def test_delete_vector_store_checks_access():
             "team_id": "team_456",
         }
     )
-    mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock(
-        return_value=mock_vector_store
-    )
+    mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=mock_vector_store)
 
     # User from different team should get 403
     user_api_key_dict = UserAPIKeyAuth(team_id="team_789")
@@ -134,9 +132,115 @@ async def test_delete_vector_store_checks_access():
     ):
         with patch("litellm.vector_store_registry", None):
             with pytest.raises(HTTPException) as exc_info:
-                await delete_vector_store(
-                    data=request, user_api_key_dict=user_api_key_dict
-                )
+                await delete_vector_store(data=request, user_api_key_dict=user_api_key_dict)
 
             assert exc_info.value.status_code == 403
             assert "Access denied" in exc_info.value.detail
+
+
+_UNSCOPED: LiteLLM_ManagedVectorStore = {
+    "vector_store_id": "vs_unscoped",
+    "custom_llm_provider": "openai",
+    "team_id": None,
+}
+_TEAM_A_OWNED: LiteLLM_ManagedVectorStore = {
+    "vector_store_id": "vs_team_a",
+    "custom_llm_provider": "openai",
+    "team_id": "team_a",
+}
+_UI_CREATED: LiteLLM_ManagedVectorStore = {
+    "vector_store_id": "vs_ui_created",
+    "custom_llm_provider": "openai",
+    "team_id": "litellm-dashboard",
+}
+
+
+async def _listed_ids(user_api_key_dict: UserAPIKeyAuth) -> list[str]:
+    from litellm.proxy.vector_store_endpoints.management_endpoints import (
+        list_vector_stores,
+    )
+
+    with patch(  # test-quality-ok: the list route reads rows through this module-level DB helper, no injection seam
+        "litellm.proxy.vector_store_endpoints.management_endpoints.VectorStoreRegistry._get_vector_stores_from_db",
+        new=AsyncMock(return_value=[_UNSCOPED, _TEAM_A_OWNED, _UI_CREATED]),
+    ):
+        response = await list_vector_stores(user_api_key_dict=user_api_key_dict)
+    return sorted(vs["vector_store_id"] for vs in response["data"])
+
+
+@pytest.mark.asyncio
+async def test_list_vector_stores_hides_ungranted_stores_from_non_admin_keys():
+    """A store with no team_id and no allowlist entry is not listed for a key it was never granted to;
+    only team ownership or an explicit object_permission grant makes a store visible."""
+    assert await _listed_ids(UserAPIKeyAuth()) == []
+    assert await _listed_ids(UserAPIKeyAuth(team_id="team_a")) == ["vs_team_a"]
+    assert await _listed_ids(
+        UserAPIKeyAuth(
+            team_id="team_b",
+            object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-1", vector_stores=["vs_unscoped"]),
+        )
+    ) == ["vs_unscoped"]
+    assert await _listed_ids(
+        UserAPIKeyAuth(
+            team_id="team_b",
+            team_object_permission=LiteLLM_ObjectPermissionTable(
+                object_permission_id="op-2", vector_stores=["vs_unscoped"]
+            ),
+        )
+    ) == ["vs_unscoped"]
+    assert await _listed_ids(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)) == [
+        "vs_team_a",
+        "vs_ui_created",
+        "vs_unscoped",
+    ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("user_team_ids", "session_key_grants", "expected"),
+    [
+        ([], None, []),
+        ([], ["vs_unscoped"], ["vs_unscoped"]),
+        (["team_a"], None, ["vs_team_a"]),
+        (["team_a", "team_granted"], None, ["vs_team_a", "vs_unscoped"]),
+    ],
+)
+async def test_list_vector_stores_dashboard_session_resolves_real_teams(
+    user_team_ids: list[str], session_key_grants: list[str] | None, expected: list[str]
+):
+    """A dashboard session lists through the user's real teams plus the session key's own grants: stores created
+    from the dashboard (team_id litellm-dashboard) are not visible just because every session shares that team id,
+    while stores owned by or granted to one of the user's teams, or granted to the session key itself, are."""
+    from litellm.models.team import LiteLLM_TeamTableCachedObj
+
+    alice = UserAPIKeyAuth(
+        team_id="litellm-dashboard",
+        user_id="alice",
+        user_role=LitellmUserRoles.INTERNAL_USER,
+        object_permission=(
+            LiteLLM_ObjectPermissionTable(object_permission_id="op-4", vector_stores=session_key_grants)
+            if session_key_grants is not None
+            else None
+        ),
+    )
+    teams = {
+        "team_a": LiteLLM_TeamTableCachedObj(team_id="team_a"),
+        "team_granted": LiteLLM_TeamTableCachedObj(
+            team_id="team_granted",
+            object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-3", vector_stores=["vs_unscoped"]),
+        ),
+    }
+
+    async def fake_get_team_object(team_id: str, **_kwargs: object) -> LiteLLM_TeamTableCachedObj:
+        return teams[team_id]
+
+    with (
+        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 _listed_ids(alice) == expected

From ab515dbc90ef6aad53f4d064ee61348db37a3410 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 13:20:13 -0700
Subject: [PATCH 37/43] fix: treat gpt-6 names as the gpt-5 request family in
 OpenAI and Azure configs

---
 .../llms/azure/chat/gpt_5_transformation.py   | 22 ++--------------
 litellm/llms/azure/chat/gpt_transformation.py |  3 ++-
 .../llms/openai/chat/gpt_5_transformation.py  | 26 ++++++++-----------
 .../llms/openai/responses/transformation.py   |  3 ++-
 .../chat/test_azure_gpt5_transformation.py    | 12 +++++++++
 .../test_openai_responses_transformation.py   |  2 ++
 .../llms/openai/test_gpt5_transformation.py   | 14 ++++++++++
 .../llms/openai/test_is_model_gpt_5_model.py  |  4 +++
 tests/test_litellm/test_main.py               | 15 +++++++++++
 9 files changed, 64 insertions(+), 37 deletions(-)

diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_5_transformation.py
index 6fdd277a04f..3189f5b57ac 100644
--- a/litellm/llms/azure/chat/gpt_5_transformation.py
+++ b/litellm/llms/azure/chat/gpt_5_transformation.py
@@ -7,6 +7,7 @@ from litellm.exceptions import UnsupportedParamsError
 from litellm.llms.openai.chat.gpt_5_transformation import (
     OpenAIGPT5Config,
     _get_effort_level,
+    is_gpt_reasoning_series_name,
 )
 from litellm.types.llms.openai import AllMessageValues
 
@@ -35,26 +36,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
 
     @classmethod
     def is_model_gpt_5_model(cls, model: str) -> bool:
-        """Check if the Azure model string refers to a gpt-5 variant.
-
-        Accepts both explicit gpt-5 model names and the ``gpt5_series/`` prefix
-        used for manual routing.
-        """
-        # The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07,
-        # …) are regular chat models: they support temperature and tool_choice but NOT
-        # reasoning_effort.  They must NOT be routed through the GPT-5 reasoning path.
-        #
-        # Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning
-        # models and must stay on the GPT-5 path.  The distinguishing feature is that
-        # the gpt-5-chat family has a literal "-chat" immediately after "gpt-5"
-        # (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version
-        # number (i.e. "gpt-5.<digit>-chat").
-        #
-        # Using a startswith("gpt-5-chat") prefix check on the normalized name (rather
-        # than a substring check) makes this boundary explicit and avoids any ambiguity
-        # if future model names coincidentally contain "gpt-5-chat" as an interior run.
-        _normalized: Final = model.split("/")[-1]  # strip provider prefix, e.g. "azure/"
-        return ("gpt-5" in model and not _normalized.startswith("gpt-5-chat")) or "gpt5_series" in model
+        return is_gpt_reasoning_series_name(model) or "gpt5_series" in model
 
     def get_supported_openai_params(self, model: str) -> list[str]:
         """Get supported parameters for Azure OpenAI GPT-5 models.
diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py
index 0ac0662205a..880a51eb584 100644
--- a/litellm/llms/azure/chat/gpt_transformation.py
+++ b/litellm/llms/azure/chat/gpt_transformation.py
@@ -14,6 +14,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
     convert_to_azure_openai_messages,
 )
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.llms.openai.chat.gpt_5_transformation import GPT_REASONING_SERIES_MARKERS
 from litellm.types.llms.azure import (
     API_VERSION_MONTH_SUPPORTED_RESPONSE_FORMAT,
     API_VERSION_YEAR_SUPPORTED_RESPONSE_FORMAT,
@@ -139,7 +140,7 @@ class AzureOpenAIConfig(BaseConfig):
         name family needs the rename, including the ``gpt-5-chat*`` models that are excluded from
         the reasoning path by https://github.com/BerriAI/litellm/issues/13781.
         """
-        return "gpt-5" in model or "gpt5_series" in model
+        return any(marker in model for marker in GPT_REASONING_SERIES_MARKERS) or "gpt5_series" in model
 
     def _is_response_format_supported_model(self, model: str) -> bool:
         """
diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py
index 0223be300b0..b02f953425d 100644
--- a/litellm/llms/openai/chat/gpt_5_transformation.py
+++ b/litellm/llms/openai/chat/gpt_5_transformation.py
@@ -61,6 +61,14 @@ def _get_effort_level(value: str | dict | None) -> str | None:
     return None
 
 
+GPT_REASONING_SERIES_MARKERS: Final = ("gpt-5", "gpt-6")
+
+
+def is_gpt_reasoning_series_name(model: str) -> bool:
+    normalized: Final = model.split("/")[-1]
+    return any(marker in model for marker in GPT_REASONING_SERIES_MARKERS) and not normalized.startswith("gpt-5-chat")
+
+
 class OpenAIGPT5Config(OpenAIGPTConfig):
     """Configuration for gpt-5 models including GPT-5-Codex variants.
 
@@ -73,21 +81,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
 
     @classmethod
     def is_model_gpt_5_model(cls, model: str) -> bool:
-        # The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07,
-        # …) are regular chat models: they support temperature and tool_choice but NOT
-        # reasoning_effort.  They must NOT be routed through the GPT-5 reasoning path.
-        #
-        # Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning
-        # models and must stay on the GPT-5 path.  The distinguishing feature is that
-        # the gpt-5-chat family has a literal "-chat" immediately after "gpt-5"
-        # (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version
-        # number (i.e. "gpt-5.<digit>-chat").
-        #
-        # Using a startswith("gpt-5-chat") prefix check on the normalized name (rather
-        # than a substring check) makes this boundary explicit and avoids any ambiguity
-        # if future model names coincidentally contain "gpt-5-chat" as an interior run.
-        _normalized: Final = model.split("/")[-1]  # strip provider prefix, e.g. "openai/"
-        return "gpt-5" in model and not _normalized.startswith("gpt-5-chat")
+        return is_gpt_reasoning_series_name(model)
 
     @classmethod
     def is_model_gpt_5_search_model(cls, model: str) -> bool:
@@ -122,6 +116,8 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
     def is_model_gpt_5_4_plus_model(cls, model: str) -> bool:
         """Check if the model is gpt-5.4 or newer (5.4, 5.5, 5.6, etc., including pro)."""
         model_name: Final = model.split("/")[-1]
+        if model_name.startswith("gpt-6"):
+            return True
         if not model_name.startswith("gpt-5."):
             return False
         try:
diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py
index 01313e95878..b97521b90c2 100644
--- a/litellm/llms/openai/responses/transformation.py
+++ b/litellm/llms/openai/responses/transformation.py
@@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
 )
 from litellm.litellm_core_utils.url_utils import encode_url_path_segment
 from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
+from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name
 from litellm.responses.litellm_completion_transformation.custom_tools import TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import *
@@ -88,7 +89,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
         parts: Final = model.split("/")
         if len(parts) > 1 and parts[0] not in ("openai",):
             return False
-        return "gpt-5" in model and "gpt-5-chat" not in model
+        return is_gpt_reasoning_series_name(model)
 
     @staticmethod
     def _supports_reasoning_effort_none(model: str) -> bool:
diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py
index e06cae97283..bd0f16a695b 100644
--- a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py
+++ b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py
@@ -336,3 +336,15 @@ class TestAzureResolvesTheDeclaredDefaultEffort:
             drop_params=True,
         )
         assert ("temperature" in mapped) is temperature_survives
+
+
+def test_azure_gpt_6_astra_takes_the_reasoning_series_request_shape():
+    params = litellm.get_optional_params(
+        model="gpt-6-astra",
+        custom_llm_provider="azure",
+        max_tokens=100,
+        reasoning_effort="max",
+    )
+    assert params["max_completion_tokens"] == 100
+    assert "max_tokens" not in params
+    assert params["reasoning_effort"] == "max"
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 b0ffd1845fe..4ac072d0ca6 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
@@ -1718,6 +1718,8 @@ class TestResponsesSurfaceSharesTheEffortRule:
             ("gpt-5.6-sol", None, False),
             ("gpt-5.6-terra", "none", True),
             ("gpt-5.6-terra", "medium", False),
+            ("gpt-6-astra", None, False),
+            ("gpt-6-astra", "low", False),
         ],
     )
     def test_temperature_follows_the_resolved_effort(
diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py
index 9c5bd34d59a..c86ce4df2ac 100644
--- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py
+++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py
@@ -1505,3 +1505,17 @@ class TestACatalogueOlderThanTheCodeDoesNotStripTemperature:
             drop_params=True,
         )
         assert "temperature" not in mapped
+
+
+def test_gpt_6_astra_takes_the_reasoning_series_request_shape():
+    params = litellm.get_optional_params(
+        model="gpt-6-astra",
+        custom_llm_provider="openai",
+        max_tokens=100,
+        reasoning_effort="max",
+        verbosity="low",
+    )
+    assert params["max_completion_tokens"] == 100
+    assert "max_tokens" not in params
+    assert params["reasoning_effort"] == "max"
+    assert params["verbosity"] == "low"
diff --git a/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py b/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py
index 1095819c98c..107a1afb2c6 100644
--- a/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py
+++ b/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py
@@ -41,6 +41,8 @@ from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config
 
 # Models that MUST be classified as GPT-5 (routed through GPT-5 reasoning path)
 GPT5_MODELS = [
+    "gpt-6-astra",
+    "openai/gpt-6-astra",
     "gpt-5",
     "gpt-5.1",
     "gpt-5.2",
@@ -120,6 +122,8 @@ class TestOpenAIGPT5ConfigIsModelGpt5Model:
 # /v1/responses bridge (when reasoning_effort is set and tools are passed) on
 # is_model_gpt_5_4_plus_model, so the gpt-5.6 family must land on the True side.
 GPT5_4_PLUS_MODELS = [
+    "gpt-6-astra",
+    "openai/gpt-6-astra",
     "gpt-5.4",
     "gpt-5.5",
     "gpt-5.5-pro",
diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py
index 3c8bf142835..c9acbe2d884 100644
--- a/tests/test_litellm/test_main.py
+++ b/tests/test_litellm/test_main.py
@@ -784,6 +784,21 @@ def test_responses_api_bridge_check_gpt_5_4_tools_plus_reasoning_routes_to_respo
     assert model_info.get("mode") == "responses"
 
 
+def test_responses_api_bridge_check_gpt_6_astra_tools_with_default_reasoning_routes_to_responses():
+    from litellm.main import responses_api_bridge_check
+
+    with patch("litellm.main._get_model_info_helper") as mock_get_model_info:
+        mock_get_model_info.return_value = {"max_tokens": 128000}
+        model_info, model = responses_api_bridge_check(
+            model="gpt-6-astra",
+            custom_llm_provider="openai",
+            tools=[{"type": "function", "function": {"name": "get_capital"}}],
+        )
+
+    assert model == "gpt-6-astra"
+    assert model_info.get("mode") == "responses"
+
+
 def test_responses_api_bridge_check_gpt_5_5_tools_plus_reasoning_routes_to_responses():
     """gpt-5.5+ with both tools and reasoning_effort should route to Responses API."""
     from litellm.main import responses_api_bridge_check

From b86a0b5562a1a5fbedcd1a5c2b6a02e7f89e64ee Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Thu, 3 Sep 2026 16:26:51 -0400
Subject: [PATCH 38/43] test(router): cover get_configured_mode so
 router_code_coverage passes

425e3069b9 added Router.get_configured_mode but only exercised it
through create_model_info_response, which the router coverage gate does
not count. The code-quality workflow has been failing on staging and on
every open PR since.
---
 tests/test_litellm/test_router.py | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 3b4c80b6b4f..f91920a636b 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -12553,3 +12553,33 @@ async def test_prompt_management_factory_marks_injection_for_every_deployment(mo
     bucket = captured.get("litellm_metadata") or captured["metadata"]
     assert captured["model_info"]["id"] == "provisional-dep"
     assert bucket["litellm_gateway_injected_cache"] == ""
+
+
+def test_get_configured_mode_reads_deployment_model_info():
+    router = Router(
+        model_list=[
+            {
+                "model_name": "my-tts",
+                "litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
+                "model_info": {"mode": "audio_speech"},
+            }
+        ]
+    )
+
+    assert router.get_configured_mode("my-tts") == "audio_speech"
+
+
+@pytest.mark.parametrize("model_info", [{}, {"mode": ""}, {"mode": "   "}, {"mode": 123}])
+def test_get_configured_mode_returns_none_for_unset_blank_or_unknown(model_info):
+    router = Router(
+        model_list=[
+            {
+                "model_name": "plain-model",
+                "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"},
+                "model_info": model_info,
+            }
+        ]
+    )
+
+    assert router.get_configured_mode("plain-model") is None
+    assert router.get_configured_mode("unknown-model") is None

From df73c623b231b68d690f349a7bb70a05b4c82333 Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Thu, 3 Sep 2026 13:39:58 -0700
Subject: [PATCH 39/43] feat(router): limit heuristic_v2 auto-routers to one
 without the auto_router license feature (#39468)

Without the auto_router feature in the signed enterprise license a proxy may hold
one complexity router with classifier_type heuristic_v2 across config.yaml and the
DB; with it the limit is lifted. The ceiling is derived once from LicenseCheck and
handed to the Router, which refuses the extra router at registration. config.yaml
over the limit refuses to start, and /model/new, /model/update and
PATCH /model/{id}/update refuse the write with a 403 before touching the DB.
Expiry follows the existing max_users/max_teams pattern: judged when the
license is verified, not on every call, and a verify that rejects the license
(expired or unreadable) leaves no signed payload behind. The rollback after a
failed upsert re-admits state that was already serving, so it is exempt from the
ceiling: an edit that fails, including one refused by a ceiling that has since
tightened, leaves the router serving its previous configuration.
A write that leaves a row on heuristic_v2 under a limited license runs in one
transaction that takes a Postgres advisory lock before counting the DB rows plus
this proxy's config.yaml routers, so concurrent writes on any pod cannot both
claim the sole slot and no surplus row is ever persisted.
Only the row insert runs under that lock: the team model bookkeeping, which
needs a second pool connection, runs after the transaction has committed.
PATCH /model/{id}/update follows the same order as create: the row is written
through the slot first and the team's model list is updated only afterwards, so
a refused write leaves the team as it was.
The slot transaction bypasses the repository's publish-on-write, so it
publishes the config change once after commit, as delete_team_models does.
---
 litellm/constants.py                          |   1 +
 litellm/proxy/auth/litellm_license.py         |  23 +-
 .../model_management_endpoints.py             | 198 +++++---
 litellm/proxy/proxy_server.py                 |  20 +-
 litellm/router.py                             |  46 +-
 .../router_utils/auto_router_model_naming.py  |  34 +-
 litellm/types/router.py                       |  12 +
 .../proxy/auth/test_litellm_license.py        |  69 +++
 .../test_model_management_endpoints.py        | 421 ++++++++++++++++--
 .../test_ptu_model_settings.py                |   6 +
 .../proxy/proxy_server/test_proxy_config.py   | 115 +++++
 .../router_strategy/test_complexity_router.py | 160 +++++++
 .../test_auto_router_model_naming.py          |  57 +++
 13 files changed, 1067 insertions(+), 95 deletions(-)

diff --git a/litellm/constants.py b/litellm/constants.py
index 1063b6ddeeb..f5acadc32ab 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -40,6 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset(
         "router_general_settings",
         "ignore_invalid_deployments",
         "fallback_access_check",
+        "heuristic_v2_router_limit",
     }
 )
 DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py
index 677f1a0fdda..55bb1e3925a 100644
--- a/litellm/proxy/auth/litellm_license.py
+++ b/litellm/proxy/auth/litellm_license.py
@@ -16,6 +16,10 @@ if TYPE_CHECKING:
     from litellm.proxy._types import EnterpriseLicenseData
 
 
+AUTO_ROUTER_LICENSE_FEATURE: Final = "auto_router"
+HEURISTIC_V2_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit."
+
+
 class LicenseCheck:
     """
     - Check if license in env
@@ -149,6 +153,19 @@ class LicenseCheck:
             return False
         return team_count > _max_teams_in_license
 
+    def heuristic_v2_router_limit(self) -> int | None:
+        """
+        How many heuristic_v2 auto-routers this proxy may hold: unlimited (None) only when the
+        signed license lists the auto_router feature, otherwise one. A license verified through
+        the API carries no feature list, so it does not lift the limit either.
+        """
+        if self.airgapped_license_data is None:
+            return 1
+        allowed_features: Final = self.airgapped_license_data.get("allowed_features")
+        if isinstance(allowed_features, list) and AUTO_ROUTER_LICENSE_FEATURE in allowed_features:
+            return None
+        return 1
+
     def verify_license_without_api_request(self, public_key, license_key):
         try:
             from cryptography.hazmat.primitives import hashes
@@ -179,19 +196,21 @@ class LicenseCheck:
             # Decode and parse the data
             license_data: Final = json.loads(message.decode())
 
-            self.airgapped_license_data = EnterpriseLicenseData(**license_data)
-
             # debug information provided in license data
             verbose_proxy_logger.debug("License data: %s", license_data)
 
             # Check expiration date
             expiration_date: Final = datetime.strptime(license_data["expiration_date"], "%Y-%m-%d")
             if expiration_date < datetime.now():
+                self.airgapped_license_data = None
                 return False, "License has expired"
 
+            self.airgapped_license_data = EnterpriseLicenseData(**license_data)
+
             return True
 
         except Exception as e:
+            self.airgapped_license_data = None
             verbose_proxy_logger.debug(
                 "litellm.proxy.auth.litellm_license.py::verify_license_without_api_request - Unable to verify License locally. - %s",
                 e,
diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py
index 613e726f89d..82ee33cbc39 100644
--- a/litellm/proxy/management_endpoints/model_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/model_management_endpoints.py
@@ -13,10 +13,11 @@ model/{model_id}/update - PATCH endpoint for model update.
 import asyncio
 import datetime
 import json
-from collections.abc import Awaitable, Mapping, Sequence
+from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
 from json import JSONDecodeError
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast
+from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeVar, cast
 
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
 from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator
@@ -49,6 +50,7 @@ from litellm.proxy._types import (
     TeamModelDeleteRequest,
     UserAPIKeyAuth,
 )
+from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_utils.config_sync_pubsub import (
     coordination_redis_cache,
@@ -96,6 +98,9 @@ from litellm.router_strategy.complexity_router import (
 from litellm.router_utils.auto_router_model_naming import (
     STRATEGY_ROUTER_PARAM_FIELDS,
     carries_complexity_router_settings,
+    count_heuristic_v2_routers,
+    heuristic_v2_limit_violation,
+    uses_heuristic_v2_classifier,
     validate_complexity_router_config_placement,
     validate_complexity_router_config_write,
     validate_strategy_router_model_write,
@@ -153,6 +158,8 @@ class _ProxyModelTable(Protocol):
 
     def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ...
 
+    def create(self, *, data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ...
+
     def update(
         self, *, where: Mapping[str, object], data: Mapping[str, object]
     ) -> Awaitable[_ProxyModelRow | None]: ...
@@ -166,6 +173,9 @@ class _TxModelTables(Protocol):
     litellm_proxymodeltable: _ProxyModelTable
 
 
+_RowT = TypeVar("_RowT")
+
+
 class _ExistingModelRow(Protocol):
     @property
     def litellm_params(self) -> Mapping[str, object]: ...
@@ -269,6 +279,66 @@ def _raise_on_strategy_router_write_violation(
     )
 
 
+HEURISTIC_V2_SLOT_LOCK_KEY: Final = 5_872_301
+_HEURISTIC_V2_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)"
+_HEURISTIC_V2_DB_ROWS_SQL: Final = """
+SELECT count(*)::int AS held FROM "LiteLLM_ProxyModelTable"
+WHERE model_id <> $1
+  AND (CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END)
+      -> 'complexity_router_config' ->> 'classifier_type' = 'heuristic_v2'
+"""
+
+
+def _effective_complexity_router_config(
+    incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None
+) -> object:
+    """The complexity config a write leaves on the row: the incoming one when the write carries it, else the stored one."""
+    incoming: Final = None if incoming_params is None else incoming_params.complexity_router_config
+    if incoming is not None or existing_params is None:
+        return incoming
+    return existing_params.complexity_router_config
+
+
+@asynccontextmanager
+async def _heuristic_v2_slot(
+    prisma_client: PrismaClient, *, effective_config: object, model_id: str | None
+) -> AsyncGenerator[_ProxyModelTable, None]:
+    """Hand out the model table to write through while the row's claim on a heuristic_v2 slot is settled.
+
+    A write that leaves the row on classifier_type heuristic_v2 under a limited license runs
+    inside one transaction that takes an advisory lock in its own statement before counting
+    (a statement's snapshot predates anything it locks), so pods cannot both pass the count:
+    the DB rows (any pod, either JSON shape) plus this proxy's config.yaml routers are judged
+    against the license limit and the write is refused with a 403 before it happens. The row
+    being edited keeps its own slot through ``model_id``. Every other write, and every write on
+    an unlimited license, goes through the repository table with no lock. Only the row write
+    itself may run inside: anything that needs a second connection (the team model bookkeeping)
+    must wait until the transaction has committed and the lock is released. The transaction
+    writes bypass the repository's publish-on-write, so the config change is published once
+    after commit, the way delete_team_models does.
+    """
+    from litellm.proxy.proxy_server import _license_check, llm_router
+
+    limit: Final = _license_check.heuristic_v2_router_limit()
+    if limit is None or not uses_heuristic_v2_classifier(effective_config):
+        yield _proxy_model_table(prisma_client)
+        return
+    async with prisma_client.db.tx() as tx_ctx:
+        tables: Final[_TxModelTables] = tx_ctx
+        await tx_ctx.query_raw(_HEURISTIC_V2_LOCK_SQL, HEURISTIC_V2_SLOT_LOCK_KEY)
+        rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw(_HEURISTIC_V2_DB_ROWS_SQL, model_id or "")
+        db_held: Final = rows[0].get("held") if rows else 0
+        config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments())
+        held: Final = (db_held if isinstance(db_held, int) else 0) + count_heuristic_v2_routers(config_rows)
+        violation: Final = heuristic_v2_limit_violation(held=held + 1, limit=limit)
+        if violation is not None:
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {HEURISTIC_V2_LICENSE_REMEDY}"
+            )
+        yield tables.litellm_proxymodeltable
+    await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable")
+
+
 ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add"
 _REQUIRED_RATE_LIMIT_FIELDS: Final = ("rpm", "tpm")
 
@@ -720,22 +790,29 @@ async def patch_model(
         )
 
         requested_model_name: Final = patch_data.model_name
+        stored_model_name: str | None = None
+
+        async def write_row(update_data: PrismaCompatibleUpdateDBModel) -> _ProxyModelRow | None:
+            nonlocal stored_model_name
+            stored_model_name = update_data.get("model_name")
+            update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name
+            update_data["updated_at"] = cast(str, get_utc_datetime())
+            async with _heuristic_v2_slot(
+                prisma_client,
+                effective_config=_effective_complexity_router_config(
+                    patch_data.litellm_params, db_model.litellm_params
+                ),
+                model_id=model_id,
+            ) as table:
+                return await table.update(where={"model_id": model_id}, data=update_data)
+
         # Handle team model updates with proper alias management
-        update_data: Final = await _update_team_model_in_db(
+        updated_model: Final = await _update_team_model_in_db(
             db_model=db_model,
             patch_data=patch_data,
             user_api_key_dict=user_api_key_dict,
             prisma_client=prisma_client,
-        )
-
-        # Add metadata about update
-        update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name
-        update_data["updated_at"] = cast(str, get_utc_datetime())
-
-        # Perform partial update
-        updated_model: Final = await _proxy_model_table(prisma_client).update(
-            where={"model_id": model_id},
-            data=update_data,
+            write_row=write_row,
         )
 
         if updated_model is None:
@@ -746,7 +823,6 @@ async def patch_model(
                 param=None,
             )
 
-        stored_model_name: Final = update_data.get("model_name")
         if (
             stored_model_name is not None
             and stored_model_name == requested_model_name
@@ -980,7 +1056,8 @@ async def _add_model_to_db(
     prisma_client: PrismaClient,
     new_encryption_key: str | None = None,
     should_create_model_in_db: bool = True,
-) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None":
+    slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None,
+) -> "_ProxyModelRow | LiteLLM_ProxyModelTable":
     # encrypt litellm params #
     _litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True)
     _original_litellm_model_name: Final = model_params.litellm_params.model
@@ -998,18 +1075,20 @@ async def _add_model_to_db(
     if model_params.model_info.id is not None:
         _data["model_id"] = model_params.model_info.id
     _create_data: Final = cast("Mapping[str, object]", _data)  # cast-ok: str-keyed json payload built just above
-    if should_create_model_in_db:
-        model_response = await ModelRepository(prisma_client).table.create(data=_create_data)
-    else:
-        model_response = LiteLLM_ProxyModelTable(**_data)
-    return model_response
+    if not should_create_model_in_db:
+        return LiteLLM_ProxyModelTable(**_data)
+    if slot is None:
+        return await _proxy_model_table(prisma_client).create(data=_create_data)
+    async with slot as table:
+        return await table.create(data=_create_data)
 
 
 async def _add_team_model_to_db(
     model_params: Deployment,
     user_api_key_dict: UserAPIKeyAuth,
     prisma_client: PrismaClient,
-) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None":
+    slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None,
+) -> "_ProxyModelRow | LiteLLM_ProxyModelTable":
     """
     If 'team_id' is provided,
 
@@ -1040,6 +1119,7 @@ async def _add_team_model_to_db(
         model_params=model_params,
         user_api_key_dict=user_api_key_dict,
         prisma_client=prisma_client,
+        slot=slot,
     )
 
     if original_model_name:
@@ -1060,7 +1140,8 @@ async def _update_team_model_in_db(
     patch_data: updateDeployment,
     user_api_key_dict: UserAPIKeyAuth,
     prisma_client: PrismaClient,
-) -> PrismaCompatibleUpdateDBModel:
+    write_row: Callable[[PrismaCompatibleUpdateDBModel], Awaitable[_RowT]],
+) -> _RowT:
     """
     Handle team model updates with proper alias management.
 
@@ -1068,6 +1149,9 @@ async def _update_team_model_in_db(
     - Creates unique internal model_name and team alias
     - Adds model to team object
     - Preserves team_public_model_name for external reference
+
+    The row is written through ``write_row`` before the team's model list is touched, so a
+    refused or failed write leaves the team as it was (the create path orders itself the same way).
     """
     # Validate team_id if present in patch_data
     from litellm.proxy.proxy_server import premium_user
@@ -1079,9 +1163,7 @@ async def _update_team_model_in_db(
         premium_user=premium_user,
     )
 
-    # Validated before any write, beside the premium check the create path already runs
-    # here. The team ACL is updated below and autocommits, so a validator that raises
-    # further down would leave the team mutated and the deployment row never written.
+    # Validated before the row write, beside the premium check the create path already runs here.
     #
     # The merged view is what gets stored, so that is what has to satisfy the invariants.
     # Validating the patch alone rejected a partial edit of an already valid deployment:
@@ -1101,7 +1183,7 @@ async def _update_team_model_in_db(
 
     # No team_id in patch, proceed with standard update
     if patch_team_id is None:
-        return update_db_model(db_model=db_model, updated_patch=patch_data)
+        return await write_row(update_db_model(db_model=db_model, updated_patch=patch_data))
 
     # Determine public model name
     public_model_name: Final = _get_public_model_name(
@@ -1120,11 +1202,14 @@ async def _update_team_model_in_db(
     db_team_id: Final = db_model.model_info.team_id if db_model.model_info else None
     is_new_team_assignment: Final = db_team_id != patch_team_id
 
+    # Team rows keep their internal UUID-based model_name; the public name lives in model_info
+    patch_data.model_name = f"model_name_{patch_team_id}_{uuid.uuid4()}" if is_new_team_assignment else None
+    row: Final = await write_row(update_db_model(db_model=db_model, updated_patch=patch_data))
+
     if is_new_team_assignment:
         await _setup_new_team_model_assignment(
             team_id=patch_team_id,
             public_model_name=public_model_name,
-            patch_data=patch_data,
             user_api_key_dict=user_api_key_dict,
         )
     else:
@@ -1132,12 +1217,11 @@ async def _update_team_model_in_db(
             team_id=patch_team_id,
             public_model_name=public_model_name,
             db_model=db_model,
-            patch_data=patch_data,
             user_api_key_dict=user_api_key_dict,
             prisma_client=prisma_client,
         )
 
-    return update_db_model(db_model=db_model, updated_patch=patch_data)
+    return row
 
 
 def _get_public_model_name(
@@ -1189,13 +1273,9 @@ def _get_public_model_name(
 async def _setup_new_team_model_assignment(
     team_id: str,
     public_model_name: str,
-    patch_data: updateDeployment,
     user_api_key_dict: UserAPIKeyAuth,
 ) -> None:
-    """Set up a new team model with unique name and team membership."""
-    unique_model_name: Final = f"model_name_{team_id}_{uuid.uuid4()}"
-    patch_data.model_name = unique_model_name
-
+    """Register a newly team-assigned model's public name on the team."""
     await team_model_add(
         data=TeamModelAddRequest(
             team_id=team_id,
@@ -1385,7 +1465,6 @@ async def _update_existing_team_model_assignment(
     team_id: str,
     public_model_name: str,
     db_model: Deployment,
-    patch_data: updateDeployment,
     user_api_key_dict: UserAPIKeyAuth,
     prisma_client: PrismaClient | None,
 ) -> None:
@@ -1409,9 +1488,6 @@ async def _update_existing_team_model_assignment(
     old_public_name: Final = db_model.model_info.team_public_model_name if db_model.model_info else None
 
     if old_public_name and public_model_name != old_public_name:
-        # Clear user-supplied public name from patch before any early return so the
-        # caller does not overwrite the internal UUID-based model_name in the DB.
-        patch_data.model_name = None
         if prisma_client is None:
             verbose_proxy_logger.warning(
                 "prisma_client not initialized; skipping public name update entirely to avoid orphaned entries"
@@ -1459,10 +1535,6 @@ async def _update_existing_team_model_assignment(
     # else: old_public_name == public_model_name (no rename needed)
     # No team_model_add/delete calls required; public name is already registered
 
-    # Always clear patch_data.model_name to prevent caller from overwriting
-    # the internal UUID-based model_name in the DB with the user-supplied public name
-    patch_data.model_name = None
-
 
 class ModelManagementAuthChecks:
     """
@@ -1878,18 +1950,19 @@ async def add_new_model(
             reload_outcome: ReconcileOutcome = ReconcileOutcome(still_desired=None, live_after=None)
             try:
                 _original_litellm_model_name: Final = model_params.model_name
-                if model_params.model_info.team_id is None:
-                    model_response = await _add_model_to_db(
-                        model_params=priced_model_params,
-                        user_api_key_dict=user_api_key_dict,
-                        prisma_client=prisma_client,
-                    )
-                else:
-                    model_response = await _add_team_model_to_db(
-                        model_params=priced_model_params,
-                        user_api_key_dict=user_api_key_dict,
-                        prisma_client=prisma_client,
-                    )
+                add_model: Final = (
+                    _add_model_to_db if model_params.model_info.team_id is None else _add_team_model_to_db
+                )
+                model_response = await add_model(
+                    model_params=priced_model_params,
+                    user_api_key_dict=user_api_key_dict,
+                    prisma_client=prisma_client,
+                    slot=_heuristic_v2_slot(
+                        prisma_client,
+                        effective_config=priced_model_params.litellm_params.complexity_router_config,
+                        model_id=priced_model_params.model_info.id,
+                    ),
+                )
                 reload_outcome = await proxy_config.add_deployment(
                     prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
                 )
@@ -1903,6 +1976,8 @@ async def add_new_model(
                         passed_model_info=priced_model_params.model_info,
                     )
             except Exception as e:
+                if isinstance(e, HTTPException):
+                    raise
                 verbose_proxy_logger.exception("Exception in add_new_model: %s", e)
 
         else:
@@ -2070,10 +2145,17 @@ async def update_model(
                 "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
                 **({} if renamed_to is None else {"model_name": renamed_to}),
             }
-            model_response: Final = await _proxy_model_table(prisma_client).update(
-                where={"model_id": _model_id},
-                data=_data,
-            )
+            async with _heuristic_v2_slot(
+                prisma_client,
+                effective_config=_effective_complexity_router_config(
+                    model_params.litellm_params, deployment.litellm_params
+                ),
+                model_id=_model_id,
+            ) as table:
+                model_response: Final = await table.update(
+                    where={"model_id": _model_id},
+                    data=_data,
+                )
             if renamed_to is not None:
                 await sync_access_groups_for_renamed_model(
                     prisma_client=prisma_client,
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 27132c90e05..96b425f2a24 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -120,6 +120,8 @@ from litellm.router_utils.add_retry_fallback_headers import (
 from litellm.router_utils.auto_router_model_naming import (
     STRATEGY_ROUTER_PARAM_FIELDS,
     carries_complexity_router_settings,
+    count_heuristic_v2_routers,
+    heuristic_v2_limit_violation,
     validate_complexity_router_config_placement,
 )
 from litellm.types.utils import (
@@ -301,7 +303,7 @@ from litellm.proxy.auth.auth_utils import (
 )
 from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
 from litellm.proxy.auth.handle_jwt import JWTHandler
-from litellm.proxy.auth.litellm_license import LicenseCheck
+from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY, LicenseCheck
 from litellm.proxy.auth.model_checks import (
     expand_wildcard_deployments_for_model_info,
     get_all_fallbacks,
@@ -4316,6 +4318,19 @@ def validate_deployment_complexity_router_placement(model: Mapping[str, object])
         raise ValueError(f"model {model.get('model_name', '')!r}: {violation}")
 
 
+def validate_heuristic_v2_router_limit(model_list: Sequence[Mapping[str, object]], *, limit: int | None) -> None:
+    """
+    Refuse to start when config.yaml defines more heuristic_v2 auto-routers than the license allows.
+
+    Checked here rather than left to router registration for the same reason as the two
+    validators above: the proxy builds its router with `ignore_invalid_deployments=True`, so
+    the router's own refusal would turn the extra router into a silently missing model.
+    """
+    violation: Final = heuristic_v2_limit_violation(held=count_heuristic_v2_routers(model_list), limit=limit)
+    if violation is not None:
+        raise ValueError(f"config.yaml model_list: {violation} {HEURISTIC_V2_LICENSE_REMEDY}")
+
+
 def pin_complexity_router_model_id(model: dict) -> None:  # mutable-ok: out-param, model_info is stamped in place
     """
     Stamps `model_info.id` from the raw litellm_params before plugin resolution swaps
@@ -5721,6 +5736,7 @@ class ProxyConfig:
         model_list: Final = config.get("model_list", None)
         if model_list:
             router_params["model_list"] = model_list
+            validate_heuristic_v2_router_limit(model_list, limit=_license_check.heuristic_v2_router_limit())
             print(  # noqa: T201
                 "\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m"
             )
@@ -5810,6 +5826,7 @@ class ProxyConfig:
             ),
             ignore_invalid_deployments=True,  # don't raise an error if a deployment is invalid
             fallback_access_check=router_fallback_access_check,
+            heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit,
         )
 
         if redis_usage_cache is not None and router.cache.redis_cache is None:
@@ -6270,6 +6287,7 @@ class ProxyConfig:
                         search_tools=search_tools,
                         ignore_invalid_deployments=True,
                         fallback_access_check=router_fallback_access_check,
+                        heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit,
                     )
                     verbose_proxy_logger.debug("updated llm_router: %s", llm_router)
             else:
diff --git a/litellm/router.py b/litellm/router.py
index cfb080a24a0..dc750941559 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -21,7 +21,7 @@ import time
 import traceback
 import weakref
 from collections import defaultdict
-from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence
+from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterator, Mapping, Sequence
 from functools import lru_cache, partial
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast
@@ -117,6 +117,9 @@ from litellm.router_utils.add_retry_fallback_headers import (
 from litellm.router_utils.auto_router_model_naming import (
     AUTO_ROUTER_MODEL_PREFIX,
     classify_strategy_router_model,
+    count_heuristic_v2_routers,
+    heuristic_v2_limit_violation,
+    uses_heuristic_v2_classifier,
 )
 from litellm.router_utils.batch_utils import (
     _get_router_metadata_variable_name,
@@ -211,6 +214,7 @@ from litellm.types.router import (
     DeploymentTypedDict,
     FallbackAccessCheck,
     GuardrailTypedDict,
+    HeuristicV2RouterLimit,
     LiteLLM_Params,
     MockRouterTestingParams,
     ModelGroupInfo,
@@ -683,6 +687,7 @@ class Router:
         background_health_check_model_groups: Sequence[str] | None = None,
         enable_weighted_failover: bool = False,
         fallback_access_check: FallbackAccessCheck | None = None,
+        heuristic_v2_router_limit: HeuristicV2RouterLimit | None = None,
     ) -> None:
         """
         Initialize the Router class with the given parameters for caching, reliability, and routing strategy.
@@ -759,6 +764,7 @@ class Router:
 
         self.set_verbose = set_verbose
         self.ignore_invalid_deployments = ignore_invalid_deployments
+        self.heuristic_v2_router_limit = heuristic_v2_router_limit
         self.fallback_access_check: Final = fallback_access_check
         self.debug_level = debug_level
         self.enable_pre_call_checks = enable_pre_call_checks
@@ -8796,6 +8802,30 @@ class Router:
         """
         return classify_strategy_router_model(litellm_params.model) == "complexity"
 
+    def config_deployments(self) -> Iterator[Mapping[str, object]]:
+        """The model_list rows that came from config.yaml rather than the DB (``model_info.db_model`` unset)."""
+        for deployment in self.model_list:
+            if not isinstance(deployment, Mapping):
+                continue
+            model_info = deployment.get("model_info")
+            if not (isinstance(model_info, Mapping) and model_info.get("db_model")):
+                yield deployment
+
+    def heuristic_v2_router_limit_violation(self) -> str | None:
+        """
+        Why one more heuristic_v2 router cannot join this router, or None when it can.
+
+        Judged against every deployment currently on the model_list; an upsert pops the row being
+        edited first, so an edit of an existing heuristic_v2 router keeps its own slot. The limit is
+        resolved on every call through ``heuristic_v2_router_limit``; unset means unlimited, which
+        is the SDK default, and the proxy injects a resolver backed by its license.
+        """
+        limit: Final = self.heuristic_v2_router_limit() if self.heuristic_v2_router_limit is not None else None
+        others: Final = count_heuristic_v2_routers(
+            deployment for deployment in self.model_list if isinstance(deployment, Mapping)
+        )
+        return heuristic_v2_limit_violation(held=others + 1, limit=limit)
+
     def init_complexity_router_deployment(self, deployment: Deployment):
         """
         Initialize the complexity-router deployment.
@@ -8813,6 +8843,10 @@ class Router:
         )
 
         complexity_router_config: Final[dict | None] = deployment.litellm_params.complexity_router_config
+        if uses_heuristic_v2_classifier(complexity_router_config):
+            limit_violation: Final = self.heuristic_v2_router_limit_violation()
+            if limit_violation is not None:
+                raise ValueError(limit_violation)
 
         default_model: str | None = deployment.litellm_params.complexity_router_default_model
 
@@ -9636,8 +9670,16 @@ class Router:
                 raise e
 
     def _restore_deployment_after_failed_upsert(self, previous_deployment: Deployment | None, model_id: str) -> None:
+        """Put a deployment back the way it was before a failed upsert popped it.
+
+        A rollback re-admits state that was already serving, so it does not go through the
+        heuristic_v2 ceiling a newcomer gets: with the ceiling tightened since the deployment first
+        registered, judging the rollback would drop a serving router over an unrelated failed edit.
+        """
         if previous_deployment is None or self.has_model_id(model_id):
             return
+        limit_resolver: Final = self.heuristic_v2_router_limit
+        self.heuristic_v2_router_limit = None
         try:
             self.add_deployment(deployment=previous_deployment)
             verbose_router_logger.info(
@@ -9652,6 +9694,8 @@ class Router:
                 model_id,
                 restore_error,
             )
+        finally:
+            self.heuristic_v2_router_limit = limit_resolver
 
     @staticmethod
     def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]:
diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_utils/auto_router_model_naming.py
index a8aa543d735..2efbfb5782e 100644
--- a/litellm/router_utils/auto_router_model_naming.py
+++ b/litellm/router_utils/auto_router_model_naming.py
@@ -10,7 +10,7 @@ the router silently dropping the deployment at load time under
 ``ignore_invalid_deployments``.
 """
 
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from dataclasses import dataclass
 from types import MappingProxyType
 from typing import Final, Literal, TypeAlias
@@ -163,6 +163,38 @@ def strategy_router_dependencies(
     )
 
 
+def uses_heuristic_v2_classifier(complexity_router_config: object) -> bool:
+    """Whether this complexity config classifies with the bundled heuristic_v2 model."""
+    return _mapping(complexity_router_config).get("classifier_type") == "heuristic_v2"
+
+
+def is_heuristic_v2_router(litellm_params: Mapping[str, object]) -> bool:
+    """Whether this deployment is a complexity router that classifies with heuristic_v2."""
+    return classify_strategy_router_model(str(litellm_params.get("model") or "")) == "complexity" and (
+        uses_heuristic_v2_classifier(litellm_params.get("complexity_router_config"))
+    )
+
+
+def count_heuristic_v2_routers(deployments: Iterable[Mapping[str, object]]) -> int:
+    """How many of ``deployments`` (router model_list entries or config.yaml rows) are heuristic_v2 routers."""
+    return sum(1 for deployment in deployments if is_heuristic_v2_router(_mapping(deployment.get("litellm_params"))))
+
+
+def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None:
+    """Why holding ``held`` heuristic_v2 routers exceeds ``limit``, or None when it fits.
+
+    ``limit`` None means unlimited. The message is shared by every enforcement point (config
+    load, model writes, router registration) and stays SDK-neutral: it names the cap and what
+    the caller can change; the proxy appends how its license lifts the cap.
+    """
+    if limit is None or held <= limit:
+        return None
+    return (
+        f"At most {limit} auto-router(s) with classifier_type 'heuristic_v2' can be registered but this would make "
+        f"{held}. Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router."
+    )
+
+
 def validate_complexity_router_config_write(complexity_router_config: Mapping[str, object] | None) -> str | None:
     """Reject a complexity config the router would refuse to build a deployment from.
 
diff --git a/litellm/types/router.py b/litellm/types/router.py
index 4f4df1a8d2e..7ebd50f1328 100644
--- a/litellm/types/router.py
+++ b/litellm/types/router.py
@@ -885,6 +885,18 @@ class FallbackAccessCheck(Protocol):
     async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ...
 
 
+class HeuristicV2RouterLimit(Protocol):
+    """
+    Resolves how many heuristic_v2 complexity routers the Router may hold right now; None means unlimited.
+
+    The Router calls it on every registration and limit query instead of caching the answer, so the
+    proxy can keep the limit on its license object (re-verified on config load) rather than hand
+    over a snapshot.
+    """
+
+    def __call__(self) -> int | None: ...
+
+
 class LiteLLM_RouterFileObject(TypedDict, total=False):
     """
     Tracking the litellm params hash, used for mapping the file id to the right model
diff --git a/tests/test_litellm/proxy/auth/test_litellm_license.py b/tests/test_litellm/proxy/auth/test_litellm_license.py
index 8da365cb587..1db53638070 100644
--- a/tests/test_litellm/proxy/auth/test_litellm_license.py
+++ b/tests/test_litellm/proxy/auth/test_litellm_license.py
@@ -2,6 +2,8 @@ import asyncio
 import json
 from unittest.mock import AsyncMock, MagicMock, patch
 
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
+
 
 from litellm.proxy.auth.litellm_license import LicenseCheck
 
@@ -30,3 +32,70 @@ def test_is_over_limit():
     assert license_check.is_over_limit(101) is False
     assert license_check.is_over_limit(100) is False
     assert license_check.is_over_limit(99) is False
+
+
+def test_heuristic_v2_router_limit() -> None:
+    """Only the signed license's auto_router feature lifts the one-router limit; an API-verified
+    license (no airgapped data) and an airgapped license without the feature keep it."""
+    license_check = LicenseCheck()
+    license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["auto_router"]}
+    assert license_check.heuristic_v2_router_limit() is None
+
+    license_check.airgapped_license_data = {
+        "expiration_date": "2999-01-01",
+        "allowed_features": ["sso", "auto_router", "audit_logs"],
+    }
+    assert license_check.heuristic_v2_router_limit() is None
+
+    license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["sso"]}
+    assert license_check.heuristic_v2_router_limit() == 1
+
+    license_check.airgapped_license_data = {"expiration_date": "2999-01-01"}
+    assert license_check.heuristic_v2_router_limit() == 1
+
+    license_check.airgapped_license_data = None
+    assert license_check.heuristic_v2_router_limit() == 1
+
+
+def _signed_license(expiration_date: str) -> tuple[RSAPublicKey, str]:
+    import base64
+
+    from cryptography.hazmat.primitives import hashes
+    from cryptography.hazmat.primitives.asymmetric import padding, rsa
+
+    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    message = json.dumps(
+        {"expiration_date": expiration_date, "user_id": "u", "allowed_features": ["auto_router"]}
+    ).encode()
+    signature = private_key.sign(
+        message,
+        padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
+        hashes.SHA256(),
+    )
+    return private_key.public_key(), base64.b64encode(message + b"." + signature).decode()
+
+
+def test_expired_or_unreadable_license_grants_no_features() -> None:
+    """The verifier stores the signed payload only after the expiry check passes and clears it when a
+    later verify rejects the license, so a stale payload cannot keep lifting the heuristic_v2 limit."""
+    license_check = LicenseCheck()
+    public_key, valid_key = _signed_license("2999-01-01")
+    assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True
+    assert license_check.heuristic_v2_router_limit() is None
+
+    _, expired_key = _signed_license("2000-01-01")
+    assert license_check.verify_license_without_api_request(public_key=public_key, license_key=expired_key) is not True
+    assert license_check.airgapped_license_data is None
+    assert license_check.heuristic_v2_router_limit() == 1
+
+    assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True
+    assert license_check.verify_license_without_api_request(public_key=public_key, license_key="not-a-license") is not True
+    assert license_check.airgapped_license_data is None
+
+
+def test_valid_signed_license_with_auto_router_lifts_the_limit() -> None:
+    license_check = LicenseCheck()
+    public_key, license_key = _signed_license("2999-01-01")
+
+    assert license_check.verify_license_without_api_request(public_key=public_key, license_key=license_key) is True
+    assert license_check.heuristic_v2_router_limit() is None
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 5fa59a85c9d..c69f8f20a13 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
@@ -28,9 +28,18 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
     delete_team_models,
 )
 from litellm.proxy.utils import PrismaClient
+from litellm.router import Router
 from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment
 
 
+async def _passthrough_row(update_data):
+    return update_data
+
+
+async def _write_empty_row(**kwargs):
+    return await kwargs["write_row"]({})
+
+
 class MockPrismaClient:
     def __init__(
         self,
@@ -1191,7 +1200,7 @@ class TestTeamModelSiblingRouting:
         team_id = "team_no_alias"
         public_name = "gpt-4.1-mini"
 
-        async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client):
+        async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client, slot=None):
             return MagicMock(model_id=str(uuid.uuid4()))
 
         mock_team_model_add = AsyncMock()
@@ -1372,7 +1381,8 @@ class TestTeamModelUpdate:
                 db_model=db_model,
                 patch_data=patch_data,
                 user_api_key_dict=user_api_key_dict,
-                prisma_client=prisma_client,  # type: ignore
+                prisma_client=prisma_client,  # type: ignore,
+                write_row=_passthrough_row,
             )
 
             assert result.get("model_name", "").startswith("model_name_test_team_123_")
@@ -1435,7 +1445,6 @@ class TestTeamModelUpdate:
                 team_id="team_123",
                 public_model_name="new-public-name",
                 db_model=db_model,
-                patch_data=patch_data,
                 user_api_key_dict=user_api_key_dict,
                 prisma_client=prisma_client,  # type: ignore
             )
@@ -1481,7 +1490,6 @@ class TestTeamModelUpdate:
                 team_id="team_123",
                 public_model_name="new-public-name",
                 db_model=db_model,
-                patch_data=patch_data,
                 user_api_key_dict=user_api_key_dict,
                 prisma_client=None,
             )
@@ -1490,39 +1498,72 @@ class TestTeamModelUpdate:
             mock_delete.assert_not_called()
 
     @pytest.mark.asyncio
-    async def test_rename_with_prisma_none_clears_patch_model_name(self):
-        """Rename path must clear patch_data.model_name even when prisma is unavailable (P1)."""
+    async def test_a_refused_row_write_leaves_the_team_untouched(self):
+        """The team's model list autocommits, so it is written only after the row write succeeded: a
+        refused write (the heuristic_v2 slot 403, a DB error) must not leave the team listing a name
+        whose row never changed."""
+        from fastapi import HTTPException
+
         from litellm.proxy.management_endpoints.model_management_endpoints import (
-            _update_existing_team_model_assignment,
+            _update_team_model_in_db,
         )
         from litellm.types.router import ModelInfo
 
         db_model = Deployment(
-            model_name="model_name_team_123_uuid1",
+            model_name="gpt-4o",
             litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"),
-            model_info=ModelInfo(
-                team_id="team_123", team_public_model_name="old-public-name"
+            model_info=ModelInfo(),
+        )
+        user_api_key_dict = UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN)
+        events: list[str] = []
+        written: dict[str, object] = {}
+
+        def patch_data() -> updateDeployment:
+            return updateDeployment(model_name="team-public", model_info=ModelInfo(team_id="team_123"))
+
+        async def refuse_row(update_data):
+            events.append("row")
+            raise HTTPException(status_code=403, detail="slot held")
+
+        async def accept_row(update_data):
+            events.append("row")
+            written.update(update_data)
+            return update_data
+
+        async def team_add(**_):
+            events.append("team_model_add")
+
+        with (
+            patch(  # test-quality-ok: the team auth check needs a live DB; the write order is what is under test
+                "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.allow_team_model_action",
+                AsyncMock(return_value=True),
             ),
-        )
-        patch_data = updateDeployment(
-            model_name="new-public-name",
-            model_info=ModelInfo(team_id="team_123"),
-        )
-        user_api_key_dict = UserAPIKeyAuth(
-            user_id="test_user",
-            user_role=LitellmUserRoles.PROXY_ADMIN,
-        )
+            patch("litellm.proxy.proxy_server.premium_user", True),  # test-quality-ok: team models are premium-gated through a proxy global with no injection seam
+            patch(  # test-quality-ok: the team list write is the collaborator whose ordering is asserted
+                "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add",
+                side_effect=team_add,
+            ),
+        ):
+            with pytest.raises(HTTPException):
+                await _update_team_model_in_db(
+                    db_model=db_model,
+                    patch_data=patch_data(),
+                    user_api_key_dict=user_api_key_dict,
+                    prisma_client=MockPrismaClient(team_exists=True),  # type: ignore
+                    write_row=refuse_row,
+                )
+            assert events == ["row"]
 
-        await _update_existing_team_model_assignment(
-            team_id="team_123",
-            public_model_name="new-public-name",
-            db_model=db_model,
-            patch_data=patch_data,
-            user_api_key_dict=user_api_key_dict,
-            prisma_client=None,
-        )
-
-        assert patch_data.model_name is None
+            await _update_team_model_in_db(
+                db_model=db_model,
+                patch_data=patch_data(),
+                user_api_key_dict=user_api_key_dict,
+                prisma_client=MockPrismaClient(team_exists=True),  # type: ignore
+                write_row=accept_row,
+            )
+            assert events == ["row", "row", "team_model_add"]
+            assert str(written["model_name"]).startswith("model_name_team_123_")
+            assert "team-public" in str(written["model_info"])
 
     @pytest.mark.asyncio
     async def test_rename_handles_legacy_string_model_info(self):
@@ -1574,7 +1615,6 @@ class TestTeamModelUpdate:
                 team_id="team_123",
                 public_model_name="new-public-name",
                 db_model=db_model,
-                patch_data=patch_data,
                 user_api_key_dict=user_api_key_dict,
                 prisma_client=prisma_client,  # type: ignore
             )
@@ -1614,7 +1654,8 @@ class TestTeamModelUpdate:
                     db_model=db_model,
                     patch_data=patch_data,
                     user_api_key_dict=user_api_key_dict,
-                    prisma_client=prisma_client,  # type: ignore
+                    prisma_client=prisma_client,  # type: ignore,
+                    write_row=_passthrough_row,
                 )
             assert "403" in str(exc_info.value)
 
@@ -1900,7 +1941,8 @@ class TestTeamModelUpdate:
                 db_model=db_model,
                 patch_data=patch_data,
                 user_api_key_dict=user_api_key_dict,
-                prisma_client=prisma_client,  # type: ignore
+                prisma_client=prisma_client,  # type: ignore,
+                write_row=_passthrough_row,
             )
 
         # team ACL must not be touched on a no-op edit
@@ -4311,6 +4353,321 @@ class TestStrategyRouterWriteValidation:
             is None
         )
 
+    @staticmethod
+    def _live_router_holding_one_heuristic_v2(limit: int | None) -> Router:
+        return Router(
+            model_list=[
+                {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}},
+                {
+                    "model_name": "held-v2",
+                    "litellm_params": {
+                        "model": "auto_router/complexity_router",
+                        "complexity_router_config": {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}},
+                    },
+                    "model_info": {"id": "held-id"},
+                },
+            ],
+            heuristic_v2_router_limit=lambda: limit,
+        )
+
+    class _FakeTx:
+        """Stands in for a prisma transaction: records the raw statements and exposes the model table."""
+
+        def __init__(self, db_held: int) -> None:
+            self.db_held = db_held
+            self.raw_calls: list[tuple[str, tuple[object, ...]]] = []
+            self.litellm_proxymodeltable = MagicMock(create=AsyncMock(), update=AsyncMock())
+
+        async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]:
+            self.raw_calls.append((sql, args))
+            return [{"held": self.db_held}] if "count(*)" in sql else []
+
+        async def __aenter__(self) -> "TestStrategyRouterWriteValidation._FakeTx":
+            return self
+
+        async def __aexit__(self, *exc: object) -> None:
+            return None
+
+    class _FakeDb:
+        """Stands in for prisma_client: the plain client and the transaction it opens are told apart by identity."""
+
+        def __init__(self, db_held: int, existing_row: object = None) -> None:
+            self.db = self
+            self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_held)
+            self.litellm_proxymodeltable = MagicMock(
+                create=AsyncMock(), update=AsyncMock(), find_unique=AsyncMock(return_value=existing_row)
+            )
+
+        def tx(self) -> "TestStrategyRouterWriteValidation._FakeTx":
+            return self.tx_obj
+
+    _V2 = {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}
+    _V1 = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini"}}
+
+    @pytest.mark.parametrize(
+        "incoming,existing,expected",
+        [
+            (_V2, None, _V2),
+            (_V2, _V1, _V2),
+            (None, _V1, _V1),
+            (None, None, None),
+            ("no-config", _V2, _V2),
+        ],
+    )
+    def test_effective_complexity_router_config(
+        self, incoming: object, existing: object, expected: object
+    ) -> None:
+        """A write is judged on the config it leaves on the row: the incoming one when it carries one, else the stored one."""
+        from litellm.proxy.management_endpoints.model_management_endpoints import (
+            _effective_complexity_router_config,
+        )
+        from litellm.types.router import updateLiteLLMParams
+
+        incoming_params = None if incoming is None else updateLiteLLMParams(
+            complexity_router_config=None if incoming == "no-config" else incoming
+        )
+        existing_params = None if existing is None else updateLiteLLMParams(complexity_router_config=existing)
+        assert _effective_complexity_router_config(incoming_params, existing_params) == expected
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "limit,effective_config,db_held,config_holds_one,model_id,expected",
+        [
+            (1, _V2, 1, False, None, "refused"),
+            (1, _V2, 0, True, None, "refused"),
+            (1, _V2, 0, False, None, "reserved"),
+            (1, _V2, 0, False, "held-id", "reserved"),
+            (2, _V2, 1, False, None, "reserved"),
+            (1, _V1, 5, True, None, "plain"),
+            (1, None, 5, True, None, "plain"),
+            (None, _V2, 5, True, None, "plain"),
+        ],
+    )
+    async def test_heuristic_v2_slot_matrix(
+        self,
+        limit: int | None,
+        effective_config: object,
+        db_held: int,
+        config_holds_one: bool,
+        model_id: str | None,
+        expected: str,
+    ) -> None:
+        """The slot is claimed inside a locked transaction only for a heuristic_v2 write under a limit; the DB rows
+        (other pods included) plus config.yaml routers decide, the row being edited is excluded through the SQL
+        parameter, and every other write runs on the plain client with no lock."""
+        from fastapi import HTTPException
+
+        from litellm.proxy.management_endpoints.model_management_endpoints import (
+            HEURISTIC_V2_SLOT_LOCK_KEY,
+            _heuristic_v2_slot,
+        )
+
+        fake = self._FakeDb(db_held)
+        live_router = self._live_router_holding_one_heuristic_v2(limit) if config_holds_one else None
+        with (
+            patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: limit),  # test-quality-ok: the guard reads the proxy license singleton with no injection seam
+            patch("litellm.proxy.proxy_server.llm_router", live_router),  # test-quality-ok: the guard reads the proxy router global with no injection seam
+            patch(  # test-quality-ok: the cross-pod publish is the side effect under test; redis is not configured here
+                "litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change",
+                new=AsyncMock(),
+            ) as published,
+        ):
+            if expected == "refused":
+                with pytest.raises(HTTPException) as exc_info:
+                    async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id):
+                        pass
+                assert exc_info.value.status_code == 403
+                assert "At most 1 auto-router" in str(exc_info.value.detail)
+                assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail)
+                return
+            async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id) as tables:
+                handle = tables
+        if expected == "plain":
+            await handle.create(data={})
+            fake.litellm_proxymodeltable.create.assert_awaited_once_with(data={})
+            assert fake.tx_obj.raw_calls == []
+            return
+        assert handle is fake.tx_obj.litellm_proxymodeltable
+        published.assert_awaited_once_with(redis_cache=None, object_type="litellm_proxymodeltable")
+        (lock_sql, lock_params), (_count_sql, count_params) = fake.tx_obj.raw_calls
+        assert "pg_advisory_xact_lock($1)" in lock_sql and "count" not in lock_sql
+        assert lock_params == (HEURISTIC_V2_SLOT_LOCK_KEY,)
+        assert count_params == (model_id or "",)
+
+    @pytest.mark.asyncio
+    async def test_team_model_bookkeeping_runs_after_the_slot_is_released(self) -> None:
+        """team_model_add needs a second pool connection, so it must run only after the slot transaction
+        (and its advisory lock) has closed; a pool-sized burst of team creates would otherwise stall on the
+        lock holder waiting for a connection the waiters are occupying."""
+        from contextlib import asynccontextmanager
+
+        from litellm.proxy.management_endpoints.model_management_endpoints import _add_team_model_to_db
+        from litellm.types.router import ModelInfo
+
+        events: list[str] = []
+        created = MagicMock(model_id="row-1")
+
+        @asynccontextmanager
+        async def slot():
+            events.append("slot-enter")
+            yield MagicMock(create=AsyncMock(return_value=created))
+            events.append("slot-exit")
+
+        async def team_model_add(**_: object) -> None:
+            events.append("team_model_add")
+
+        deployment = Deployment(
+            model_name="public-v2",
+            litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2),
+            model_info=ModelInfo(id="row-1", team_id="team-1"),
+        )
+        with (
+            patch(  # test-quality-ok: params are encrypted with the proxy master key, which this test does not configure
+                "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
+                lambda value, new_encryption_key=None: value,
+            ),
+            patch(  # test-quality-ok: the team list write is the collaborator whose ordering is asserted
+                "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add",
+                side_effect=team_model_add,
+            ),
+        ):
+            result = await _add_team_model_to_db(
+                model_params=deployment,
+                user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
+                prisma_client=MagicMock(),
+                slot=slot(),
+            )
+
+        assert result is created
+        assert events == ["slot-enter", "slot-exit", "team_model_add"]
+
+    @pytest.mark.asyncio
+    async def test_add_new_model_refuses_a_second_heuristic_v2_router_before_the_db_write(self) -> None:
+        from litellm.proxy._types import ProxyException
+        from litellm.proxy.management_endpoints.model_management_endpoints import (
+            add_new_model,
+        )
+
+        admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
+        fake = self._FakeDb(db_held=1)
+
+        with (
+            patch("litellm.proxy.proxy_server.prisma_client", fake),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch("litellm.proxy.proxy_server.llm_router", None),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch("litellm.proxy.proxy_server.store_model_in_db", True),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch("litellm.proxy.proxy_server.premium_user", True),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch(  # test-quality-ok: prior auth check needs a live DB; only the license limit is under test
+                "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
+                new=AsyncMock(return_value=None),
+            ),
+            patch(  # test-quality-ok: params are encrypted before the slot is entered; no master key in this test
+                "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
+                lambda value, new_encryption_key=None: value,
+            ),
+        ):
+            with pytest.raises(ProxyException) as exc_info:
+                await add_new_model(
+                    model_params=Deployment(
+                        model_name="second-v2",
+                        litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2),
+                    ),
+                    user_api_key_dict=admin,
+                )
+            assert exc_info.value.code == "403"
+            assert "At most 1 auto-router" in str(exc_info.value.message)
+            fake.tx_obj.litellm_proxymodeltable.create.assert_not_awaited()
+            fake.litellm_proxymodeltable.create.assert_not_awaited()
+
+    @pytest.mark.asyncio
+    async def test_patch_model_refuses_switching_another_router_to_heuristic_v2(self) -> None:
+        """patch_model relays HTTPException as-is, so the license refusal reaches the client as a plain 403."""
+        from fastapi import HTTPException
+
+        from litellm.proxy.management_endpoints.model_management_endpoints import (
+            patch_model,
+        )
+        from litellm.types.router import updateLiteLLMParams
+
+        model_id = "other-id"
+        admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
+        fake = self._FakeDb(db_held=1)
+
+        with (
+            patch("litellm.proxy.proxy_server.prisma_client", fake),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch("litellm.proxy.proxy_server.llm_router", None),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch("litellm.proxy.proxy_server.store_model_in_db", True),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch("litellm.proxy.proxy_server.premium_user", True),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch(  # test-quality-ok: the write must be refused before this DB step runs
+                "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model",
+                new=AsyncMock(return_value=self._db_complexity_router(model_id)),
+            ),
+            patch(  # test-quality-ok: prior auth check needs a live DB; only the license limit is under test
+                "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
+                new=AsyncMock(return_value=None),
+            ),
+            patch(  # test-quality-ok: the helper's team bookkeeping needs a live DB; the row writer it is handed is what is under test
+                "litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db",
+                new=AsyncMock(side_effect=_write_empty_row),
+            ),
+        ):
+            with pytest.raises(HTTPException) as exc_info:
+                await patch_model(
+                    model_id=model_id,
+                    patch_data=updateDeployment(litellm_params=updateLiteLLMParams(complexity_router_config=self._V2)),
+                    user_api_key_dict=admin,
+                )
+            assert exc_info.value.status_code == 403
+            fake.tx_obj.litellm_proxymodeltable.update.assert_not_awaited()
+            fake.litellm_proxymodeltable.update.assert_not_awaited()
+
+    @pytest.mark.asyncio
+    async def test_update_model_refuses_switching_another_router_to_heuristic_v2(self) -> None:
+        from litellm.proxy._types import ProxyException
+        from litellm.proxy.management_endpoints.model_management_endpoints import (
+            update_model,
+        )
+        from litellm.types.router import ModelInfo, updateLiteLLMParams
+
+        model_id = "other-id"
+        admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
+
+        existing_row = MagicMock()
+        existing_row.model_dump.return_value = {
+            "model_name": "my-auto-router",
+            "litellm_params": {
+                "model": "auto_router/complexity_router",
+                "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}},
+            },
+            "model_info": {"id": model_id},
+        }
+        existing_row.litellm_params = existing_row.model_dump.return_value["litellm_params"]
+        fake = self._FakeDb(db_held=1, existing_row=existing_row)
+
+        with (
+            patch("litellm.proxy.proxy_server.prisma_client", fake),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch("litellm.proxy.proxy_server.llm_router", None),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch("litellm.proxy.proxy_server.store_model_in_db", True),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch("litellm.proxy.proxy_server.premium_user", True),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1),  # test-quality-ok: endpoint reads proxy server globals with no injection seam
+            patch(  # test-quality-ok: prior auth check needs a live DB; only the license limit is under test
+                "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
+                new=AsyncMock(return_value=None),
+            ),
+        ):
+            with pytest.raises(ProxyException) as exc_info:
+                await update_model(
+                    model_params=updateDeployment(
+                        litellm_params=updateLiteLLMParams(complexity_router_config=self._V2),
+                        model_info=ModelInfo(id=model_id),
+                    ),
+                    user_api_key_dict=admin,
+                )
+            assert exc_info.value.code == "403"
+            fake.tx_obj.litellm_proxymodeltable.update.assert_not_awaited()
+            fake.litellm_proxymodeltable.update.assert_not_awaited()
+
     @pytest.mark.asyncio
     async def test_update_model_rejects_prefix_strip(self):
         from litellm.proxy._types import ProxyException
diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py
index a1c38d26b9d..d35b77f732c 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py
@@ -45,6 +45,10 @@ from litellm.types.router import (
 from litellm.types.utils import Usage
 
 
+async def _passthrough_row(update_data):
+    return update_data
+
+
 def test_model_info_accepts_valid_ptu_fields():
     info = ModelInfo(
         id="x",
@@ -385,6 +389,7 @@ class TestTeamModelUpdateValidatesBeforeWriting:
             patch_data=patch_data,
             user_api_key_dict=MagicMock(),
             prisma_client=MagicMock(),
+            write_row=_passthrough_row,
         )
         return result, touched
 
@@ -914,6 +919,7 @@ class TestPtuDeploymentsAreNotBilledPerToken:
                     patch_data=patch,
                     user_api_key_dict=UserAPIKeyAuth(user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN),
                     prisma_client=MagicMock(),
+                    write_row=_passthrough_row,
                 )
 
         assert exc.value.status_code == 400
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 1ab18639fff..dcfad8f6815 100644
--- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
+++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
@@ -28,6 +28,7 @@ from litellm.proxy.proxy_server import (
     resolve_routing_plugins,
     validate_deployment_complexity_router_placement,
     validate_deployment_max_agentic_loops,
+    validate_heuristic_v2_router_limit,
 )
 
 from .conftest import normalize
@@ -193,6 +194,120 @@ def test_validate_deployment_complexity_router_placement_leaves_valid_deployment
     assert model["litellm_params"] == litellm_params
 
 
+def _heuristic_v2_row(model_name: str, classifier_type: str = "heuristic_v2") -> dict[str, object]:
+    return {
+        "model_name": model_name,
+        "litellm_params": {
+            "model": "auto_router/complexity_router",
+            "complexity_router_config": {"classifier_type": classifier_type, "tiers": {"SIMPLE": "gpt-4o-mini"}},
+        },
+    }
+
+
+def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() -> None:
+    """Same reason as the two validators above: the proxy router swallows registration errors, so
+    an over-limit config.yaml must fail here instead of booting with a silently missing router."""
+    with pytest.raises(ValueError, match=re.escape("At most 1 auto-router")) as exc_info:
+        validate_heuristic_v2_router_limit(
+            [_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], limit=1
+        )
+    assert "'auto_router' feature lifts the limit" in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+    "model_list,limit",
+    [
+        ([_heuristic_v2_row("a"), _heuristic_v2_row("b")], None),
+        ([_heuristic_v2_row("a"), _heuristic_v2_row("c", "heuristic")], 1),
+        ([{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}], 1),
+    ],
+)
+def test_validate_heuristic_v2_router_limit_leaves_configs_within_the_limit_alone(
+    model_list: list[dict[str, object]], limit: int | None
+) -> None:
+    assert validate_heuristic_v2_router_limit(model_list, limit=limit) is None
+
+
+_TWO_HEURISTIC_V2_ROUTERS_YAML = (
+    "model_list:\n"
+    "  - model_name: gpt-4o-mini\n"
+    "    litellm_params:\n"
+    "      model: openai/gpt-4o-mini\n"
+    "      api_key: k\n"
+    "  - model_name: v2-a\n"
+    "    litellm_params:\n"
+    "      model: auto_router/complexity_router\n"
+    "      complexity_router_config:\n"
+    "        classifier_type: heuristic_v2\n"
+    "        tiers: {SIMPLE: gpt-4o-mini}\n"
+    "  - model_name: v2-b\n"
+    "    litellm_params:\n"
+    "      model: auto_router/complexity_router\n"
+    "      complexity_router_config:\n"
+    "        classifier_type: heuristic_v2\n"
+    "        tiers: {SIMPLE: gpt-4o-mini}\n"
+    "router_settings:\n"
+    "  heuristic_v2_router_limit: 99\n"
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("license_limit", [1, None])
+async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_license_only(
+    tmp_path, monkeypatch, license_limit: int | None
+) -> None:
+    """`router_settings.heuristic_v2_router_limit` is managed outside config.yaml: an operator
+    cannot grant the entitlement by editing the config, and a licensed proxy boots both routers."""
+    f = tmp_path / "c.yaml"
+    f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
+    monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: license_limit
+    )
+
+    if license_limit is None:
+        router, _model_list, _general_settings = await ProxyConfig().load_config(
+            router=None, config_file_path=str(f)
+        )
+        assert router.heuristic_v2_router_limit is not None
+        assert router.heuristic_v2_router_limit() is None
+        assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
+        return
+
+    with pytest.raises(ValueError, match=re.escape("config.yaml model_list: At most 1 auto-router")):
+        await ProxyConfig().load_config(router=None, config_file_path=str(f))
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig_load_config_router_refuses_a_db_heuristic_v2_router_beyond_the_license(
+    tmp_path, monkeypatch
+) -> None:
+    """config.yaml holds the one allowed heuristic_v2 router; a second one arriving later from the DB
+    is refused at registration because the router was built with the license's ceiling."""
+    from litellm.types.router import Deployment
+
+    f = tmp_path / "c.yaml"
+    f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML.replace("  - model_name: v2-b\n", "  - model_name: v1-b\n", 1).replace(
+        "classifier_type: heuristic_v2\n        tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings",
+        "classifier_type: heuristic\n        tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings",
+    ))
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
+    monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
+    monkeypatch.setattr("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1)
+
+    router, _model_list, _general_settings = await ProxyConfig().load_config(router=None, config_file_path=str(f))
+
+    assert router.heuristic_v2_router_limit is not None
+    assert router.heuristic_v2_router_limit() == 1
+    assert sorted(router.complexity_routers) == ["v1-b", "v2-a"]
+    db_row = Deployment(**_heuristic_v2_row("v2-from-db"), model_info={"id": "db-id"})
+    assert router.upsert_deployment(db_row) is None
+    assert sorted(router.complexity_routers) == ["v1-b", "v2-a"]
+
+
 def test_validate_deployment_max_agentic_loops_allows_a_deployment_without_the_key():
     model = {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}
 
diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py
index aa1b51afe10..da3791da39a 100644
--- a/tests/test_litellm/router_strategy/test_complexity_router.py
+++ b/tests/test_litellm/router_strategy/test_complexity_router.py
@@ -14,6 +14,7 @@ from pydantic import ValidationError
 
 import litellm
 from litellm import Router
+from litellm.router_utils.auto_router_model_naming import count_heuristic_v2_routers
 from litellm._logging import verbose_router_logger
 from litellm.caching.dual_cache import DualCache
 from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
@@ -1085,6 +1086,165 @@ class TestRouterComplexityDeploymentMethods:
         router.init_complexity_router_deployment(deployment)
         assert "auto_router/complexity_router/test-router" in router.complexity_routers
 
+    @staticmethod
+    def _router_row(model_name: str, model_id: str, classifier_type: str) -> dict[str, object]:
+        return {
+            "model_name": model_name,
+            "litellm_params": {
+                "model": "auto_router/complexity_router",
+                "complexity_router_config": {
+                    "classifier_type": classifier_type,
+                    "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
+                },
+            },
+            "model_info": {"id": model_id},
+        }
+
+    _POOL: dict[str, object] = {
+        "model_name": "gpt-4o-mini",
+        "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"},
+    }
+
+    def test_heuristic_v2_ceiling_keeps_the_first_router_and_drops_the_rest(self) -> None:
+        """The proxy runs with ignore_invalid_deployments, so the second heuristic_v2 router is dropped
+        at registration while a heuristic (v1) sibling and the first v2 router stay routable."""
+        router = Router(
+            model_list=[
+                self._POOL,
+                self._router_row("v2-a", "id-a", "heuristic_v2"),
+                self._router_row("v2-b", "id-b", "heuristic_v2"),
+                self._router_row("v1-c", "id-c", "heuristic"),
+            ],
+            heuristic_v2_router_limit=lambda: 1,
+            ignore_invalid_deployments=True,
+        )
+
+        assert sorted(router.complexity_routers) == ["v1-c", "v2-a"]
+        assert router.get_deployment(model_id="id-b") is None
+
+    def test_heuristic_v2_ceiling_raises_without_ignore_invalid_deployments(self) -> None:
+        with pytest.raises(ValueError, match="At most 1 auto-router"):
+            Router(
+                model_list=[
+                    self._POOL,
+                    self._router_row("v2-a", "id-a", "heuristic_v2"),
+                    self._router_row("v2-b", "id-b", "heuristic_v2"),
+                ],
+                heuristic_v2_router_limit=lambda: 1,
+            )
+
+    def test_heuristic_v2_limit_is_resolved_on_every_registration(self) -> None:
+        """The Router never caches the limit: when the resolver's answer moves (the proxy re-verified
+        its license), the next registration and the next limit query see the new value."""
+        limits = {"value": None}
+        router = Router(
+            model_list=[
+                self._POOL,
+                self._router_row("v2-a", "id-a", "heuristic_v2"),
+                self._router_row("v2-b", "id-b", "heuristic_v2"),
+            ],
+            heuristic_v2_router_limit=lambda: limits["value"],
+            ignore_invalid_deployments=True,
+        )
+        assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
+        assert router.heuristic_v2_router_limit_violation() is None
+
+        limits["value"] = 1
+        assert router.heuristic_v2_router_limit_violation() is not None
+        assert router.upsert_deployment(Deployment(**self._router_row("v2-c", "id-c", "heuristic_v2"))) is None
+        assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
+
+    def test_heuristic_v2_ceiling_tightening_refuses_the_edit_and_keeps_the_live_router(self) -> None:
+        """Two heuristic_v2 routers registered under an unlimited ceiling, then the ceiling drops to one:
+        an edit to either must be refused before its live row is popped, or the failed re-add and
+        the failed restore would drop a serving router while the write reports success."""
+        limits = {"value": None}
+        router = Router(
+            model_list=[
+                self._POOL,
+                self._router_row("v2-a", "id-a", "heuristic_v2"),
+                self._router_row("v2-b", "id-b", "heuristic_v2"),
+            ],
+            heuristic_v2_router_limit=lambda: limits["value"],
+            ignore_invalid_deployments=True,
+        )
+        limits["value"] = 1
+
+        assert router.upsert_deployment(Deployment(**self._router_row("v2-a-renamed", "id-a", "heuristic_v2"))) is None
+        assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
+        assert router.get_deployment(model_id="id-a") is not None
+
+        assert router.upsert_deployment(Deployment(**self._router_row("v1-a", "id-a", "heuristic"))) is not None
+        assert sorted(router.complexity_routers) == ["v1-a", "v2-b"]
+
+    def test_config_deployments_excludes_db_rows(self) -> None:
+        """The proxy counts config.yaml routers from here and DB rows from the database, so a DB-loaded
+        row (``model_info.db_model``) must not show up twice."""
+        router = Router(model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")])
+        db_row = self._router_row("v2-db", "id-db", "heuristic_v2")
+        db_row["model_info"] = {"id": "id-db", "db_model": True}
+        assert router.upsert_deployment(Deployment(**db_row)) is not None
+
+        assert sorted(str(row["model_name"]) for row in router.config_deployments()) == ["gpt-4o-mini", "v2-a"]
+        assert count_heuristic_v2_routers(router.config_deployments()) == 1
+
+    def test_failed_edit_of_a_live_v2_router_rolls_back_without_the_ceiling(self) -> None:
+        """A rollback after a failed upsert re-admits state that was already serving, so it must not be
+        judged by a ceiling that tightened since: converting one of two live heuristic_v2 routers to a
+        config whose registration fails must leave it serving its previous v2 configuration."""
+        limits = {"value": None}
+        router = Router(
+            model_list=[
+                self._POOL,
+                self._router_row("v2-a", "id-a", "heuristic_v2"),
+                self._router_row("v2-b", "id-b", "heuristic_v2"),
+            ],
+            heuristic_v2_router_limit=lambda: limits["value"],
+            ignore_invalid_deployments=True,
+        )
+        limits["value"] = 1
+
+        broken = self._router_row("v1-a", "id-a", "heuristic")
+        broken["litellm_params"]["complexity_router_config"]["tiers"] = {}
+        assert router.upsert_deployment(Deployment(**broken)) is None
+
+        assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
+        live = router.get_deployment(model_id="id-a")
+        assert live is not None and live.litellm_params.complexity_router_config["classifier_type"] == "heuristic_v2"
+        assert router.heuristic_v2_router_limit_violation() is not None
+
+    def test_heuristic_v2_routers_are_unlimited_by_default(self) -> None:
+        router = Router(
+            model_list=[
+                self._POOL,
+                self._router_row("v2-a", "id-a", "heuristic_v2"),
+                self._router_row("v2-b", "id-b", "heuristic_v2"),
+            ]
+        )
+
+        assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
+        assert router.heuristic_v2_router_limit_violation() is None
+
+    def test_heuristic_v2_router_limit_violation_frees_the_slot_of_the_router_being_edited(self) -> None:
+        """A DB reload upserts the existing heuristic_v2 router again; that edit must keep its own slot
+        while a different deployment switching to heuristic_v2 is refused."""
+        router = Router(
+            model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")],
+            heuristic_v2_router_limit=lambda: 1,
+            ignore_invalid_deployments=True,
+        )
+
+        assert router.heuristic_v2_router_limit_violation() is not None
+
+        edited = self._router_row("v2-a-renamed", "id-a", "heuristic_v2")
+        assert router.upsert_deployment(Deployment(**edited)) is not None
+        assert sorted(router.complexity_routers) == ["v2-a-renamed"]
+
+        assert router.upsert_deployment(Deployment(**self._router_row("v2-b", "id-b", "heuristic_v2"))) is None
+        assert sorted(router.complexity_routers) == ["v2-a-renamed"]
+        assert router.upsert_deployment(Deployment(**self._router_row("v1-c", "id-c", "heuristic"))) is not None
+        assert sorted(router.complexity_routers) == ["v1-c", "v2-a-renamed"]
+
     def test_hybrid_initialization_waits_for_later_pool_deployments(self):
         router = Router(
             model_list=[
diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py
index 0007f09896a..238d0546518 100644
--- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py
+++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py
@@ -1,8 +1,13 @@
+from collections.abc import Mapping
+
 import pytest
 
 from litellm.router_utils.auto_router_model_naming import (
     carries_complexity_router_settings,
     classify_strategy_router_model,
+    count_heuristic_v2_routers,
+    heuristic_v2_limit_violation,
+    is_heuristic_v2_router,
     strategy_router_dependencies,
     validate_complexity_router_config_placement,
     validate_complexity_router_config_write,
@@ -369,3 +374,55 @@ def test_placement_is_scoped_to_complexity_router_deployments(model, present_fie
     flat param on an s3_vectors vector store, so an unscoped gate would reject a valid deployment.
     Either complexity field names one on its own, which is what the load itself requires."""
     assert carries_complexity_router_settings(model, present_fields) is scoped
+
+
+@pytest.mark.parametrize(
+    "litellm_params,expected",
+    [
+        ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True),
+        ({"model": "auto_router/complexity_router-eu", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True),
+        ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, False),
+        ({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, False),
+        ({"model": "auto_router/complexity_router"}, False),
+        ({"model": "auto_router/quality_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False),
+        ({"model": "openai/gpt-4o", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False),
+        ({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, False),
+        ({}, False),
+    ],
+)
+def test_is_heuristic_v2_router(litellm_params: Mapping[str, object], expected: bool) -> None:
+    """Only a complexity router whose config selects heuristic_v2 counts toward the license limit."""
+    assert is_heuristic_v2_router(litellm_params) is expected
+
+
+def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_ones() -> None:
+    v2 = {"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}
+    rows: list[Mapping[str, object]] = [
+        {"model_name": "a", "litellm_params": v2},
+        {"model_name": "b", "litellm_params": {"model": "openai/gpt-4o"}},
+        {"model_name": "c", "litellm_params": v2},
+        {"model_name": "d"},
+        {"model_name": "e", "litellm_params": "not a mapping"},
+    ]
+    assert count_heuristic_v2_routers(rows) == 2
+    assert count_heuristic_v2_routers(()) == 0
+
+
+@pytest.mark.parametrize(
+    "held,limit,violates",
+    [
+        (1, 1, False),
+        (2, 1, True),
+        (0, 1, False),
+        (5, None, False),
+        (3, 3, False),
+        (4, 3, True),
+    ],
+)
+def test_heuristic_v2_limit_violation(held: int, limit: int | None, violates: bool) -> None:
+    violation = heuristic_v2_limit_violation(held=held, limit=limit)
+    assert (violation is not None) is violates
+    if violation is not None:
+        assert f"At most {limit} auto-router" in violation
+        assert f"would make {held}" in violation
+        assert "license" not in violation

From 4e0907fb2dd9f0656f2881cecfabff336751feaf Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Thu, 3 Sep 2026 16:40:58 -0400
Subject: [PATCH 40/43] test(router): use an unmapped model so
 get_configured_mode tests do not write into the global cost map

---
 tests/test_litellm/test_router.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index f91920a636b..26c0209bc05 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -12560,7 +12560,7 @@ def test_get_configured_mode_reads_deployment_model_info():
         model_list=[
             {
                 "model_name": "my-tts",
-                "litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
+                "litellm_params": {"model": "openai/some-unmapped-mode-model"},
                 "model_info": {"mode": "audio_speech"},
             }
         ]
@@ -12575,7 +12575,7 @@ def test_get_configured_mode_returns_none_for_unset_blank_or_unknown(model_info)
         model_list=[
             {
                 "model_name": "plain-model",
-                "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"},
+                "litellm_params": {"model": "openai/some-unmapped-mode-model"},
                 "model_info": model_info,
             }
         ]

From 108f55894652b1a995e86a928d8ceb6a66c4c673 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 13:46:33 -0700
Subject: [PATCH 41/43] test: drop the internal patch from the gpt-6-astra
 bridge test

---
 tests/test_litellm/test_main.py | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py
index c9acbe2d884..2df1c2f4ca5 100644
--- a/tests/test_litellm/test_main.py
+++ b/tests/test_litellm/test_main.py
@@ -787,13 +787,11 @@ def test_responses_api_bridge_check_gpt_5_4_tools_plus_reasoning_routes_to_respo
 def test_responses_api_bridge_check_gpt_6_astra_tools_with_default_reasoning_routes_to_responses():
     from litellm.main import responses_api_bridge_check
 
-    with patch("litellm.main._get_model_info_helper") as mock_get_model_info:
-        mock_get_model_info.return_value = {"max_tokens": 128000}
-        model_info, model = responses_api_bridge_check(
-            model="gpt-6-astra",
-            custom_llm_provider="openai",
-            tools=[{"type": "function", "function": {"name": "get_capital"}}],
-        )
+    model_info, model = responses_api_bridge_check(
+        model="gpt-6-astra",
+        custom_llm_provider="openai",
+        tools=[{"type": "function", "function": {"name": "get_capital"}}],
+    )
 
     assert model == "gpt-6-astra"
     assert model_info.get("mode") == "responses"

From aff6b7e21296f4b6b97883ce8302b640992ec0b1 Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Thu, 3 Sep 2026 13:51:04 -0700
Subject: [PATCH 42/43] fix(ui): clear agents when updating team permissions
 (#39600)

Always serialize object_permission.agents and agent_access_groups in
the team update payload so removing the last agent in the dashboard
sends an explicit empty array instead of omitting the key, which the
backend merge treats as no change

Resolves LIT-6861

Co-authored-by: yassin <yassin@berri.ai>
---
 .../src/components/team/TeamInfo.test.tsx     | 45 +++++++++++++++++++
 .../src/components/team/TeamInfo.tsx          |  8 +---
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
index a9c1077e96b..aedc04283f5 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
@@ -1888,6 +1888,8 @@ describe("TeamInfoView - the exact bytes the update call sends", () => {
     mcp_access_groups: [],
     mcp_tool_permissions: {},
     mcp_toolsets: [],
+    agents: [],
+    agent_access_groups: [],
     vector_stores: ["vs-1"],
   };
 
@@ -1908,6 +1910,49 @@ describe("TeamInfoView - the exact bytes the update call sends", () => {
     });
   });
 
+  const openEditorWithAgents = async (user: ReturnType<typeof userEvent.setup>) => {
+    vi.mocked(networking.teamInfoCall).mockResolvedValue(
+      createMockTeamData({
+        models: ["gpt-4"],
+        object_permission: { agents: ["agent-1"], agent_access_groups: ["group-a"] },
+      }),
+    );
+    vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any);
+
+    renderWithProviders(<TeamInfoView {...props} />);
+    await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
+    await user.click(screen.getByRole("tab", { name: "Settings" }));
+    await user.click(await screen.findByRole("button", { name: /edit settings/i }));
+    await screen.findByLabelText("Team Name");
+  };
+
+  it("resends the stored agents and agent_access_groups when the selector is left untouched", async () => {
+    const user = userEvent.setup({ delay: null });
+    await openEditorWithAgents(user);
+
+    const payload = await save(user);
+
+    const objectPermission = wireBody(payload).object_permission as Record<string, unknown>;
+    expect(objectPermission.agents).toStrictEqual(["agent-1"]);
+    expect(objectPermission.agent_access_groups).toStrictEqual(["group-a"]);
+  });
+
+  it("sends empty agents and agent_access_groups arrays after the last agent chip is removed", async () => {
+    const user = userEvent.setup({ delay: null });
+    await openEditorWithAgents(user);
+
+    await user.click(within(screen.getByLabelText("agent-1")).getByRole("button"));
+    await user.click(within(screen.getByLabelText("group:group-a")).getByRole("button"));
+    expect(screen.queryByLabelText("agent-1")).not.toBeInTheDocument();
+    expect(screen.queryByLabelText("group:group-a")).not.toBeInTheDocument();
+
+    const payload = await save(user);
+
+    const objectPermission = wireBody(payload).object_permission as Record<string, unknown>;
+    expect(objectPermission.agents).toStrictEqual([]);
+    expect(objectPermission.agent_access_groups).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 3f6d6a96972..c2b8cd3cc56 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
@@ -863,12 +863,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
         agents: [],
         accessGroups: [],
       };
-      if (agents && agents.length > 0) {
-        updateData.object_permission.agents = agents;
-      }
-      if (agentAccessGroups && agentAccessGroups.length > 0) {
-        updateData.object_permission.agent_access_groups = agentAccessGroups;
-      }
+      updateData.object_permission.agents = agents;
+      updateData.object_permission.agent_access_groups = agentAccessGroups;
       delete values.agents_and_groups;
 
       // Handle vector stores permissions

From eb6c24a2a036ee28aa557e20673d4cf603c5487f Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Thu, 3 Sep 2026 13:53:30 -0700
Subject: [PATCH 43/43] fix(auto_router): bill the routing embedding to the
 caller's key and team (#39532)

* fix(auto_router): bill the routing embedding to the caller's key and team

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

* refactor(auto_router): validate the forwarded caller metadata with a pydantic model

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>
---
 .../internal_call_metadata.py                 |  15 +++
 .../auto_router/auto_router.py                |  59 +++++++--
 .../router_strategy/test_auto_router.py       | 120 ++++++++++++------
 3 files changed, 146 insertions(+), 48 deletions(-)

diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py
index 4d043701f40..87f007ca1d5 100644
--- a/litellm/litellm_core_utils/internal_call_metadata.py
+++ b/litellm/litellm_core_utils/internal_call_metadata.py
@@ -18,9 +18,11 @@ caller's identity metadata, minus two things that must never be forwarded as-is:
 from __future__ import annotations
 
 from collections.abc import Mapping
+from types import MappingProxyType
 from typing import Final
 
 from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, NON_INFERENCE_CALL_TYPES
+from litellm.litellm_core_utils.initialize_dynamic_callback_params import initialize_standard_callback_dynamic_params
 from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, InternalCallOrigin
 
 BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"})
@@ -142,6 +144,19 @@ def forwarded_internal_call_metadata(
     }
 
 
+def parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> Mapping[str, str]:
+    kwargs: Final = request_kwargs or MappingProxyType({})
+    return MappingProxyType(
+        {k: v for k in ("litellm_session_id", "litellm_trace_id") if isinstance(v := kwargs.get(k), str)}
+    )
+
+
+def effective_turn_off_message_logging(request_kwargs: Mapping[str, object] | None) -> bool | None:
+    return initialize_standard_callback_dynamic_params(dict(request_kwargs) if request_kwargs else None).get(
+        "turn_off_message_logging"
+    )
+
+
 def sanitized_forwardable_call_metadata(
     parent_metadata: Mapping[str, object],
     call_origin: InternalCallOrigin,
diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py
index c77745a498d..6b443026f61 100644
--- a/litellm/router_strategy/auto_router/auto_router.py
+++ b/litellm/router_strategy/auto_router/auto_router.py
@@ -2,23 +2,41 @@
 Auto-Routing Strategy that works with a Semantic Router Config
 """
 
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, Final, Optional
 
+from pydantic import BaseModel, ConfigDict
+
 from litellm._logging import verbose_router_logger
 from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.internal_call_metadata import (
+    effective_turn_off_message_logging,
+    forwarded_internal_call_metadata,
+    parent_session_kwargs,
+)
+from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN
 
 if TYPE_CHECKING:
     from semantic_router.routers import SemanticRouter
     from semantic_router.routers.base import Route
 
     from litellm.router import Router
+    from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder
     from litellm.types.router import PreRoutingHookResponse
 else:
     Router = Any
     PreRoutingHookResponse = Any
     Route = Any
     SemanticRouter = Any
+    LiteLLMRouterEncoder = Any
+
+
+class _CallerMetadata(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+
+    metadata: Mapping[str, object] | None = None
+    litellm_metadata: Mapping[str, object] | None = None
 
 
 class AutoRouter(CustomLogger):
@@ -50,6 +68,8 @@ class AutoRouter(CustomLogger):
         """
         from semantic_router.routers import SemanticRouter
 
+        from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder
+
         self.auto_router_config_path: str | None = auto_router_config_path
         self.auto_router_config: str | None = auto_router_config
         self.auto_sync_value = self.DEFAULT_AUTO_SYNC_VALUE
@@ -59,6 +79,11 @@ class AutoRouter(CustomLogger):
         self.embedding_model: str = embedding_model
         self.max_input_chars: int = max_input_chars
         self.litellm_router_instance: Router = litellm_router_instance
+        self.encoder: LiteLLMRouterEncoder = LiteLLMRouterEncoder(
+            litellm_router_instance=litellm_router_instance,
+            model_name=embedding_model,
+            max_input_chars=max_input_chars,
+        )
 
     def _load_semantic_routing_routes(self) -> list[Route]:
         from semantic_router.routers import SemanticRouter
@@ -129,9 +154,6 @@ class AutoRouter(CustomLogger):
         from semantic_router.routers import SemanticRouter
 
         from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages
-        from litellm.router_strategy.auto_router.litellm_encoder import (
-            LiteLLMRouterEncoder,
-        )
         from litellm.types.router import PreRoutingHookResponse
 
         resolved_messages: Final = (
@@ -149,34 +171,47 @@ class AutoRouter(CustomLogger):
             #######################
             routelayer = SemanticRouter(
                 routes=self.loaded_routes,
-                encoder=LiteLLMRouterEncoder(
-                    litellm_router_instance=self.litellm_router_instance,
-                    model_name=self.embedding_model,
-                    max_input_chars=self.max_input_chars,
-                ),
+                encoder=self.encoder,
                 auto_sync=self.auto_sync_value,
             )
             self.routelayer = routelayer
 
         message_content: Final = self._extract_text_from_messages(resolved_messages)
-        route_name: Final = self._matched_route_name(routelayer, message_content)
+        route_name: Final = await self._matched_route_name(routelayer, message_content, request_kwargs)
 
         return PreRoutingHookResponse(
             model=route_name or self.default_model,
             messages=messages,
         )
 
-    def _matched_route_name(self, routelayer: "SemanticRouter", text: str) -> str | None:
+    async def _matched_route_name(
+        self, routelayer: "SemanticRouter", text: str, request_kwargs: Mapping[str, object]
+    ) -> str | None:
         """Name of the route `text` matches, or None when nothing matched or the match failed.
 
-        The route layer embeds `text` to compare it against the routes, and that embedding call can
+        `text` is embedded here rather than by `routelayer(text=...)` so the caller's metadata reaches
+        `aembedding()` and the embedding's spend lands on the key/team that sent the request;
+        SemanticRouter has no way to pass kwargs through to its encoder. That embedding call can
         fail (context limit, timeout, provider error). Choosing a model is a routing decision, so a
         failure here falls back to the default model rather than failing the user's request.
         """
         from semantic_router.schema import RouteChoice
 
         try:
-            route_choice: Final = routelayer(text=text)
+            caller: Final = _CallerMetadata.model_validate(request_kwargs)
+            query_vector: Final = (
+                await self.encoder.aencode_queries(
+                    [text],
+                    metadata=forwarded_internal_call_metadata(caller.metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN),
+                    litellm_metadata=forwarded_internal_call_metadata(
+                        caller.litellm_metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN
+                    ),
+                    proxy_server_request={"body": {"model": self.embedding_model, "input": [text]}},
+                    turn_off_message_logging=effective_turn_off_message_logging(request_kwargs),
+                    **parent_session_kwargs(request_kwargs),
+                )
+            )[0]
+            route_choice: Final = await routelayer.acall(vector=query_vector)
         except Exception as e:  # noqa: BLE001 -- the embedding call behind the route layer can fail many ways (context limit, timeout, provider/network error); none of them may fail the request
             verbose_router_logger.warning(
                 "AutoRouter: semantic routing failed (%s), falling back to default model %s", e, self.default_model
diff --git a/tests/test_litellm/router_strategy/test_auto_router.py b/tests/test_litellm/router_strategy/test_auto_router.py
index 36199b45847..123ada83ca4 100644
--- a/tests/test_litellm/router_strategy/test_auto_router.py
+++ b/tests/test_litellm/router_strategy/test_auto_router.py
@@ -330,36 +330,46 @@ ROUTER_CONFIG: Final = json.dumps(
 )
 
 
-class FailingRouteLayer:
-    """Route layer whose embedding call fails, as it does when the prompt exceeds the encoder's window."""
-
-    def __call__(self, text: str) -> Any:
-        raise ValueError(
-            "Internal_litellm_router API call failed. Error: litellm.InternalServerError: "
-            "input is too large to process. increase the physical batch size"
-        )
-
-
 class FixedRouteLayer:
-    """Route layer that returns whatever the test tells it to, recording the text it was asked about."""
+    """Route layer that returns whatever the test tells it to for the query vector it is handed."""
 
     def __init__(self, route_choice: Any) -> None:
         self.route_choice = route_choice
-        self.seen_text: str | None = None
 
-    def __call__(self, text: str) -> Any:
-        self.seen_text = text
+    async def acall(self, vector: Any) -> Any:
         return self.route_choice
 
 
+def _embedding_response(input: List[str]) -> Any:
+    import litellm
+
+    return litellm.EmbeddingResponse(
+        data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))]
+    )
+
+
 class StubEmbeddingRouter:
-    """Stands in for the LiteLLM Router when the route index has to be built for real."""
+    """Stands in for the LiteLLM Router, recording the text and kwargs each query embedding was made with."""
+
+    def __init__(self) -> None:
+        self.seen_text: str | None = None
+        self.aembedding_kwargs: Dict[str, Any] | None = None
 
     def embedding(self, input: List[str], model: str, **kwargs: Any) -> Any:
-        import litellm
+        return _embedding_response(input)
 
-        return litellm.EmbeddingResponse(
-            data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))]
+    async def aembedding(self, input: List[str], model: str, **kwargs: Any) -> Any:
+        self.seen_text = input[0]
+        self.aembedding_kwargs = kwargs
+        return _embedding_response(input)
+
+
+class FailingEmbeddingRouter(StubEmbeddingRouter):
+    """Router whose query embedding fails, as it does when the prompt exceeds the encoder's window."""
+
+    async def aembedding(self, input: List[str], model: str, **kwargs: Any) -> Any:
+        raise ValueError(
+            "litellm.InternalServerError: input is too large to process. increase the physical batch size"
         )
 
 
@@ -369,7 +379,7 @@ def _auto_router(routelayer: Any, litellm_router_instance: Any = None, **kwargs:
         auto_router_config=ROUTER_CONFIG,
         default_model="fallback-model",
         embedding_model="text-embedding-3-small",
-        litellm_router_instance=litellm_router_instance or MagicMock(),
+        litellm_router_instance=litellm_router_instance or StubEmbeddingRouter(),
         **kwargs,
     )
     auto_router.routelayer = routelayer
@@ -381,7 +391,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel:
 
     @pytest.mark.asyncio
     async def test_should_fall_back_to_default_model_when_the_embedding_call_fails(self):
-        auto_router: Final = _auto_router(FailingRouteLayer())
+        auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=FailingEmbeddingRouter())
 
         result: Final = await auto_router.async_pre_routing_hook(
             model="my-auto-router",
@@ -440,8 +450,8 @@ class TestAutoRouterAlwaysResolvesARoutableModel:
     async def test_should_still_route_to_the_matched_route_when_one_matches(self):
         from semantic_router.schema import RouteChoice
 
-        layer: Final = FixedRouteLayer(RouteChoice(name="code-model"))
-        auto_router: Final = _auto_router(layer)
+        router: Final = StubEmbeddingRouter()
+        auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router)
 
         result: Final = await auto_router.async_pre_routing_hook(
             model="my-auto-router",
@@ -451,7 +461,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel:
 
         assert result is not None
         assert result.model == "code-model"
-        assert layer.seen_text == "fix this stack trace"
+        assert router.seen_text == "fix this stack trace"
 
 
 class TestAutoRouterEmbeddingInputCap:
@@ -483,8 +493,8 @@ class TestAutoRouterRoutesResponsesApiInput:
     async def test_should_route_a_string_input_when_messages_is_none(self):
         from semantic_router.schema import RouteChoice
 
-        layer: Final = FixedRouteLayer(RouteChoice(name="code-model"))
-        auto_router: Final = _auto_router(layer)
+        router: Final = StubEmbeddingRouter()
+        auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router)
 
         result: Final = await auto_router.async_pre_routing_hook(
             model="my-auto-router",
@@ -498,14 +508,14 @@ class TestAutoRouterRoutesResponsesApiInput:
         assert result is not None
         assert result.model == "code-model"
         assert result.messages is None
-        assert layer.seen_text == "fix this stack trace"
+        assert router.seen_text == "fix this stack trace"
 
     @pytest.mark.asyncio
     async def test_should_route_a_list_input_with_instructions_when_messages_is_none(self):
         from semantic_router.schema import RouteChoice
 
-        layer: Final = FixedRouteLayer(RouteChoice(name="code-model"))
-        auto_router: Final = _auto_router(layer)
+        router: Final = StubEmbeddingRouter()
+        auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router)
 
         result: Final = await auto_router.async_pre_routing_hook(
             model="my-auto-router",
@@ -525,13 +535,13 @@ class TestAutoRouterRoutesResponsesApiInput:
 
         assert result is not None
         assert result.model == "code-model"
-        assert layer.seen_text is not None
-        assert "fix this stack trace" in layer.seen_text
+        assert router.seen_text is not None
+        assert "fix this stack trace" in router.seen_text
 
     @pytest.mark.asyncio
     async def test_should_skip_routing_when_neither_messages_nor_input_is_present(self):
-        layer: Final = FixedRouteLayer(None)
-        auto_router: Final = _auto_router(layer)
+        router: Final = StubEmbeddingRouter()
+        auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router)
 
         result: Final = await auto_router.async_pre_routing_hook(
             model="my-auto-router",
@@ -540,12 +550,12 @@ class TestAutoRouterRoutesResponsesApiInput:
         )
 
         assert result is None
-        assert layer.seen_text is None
+        assert router.seen_text is None
 
     @pytest.mark.asyncio
     async def test_should_keep_routing_an_empty_messages_list_to_the_default_model(self):
-        layer: Final = FixedRouteLayer(None)
-        auto_router: Final = _auto_router(layer)
+        router: Final = StubEmbeddingRouter()
+        auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router)
 
         result: Final = await auto_router.async_pre_routing_hook(
             model="my-auto-router",
@@ -555,4 +565,42 @@ class TestAutoRouterRoutesResponsesApiInput:
 
         assert result is not None
         assert result.model == "fallback-model"
-        assert layer.seen_text == ""
+        assert router.seen_text == ""
+
+
+class TestAutoRouterAttributesItsEmbeddingSpend:
+    """The query embedding is billed to the key that sent the request, like any other call it made."""
+
+    @pytest.mark.asyncio
+    async def test_should_forward_the_callers_identity_to_the_query_embedding_minus_its_budget_reservation(self):
+        from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
+
+        router: Final = StubEmbeddingRouter()
+        auto_router: Final = _auto_router(None, litellm_router_instance=router)
+        request_kwargs: Final = {
+            "metadata": {
+                "user_api_key": "hashed-key",
+                "user_api_key_team_id": "team-1",
+                "user_api_key_budget_reservation": {"reservation_id": "r-1"},
+            },
+            "litellm_session_id": "session-1",
+        }
+
+        result: Final = await auto_router.async_pre_routing_hook(
+            model="my-auto-router",
+            request_kwargs=request_kwargs,
+            messages=[{"role": "user", "content": "fix this stack trace"}],
+        )
+
+        assert result is not None
+        assert router.seen_text == "fix this stack trace"
+        assert router.aembedding_kwargs is not None
+        forwarded: Final = router.aembedding_kwargs["metadata"]
+        assert forwarded["user_api_key"] == "hashed-key"
+        assert forwarded["user_api_key_team_id"] == "team-1"
+        assert forwarded[INTERNAL_CALL_ORIGIN_METADATA_KEY] == "autorouter_classifier"
+        assert "user_api_key_budget_reservation" not in forwarded
+        assert router.aembedding_kwargs["litellm_session_id"] == "session-1"
+        assert router.aembedding_kwargs["proxy_server_request"] == {
+            "body": {"model": "text-embedding-3-small", "input": ["fix this stack trace"]}
+        }